From 34b549cfe2e2e060561eabe369a99cfb4d9c9568 Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Wed, 29 Apr 2026 15:58:56 +0200 Subject: [PATCH] perf: optimize datatable app chat schemas (#8960) * perf: optimize datatable app chat schemas Co-Authored-By: Claude Opus 4.5 * perf: optimize datatable catalog queries Co-Authored-By: Claude Opus 4.5 * refactor: narrow datatable chat optimization Co-Authored-By: Claude Opus 4.5 * fix: restrict datatable schema lookups Co-Authored-By: Claude Opus 4.5 * fix: block system datatable schema lookups Co-Authored-By: Claude Opus 4.5 * fix: handle datatable context edge cases Co-Authored-By: Claude Opus 4.5 * fix: handle datatable schema edge cases Co-Authored-By: Claude Opus 4.5 --------- Co-authored-by: Claude Opus 4.5 --- .../windmill-api-workspaces/src/workspaces.rs | 313 ++++++++++++++++-- backend/windmill-api/openapi.yaml | 82 +++++ .../copilot/chat/AIChatInput.svelte | 115 ++++--- .../copilot/chat/AIChatManager.svelte.ts | 2 +- .../components/copilot/chat/app/core.test.ts | 34 +- .../lib/components/copilot/chat/app/core.ts | 98 ++---- .../lib/components/copilot/chat/context.ts | 19 +- .../components/raw_apps/RawAppEditor.svelte | 185 +++++++---- .../raw_apps/dataTableRefUtils.test.ts | 37 +++ .../components/raw_apps/dataTableRefUtils.ts | 54 +++ 10 files changed, 692 insertions(+), 247 deletions(-) create mode 100644 frontend/src/lib/components/raw_apps/dataTableRefUtils.test.ts diff --git a/backend/windmill-api-workspaces/src/workspaces.rs b/backend/windmill-api-workspaces/src/workspaces.rs index 2b33d99b3f..cd9ce277c5 100644 --- a/backend/windmill-api-workspaces/src/workspaces.rs +++ b/backend/windmill-api-workspaces/src/workspaces.rs @@ -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; /// Schemas mapped by name to their tables type SchemaMap = HashMap; +/// Schemas mapped by name to their table names +type TableListMap = HashMap>; + #[derive(Serialize, Debug)] struct DataTableSchema { datatable_name: String, @@ -1374,26 +1382,36 @@ struct DataTableSchema { error: Option, } +#[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, +} + +#[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, Path(w_id): Path, ) -> JsonResult> { - // Get all datatable names for this workspace - let datatable_names: Vec = 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, + Path(w_id): Path, +) -> JsonResult> { + 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, + Path(w_id): Path, + Query(query): Query, +) -> JsonResult { + 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> { + 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 { // 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 = 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 { + 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 = 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 { + 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 = 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 { + 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::() + ) + } 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, diff --git a/backend/windmill-api/openapi.yaml b/backend/windmill-api/openapi.yaml index fffdbeb910..dbc74330d3 100644 --- a/backend/windmill-api/openapi.yaml +++ b/backend/windmill-api/openapi.yaml @@ -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: diff --git a/frontend/src/lib/components/copilot/chat/AIChatInput.svelte b/frontend/src/lib/components/copilot/chat/AIChatInput.svelte index b496dbb6c8..1a4169f677 100644 --- a/frontend/src/lib/components/copilot/chat/AIChatInput.svelte +++ b/frontend/src/lib/components/copilot/chat/AIChatInput.svelte @@ -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 @@
{#snippet trigger()} - -
@
- - {/snippet} +
@
+ {/snippet} {#snippet content({ close })} - - { - addContextToSelection(element) - close() - }} - onSelectWorkspaceItem={(element) => { - addContextToSelection(element) - close() - }} - /> - - {/snippet} + { + void addContextToSelection(element) + close() + }} + onSelectWorkspaceItem={(element) => { + void addContextToSelection(element) + close() + }} + /> + {/snippet} - {#each selectedContext as element} + {#each selectedContext as element (element.type + '-' + element.title)} addContextToSelection(contextElement)} + onAddContext={(contextElement) => void addContextToSelection(contextElement)} onSendRequest={() => { if (disabled) { return @@ -423,25 +448,21 @@
{#snippet trigger()} - -
@
- - {/snippet} +
@
+ {/snippet} {#snippet content({ close })} - - { - addContextToSelection(element) - close() - }} - /> - - {/snippet} + { + void addContextToSelection(element) + close() + }} + /> + {/snippet} {#each selectedContext as element (element.type + '-' + element.title)} = {}): 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.") }) diff --git a/frontend/src/lib/components/copilot/chat/app/core.ts b/frontend/src/lib/components/copilot/chat/app/core.ts index 9eaf653149..4cd0110a91 100644 --- a/frontend/src/lib/components/copilot/chat/app/core.ts +++ b/frontend/src/lib/components/copilot/chat/app/core.ts @@ -103,6 +103,13 @@ export interface DataTableSchema { error?: string } +export interface AppDatatableMetadata { + datatable_name: string + schemas: Record + 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 + /** List configured datatables with schema/table names only. */ + listDatatableTables: () => Promise + /** Get columns for one datatable table. */ + getDatatableTableSchema: ( + datatableName: string, + schemaName: string, + tableName: string + ) => Promise> /** 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 - tableCount: number - error?: string -} - -function summarizeDatatables(datatables: DataTableSchema[]): AppDatatableMetadata[] { - return datatables.map((datatable) => { - const schemas: Record = {} - 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[] => [ 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[] => [ 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[] => [ 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` + } } } } diff --git a/frontend/src/lib/components/copilot/chat/context.ts b/frontend/src/lib/components/copilot/chat/context.ts index f110ba2c70..f077b4a46f 100644 --- a/frontend/src/lib/components/copilot/chat/context.ts +++ b/frontend/src/lib/components/copilot/chat/context.ts @@ -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 + /** The table columns: column_name -> compact_type. Loaded only when a table is explicitly selected. */ + columns?: Record } export function createAppSelectedContext(options: SelectedContext = {}): SelectedContext { @@ -184,7 +184,7 @@ export function createAppDatatableContextElement( datatableName: string, schemaName: string, tableName: string, - columns: Record + columns?: Record ): 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) ) ) }) diff --git a/frontend/src/lib/components/raw_apps/RawAppEditor.svelte b/frontend/src/lib/components/raw_apps/RawAppEditor.svelte index 3ad55f4f4c..42fab104c1 100644 --- a/frontend/src/lib/components/raw_apps/RawAppEditor.svelte +++ b/frontend/src/lib/components/raw_apps/RawAppEditor.svelte @@ -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 + error?: string + } + + function countTables(schemas: Record): 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 = {} + 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 => { + listDatatableTables: async (): Promise => { 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> => { + if (!$workspaceStore) { + return {} } - // Build a map of whitelisted tables per datatable: datatable -> schema -> Set - const whitelistedTables = new Map>>() - 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>> = {} - for (const [schemaName, tables] of Object.entries(schema.schemas)) { - const allowedTablesInSchema = allowedTables.get(schemaName) - if (!allowedTablesInSchema) continue - - const filteredTables: Record> = {} - 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() } } }) diff --git a/frontend/src/lib/components/raw_apps/dataTableRefUtils.test.ts b/frontend/src/lib/components/raw_apps/dataTableRefUtils.test.ts new file mode 100644 index 0000000000..500311aecb --- /dev/null +++ b/frontend/src/lib/components/raw_apps/dataTableRefUtils.test.ts @@ -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) + }) +}) diff --git a/frontend/src/lib/components/raw_apps/dataTableRefUtils.ts b/frontend/src/lib/components/raw_apps/dataTableRefUtils.ts index 89c727ee91..b6f8c636d3 100644 --- a/frontend/src/lib/components/raw_apps/dataTableRefUtils.ts +++ b/frontend/src/lib/components/raw_apps/dataTableRefUtils.ts @@ -25,6 +25,60 @@ export const DEFAULT_DATA: RawAppData = { schema: undefined } +export type DataTableWhitelist = { + datatables: Set + allTablesDatatables: Set + tables: Map>> +} + +export function buildDataTableWhitelist(refs: DataTableRef[]): DataTableWhitelist { + const datatables = new Set() + const allTablesDatatables = new Set() + const tables = new Map>>() + + 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: /:
or /
(for public schema)