Merge pull request #109 from buster-so/dal/database_identifier_lineage

Fix: database_identifier_fix

api/migrations/2025-02-06-220857_add_database_identifier_to_datasets/down.sql

@@ -0,0 +1,2 @@
+-- This file should undo anything in `up.sql`
+ALTER TABLE datasets DROP COLUMN database_identifier;

api/migrations/2025-02-06-220857_add_database_identifier_to_datasets/up.sql

@@ -0,0 +1,2 @@
+-- Your SQL goes here
+ALTER TABLE datasets ADD COLUMN database_identifier TEXT NULL;

api/src/database/models.rs

@@ -189,6 +189,7 @@ pub struct Dataset {
     pub deleted_at: Option<DateTime<Utc>>,
     pub model: Option<String>,
     pub yml_file: Option<String>,
+    pub database_identifier: Option<String>,
 }
 
 #[derive(Insertable, Queryable, Associations, Debug)]

api/src/database/schema.rs

@@ -252,6 +252,7 @@ diesel::table! {
         deleted_at -> Nullable<Timestamptz>,
         model -> Nullable<Text>,
         yml_file -> Nullable<Text>,
+        database_identifier -> Nullable<Text>,
     }
 }
 

api/src/routes/rest/routes/datasets/deploy_datasets.rs

@@ -48,6 +48,7 @@ pub struct FullDeployDatasetsRequest {
     pub entity_relationships: Option<Vec<DeployDatasetsEntityRelationshipsRequest>>,
     pub columns: Vec<DeployDatasetsColumnsRequest>,
     pub yml_file: Option<String>,
+    pub database_identifier: Option<String>,
 }
 
 #[derive(Debug, Deserialize)]
@@ -238,6 +239,7 @@ async fn process_deploy_request(
                     entity_relationships: Some(entity_relationships),
                     columns,
                     yml_file: Some(yml.clone()),
+                    database_identifier: None,
                 });
             }
 
@@ -323,6 +325,7 @@ async fn deploy_datasets_handler(
             organization_id,
             model: req.model.clone(),
             yml_file: req.yml_file.clone(),
+            database_identifier: req.database_identifier.clone(),
         };
 
         match req.id {

api/src/routes/rest/routes/datasets/post_dataset.rs

@@ -127,6 +127,7 @@ async fn post_dataset_handler(
         deleted_at: None,
         model: None,
         yml_file: None,
+        database_identifier: None,
     };
 
     diesel::insert_into(datasets::table)

api/src/routes/ws/datasets/post_dataset.rs

@@ -235,6 +235,7 @@ async fn create_dataset(user_id: &Uuid, name: &String, data_source_id: &Uuid) ->
         organization_id: user_org_id,
         yml_file: None,
         model: None,
+        database_identifier: None,
     };
 
     let mut conn = match get_pg_pool().get().await {

api/src/routes/ws/threads_and_messages/post_thread/post_thread.rs

@@ -929,9 +929,16 @@ fn create_dataset_ddl(dataset: &Dataset, dataset_columns: &Vec<DatasetColumn>) -
     if let Some(when_to_use) = &dataset.when_to_use {
         ddl.push_str(&format!("    -- Description: {}\n", when_to_use));
     }
+
+    let schema_identifier = if let Some(db_id) = &dataset.database_identifier {
+        format!("{}.{}", db_id, dataset.schema)
+    } else {
+        dataset.schema.clone()
+    };
+
     ddl.push_str(&format!(
         "    CREATE TABLE {}.{} (\n",
-        dataset.schema, dataset.database_name
+        schema_identifier, dataset.database_name
     ));
 
     // Add columns

api/src/utils/prompts/generate_sql_prompts/sql_gen_prompt.rs

@@ -39,6 +39,7 @@ Please output the SQL delimited in ```sql tags.
 - Always order dates in ascending order.
 - When working with time series data, always return a date field.
 - You must use the schema when referencing tables. Like this pattern <SCHEMA_NAME>.<TABLE_NAME>
+- Pay attention to the database identifier. It may be used to reference across multiple databases. 
 - Never use the 'SELECT *'  or 'SELECT COUNT(*)' command.  You must select the columns you want to see/use.
 - Users may mention formatting or charting.  Although this task is specific to SQL generation, the user is referring to future steps for visualization.
 - A request for a line chart should default to using a date-related field unless the user specifies otherwise or it is not available.

api/src/utils/query_engine/import_datasets.rs

@@ -82,6 +82,7 @@ async fn create_datasets(
             deleted_at: None,
             yml_file: None,
             model: None,
+            database_identifier: None,
         })
         .collect::<Vec<Dataset>>();