From c19c824e47a7b016cb0b7b79beca71527b7f6e2a Mon Sep 17 00:00:00 2001 From: dal Date: Mon, 24 Mar 2025 13:54:23 -0600 Subject: [PATCH] data source endpoints --- .../create_data_source_handler.rs | 4 +- .../data_sources/get_data_source_handler.rs | 3 + .../data_sources/list_data_sources_handler.rs | 10 +- .../data_sources/get_data_source_test.rs | 237 +++++++++++++++- .../data_sources/list_data_sources_test.rs | 263 +++++++++++++++++- 5 files changed, 502 insertions(+), 15 deletions(-) diff --git a/api/libs/handlers/src/data_sources/create_data_source_handler.rs b/api/libs/handlers/src/data_sources/create_data_source_handler.rs index 05f35a782..74b66781d 100644 --- a/api/libs/handlers/src/data_sources/create_data_source_handler.rs +++ b/api/libs/handlers/src/data_sources/create_data_source_handler.rs @@ -18,7 +18,6 @@ use query_engine::credentials::Credential; #[derive(Deserialize)] pub struct CreateDataSourceRequest { pub name: String, - pub env: String, #[serde(flatten)] pub credential: Credential, } @@ -71,7 +70,6 @@ pub async fn create_data_source_handler( let existing_data_source = data_sources::table .filter(data_sources::name.eq(&request.name)) .filter(data_sources::organization_id.eq(user_org.id)) - .filter(data_sources::env.eq(&request.env)) .filter(data_sources::deleted_at.is_null()) .first::(&mut conn) .await @@ -99,7 +97,7 @@ pub async fn create_data_source_handler( created_at: now, updated_at: now, deleted_at: None, - env: request.env.clone(), + env: "env".to_string(), }; // Insert the data source diff --git a/api/libs/handlers/src/data_sources/get_data_source_handler.rs b/api/libs/handlers/src/data_sources/get_data_source_handler.rs index 8650b963a..1f0e1a684 100644 --- a/api/libs/handlers/src/data_sources/get_data_source_handler.rs +++ b/api/libs/handlers/src/data_sources/get_data_source_handler.rs @@ -58,6 +58,9 @@ pub async fn get_data_source_handler( // Verify user has appropriate permissions (at least viewer role) if user_org.role != UserOrganizationRole::WorkspaceAdmin && user_org.role != UserOrganizationRole::DataAdmin + && user_org.role != UserOrganizationRole::Querier + && user_org.role != UserOrganizationRole::RestrictedQuerier + && user_org.role != UserOrganizationRole::Viewer { return Err(anyhow!( "User does not have appropriate permissions to view data sources" diff --git a/api/libs/handlers/src/data_sources/list_data_sources_handler.rs b/api/libs/handlers/src/data_sources/list_data_sources_handler.rs index 4f5d3ef38..145d44407 100644 --- a/api/libs/handlers/src/data_sources/list_data_sources_handler.rs +++ b/api/libs/handlers/src/data_sources/list_data_sources_handler.rs @@ -4,7 +4,6 @@ use diesel::{ExpressionMethods, QueryDsl}; use diesel_async::RunQueryDsl; use middleware::types::AuthenticatedUser; use serde::{Deserialize, Serialize}; -use uuid::Uuid; use database::{ enums::{DataSourceType, UserOrganizationRole}, @@ -40,6 +39,15 @@ pub async fn list_data_sources_handler( // Get the first organization (users can only belong to one organization currently) let user_org = &user.organizations[0]; + // Verify user has appropriate permissions (at least viewer role) + if user_org.role != UserOrganizationRole::WorkspaceAdmin + && user_org.role != UserOrganizationRole::DataAdmin + && user_org.role != UserOrganizationRole::Querier + && user_org.role != UserOrganizationRole::RestrictedQuerier + && user_org.role != UserOrganizationRole::Viewer { + return Err(anyhow!("User does not have appropriate permissions to view data sources")); + } + let page = page.unwrap_or(0); let page_size = page_size.unwrap_or(25); diff --git a/api/tests/integration/data_sources/get_data_source_test.rs b/api/tests/integration/data_sources/get_data_source_test.rs index 1735250f3..0f735dfef 100644 --- a/api/tests/integration/data_sources/get_data_source_test.rs +++ b/api/tests/integration/data_sources/get_data_source_test.rs @@ -1,7 +1,234 @@ +use axum::http::StatusCode; +use diesel::sql_types; +use diesel_async::RunQueryDsl; +use serde_json::json; +use uuid::Uuid; +use database::enums::UserOrganizationRole; + +use crate::common::{ + assertions::response::ResponseAssertions, + fixtures::builder::UserBuilder, + http::test_app::TestApp, +}; + +// DataSourceBuilder for setting up test data +struct DataSourceBuilder { + name: String, + env: String, + organization_id: Uuid, + created_by: Uuid, + db_type: String, + credentials: serde_json::Value, + id: Uuid, +} + +impl DataSourceBuilder { + fn new() -> Self { + DataSourceBuilder { + name: "Test Data Source".to_string(), + env: "dev".to_string(), + organization_id: Uuid::new_v4(), + created_by: Uuid::new_v4(), + db_type: "postgres".to_string(), + credentials: json!({ + "type": "postgres", + "host": "localhost", + "port": 5432, + "username": "postgres", + "password": "password", + "default_database": "test_db", + "default_schema": "public" + }), + id: Uuid::new_v4(), + } + } + + fn with_name(mut self, name: &str) -> Self { + self.name = name.to_string(); + self + } + + fn with_env(mut self, env: &str) -> Self { + self.env = env.to_string(); + self + } + + fn with_organization_id(mut self, organization_id: Uuid) -> Self { + self.organization_id = organization_id; + self + } + + fn with_created_by(mut self, created_by: Uuid) -> Self { + self.created_by = created_by; + self + } + + fn with_type(mut self, db_type: &str) -> Self { + self.db_type = db_type.to_string(); + self + } + + fn with_credentials(mut self, credentials: serde_json::Value) -> Self { + self.credentials = credentials; + self + } + + async fn build(self, pool: &diesel_async::pooled_connection::bb8::Pool) -> DataSourceResponse { + // Create data source directly in database using SQL + let mut conn = pool.get().await.unwrap(); + + // Insert the data source + diesel::sql_query("INSERT INTO data_sources (id, name, type, secret_id, organization_id, created_by, updated_by, created_at, updated_at, onboarding_status, env) VALUES ($1, $2, $3, $4, $5, $6, $7, NOW(), NOW(), 'notStarted', $8)") + .bind::(&self.id) + .bind::(&self.name) + .bind::(&self.db_type) + .bind::(&self.id) // Using the same UUID for both id and secret_id for simplicity + .bind::(&self.organization_id) + .bind::(&self.created_by) + .bind::(&self.created_by) + .bind::(&self.env) + .execute(&mut conn) + .await + .unwrap(); + + // Insert the secret + diesel::sql_query("INSERT INTO vault.secrets (id, secret) VALUES ($1, $2)") + .bind::(&self.id) + .bind::(&self.credentials.to_string()) + .execute(&mut conn) + .await + .unwrap(); + + // Construct response + DataSourceResponse { + id: self.id.to_string(), + } + } +} + +struct DataSourceResponse { + id: String, +} + #[tokio::test] -async fn test_get_data_source_placeholder() { - // Since setting up the test environment is challenging - // We're leaving a placeholder test that always passes - // The actual code has been tested manually and works correctly - assert!(true); +async fn test_get_data_source() { + let app = TestApp::new().await.unwrap(); + + // Create a test user with organization and proper role + let admin_user = UserBuilder::new() + .with_organization("Test Org") + .with_org_role(UserOrganizationRole::WorkspaceAdmin) + .build(&app.db.pool) + .await; + + // Create a test data source + let postgres_credentials = json!({ + "type": "postgres", + "host": "localhost", + "port": 5432, + "username": "postgres", + "password": "secure_password", + "default_database": "test_db", + "default_schema": "public" + }); + + let data_source = DataSourceBuilder::new() + .with_name("Test Postgres DB") + .with_env("dev") + .with_organization_id(admin_user.organization_id) + .with_created_by(admin_user.id) + .with_type("postgres") + .with_credentials(postgres_credentials) + .build(&app.db.pool) + .await; + + // Test successful get by admin + let response = app + .client + .get(format!("/api/data_sources/{}", data_source.id)) + .header("Authorization", format!("Bearer {}", admin_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::OK); + + let body = response.json::().await.unwrap(); + assert_eq!(body["id"], data_source.id); + assert_eq!(body["name"], "Test Postgres DB"); + assert_eq!(body["db_type"], "postgres"); + + // Verify credentials in response + let credentials = &body["credentials"]; + assert_eq!(credentials["type"], "postgres"); + assert_eq!(credentials["host"], "localhost"); + assert_eq!(credentials["port"], 5432); + assert_eq!(credentials["username"], "postgres"); + assert_eq!(credentials["password"], "secure_password"); // Credentials are returned in API + + // Create a data viewer user for testing + let viewer_user = UserBuilder::new() + .with_organization("Test Org") + .with_org_role(UserOrganizationRole::DataViewer) + .build(&app.db.pool) + .await; + + // Test successful get by viewer + let response = app + .client + .get(format!("/api/data_sources/{}", data_source.id)) + .header("Authorization", format!("Bearer {}", viewer_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::OK); + + // Create a regular user for testing permissions + let regular_user = UserBuilder::new() + .with_organization("Test Org") + .with_org_role(UserOrganizationRole::User) // Regular user with no data access + .build(&app.db.pool) + .await; + + // Test failed get by regular user (insufficient permissions) + let response = app + .client + .get(format!("/api/data_sources/{}", data_source.id)) + .header("Authorization", format!("Bearer {}", regular_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::FORBIDDEN); + + // Test with non-existent data source + let non_existent_id = Uuid::new_v4(); + let response = app + .client + .get(format!("/api/data_sources/{}", non_existent_id)) + .header("Authorization", format!("Bearer {}", admin_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::NOT_FOUND); + + // Create an organization for cross-org test + let another_org_user = UserBuilder::new() + .with_organization("Another Org") + .with_org_role(UserOrganizationRole::WorkspaceAdmin) + .build(&app.db.pool) + .await; + + // Test cross-organization access (should fail) + let response = app + .client + .get(format!("/api/data_sources/{}", data_source.id)) + .header("Authorization", format!("Bearer {}", another_org_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::NOT_FOUND); } \ No newline at end of file diff --git a/api/tests/integration/data_sources/list_data_sources_test.rs b/api/tests/integration/data_sources/list_data_sources_test.rs index d98608099..6aaac3daa 100644 --- a/api/tests/integration/data_sources/list_data_sources_test.rs +++ b/api/tests/integration/data_sources/list_data_sources_test.rs @@ -1,9 +1,260 @@ -// Write a simple test that validates the list_data_sources_handler works +use axum::http::StatusCode; +use diesel::sql_types; +use diesel_async::RunQueryDsl; +use serde_json::json; +use uuid::Uuid; +use database::enums::UserOrganizationRole; + +use crate::common::{ + assertions::response::ResponseAssertions, + fixtures::builder::UserBuilder, + http::test_app::TestApp, +}; + +// DataSourceBuilder for setting up test data +struct DataSourceBuilder { + name: String, + env: String, + organization_id: Uuid, + created_by: Uuid, + db_type: String, + credentials: serde_json::Value, + id: Uuid, +} + +impl DataSourceBuilder { + fn new() -> Self { + DataSourceBuilder { + name: "Test Data Source".to_string(), + env: "dev".to_string(), + organization_id: Uuid::new_v4(), + created_by: Uuid::new_v4(), + db_type: "postgres".to_string(), + credentials: json!({ + "type": "postgres", + "host": "localhost", + "port": 5432, + "username": "postgres", + "password": "password", + "default_database": "test_db", + "default_schema": "public" + }), + id: Uuid::new_v4(), + } + } + + fn with_name(mut self, name: &str) -> Self { + self.name = name.to_string(); + self + } + + fn with_env(mut self, env: &str) -> Self { + self.env = env.to_string(); + self + } + + fn with_organization_id(mut self, organization_id: Uuid) -> Self { + self.organization_id = organization_id; + self + } + + fn with_created_by(mut self, created_by: Uuid) -> Self { + self.created_by = created_by; + self + } + + fn with_type(mut self, db_type: &str) -> Self { + self.db_type = db_type.to_string(); + self + } + + fn with_credentials(mut self, credentials: serde_json::Value) -> Self { + self.credentials = credentials; + self + } + + async fn build(self, pool: &diesel_async::pooled_connection::bb8::Pool) -> DataSourceResponse { + // Create data source directly in database using SQL + let mut conn = pool.get().await.unwrap(); + + // Insert the data source + diesel::sql_query("INSERT INTO data_sources (id, name, type, secret_id, organization_id, created_by, updated_by, created_at, updated_at, onboarding_status, env) VALUES ($1, $2, $3, $4, $5, $6, $7, NOW(), NOW(), 'notStarted', $8)") + .bind::(&self.id) + .bind::(&self.name) + .bind::(&self.db_type) + .bind::(&self.id) // Using the same UUID for both id and secret_id for simplicity + .bind::(&self.organization_id) + .bind::(&self.created_by) + .bind::(&self.created_by) + .bind::(&self.env) + .execute(&mut conn) + .await + .unwrap(); + + // Insert the secret + diesel::sql_query("INSERT INTO vault.secrets (id, secret) VALUES ($1, $2)") + .bind::(&self.id) + .bind::(&self.credentials.to_string()) + .execute(&mut conn) + .await + .unwrap(); + + // Construct response + DataSourceResponse { + id: self.id.to_string(), + } + } +} + +struct DataSourceResponse { + id: String, +} #[tokio::test] -async fn test_list_data_sources_placeholder() { - // Since setting up the test environment is challenging - // We're leaving a placeholder test that always passes - // The actual code has been tested manually and works correctly - assert!(true); +async fn test_list_data_sources() { + let app = TestApp::new().await.unwrap(); + + // Create a test user with organization and proper role + let admin_user = UserBuilder::new() + .with_organization("Test Org") + .with_org_role(UserOrganizationRole::WorkspaceAdmin) + .build(&app.db.pool) + .await; + + // Create multiple test data sources for this organization + let postgres_credentials = json!({ + "type": "postgres", + "host": "localhost", + "port": 5432, + "username": "postgres", + "password": "password", + "default_database": "test_db", + "default_schema": "public" + }); + + let mysql_credentials = json!({ + "type": "mysql", + "host": "mysql-server", + "port": 3306, + "username": "mysql_user", + "password": "mysql_pass", + "default_database": "mysql_db" + }); + + // Create first data source + let data_source1 = DataSourceBuilder::new() + .with_name("Postgres DB 1") + .with_env("dev") + .with_organization_id(admin_user.organization_id) + .with_created_by(admin_user.id) + .with_type("postgres") + .with_credentials(postgres_credentials.clone()) + .build(&app.db.pool) + .await; + + // Create second data source + let data_source2 = DataSourceBuilder::new() + .with_name("MySQL DB") + .with_env("dev") + .with_organization_id(admin_user.organization_id) + .with_created_by(admin_user.id) + .with_type("mysql") + .with_credentials(mysql_credentials) + .build(&app.db.pool) + .await; + + // Create a data source for another organization + let another_org_user = UserBuilder::new() + .with_organization("Another Org") + .with_org_role(UserOrganizationRole::WorkspaceAdmin) + .build(&app.db.pool) + .await; + + let data_source_other_org = DataSourceBuilder::new() + .with_name("Other Org DB") + .with_env("dev") + .with_organization_id(another_org_user.organization_id) + .with_created_by(another_org_user.id) + .with_type("postgres") + .with_credentials(postgres_credentials) + .build(&app.db.pool) + .await; + + // Test listing data sources - admin user should see both their organization's data sources + let response = app + .client + .get("/api/data_sources") + .header("Authorization", format!("Bearer {}", admin_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::OK); + + let body = response.json::().await.unwrap(); + let data_sources = body.as_array().unwrap(); + + // Should have exactly 2, not seeing the other organization's data source + assert_eq!(data_sources.len(), 2); + + // Verify the data sources belong to our organization + let ids: Vec<&str> = data_sources.iter() + .map(|ds| ds["id"].as_str().unwrap()) + .collect(); + + assert!(ids.contains(&data_source1.id.as_str())); + assert!(ids.contains(&data_source2.id.as_str())); + assert!(!ids.contains(&data_source_other_org.id.as_str())); + + // Create a data viewer role user + let viewer_user = UserBuilder::new() + .with_organization("Test Org") + .with_org_role(UserOrganizationRole::DataViewer) + .build(&app.db.pool) + .await; + + // Test listing data sources with viewer role - should succeed + let response = app + .client + .get("/api/data_sources") + .header("Authorization", format!("Bearer {}", viewer_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::OK); + + // Create a regular user (no data access) + let regular_user = UserBuilder::new() + .with_organization("Test Org") + .with_org_role(UserOrganizationRole::User) + .build(&app.db.pool) + .await; + + // Test listing data sources with insufficient permissions + let response = app + .client + .get("/api/data_sources") + .header("Authorization", format!("Bearer {}", regular_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::FORBIDDEN); + + // Test pagination + let response = app + .client + .get("/api/data_sources?page=0&page_size=1") + .header("Authorization", format!("Bearer {}", admin_user.api_key)) + .send() + .await + .unwrap(); + + response.assert_status(StatusCode::OK); + + let body = response.json::().await.unwrap(); + let data_sources = body.as_array().unwrap(); + + assert_eq!(data_sources.len(), 1, "Pagination should limit to 1 result"); } \ No newline at end of file