diff --git a/backend/windmill-common/src/workspaces.rs b/backend/windmill-common/src/workspaces.rs index 75d731c3e0..a1eafa0549 100644 --- a/backend/windmill-common/src/workspaces.rs +++ b/backend/windmill-common/src/workspaces.rs @@ -170,6 +170,49 @@ pub enum DataTableCatalogResourceType { Instance, } +pub async fn get_datatable_resource_from_db_unchecked( + db: &DB, + w_id: &str, + name: &str, +) -> Result { + let datatable = sqlx::query_scalar!( + r#" + SELECT ws.datatable->'datatables'->$2 AS config + FROM workspace_settings ws + WHERE ws.workspace_id = $1 + "#, + &w_id, + name + ) + .fetch_one(db) + .await + .map_err(|err| Error::internal_err(format!("getting datatable {name}: {err}")))? + .ok_or_else(|| Error::internal_err(format!("datatable {name} not found")))?; + let datatable = serde_json::from_value::(datatable)?; + + let db_resource = if datatable.database.resource_type == DataTableCatalogResourceType::Instance + { + let pg_creds = parse_postgres_url(&get_database_url().await?.as_str().await)?; + json!({ + "dbname": datatable.database.resource_path, + "host": pg_creds.host, + "port": pg_creds.port, + "user": "custom_instance_user", + "sslmode": pg_creds.ssl_mode, + "password": get_custom_pg_instance_password(&db).await?, + }) + } else { + transform_json_unchecked( + &serde_json::Value::String(format!("$res:{}", datatable.database.resource_path)), + w_id, + db, + ) + .await? + }; + + Ok(db_resource) +} + #[derive(Deserialize, Serialize, Debug)] pub struct Ducklake { pub catalog: DucklakeCatalog, diff --git a/backend/windmill-worker/src/agent_workers.rs b/backend/windmill-worker/src/agent_workers.rs index 0aaecdbe1e..9bf0ea1834 100644 --- a/backend/windmill-worker/src/agent_workers.rs +++ b/backend/windmill-worker/src/agent_workers.rs @@ -64,4 +64,18 @@ pub async fn get_ducklake_from_agent_http( .await } +#[allow(dead_code)] +pub async fn get_datatable_resource_from_agent_http( + client: &HttpClient, + name: &str, + w_id: &str, +) -> anyhow::Result { + client + .get(&format!( + "/api/w/{}/agent_workers/get_datatable_resource/{}", + w_id, &name + )) + .await +} + pub const UPDATE_PING_URL: &str = "/api/agent_workers/update_ping"; diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 15cc6ccf30..0013f24a87 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -31,6 +31,7 @@ use windmill_common::s3_helpers::convert_json_line_stream; use windmill_common::worker::{ to_raw_value, Connection, SqlResultCollectionStrategy, CLOUD_HOSTED, }; +use windmill_common::workspaces::get_datatable_resource_from_db_unchecked; use windmill_parser::{Arg, Typ}; use windmill_parser_sql::{ parse_db_resource, parse_pg_statement_arg_indices, parse_pgsql_sig, parse_s3_mode, @@ -38,6 +39,7 @@ use windmill_parser_sql::{ }; use windmill_queue::{CanceledBy, MiniPulledJob}; +use crate::agent_workers::get_datatable_resource_from_agent_http; use crate::common::{ build_args_values, get_reserved_variables, s3_mode_args_to_worker_data, sizeof_val, OccupancyMetrics, S3ModeWorkerData, @@ -214,7 +216,26 @@ pub async fn do_postgresql( .await?, ) } else { - pg_args.get("database").cloned() + println!( + "PG ARGS BEFORE DATATABLE CHECK: {:?}", + pg_args.get("database") + ); + match pg_args.get("database").cloned() { + Some(Value::String(db_str)) if db_str.starts_with("datatable://") => { + let db_str = db_str.trim_start_matches("datatable://"); + Some(match conn { + Connection::Http(client) => { + get_datatable_resource_from_agent_http(client, &db_str, &job.workspace_id) + .await? + } + Connection::Sql(db) => { + get_datatable_resource_from_db_unchecked(db, &job.workspace_id, &db_str) + .await? + } + }) + } + database => database, + } }; let database = if let Some(db) = db_arg {