fix(datatables): bind fork database copies to their workspace, and count every use before dropping one

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

backend/windmill-api-workspaces/src/workspaces.rs

@@ -3437,8 +3437,13 @@ async fn create_pg_database(
     }
 
     if is_instance_datatable_source(&db, &w_id, &req.source).await? {
-        windmill_common::create_custom_instance_database(&db, &req.target_dbname, "datatable")
-            .await?;
+        windmill_common::create_custom_instance_database(
+            &db,
+            &req.target_dbname,
+            "datatable",
+            Some(&w_id),
+        )
+        .await?;
     } else {
         let source_pg =
             resolve_pg_source_checked(&db, &user_db, &authed, &w_id, &req.source).await?;
@@ -3592,6 +3597,10 @@ async fn import_pg_database(
                         .to_string(),
                 ));
             }
+            if is_instance_datatable_source(&db, &w_id, &req.target).await? {
+                windmill_common::ensure_fork_database_available_to(&db, override_dbname, &w_id)
+                    .await?;
+            }
         }
         target_pg.dbname = override_dbname.clone();
     }
@@ -8030,6 +8039,7 @@ async fn point_kept_datatables_at_parent(
     forked_w_id: &str,
     cloned: &[ForkedDatatableInfo],
 ) -> Result<()> {
+    windmill_common::workspaces::lock_fork_datatables(tx, parent_w_id).await?;
     let settings: Option<serde_json::Value> = sqlx::query_scalar!(
         "SELECT datatable FROM workspace_settings WHERE workspace_id = $1",
         forked_w_id
@@ -8198,6 +8208,12 @@ async fn apply_forked_datatable(
         })?,
     };
 
+    if database.resource_type == DataTableCatalogResourceType::Instance
+        && !windmill_api_auth::is_super_admin_authed(db, authed).await?
+    {
+        windmill_common::ensure_fork_database_available_to(db, &fdt.new_dbname, parent_w_id)
+            .await?;
+    }
     if database.resource_type == DataTableCatalogResourceType::Instance {
         // The whole `database` object, not just its `resource_path`: a pointer entry has none to
         // patch. `reference` goes with it — exactly one of the two may be set.

backend/windmill-api-workspaces/src/workspaces_extra.rs

@@ -1420,7 +1420,31 @@ pub async fn drop_forked_datatable_databases(
                 ));
                 continue;
             }
-            if let Err(e) = windmill_common::drop_custom_instance_database(&db, db_to_drop).await {
+            // The fork's own entry is what is going away; anything else still reaching the copy,
+            // a child fork's pointer at this entry included, keeps it. The lock keeps a child fork
+            // from gaining such a pointer before the drop.
+            let dropped = async {
+                let mut tx = db.begin().await?;
+                windmill_common::workspaces::lock_fork_datatables(&mut tx, &w_id).await?;
+                let uses = windmill_common::workspaces::managed_database_uses(
+                    &mut tx,
+                    windmill_common::workspaces::DataTableCatalogResourceType::Instance,
+                    db_to_drop,
+                    Some((&w_id, dt_name)),
+                )
+                .await?;
+                if !uses.is_empty() {
+                    return Err(Error::BadRequest(format!(
+                        "it is still used by {}",
+                        uses.join(", ")
+                    )));
+                }
+                windmill_common::drop_custom_instance_database(&db, db_to_drop).await?;
+                tx.commit().await?;
+                Ok::<_, Error>(())
+            }
+            .await;
+            if let Err(e) = dropped {
                 errors.push(format!(
                     "Could not drop instance database '{}' for datatable://{}: {}",
                     db_to_drop, dt_name, e

backend/windmill-common/src/instance_config.rs

@@ -809,6 +809,9 @@ pub struct CustomInstanceDb {
     pub error: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub tag: Option<String>,
+    /// The workspace a member created this fork copy for. Absent when a superadmin created it.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub workspace_id: Option<String>,
 }
 
 /// Setup log entries for a custom instance database.

backend/windmill-common/src/lib.rs

@@ -1568,11 +1568,13 @@ pub async fn ensure_instance_db_grant_options_unchecked(
 }
 
 /// Create a custom instance database: CREATE DATABASE, grant permissions, register in global_settings.
-/// The `tag` is stored in global_settings metadata (e.g. "datatable" or "ducklake").
+/// The `tag` is stored in global_settings metadata (e.g. "datatable" or "ducklake"). `for_workspace`
+/// is the workspace a member creates a fork copy for; see [`ensure_fork_database_available_to`].
 pub async fn create_custom_instance_database(
     db: &DB,
     dbname: &str,
     tag: &str,
+    for_workspace: Option<&str>,
 ) -> error::Result<()> {
     let dbname = dbname.trim();
     validate_dbname(dbname)?;
@@ -1626,7 +1628,8 @@ pub async fn create_custom_instance_database(
         },
         "success": true,
         "error": null,
-        "tag": tag
+        "tag": tag,
+        "workspace_id": for_workspace,
     });
     sqlx::query!(
         r#"UPDATE global_settings SET value = jsonb_set(value, '{databases}', (COALESCE(value->'databases', '{}'::jsonb) || to_jsonb($1::json))) WHERE name = 'custom_instance_pg_databases'"#,
@@ -1646,6 +1649,44 @@ pub async fn create_custom_instance_database(
     Ok(())
 }
 
+/// Refuse a workspace member writing a fork copy into, or pointing a fork at, the instance database
+/// `dbname`, unless `w_id` created it for that ([`create_custom_instance_database`]) and nothing uses
+/// it yet. The `wm_fork_` prefix is no authorization: every instance database answers to the same
+/// `custom_instance_user`, so a name is all it takes to reach another workspace's copy.
+pub async fn ensure_fork_database_available_to(
+    db: &DB,
+    dbname: &str,
+    w_id: &str,
+) -> error::Result<()> {
+    let created_for = sqlx::query_scalar::<_, Option<String>>(
+        "SELECT value->'databases'->$1->>'workspace_id' FROM global_settings
+         WHERE name = 'custom_instance_pg_databases'",
+    )
+    .bind(dbname)
+    .fetch_optional(db)
+    .await?
+    .flatten();
+    if created_for.as_deref() != Some(w_id) {
+        return Err(Error::BadRequest(format!(
+            "Database '{dbname}' was not created for a fork of workspace '{w_id}'"
+        )));
+    }
+    let uses = workspaces::managed_database_uses(
+        &mut *db.acquire().await?,
+        workspaces::DataTableCatalogResourceType::Instance,
+        dbname,
+        None,
+    )
+    .await?;
+    if !uses.is_empty() {
+        return Err(Error::BadRequest(format!(
+            "Database '{dbname}' is already in use: {}",
+            uses.join(", ")
+        )));
+    }
+    Ok(())
+}
+
 /// Connection options parsed from a database URL.
 ///
 /// The only place a database URL becomes `PgConnectOptions`. Providers that mint the password

backend/windmill-common/src/workspaces.rs

@@ -1471,6 +1471,73 @@ pub struct GoverningDatatable {
     pub datatable: DataTable,
 }
 
+/// Everything still using the Windmill-managed database `dbname`, one description per use: data
+/// table entries naming it, fork entries pointing at those, Ducklake catalogs on it, and fork
+/// Ducklake metadata schemas there that cleanup has not dropped yet. `exempt` is the one data table
+/// entry, `(workspace_id, name)`, the caller is about to stop using it through; pointers at that
+/// entry still count, since dropping the database would leave them resolving to nothing.
+///
+/// Authorization: reads every workspace's settings and checks nothing. Callers MUST only turn the
+/// answer into a refusal for someone allowed to administer `dbname`.
+pub async fn managed_database_uses(
+    conn: &mut sqlx::PgConnection,
+    kind: DataTableCatalogResourceType,
+    dbname: &str,
+    exempt: Option<(&str, &str)>,
+) -> Result<Vec<String>> {
+    let (exempt_workspace, exempt_name) = exempt.unzip();
+    Ok(sqlx::query_scalar::<_, String>(
+        "WITH entries AS (
+             SELECT ws.workspace_id::text AS workspace_id, dt.key AS name, dt.value
+             FROM workspace_settings ws
+             CROSS JOIN LATERAL jsonb_each(
+                 CASE WHEN jsonb_typeof(ws.datatable->'datatables') = 'object'
+                     THEN ws.datatable->'datatables' ELSE '{}'::jsonb END) dt
+         ), naming AS (
+             SELECT workspace_id, name FROM entries
+             WHERE value->'database'->>'resource_type' = $1
+               AND value->'database'->>'resource_path' = $2
+         )
+         SELECT format('data table ''%s'' in workspace ''%s''', name, workspace_id) FROM naming
+         WHERE $3::text IS NULL OR NOT (workspace_id = $3 AND name = $4)
+         UNION ALL
+         SELECT format('data table ''%s'' in workspace ''%s'', which points at the one in ''%s''',
+                       e.name, e.workspace_id, n.workspace_id)
+         FROM entries e JOIN naming n
+           ON e.value->'reference'->>'workspace_id' = n.workspace_id
+          AND e.value->'reference'->>'datatable' = n.name
+         UNION ALL
+         SELECT format('Ducklake ''%s'' in workspace ''%s''', dl.key, ws.workspace_id)
+         FROM workspace_settings ws
+         CROSS JOIN LATERAL jsonb_each(
+             CASE WHEN jsonb_typeof(ws.ducklake->'ducklakes') = 'object'
+                 THEN ws.ducklake->'ducklakes' ELSE '{}'::jsonb END) dl
+         WHERE dl.value->'catalog'->>'resource_type' = $1
+           AND dl.value->'catalog'->>'resource_path' = $2
+         UNION ALL
+         SELECT format('the Ducklake namespace of fork ''%s'', not cleaned up yet', workspace_id)
+         FROM fork_ducklake_namespace
+         WHERE catalog = $1 || ':' || $2 AND NOT schema_dropped
+         ORDER BY 1",
+    )
+    .bind(kind.as_ref())
+    .bind(dbname)
+    .bind(exempt_workspace)
+    .bind(exempt_name)
+    .fetch_all(&mut *conn)
+    .await?)
+}
+
+/// Held by fork cleanup of `w_id`'s data tables and by forking `w_id`, which can hand the new fork
+/// pointers at them, so a pointer cannot appear between cleanup's check and its drop.
+pub async fn lock_fork_datatables(conn: &mut sqlx::PgConnection, w_id: &str) -> Result<()> {
+    sqlx::query("SELECT pg_advisory_xact_lock(hashtext('fork_datatables:' || $1))")
+        .bind(w_id)
+        .execute(&mut *conn)
+        .await?;
+    Ok(())
+}
+
 impl GoverningDatatable {
     /// Backed by the Windmill instance's own Postgres, which is the only substrate data table
     /// roles apply to.