diff --git a/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.test.ts b/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.test.ts index b7e00b41..979873a4 100644 --- a/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.test.ts +++ b/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.test.ts @@ -36,6 +36,7 @@ describe('DiscoverDataTool', () => { catalog.searchByName.mockResolvedValue([ { kind: 'table', + connectionName: 'warehouse', ref: { catalog: null, db: 'public', name: 'orders' }, display: 'public.orders', matchedOn: 'name', @@ -55,6 +56,33 @@ describe('DiscoverDataTool', () => { expect(result.structured.raw?.hits).toHaveLength(1); }); + it('includes connectionName on raw schema hits so entity_details can follow up', async () => { + const multiConnectionContext: ToolContext = { + ...context, + session: { allowedConnectionNames: new Set(['warehouse', 'analytics']) } as any, + }; + catalog.searchByName.mockImplementation(async (connectionName: string, query: string) => [ + { + kind: 'table', + connectionName, + ref: { catalog: null, db: 'public', name: `${connectionName}_${query}` }, + display: `public.${connectionName}_${query}`, + matchedOn: 'name', + }, + ]); + + const result = await tool.call({ query: 'orders', limit: 10 }, multiConnectionContext); + + expect(catalog.searchByName).toHaveBeenCalledWith('analytics', 'orders', 10); + expect(catalog.searchByName).toHaveBeenCalledWith('warehouse', 'orders', 10); + expect(result.markdown).toContain('connectionName=analytics'); + expect(result.markdown).toContain('connectionName=warehouse'); + expect(result.markdown).toContain( + 'entity_details({connectionName: "analytics", targets: [{display: "public.analytics_orders"}]})', + ); + expect(result.structured.raw?.hits.map((hit) => hit.connectionName)).toEqual(['analytics', 'warehouse']); + }); + it('refuses explicit out-of-scope connection names', async () => { const result = await tool.call({ query: 'orders', connectionName: 'billing' }, context); diff --git a/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.ts b/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.ts index 4bd56509..667d8f83 100644 --- a/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.ts +++ b/packages/context/src/ingest/tools/warehouse-verification/discover-data.tool.ts @@ -113,11 +113,18 @@ export class DiscoverDataTool extends BaseTool { rawHits.push(...(await catalog.searchByName(connectionName, query, limit))); } if (rawHits.length > 0) { - parts.push('## Raw Warehouse Schema', '> use `entity_details({connectionName, targets: [{display}]})` for full DDL + sample values'); + parts.push( + '## Raw Warehouse Schema', + '> use `entity_details({connectionName, targets: [{display}]})` for full DDL + sample values', + ); parts.push( rawHits .slice(0, limit) - .map((hit) => `- ${hit.kind}: ${hit.display} (matched on ${hit.matchedOn})`) + .map( + (hit) => + `- ${hit.kind}: ${hit.display} [connectionName=${hit.connectionName}] (matched on ${hit.matchedOn}) - ` + + `follow up with \`entity_details({connectionName: "${hit.connectionName}", targets: [{display: "${hit.display}"}]})\``, + ) .join('\n'), ); raw = { hits: rawHits.slice(0, limit) }; diff --git a/packages/context/src/ingest/tools/warehouse-verification/warehouse-catalog.service.ts b/packages/context/src/ingest/tools/warehouse-verification/warehouse-catalog.service.ts index 9989ff38..691f88e9 100644 --- a/packages/context/src/ingest/tools/warehouse-verification/warehouse-catalog.service.ts +++ b/packages/context/src/ingest/tools/warehouse-verification/warehouse-catalog.service.ts @@ -38,8 +38,20 @@ export interface TableDetail { } export type RawSchemaHit = - | { kind: 'table'; ref: KtxTableRef; display: string; matchedOn: 'name' | 'db' | 'comment' | 'description' } - | { kind: 'column'; ref: KtxTableRef & { column: string }; display: string; matchedOn: 'name' | 'comment' | 'description' }; + | { + kind: 'table'; + connectionName: string; + ref: KtxTableRef; + display: string; + matchedOn: 'name' | 'db' | 'comment' | 'description'; + } + | { + kind: 'column'; + connectionName: string; + ref: KtxTableRef & { column: string }; + display: string; + matchedOn: 'name' | 'comment' | 'description'; + }; export interface DisplayTargetResolution { resolved: (KtxTableRef & { column?: string }) | null; @@ -370,6 +382,7 @@ export class WarehouseCatalogService { if (tableMatch) { hits.push({ kind: 'table', + connectionName, ref: { catalog: table.catalog, db: table.db, name: table.name }, display: formatDisplay(catalog.driver, table), matchedOn: tableMatch, @@ -382,6 +395,7 @@ export class WarehouseCatalogService { } hits.push({ kind: 'column', + connectionName, ref: { catalog: table.catalog, db: table.db, name: table.name, column: column.name }, display: `${formatDisplay(catalog.driver, table)}.${column.name}`, matchedOn: columnMatch,