feat: add get_datatable_full_schema endpoint and snapshot schema on fork

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-09-06 08:01:35 +00:00 · 2026-03-25 13:11:58 +01:00
parent 5bb67954bd
commit 4a91e30187
5 changed files with 346 additions and 5 deletions
@@ -153,6 +153,10 @@ pub fn workspaced_service() -> Router {
        .route("/compare/:target_workspace_id", get(compare_workspaces))
        .route("/fork_pg_database", post(fork_pg_database))
        .route("/export_pg_schema", post(export_pg_schema))
+        .route(
+            "/get_datatable_full_schema",
+            post(get_datatable_full_schema),
+        )
        .route("/protection_rules", get(list_protection_rules))
        .route("/protection_rules", post(create_protection_rule))
        .route(
@@ -1571,6 +1575,233 @@ async fn export_pg_schema(
    pg_dump_database(&pg, true).await
 }

+#[derive(Deserialize)]
+struct GetDatatableFullSchemaRequest {
+    source: String,
+}
+
+#[derive(Serialize)]
+struct TableEditorValuesColumn {
+    name: String,
+    datatype: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    primary_key: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    default_value: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    nullable: Option<bool>,
+}
+
+#[derive(Serialize)]
+struct TableEditorForeignKey {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    target_table: Option<String>,
+    columns: Vec<ForeignKeyColumnPair>,
+    on_delete: String,
+    on_update: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    fk_constraint_name: Option<String>,
+}
+
+#[derive(Serialize)]
+struct ForeignKeyColumnPair {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    source_column: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    target_column: Option<String>,
+}
+
+#[derive(Serialize)]
+struct TableEditorValues {
+    name: String,
+    columns: Vec<TableEditorValuesColumn>,
+    foreign_keys: Vec<TableEditorForeignKey>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pk_constraint_name: Option<String>,
+}
+
+async fn get_datatable_full_schema(
+    ApiAuthed { is_admin, username, .. }: ApiAuthed,
+    Extension(db): Extension<DB>,
+    Path(w_id): Path<String>,
+    Json(req): Json<GetDatatableFullSchemaRequest>,
+) -> JsonResult<HashMap<String, HashMap<String, TableEditorValues>>> {
+    require_admin(is_admin, &username)?;
+
+    let pg = resolve_pg_source(&db, &w_id, &req.source).await?;
+    let (client, connection) = pg.connect().await?;
+    let join_handle = tokio::spawn(async move { connection.await });
+
+    // Single query: columns with PK info, per schema/table
+    let column_rows = client
+        .query(
+            "SELECT
+                ns.nspname AS schema_name,
+                c.relname AS table_name,
+                a.attname AS column_name,
+                pg_catalog.format_type(a.atttypid, a.atttypmod) AS datatype,
+                (SELECT substring(pg_catalog.pg_get_expr(d.adbin, d.adrelid, true) for 128)
+                 FROM pg_catalog.pg_attrdef d
+                 WHERE d.adrelid = a.attrelid AND d.adnum = a.attnum AND a.atthasdef) AS default_value,
+                CASE a.attnotnull WHEN false THEN true ELSE false END AS nullable,
+                EXISTS (
+                    SELECT 1 FROM pg_catalog.pg_index i
+                    WHERE i.indrelid = c.oid AND i.indisprimary AND a.attnum = ANY(i.indkey)
+                ) AS is_primary_key,
+                (SELECT con.conname FROM pg_catalog.pg_constraint con
+                 WHERE con.conrelid = c.oid AND con.contype = 'p' LIMIT 1) AS pk_constraint_name
+            FROM pg_catalog.pg_attribute a
+            JOIN pg_catalog.pg_class c ON a.attrelid = c.oid
+            JOIN pg_catalog.pg_namespace ns ON c.relnamespace = ns.oid
+            WHERE c.relkind = 'r'
+                AND a.attnum > 0
+                AND NOT a.attisdropped
+                AND ns.nspname NOT IN ('pg_catalog', 'information_schema')
+            ORDER BY ns.nspname, c.relname, a.attnum",
+            &[],
+        )
+        .await
+        .map_err(|e| Error::internal_err(format!("Failed to query columns: {}", e)))?;
+
+    // Foreign keys query
+    let fk_rows = client
+        .query(
+            "SELECT
+                ns.nspname AS schema_name,
+                c.relname AS table_name,
+                con.conname AS fk_constraint_name,
+                att_src.attname AS source_column,
+                ns_ref.nspname AS ref_schema,
+                c_ref.relname AS ref_table,
+                att_ref.attname AS ref_column,
+                con.confdeltype AS on_delete,
+                con.confupdtype AS on_update
+            FROM pg_catalog.pg_constraint con
+            JOIN pg_catalog.pg_class c ON con.conrelid = c.oid
+            JOIN pg_catalog.pg_namespace ns ON c.relnamespace = ns.oid
+            JOIN pg_catalog.pg_class c_ref ON con.confrelid = c_ref.oid
+            JOIN pg_catalog.pg_namespace ns_ref ON c_ref.relnamespace = ns_ref.oid
+            CROSS JOIN LATERAL unnest(con.conkey, con.confkey) WITH ORDINALITY AS u(src_attnum, ref_attnum, ord)
+            JOIN pg_catalog.pg_attribute att_src ON att_src.attrelid = c.oid AND att_src.attnum = u.src_attnum
+            JOIN pg_catalog.pg_attribute att_ref ON att_ref.attrelid = c_ref.oid AND att_ref.attnum = u.ref_attnum
+            WHERE con.contype = 'f'
+                AND ns.nspname NOT IN ('pg_catalog', 'information_schema')
+            ORDER BY ns.nspname, c.relname, con.conname, u.ord",
+            &[],
+        )
+        .await
+        .map_err(|e| Error::internal_err(format!("Failed to query foreign keys: {}", e)))?;
+
+    drop(client);
+    join_handle
+        .await
+        .map_err(|e| Error::internal_err(format!("join error: {}", e)))?
+        .map_err(|e| Error::internal_err(format!("tokio_postgres error: {}", e)))?;
+
+    // Build result: { schema_name: { table_name: TableEditorValues } }
+    let mut result: HashMap<String, HashMap<String, TableEditorValues>> = HashMap::new();
+
+    for row in &column_rows {
+        let schema_name: &str = row.get("schema_name");
+        let table_name: &str = row.get("table_name");
+        let column_name: &str = row.get("column_name");
+        let datatype: &str = row.get("datatype");
+        let default_value: Option<&str> = row.get("default_value");
+        let nullable: bool = row.get("nullable");
+        let is_primary_key: bool = row.get("is_primary_key");
+        let pk_constraint_name: Option<&str> = row.get("pk_constraint_name");
+
+        let schema_tables = result.entry(schema_name.to_string()).or_default();
+        let table = schema_tables
+            .entry(table_name.to_string())
+            .or_insert_with(|| TableEditorValues {
+                name: table_name.to_string(),
+                columns: vec![],
+                foreign_keys: vec![],
+                pk_constraint_name: pk_constraint_name.map(|s| s.to_string()),
+            });
+
+        table.columns.push(TableEditorValuesColumn {
+            name: column_name.to_string(),
+            datatype: datatype.to_string(),
+            primary_key: if is_primary_key { Some(true) } else { None },
+            default_value: default_value.map(|s| s.to_string()),
+            nullable: Some(nullable),
+        });
+    }
+
+    // Group FK rows by (schema, table, constraint_name)
+    let mut fk_map: HashMap<
+        (String, String, String),
+        (Option<String>, Vec<ForeignKeyColumnPair>, String, String),
+    > = HashMap::new();
+
+    fn pg_action_to_string(action: &str) -> String {
+        match action {
+            "a" => "NO ACTION".to_string(),
+            "r" => "RESTRICT".to_string(),
+            "c" => "CASCADE".to_string(),
+            "n" => "SET NULL".to_string(),
+            "d" => "SET DEFAULT".to_string(),
+            _ => "NO ACTION".to_string(),
+        }
+    }
+
+    for row in &fk_rows {
+        let schema_name: &str = row.get("schema_name");
+        let table_name: &str = row.get("table_name");
+        let fk_name: &str = row.get("fk_constraint_name");
+        let source_column: &str = row.get("source_column");
+        let ref_schema: &str = row.get("ref_schema");
+        let ref_table: &str = row.get("ref_table");
+        let ref_column: &str = row.get("ref_column");
+        let on_delete: &str = row.get("on_delete");
+        let on_update: &str = row.get("on_update");
+
+        let target_table = if ref_schema == schema_name {
+            ref_table.to_string()
+        } else {
+            format!("{}.{}", ref_schema, ref_table)
+        };
+
+        let key = (
+            schema_name.to_string(),
+            table_name.to_string(),
+            fk_name.to_string(),
+        );
+        let entry = fk_map.entry(key).or_insert_with(|| {
+            (
+                Some(target_table.clone()),
+                vec![],
+                pg_action_to_string(on_delete),
+                pg_action_to_string(on_update),
+            )
+        });
+        entry.1.push(ForeignKeyColumnPair {
+            source_column: Some(source_column.to_string()),
+            target_column: Some(ref_column.to_string()),
+        });
+    }
+
+    for ((schema_name, table_name, fk_name), (target_table, columns, on_delete, on_update)) in
+        fk_map
+    {
+        if let Some(schema_tables) = result.get_mut(&schema_name) {
+            if let Some(table) = schema_tables.get_mut(&table_name) {
+                table.foreign_keys.push(TableEditorForeignKey {
+                    target_table,
+                    columns,
+                    on_delete,
+                    on_update,
+                    fk_constraint_name: Some(fk_name),
+                });
+            }
+        }
+    }
+
+    Ok(Json(result))
+}
+
 async fn edit_ducklake_config(
    authed: ApiAuthed,
    Extension(db): Extension<DB>,