perf: optimize datatable app chat schemas (#8960)

* perf: optimize datatable app chat schemas Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * perf: optimize datatable catalog queries Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * refactor: narrow datatable chat optimization Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: restrict datatable schema lookups Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: block system datatable schema lookups Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: handle datatable context edge cases Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * fix: handle datatable schema edge cases Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
2026-09-12 00:06:14 +00:00 · 2026-04-29 13:58:56 +00:00
parent c0eeea9c83
commit 34b549cfe2
10 changed files with 692 additions and 247 deletions
@@ -115,6 +115,11 @@ pub fn workspaced_service() -> Router {
        .route("/list_ducklakes", get(list_ducklakes))
        .route("/list_datatables", get(list_datatables))
        .route("/list_datatable_schemas", get(list_datatable_schemas))
+        .route("/list_datatable_tables", get(list_datatable_tables))
+        .route(
+            "/get_datatable_table_schema",
+            get(get_datatable_table_schema),
+        )
        .route("/edit_datatable_config", post(edit_datatable_config))
        .route("/git_sync_enabled", get(get_git_sync_enabled))
        .route("/edit_git_sync_config", post(edit_git_sync_config))
@@ -1365,6 +1370,9 @@ type TableMap = HashMap<String, ColumnMap>;
 /// Schemas mapped by name to their tables
 type SchemaMap = HashMap<String, TableMap>;

+/// Schemas mapped by name to their table names
+type TableListMap = HashMap<String, Vec<String>>;
+
 #[derive(Serialize, Debug)]
 struct DataTableSchema {
    datatable_name: String,
@@ -1374,26 +1382,36 @@ struct DataTableSchema {
    error: Option<String>,
 }

+#[derive(Serialize, Debug)]
+struct DataTableTables {
+    datatable_name: String,
+    /// Hierarchical metadata: schema_name -> table_names
+    schemas: TableListMap,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    error: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct GetDataTableSchemaQuery {
+    datatable_name: String,
+    schema_name: String,
+    table_name: String,
+}
+
+#[derive(Serialize, Debug)]
+struct DataTableTableSchema {
+    datatable_name: String,
+    schema_name: String,
+    table_name: String,
+    columns: ColumnMap,
+}
+
 async fn list_datatable_schemas(
    _authed: ApiAuthed,
    Extension(db): Extension<DB>,
    Path(w_id): Path<String>,
 ) -> JsonResult<Vec<DataTableSchema>> {
-    // Get all datatable names for this workspace
-    let datatable_names: Vec<String> = sqlx::query_scalar!(
-        r#"
-            SELECT jsonb_object_keys(ws.datatable->'datatables') AS datatable_name
-            FROM workspace_settings ws
-            WHERE ws.workspace_id = $1
-        "#,
-        &w_id
-    )
-    .fetch_all(&db)
-    .await?
-    .into_iter()
-    .filter_map(|s| s)
-    .collect();
-
+    let datatable_names = list_datatable_names(&db, &w_id).await?;
    let mut results = Vec::new();

    for datatable_name in datatable_names {
@@ -1411,6 +1429,68 @@ async fn list_datatable_schemas(
    Ok(Json(results))
 }

+async fn list_datatable_tables(
+    _authed: ApiAuthed,
+    Extension(db): Extension<DB>,
+    Path(w_id): Path<String>,
+) -> JsonResult<Vec<DataTableTables>> {
+    let datatable_names = list_datatable_names(&db, &w_id).await?;
+    let mut results = Vec::new();
+
+    for datatable_name in datatable_names {
+        let tables = match get_datatable_tables(&db, &w_id, &datatable_name).await {
+            Ok(schemas) => DataTableTables { datatable_name, schemas, error: None },
+            Err(e) => DataTableTables {
+                datatable_name,
+                schemas: HashMap::new(),
+                error: Some(e.to_string()),
+            },
+        };
+        results.push(tables);
+    }
+
+    Ok(Json(results))
+}
+
+async fn get_datatable_table_schema(
+    _authed: ApiAuthed,
+    Extension(db): Extension<DB>,
+    Path(w_id): Path<String>,
+    Query(query): Query<GetDataTableSchemaQuery>,
+) -> JsonResult<DataTableTableSchema> {
+    let columns = get_datatable_table_columns(
+        &db,
+        &w_id,
+        &query.datatable_name,
+        &query.schema_name,
+        &query.table_name,
+    )
+    .await?;
+
+    Ok(Json(DataTableTableSchema {
+        datatable_name: query.datatable_name,
+        schema_name: query.schema_name,
+        table_name: query.table_name,
+        columns,
+    }))
+}
+
+async fn list_datatable_names(db: &DB, w_id: &str) -> Result<Vec<String>> {
+    Ok(sqlx::query_scalar!(
+        r#"
+            SELECT jsonb_object_keys(ws.datatable->'datatables') AS datatable_name
+            FROM workspace_settings ws
+            WHERE ws.workspace_id = $1
+        "#,
+        w_id
+    )
+    .fetch_all(db)
+    .await?
+    .into_iter()
+    .filter_map(|s| s)
+    .collect())
+}
+
 async fn get_datatable_schema(db: &DB, w_id: &str, datatable_name: &str) -> Result<SchemaMap> {
    // Get the datatable resource (connection credentials)
    let db_resource = get_datatable_resource_from_db_unchecked(db, w_id, datatable_name).await?;
@@ -1486,33 +1566,202 @@ async fn get_datatable_schema(db: &DB, w_id: &str, datatable_name: &str) -> Resu
        let is_nullable: String = row.get(4);
        let column_default: Option<String> = row.get(5);

-        // Build compact type representation: "type[?][=default]"
-        let mut compact = udt_name;
-        if is_nullable == "YES" {
-            compact.push('?');
-        }
-        if let Some(default) = column_default {
-            // Truncate long defaults for compactness
-            let short_default = if default.len() > 30 {
-                format!("{}...", &default[..27])
-            } else {
-                default
-            };
-            compact.push('=');
-            compact.push_str(&short_default);
-        }
-
        schema_map
            .entry(table_schema)
            .or_default()
            .entry(table_name)
            .or_default()
-            .insert(column_name, compact);
+            .insert(
+                column_name,
+                compact_column_type(udt_name, is_nullable, column_default),
+            );
    }

    Ok(schema_map)
 }

+async fn get_datatable_tables(db: &DB, w_id: &str, datatable_name: &str) -> Result<TableListMap> {
+    let db_resource = get_datatable_resource_from_db_unchecked(db, w_id, datatable_name).await?;
+    let pg_db: PgDatabase = serde_json::from_value(db_resource)
+        .map_err(|e| Error::internal_err(format!("Failed to parse database credentials: {}", e)))?;
+    let (client, connection) = pg_db.connect(Some(db)).await?;
+
+    tokio::spawn(async move {
+        if let Err(e) = connection.await {
+            tracing::error!("Datatable connection error: {}", e);
+        }
+    });
+
+    let schema_rows = client
+        .query(
+            r#"
+            SELECT nspname::text AS schema_name
+            FROM pg_namespace
+            WHERE nspname NOT IN ('information_schema', 'pg_toast', 'pg_catalog')
+              AND nspname NOT LIKE 'pg_%'
+            ORDER BY nspname
+            "#,
+            &[],
+        )
+        .await
+        .map_err(|e| Error::internal_err(format!("Failed to query schemas: {}", e)))?;
+
+    let mut table_map: TableListMap = HashMap::new();
+    let schema_names: Vec<String> = schema_rows
+        .iter()
+        .map(|row| {
+            let name: String = row.get(0);
+            table_map.entry(name.clone()).or_default();
+            name
+        })
+        .collect();
+
+    let rows = client
+        .query(
+            r#"
+            SELECT DISTINCT
+                table_schema::text,
+                table_name::text
+            FROM information_schema.columns
+            WHERE table_schema = ANY($1)
+              AND table_name IS NOT NULL
+            ORDER BY table_schema, table_name
+            "#,
+            &[&schema_names],
+        )
+        .await
+        .map_err(|e| Error::internal_err(format!("Failed to query tables: {}", e)))?;
+
+    for row in rows {
+        let table_schema: String = row.get(0);
+        let table_name: String = row.get(1);
+        table_map.entry(table_schema).or_default().push(table_name);
+    }
+
+    Ok(table_map)
+}
+
+async fn get_datatable_table_columns(
+    db: &DB,
+    w_id: &str,
+    datatable_name: &str,
+    schema_name: &str,
+    table_name: &str,
+) -> Result<ColumnMap> {
+    if is_system_pg_schema(schema_name) {
+        return Err(Error::BadRequest(format!(
+            "Schema '{}' is not available for datatable schema lookup",
+            schema_name
+        )));
+    }
+
+    let db_resource = get_datatable_resource_from_db_unchecked(db, w_id, datatable_name).await?;
+    let pg_db: PgDatabase = serde_json::from_value(db_resource)
+        .map_err(|e| Error::internal_err(format!("Failed to parse database credentials: {}", e)))?;
+    let (client, connection) = pg_db.connect(Some(db)).await?;
+
+    tokio::spawn(async move {
+        if let Err(e) = connection.await {
+            tracing::error!("Datatable connection error: {}", e);
+        }
+    });
+
+    let rows = client
+        .query(
+            r#"
+            SELECT
+                column_name::text,
+                udt_name::text,
+                is_nullable::text,
+                column_default::text
+            FROM information_schema.columns
+            WHERE table_schema = $1
+              AND table_name = $2
+            ORDER BY ordinal_position
+            "#,
+            &[&schema_name, &table_name],
+        )
+        .await
+        .map_err(|e| Error::internal_err(format!("Failed to query columns: {}", e)))?;
+
+    if rows.is_empty() {
+        return Err(Error::NotFound(format!(
+            "Table '{}.{}' not found in datatable '{}'",
+            schema_name, table_name, datatable_name
+        )));
+    }
+
+    let mut columns: ColumnMap = HashMap::new();
+    for row in rows {
+        let column_name: String = row.get(0);
+        let udt_name: String = row.get(1);
+        let is_nullable: String = row.get(2);
+        let column_default: Option<String> = row.get(3);
+        columns.insert(
+            column_name,
+            compact_column_type(udt_name, is_nullable, column_default),
+        );
+    }
+
+    Ok(columns)
+}
+
+fn is_system_pg_schema(schema_name: &str) -> bool {
+    // Match the datatable listing filter: PostgreSQL reserves pg_* schemas for system use.
+    matches!(
+        schema_name,
+        "information_schema" | "pg_toast" | "pg_catalog"
+    ) || schema_name.starts_with("pg_")
+}
+
+fn compact_column_type(
+    udt_name: String,
+    is_nullable: String,
+    column_default: Option<String>,
+) -> String {
+    let mut compact = udt_name;
+    if is_nullable == "YES" {
+        compact.push('?');
+    }
+    if let Some(default) = column_default {
+        compact.push('=');
+        compact.push_str(&truncate_column_default(default));
+    }
+    compact
+}
+
+fn truncate_column_default(default: String) -> String {
+    const MAX_DEFAULT_CHARS: usize = 30;
+    const TRUNCATED_DEFAULT_CHARS: usize = 27;
+
+    if default.chars().count() > MAX_DEFAULT_CHARS {
+        format!(
+            "{}...",
+            default
+                .chars()
+                .take(TRUNCATED_DEFAULT_CHARS)
+                .collect::<String>()
+        )
+    } else {
+        default
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn compact_column_type_truncates_multibyte_defaults_safely() {
+        let default = "é".repeat(31);
+
+        assert_eq!(
+            compact_column_type("text".to_string(), "NO".to_string(), Some(default)),
+            format!("text={}...", "é".repeat(27))
+        );
+    }
+}
+
 /// Resolve a source string to PgDatabase credentials with user-scoped permission checks.
 /// For `datatable://name`: accessible to everyone (variables are resolved internally).
 /// For `$res:path`: uses UserDB (row-level security) to verify the user can see the resource,
@@ -3996,6 +3996,55 @@ paths:
                items:
                  $ref: "#/components/schemas/DataTableSchema"

+  /w/{workspace}/workspaces/list_datatable_tables:
+    get:
+      summary: list tables of all connected Datatables
+      operationId: listDataTableTables
+      tags:
+        - workspace
+      parameters:
+        - $ref: "#/components/parameters/WorkspaceId"
+      responses:
+        "200":
+          description: table metadata of all datatables
+          content:
+            application/json:
+              schema:
+                type: array
+                items:
+                  $ref: "#/components/schemas/DataTableTables"
+
+  /w/{workspace}/workspaces/get_datatable_table_schema:
+    get:
+      summary: get one Datatable table schema
+      operationId: getDataTableTableSchema
+      tags:
+        - workspace
+      parameters:
+        - $ref: "#/components/parameters/WorkspaceId"
+        - name: datatable_name
+          in: query
+          required: true
+          schema:
+            type: string
+        - name: schema_name
+          in: query
+          required: true
+          schema:
+            type: string
+        - name: table_name
+          in: query
+          required: true
+          schema:
+            type: string
+      responses:
+        "200":
+          description: schema of one datatable table
+          content:
+            application/json:
+              schema:
+                $ref: "#/components/schemas/DataTableTableSchema"
+
  /w/{workspace}/workspaces/edit_ducklake_config:
    post:
      summary: edit ducklake settings
@@ -26477,6 +26526,39 @@ components:
        error:
          type: string

+    DataTableTables:
+      type: object
+      required: [datatable_name, schemas]
+      properties:
+        datatable_name:
+          type: string
+        schemas:
+          type: object
+          description: "Hierarchical metadata: schema_name -> table_names"
+          additionalProperties:
+            type: array
+            items:
+              type: string
+        error:
+          type: string
+
+    DataTableTableSchema:
+      type: object
+      required: [datatable_name, schema_name, table_name, columns]
+      properties:
+        datatable_name:
+          type: string
+        schema_name:
+          type: string
+        table_name:
+          type: string
+        columns:
+          type: object
+          description: "Columns in this table: column_name -> compact_type"
+          additionalProperties:
+            type: string
+            description: "Compact type: 'type[?][=default]' where ? means nullable"
+
    DynamicInputData:
      type: object
      properties:
@@ -12,6 +12,7 @@
 	import Portal from '$lib/components/Portal.svelte'
 	import { zIndexes } from '$lib/zIndexes'
 	import { tick, untrack } from 'svelte'
+	import { sendUserToast } from '$lib/toast'

 	interface Props {
 		availableContext: ContextElement[]
@@ -109,7 +110,7 @@
 		}
 	}

-	function addContextToSelection(contextElement: ContextElement) {
+	async function addContextToSelection(contextElement: ContextElement) {
 		if (!selectedContext || !availableContext) return

 		const alreadySelected = selectedContext.find(
@@ -130,20 +131,48 @@
 			return
 		}

-		selectedContext = [...selectedContext, contextElement]
+		let contextToAdd = contextElement

-		// If it's a datatable table, add it to the app's whitelisted tables
 		if (
 			contextElement.type === 'app_datatable' &&
 			aiChatManager.mode === AIMode.APP &&
 			aiChatManager.appAiChatHelpers
 		) {
-			aiChatManager.appAiChatHelpers.addTableToWhitelist(
+			const appAiChatHelpers = aiChatManager.appAiChatHelpers
+			appAiChatHelpers.addTableToWhitelist(
 				contextElement.datatableName,
 				contextElement.schemaName,
 				contextElement.tableName
 			)
+
+			if (!contextElement.columns) {
+				try {
+					contextToAdd = {
+						...contextElement,
+						columns: await appAiChatHelpers.getDatatableTableSchema(
+							contextElement.datatableName,
+							contextElement.schemaName,
+							contextElement.tableName
+						)
+					}
+				} catch (e) {
+					console.error('Failed to load datatable table schema:', e)
+					sendUserToast(
+						'Failed to load datatable table schema',
+						true,
+						[],
+						e instanceof Error ? e.message : String(e)
+					)
+				}
+			}
 		}
+
+		const duplicateAfterAwait = selectedContext.find(
+			(c) => c.type === contextToAdd.type && c.title === contextToAdd.title
+		)
+		if (duplicateAfterAwait) return
+
+		selectedContext = [...selectedContext, contextToAdd]
 	}

 	function sendRequest() {
@@ -342,7 +371,7 @@
 	}

 	function handleAppContextSelection(contextElement: ContextElement) {
-		addContextToSelection(contextElement)
+		void addContextToSelection(contextElement)
 		// Update instructions with the selected context title
 		const index = instructions.lastIndexOf('@')
 		if (index !== -1) {
@@ -364,31 +393,27 @@
 			<div class="flex flex-row gap-1 mb-1 overflow-scroll pt-2 no-scrollbar">
 				<Popover>
 					{#snippet trigger()}
-									
-							<div
-								class="border rounded-md px-1 py-0.5 font-normal text-primary text-xs hover:bg-surface-hover bg-surface"
-								>@</div
-							>
-						
-									{/snippet}
+						<div
+							class="border rounded-md px-1 py-0.5 font-normal text-primary text-xs hover:bg-surface-hover bg-surface"
+							>@</div
+						>
+					{/snippet}
 					{#snippet content({ close })}
-									
-							<AvailableContextList
-								{availableContext}
-								{selectedContext}
-								onSelect={(element) => {
-									addContextToSelection(element)
-									close()
-								}}
-								onSelectWorkspaceItem={(element) => {
-									addContextToSelection(element)
-									close()
-								}}
-							/>
-						
-									{/snippet}
+						<AvailableContextList
+							{availableContext}
+							{selectedContext}
+							onSelect={(element) => {
+								void addContextToSelection(element)
+								close()
+							}}
+							onSelectWorkspaceItem={(element) => {
+								void addContextToSelection(element)
+								close()
+							}}
+						/>
+					{/snippet}
 				</Popover>
-				{#each selectedContext as element}
+				{#each selectedContext as element (element.type + '-' + element.title)}
 					<ContextElementBadge
 						contextElement={element}
 						deletable
@@ -408,7 +433,7 @@
 			{selectedContext}
 			{isFirstMessage}
 			placeholder={modePlaceholder}
-			onAddContext={(contextElement) => addContextToSelection(contextElement)}
+			onAddContext={(contextElement) => void addContextToSelection(contextElement)}
 			onSendRequest={() => {
 				if (disabled) {
 					return
@@ -423,25 +448,21 @@
 			<div class="flex flex-row gap-1 mb-1 overflow-scroll pt-2 no-scrollbar">
 				<Popover>
 					{#snippet trigger()}
-											
-							<div
-								class="border rounded-md px-1 py-0.5 font-normal text-primary text-xs hover:bg-surface-hover bg-surface"
-								>@</div
-							>
-						
-											{/snippet}
+						<div
+							class="border rounded-md px-1 py-0.5 font-normal text-primary text-xs hover:bg-surface-hover bg-surface"
+							>@</div
+						>
+					{/snippet}
 					{#snippet content({ close })}
-											
-							<AppAvailableContextList
-								{availableContext}
-								{selectedContext}
-								onSelect={(element) => {
-									addContextToSelection(element)
-									close()
-								}}
-							/>
-						
-											{/snippet}
+						<AppAvailableContextList
+							{availableContext}
+							{selectedContext}
+							onSelect={(element) => {
+								void addContextToSelection(element)
+								close()
+							}}
+						/>
+					{/snippet}
 				</Popover>
 				{#each selectedContext as element (element.type + '-' + element.title)}
 					<ContextElementBadge
@@ -1043,7 +1043,7 @@ class AIChatManager {
 		}

 		try {
-			const datatables = await this.appAiChatHelpers.getDatatables()
+			const datatables = await this.appAiChatHelpers.listDatatableTables()
 			this.cachedDatatables = flattenDatatablesToAppContextElements(datatables)
 		} catch (err) {
 			console.error('Failed to refresh datatables:', err)
@@ -93,7 +93,8 @@ function createHelpers(overrides: Partial<AppAIChatHelpers> = {}): AppAIChatHelp
 		snapshot: () => 1,
 		revertToSnapshot: () => undefined,
 		lint: () => EMPTY_LINT_RESULT,
-		getDatatables: async () => [],
+		listDatatableTables: async () => [],
+		getDatatableTableSchema: async () => ({}),
 		getAvailableDatatableNames: () => [],
 		execDatatableSql: async () => ({ success: true }),
 		addTableToWhitelist: () => undefined,
@@ -195,18 +196,14 @@ describe('app datatable tools', () => {
 			args: {},
 			workspace: 'test-workspace',
 			helpers: createHelpers({
-				getDatatables: async () => [
+				listDatatableTables: async () => [
 					{
 						datatable_name: 'main',
 						schemas: {
-							public: {
-								users: { id: 'int4', email: 'text' },
-								orders: { id: 'int4', total: 'numeric' }
-							},
-							analytics: {
-								events: { id: 'int4', payload: 'jsonb' }
-							}
-						}
+							public: ['users', 'orders'],
+							analytics: ['events']
+						},
+						tableCount: 3
 					}
 				]
 			}),
@@ -240,17 +237,7 @@ describe('app datatable tools', () => {
 			},
 			workspace: 'test-workspace',
 			helpers: createHelpers({
-				getDatatables: async () => [
-					{
-						datatable_name: 'main',
-						schemas: {
-							public: {
-								users: { id: 'int4', email: 'text' },
-								orders: { id: 'int4', total: 'numeric' }
-							}
-						}
-					}
-				]
+				getDatatableTableSchema: async () => ({ id: 'int4', email: 'text' })
 			}),
 			toolCallbacks: createToolCallbacks(),
 			toolId: 'tool-get-table-schema'
@@ -287,7 +274,10 @@ describe('app patch_file tool', () => {
 			toolId: 'tool-1'
 		})

-		expect(setFrontendFile).toHaveBeenCalledWith('/index.tsx', 'export const title = "Hello cookbook"\n')
+		expect(setFrontendFile).toHaveBeenCalledWith(
+			'/index.tsx',
+			'export const title = "Hello cookbook"\n'
+		)
 		expect(result).toContain("Patched '/index.tsx' successfully.")
 	})

@@ -103,6 +103,13 @@ export interface DataTableSchema {
 	error?: string
 }

+export interface AppDatatableMetadata {
+	datatable_name: string
+	schemas: Record<string, string[]>
+	tableCount: number
+	error?: string
+}
+
 export interface AppAIChatHelpers {
 	// Frontend file operations
 	listFrontendFiles: () => string[]
@@ -127,8 +134,14 @@ export interface AppAIChatHelpers {
 	/** Lint all frontend files and backend runnables, returns errors and warnings */
 	lint: () => LintResult
 	// Data table operations
-	/** Get all datatables configured in the app with their schemas */
-	getDatatables: () => Promise<DataTableSchema[]>
+	/** List configured datatables with schema/table names only. */
+	listDatatableTables: () => Promise<AppDatatableMetadata[]>
+	/** Get columns for one datatable table. */
+	getDatatableTableSchema: (
+		datatableName: string,
+		schemaName: string,
+		tableName: string
+	) => Promise<Record<string, string>>
 	/** Get unique datatable names configured in the app (for UI policy selector) */
 	getAvailableDatatableNames: () => string[]
 	/** Execute a SQL query on a datatable. Optionally specify newTable to register a newly created table. */
@@ -406,33 +419,6 @@ const getListFilesToolDef = memo(() =>

 // ============= Data Table Tools =============

-interface AppDatatableMetadata {
-	datatable_name: string
-	schemas: Record<string, string[]>
-	tableCount: number
-	error?: string
-}
-
-function summarizeDatatables(datatables: DataTableSchema[]): AppDatatableMetadata[] {
-	return datatables.map((datatable) => {
-		const schemas: Record<string, string[]> = {}
-		let tableCount = 0
-
-		for (const [schemaName, tables] of Object.entries(datatable.schemas)) {
-			const tableNames = Object.keys(tables)
-			schemas[schemaName] = tableNames
-			tableCount += tableNames.length
-		}
-
-		return {
-			datatable_name: datatable.datatable_name,
-			schemas,
-			tableCount,
-			...(datatable.error && { error: datatable.error })
-		}
-	})
-}
-
 const getListDatatablesSchema = memo(() => z.object({}))
 const getListDatatablesToolDef = memo(() =>
 	createToolDef(
@@ -641,7 +627,9 @@ export const getAppTools = memo((): Tool<AppAIChatHelpers>[] => [

 			if (target.type === 'frontend') {
 				if (target.path === '/wmill.d.ts') {
-					throw new Error("'/wmill.d.ts' is generated automatically. Edit backend runnables instead.")
+					throw new Error(
+						"'/wmill.d.ts' is generated automatically. Edit backend runnables instead."
+					)
 				}

 				const frontendContent = helpers.getFrontendFile(target.path)
@@ -824,12 +812,11 @@ export const getAppTools = memo((): Tool<AppAIChatHelpers>[] => [
 		fn: async ({ helpers, toolId, toolCallbacks }) => {
 			toolCallbacks.setToolStatus(toolId, { content: 'Listing datatables...' })
 			try {
-				const datatables = await helpers.getDatatables()
-				if (datatables.length === 0) {
+				const metadata = await helpers.listDatatableTables()
+				if (metadata.length === 0) {
 					toolCallbacks.setToolStatus(toolId, { content: 'No datatables configured' })
 					return 'No datatables are configured in this app. Use the Data panel in the sidebar to add datatable references.'
 				}
-				const metadata = summarizeDatatables(datatables)
 				const totalTables = metadata.reduce((acc, datatable) => acc + datatable.tableCount, 0)
 				toolCallbacks.setToolStatus(toolId, {
 					content: `Listed ${metadata.length} datatable(s) with ${totalTables} table(s)`
@@ -850,32 +837,11 @@ export const getAppTools = memo((): Tool<AppAIChatHelpers>[] => [
 				content: `Getting schema for ${parsedArgs.datatable_name}.${parsedArgs.schema_name}.${parsedArgs.table_name}...`
 			})
 			try {
-				const datatables = await helpers.getDatatables()
-				const datatable = datatables.find(
-					(candidate) => candidate.datatable_name === parsedArgs.datatable_name
+				const columns = await helpers.getDatatableTableSchema(
+					parsedArgs.datatable_name,
+					parsedArgs.schema_name,
+					parsedArgs.table_name
 				)
-				if (!datatable) {
-					const availableDatatables = datatables.map((candidate) => candidate.datatable_name)
-					const errorMsg = `Datatable '${parsedArgs.datatable_name}' not found. Available datatables: ${availableDatatables.join(', ') || 'none'}`
-					toolCallbacks.setToolStatus(toolId, { content: errorMsg, error: errorMsg })
-					return errorMsg
-				}
-
-				const schema = datatable.schemas[parsedArgs.schema_name]
-				if (!schema) {
-					const availableSchemas = Object.keys(datatable.schemas)
-					const errorMsg = `Schema '${parsedArgs.schema_name}' not found in datatable '${parsedArgs.datatable_name}'. Available schemas: ${availableSchemas.join(', ') || 'none'}`
-					toolCallbacks.setToolStatus(toolId, { content: errorMsg, error: errorMsg })
-					return errorMsg
-				}
-
-				const columns = schema[parsedArgs.table_name]
-				if (!columns) {
-					const availableTables = Object.keys(schema)
-					const errorMsg = `Table '${parsedArgs.table_name}' not found in '${parsedArgs.datatable_name}.${parsedArgs.schema_name}'. Available tables: ${availableTables.join(', ') || 'none'}`
-					toolCallbacks.setToolStatus(toolId, { content: errorMsg, error: errorMsg })
-					return errorMsg
-				}

 				toolCallbacks.setToolStatus(toolId, {
 					content: `Retrieved schema for ${parsedArgs.schema_name}.${parsedArgs.table_name}`
@@ -1271,12 +1237,16 @@ export function prepareAppUserMessage(
 					content += `- **Schema**: ${datatableCtx.schemaName}\n`
 					content += `- **Table**: ${datatableCtx.tableName}\n`
 					// Format columns as column_name: type
-					const columnsStr = JSON.stringify(datatableCtx.columns, null, 2)
-					const truncatedColumns =
-						columnsStr.length > MAX_CONTEXT_CONTENT_LENGTH
-							? columnsStr.slice(0, MAX_CONTEXT_CONTENT_LENGTH) + '\n... [TRUNCATED]'
-							: columnsStr
-					content += `- **Columns** (column_name -> type):\n\`\`\`json\n${truncatedColumns}\n\`\`\`\n`
+					if (datatableCtx.columns) {
+						const columnsStr = JSON.stringify(datatableCtx.columns, null, 2)
+						const truncatedColumns =
+							columnsStr.length > MAX_CONTEXT_CONTENT_LENGTH
+								? columnsStr.slice(0, MAX_CONTEXT_CONTENT_LENGTH) + '\n... [TRUNCATED]'
+								: columnsStr
+						content += `- **Columns** (column_name -> type):\n\`\`\`json\n${truncatedColumns}\n\`\`\`\n`
+					} else {
+						content += `- **Columns**: not loaded. Use get_datatable_table_schema() if column names or types are needed.\n`
+					}
 				}
 			}
 		}
@@ -12,7 +12,7 @@ import BarsStaggered from '$lib/components/icons/BarsStaggered.svelte'
 import type { ScriptLang, Script, OpenFlow } from '$lib/gen/types.gen'
 import { type DBSchema } from '$lib/stores'
 import { type Change } from 'diff'
-import type { BackendRunnable, DataTableSchema, SelectedContext } from './app/core'
+import type { AppDatatableMetadata, BackendRunnable, SelectedContext } from './app/core'

 export const ContextIconMap = {
 	code: Code,
@@ -136,8 +136,8 @@ export interface AppDatatableElement {
 	tableName: string
 	/** Title for display (e.g., "main/public:users" or "main/users") */
 	title: string
-	/** The table columns: column_name -> compact_type */
-	columns: Record<string, string>
+	/** The table columns: column_name -> compact_type. Loaded only when a table is explicitly selected. */
+	columns?: Record<string, string>
 }

 export function createAppSelectedContext(options: SelectedContext = {}): SelectedContext {
@@ -184,7 +184,7 @@ export function createAppDatatableContextElement(
 	datatableName: string,
 	schemaName: string,
 	tableName: string,
-	columns: Record<string, string>
+	columns?: Record<string, string>
 ): AppDatatableElement {
 	return {
 		type: 'app_datatable',
@@ -197,7 +197,7 @@ export function createAppDatatableContextElement(
 }

 export function flattenDatatablesToAppContextElements(
-	datatables: DataTableSchema[]
+	datatables: AppDatatableMetadata[]
 ): AppDatatableElement[] {
 	return datatables.flatMap((datatable) => {
 		if (datatable.error) {
@@ -205,13 +205,8 @@ export function flattenDatatablesToAppContextElements(
 		}

 		return Object.entries(datatable.schemas).flatMap(([schemaName, tables]) =>
-			Object.entries(tables).map(([tableName, columns]) =>
-				createAppDatatableContextElement(
-					datatable.datatable_name,
-					schemaName,
-					tableName,
-					columns
-				)
+			tables.map((tableName) =>
+				createAppDatatableContextElement(datatable.datatable_name, schemaName, tableName)
 			)
 		)
 	})
@@ -20,7 +20,11 @@
 	import { isRunnableByName, isRunnableByPath } from '../apps/inputType'
 	import { aiChatManager, AIMode } from '../copilot/chat/AIChatManager.svelte'
 	import { onMount, untrack } from 'svelte'
-	import type { LintResult, DataTableSchema, InspectorElementInfo } from '../copilot/chat/app/core'
+	import type {
+		AppDatatableMetadata,
+		LintResult,
+		InspectorElementInfo
+	} from '../copilot/chat/app/core'
 	import { createAppSelectedContext, type AppCodeSelectionElement } from '../copilot/chat/context'
 	import { rawAppLintStore } from './lintStore'
 	import { dbSchemas } from '$lib/stores'
@@ -28,8 +32,10 @@
 	import { RawAppHistoryManager } from './RawAppHistoryManager.svelte'
 	import { sendUserToast } from '$lib/utils'
 	import {
+		buildDataTableWhitelist,
 		parseDataTableRef,
 		formatDataTableRef,
+		isDatatableTableAllowed,
 		type RawAppData,
 		DEFAULT_DATA
 	} from './dataTableRefUtils'
@@ -74,6 +80,87 @@

 	// Convert to object format for child components
 	let dataTableRefsObjects = $derived(data.tables.map(parseDataTableRef))
+	let dataTableWhitelist = $derived(buildDataTableWhitelist(dataTableRefsObjects))
+
+	type DataTableTablesMetadata = {
+		datatable_name: string
+		schemas: Record<string, string[]>
+		error?: string
+	}
+
+	function countTables(schemas: Record<string, string[]>): number {
+		return Object.values(schemas).reduce((acc, tables) => acc + tables.length, 0)
+	}
+
+	function withTableCount(datatable: DataTableTablesMetadata): AppDatatableMetadata {
+		return {
+			datatable_name: datatable.datatable_name,
+			schemas: datatable.schemas,
+			tableCount: countTables(datatable.schemas),
+			...(datatable.error && { error: datatable.error })
+		}
+	}
+
+	function isDatatableTableWhitelisted(
+		datatableName: string,
+		schemaName: string,
+		tableName: string
+	): boolean {
+		return isDatatableTableAllowed(dataTableWhitelist, datatableName, schemaName, tableName)
+	}
+
+	function filterDatatableTables(allTables: DataTableTablesMetadata[]): AppDatatableMetadata[] {
+		if (dataTableWhitelist.datatables.size === 0) {
+			return allTables.map(withTableCount)
+		}
+
+		const results: AppDatatableMetadata[] = []
+		for (const datatable of allTables) {
+			if (!dataTableWhitelist.datatables.has(datatable.datatable_name)) {
+				continue
+			}
+
+			if (dataTableWhitelist.allTablesDatatables.has(datatable.datatable_name)) {
+				results.push(withTableCount(datatable))
+				continue
+			}
+
+			const allowedTables = dataTableWhitelist.tables.get(datatable.datatable_name)
+			if (!allowedTables) {
+				results.push(
+					withTableCount({
+						datatable_name: datatable.datatable_name,
+						schemas: {},
+						error: datatable.error
+					})
+				)
+				continue
+			}
+
+			const filteredSchemas: Record<string, string[]> = {}
+			for (const [schemaName, tableNames] of Object.entries(datatable.schemas)) {
+				const allowedTablesInSchema = allowedTables.get(schemaName)
+				if (!allowedTablesInSchema) continue
+
+				const filteredTables = tableNames.filter((tableName) =>
+					allowedTablesInSchema.has(tableName)
+				)
+				if (filteredTables.length > 0) {
+					filteredSchemas[schemaName] = filteredTables
+				}
+			}
+
+			results.push(
+				withTableCount({
+					datatable_name: datatable.datatable_name,
+					schemas: filteredSchemas,
+					error: datatable.error
+				})
+			)
+		}
+
+		return results
+	}

 	// Initialize history manager
 	const historyManager = new RawAppHistoryManager({
@@ -445,88 +532,45 @@
 				console.log('reverting to snapshot', id)
 				handleHistorySelect(id)
 			},
-			getDatatables: async (): Promise<DataTableSchema[]> => {
+			listDatatableTables: async (): Promise<AppDatatableMetadata[]> => {
 				if (!$workspaceStore) {
 					return []
 				}

-				// Get all datatable schemas from the backend
-				const allSchemas = await WorkspaceService.listDataTableSchemas({
+				const tables = await WorkspaceService.listDataTableTables({
 					workspace: $workspaceStore
 				})
-
-				// Get unique datatable names from dataTableRefs (the whitelisted tables)
-				const whitelistedDatatables = new Set(dataTableRefsObjects.map((ref) => ref.datatable))
-
-				// If no datatables are configured, return all available datatables
-				// This allows users to see all datatables in the @ context menu
-				if (whitelistedDatatables.size === 0) {
-					return allSchemas
+				return filterDatatableTables(tables)
+			},
+			getDatatableTableSchema: async (
+				datatableName: string,
+				schemaName: string,
+				tableName: string
+			): Promise<Record<string, string>> => {
+				if (!$workspaceStore) {
+					return {}
 				}

-				// Build a map of whitelisted tables per datatable: datatable -> schema -> Set<table>
-				const whitelistedTables = new Map<string, Map<string, Set<string>>>()
-				for (const ref of dataTableRefsObjects) {
-					if (!ref.table) continue
-					if (!whitelistedTables.has(ref.datatable)) {
-						whitelistedTables.set(ref.datatable, new Map())
-					}
-					const schemaKey = ref.schema || 'public'
-					const schemaMap = whitelistedTables.get(ref.datatable)!
-					if (!schemaMap.has(schemaKey)) {
-						schemaMap.set(schemaKey, new Set())
-					}
-					schemaMap.get(schemaKey)!.add(ref.table)
-				}
-
-				// Filter schemas to only include whitelisted datatables and tables
-				const results: DataTableSchema[] = []
-				for (const schema of allSchemas) {
-					if (!whitelistedDatatables.has(schema.datatable_name)) {
-						continue
-					}
-
-					const allowedTables = whitelistedTables.get(schema.datatable_name)
-					if (!allowedTables) {
-						// Include the datatable but with empty schemas
-						results.push({
-							datatable_name: schema.datatable_name,
-							schemas: {},
-							error: schema.error
-						})
-						continue
-					}
-
-					// Filter schemas to only include whitelisted tables
-					const filteredSchemas: Record<string, Record<string, Record<string, string>>> = {}
-					for (const [schemaName, tables] of Object.entries(schema.schemas)) {
-						const allowedTablesInSchema = allowedTables.get(schemaName)
-						if (!allowedTablesInSchema) continue
-
-						const filteredTables: Record<string, Record<string, string>> = {}
-						for (const [tableName, columns] of Object.entries(tables)) {
-							if (allowedTablesInSchema.has(tableName)) {
-								filteredTables[tableName] = columns
-							}
-						}
-
-						if (Object.keys(filteredTables).length > 0) {
-							filteredSchemas[schemaName] = filteredTables
-						}
-					}
-
-					results.push({
-						datatable_name: schema.datatable_name,
-						schemas: filteredSchemas,
-						error: schema.error
+				if (!isDatatableTableWhitelisted(datatableName, schemaName, tableName)) {
+					const tableRef = formatDataTableRef({
+						datatable: datatableName,
+						schema: schemaName === 'public' ? undefined : schemaName,
+						table: tableName
 					})
+					throw new Error(`Table '${tableRef}' is not configured in this app`)
 				}

-				return results
+				const schema = await WorkspaceService.getDataTableTableSchema({
+					workspace: $workspaceStore,
+					datatableName,
+					schemaName,
+					tableName
+				})
+				return schema.columns
 			},
 			getAvailableDatatableNames: (): string[] => {
 				// Get unique datatable names from dataTableRefs
-				return [...new Set(dataTableRefsObjects.map((ref) => ref.datatable))]
+				return [...dataTableWhitelist.datatables]
 			},
 			execDatatableSql: async (
 				datatableName: string,
@@ -564,6 +608,8 @@
 						}
 					}

+					void aiChatManager.refreshDatatables()
+
 					// Check if result is an array (SELECT) or something else
 					if (Array.isArray(result)) {
 						return { success: true, result }
@@ -586,6 +632,7 @@
 				if (!data.tables.includes(newRef)) {
 					data.tables = [...data.tables, newRef]
 					saveFrontendDraft()
+					void aiChatManager.refreshDatatables()
 				}
 			}
 		})
@@ -0,0 +1,37 @@
+import { describe, expect, it } from 'vitest'
+import { buildDataTableWhitelist, isDatatableTableAllowed } from './dataTableRefUtils'
+
+describe('datatable whitelist helpers', () => {
+	it('allows every datatable table when no refs are configured', () => {
+		const whitelist = buildDataTableWhitelist([])
+
+		expect(isDatatableTableAllowed(whitelist, 'main', 'public', 'users')).toBe(true)
+		expect(isDatatableTableAllowed(whitelist, 'analytics', 'events', 'clicks')).toBe(true)
+	})
+
+	it('treats datatable-level refs as all tables in that datatable', () => {
+		const whitelist = buildDataTableWhitelist([{ datatable: 'main' }])
+
+		expect(isDatatableTableAllowed(whitelist, 'main', 'public', 'users')).toBe(true)
+		expect(isDatatableTableAllowed(whitelist, 'main', 'analytics', 'events')).toBe(true)
+		expect(isDatatableTableAllowed(whitelist, 'other', 'public', 'users')).toBe(false)
+	})
+
+	it('allows only explicitly listed tables for table-level refs', () => {
+		const whitelist = buildDataTableWhitelist([
+			{ datatable: 'main', schema: 'public', table: 'users' },
+			{ datatable: 'main', schema: 'analytics', table: 'events' }
+		])
+
+		expect(isDatatableTableAllowed(whitelist, 'main', 'public', 'users')).toBe(true)
+		expect(isDatatableTableAllowed(whitelist, 'main', 'analytics', 'events')).toBe(true)
+		expect(isDatatableTableAllowed(whitelist, 'main', 'public', 'orders')).toBe(false)
+	})
+
+	it('does not treat an explicit empty schema as public', () => {
+		const whitelist = buildDataTableWhitelist([{ datatable: 'main', schema: '', table: 'users' }])
+
+		expect(isDatatableTableAllowed(whitelist, 'main', '', 'users')).toBe(true)
+		expect(isDatatableTableAllowed(whitelist, 'main', 'public', 'users')).toBe(false)
+	})
+})
@@ -25,6 +25,60 @@ export const DEFAULT_DATA: RawAppData = {
 	schema: undefined
 }

+export type DataTableWhitelist = {
+	datatables: Set<string>
+	allTablesDatatables: Set<string>
+	tables: Map<string, Map<string, Set<string>>>
+}
+
+export function buildDataTableWhitelist(refs: DataTableRef[]): DataTableWhitelist {
+	const datatables = new Set<string>()
+	const allTablesDatatables = new Set<string>()
+	const tables = new Map<string, Map<string, Set<string>>>()
+
+	for (const ref of refs) {
+		datatables.add(ref.datatable)
+
+		if (!ref.table) {
+			allTablesDatatables.add(ref.datatable)
+			continue
+		}
+
+		if (!tables.has(ref.datatable)) {
+			tables.set(ref.datatable, new Map())
+		}
+		const schemaKey = ref.schema ?? 'public'
+		const schemaMap = tables.get(ref.datatable)!
+		if (!schemaMap.has(schemaKey)) {
+			schemaMap.set(schemaKey, new Set())
+		}
+		schemaMap.get(schemaKey)!.add(ref.table)
+	}
+
+	return { datatables, allTablesDatatables, tables }
+}
+
+export function isDatatableTableAllowed(
+	whitelist: DataTableWhitelist,
+	datatableName: string,
+	schemaName: string,
+	tableName: string
+): boolean {
+	if (whitelist.datatables.size === 0) {
+		return true
+	}
+
+	if (!whitelist.datatables.has(datatableName)) {
+		return false
+	}
+
+	if (whitelist.allTablesDatatables.has(datatableName)) {
+		return true
+	}
+
+	return whitelist.tables.get(datatableName)?.get(schemaName ?? 'public')?.has(tableName) ?? false
+}
+
 /**
 * Parse a string ref into a DataTableRef object
 * Format: <datatableName>/<schema>:<table> or <datatableName>/<table> (for public schema)