Enhance TableSourceCollector to include query and referenced columns

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<void> {
         // 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<void> {
         expr.input.accept(this);
         expr.castType.accept(this);
     }
-}
+}

tests/transformers/TableSourceCollector.test.ts

@@ -436,4 +436,52 @@ order by
         // Assert
         expect(tableNames.length).toBe(0);
     });
-});
+
+    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([]);
+    });
+});