diff --git a/src/transformers/TableSourceCollector.ts b/src/transformers/TableSourceCollector.ts index 670c646..b7245f2 100644 --- a/src/transformers/TableSourceCollector.ts +++ b/src/transformers/TableSourceCollector.ts @@ -11,6 +11,7 @@ import { TypeValue } from "../models/ValueComponent"; import { CTECollector } from "./CTECollector"; +import { SelectableColumnCollector } from "./SelectableColumnCollector"; /** * A visitor that collects all table source names from a SQL query structure. @@ -287,6 +288,15 @@ export class TableSourceCollector implements SqlComponentVisitor { // Check if this is a table managed by a CTE if (!this.tableNameMap.has(identifier) && !this.isCTETable(source.table.name)) { this.tableNameMap.set(identifier, true); + + // Collect referenced columns using SelectableColumnCollector + const columnCollector = new SelectableColumnCollector(); + const referencedColumns = columnCollector.collect(source); + + // Add the query and referenced columns to the table source + source.query = source; + source.referencedColumns = referencedColumns.map(col => col.name); + this.tableSources.push(source); } } @@ -458,4 +468,4 @@ export class TableSourceCollector implements SqlComponentVisitor { expr.input.accept(this); expr.castType.accept(this); } -} \ No newline at end of file +} diff --git a/tests/transformers/TableSourceCollector.test.ts b/tests/transformers/TableSourceCollector.test.ts index 1057c6d..a6330da 100644 --- a/tests/transformers/TableSourceCollector.test.ts +++ b/tests/transformers/TableSourceCollector.test.ts @@ -436,4 +436,52 @@ order by // Assert expect(tableNames.length).toBe(0); }); -}); \ No newline at end of file + + test('includes the query the table belongs to in the return value', () => { + // Arrange + const sql = `SELECT * FROM users`; + const query = SelectQueryParser.parse(sql); + const collector = new TableSourceCollector(); + + // Act + collector.visit(query); + const tableSources = collector.getTableSources(); + + // Assert + expect(tableSources.length).toBe(1); + expect(tableSources[0].table.name).toBe('users'); + expect(tableSources[0].query).toBeDefined(); + }); + + test('includes the list of column names referenced from the table in the return value', () => { + // Arrange + const sql = `SELECT id, name FROM users`; + const query = SelectQueryParser.parse(sql); + const collector = new TableSourceCollector(); + + // Act + collector.visit(query); + const tableSources = collector.getTableSources(); + + // Assert + expect(tableSources.length).toBe(1); + expect(tableSources[0].table.name).toBe('users'); + expect(tableSources[0].referencedColumns).toEqual(['id', 'name']); + }); + + test('handles wildcard columns correctly', () => { + // Arrange + const sql = `SELECT * FROM users`; + const query = SelectQueryParser.parse(sql); + const collector = new TableSourceCollector(); + + // Act + collector.visit(query); + const tableSources = collector.getTableSources(); + + // Assert + expect(tableSources.length).toBe(1); + expect(tableSources[0].table.name).toBe('users'); + expect(tableSources[0].referencedColumns).toEqual([]); + }); +});