Add initial support for table identification (#76)

Co-authored-by: Matthew Rathbone <[email protected]>

Author: azmy60

package.json

@@ -34,12 +34,11 @@
     "prettier": "^2.3.2",
     "terser-webpack-plugin": "^5.1.1",
     "ts-loader": "^8.0.17",
-    "ts-node": "^10.9.1",
+    "ts-node": "^10.9.2",
     "typescript": "^4.1.5",
     "webpack": "^5.11.1",
     "webpack-cli": "^4.3.1"
   },
-  "dependencies": {},
   "engines": {
     "node": ">= 10.13"
   }

src/defines.ts

@@ -79,6 +79,7 @@ export type ExecutionType = 'LISTING' | 'MODIFICATION' | 'INFORMATION' | 'ANON_B
 export interface IdentifyOptions {
   strict?: boolean;
   dialect?: Dialect;
+  identifyTables?: boolean;
 }
 
 export interface IdentifyResult {
@@ -88,6 +89,7 @@ export interface IdentifyResult {
   type: StatementType;
   executionType: ExecutionType;
   parameters: string[];
+  tables: string[];
 }
 
 export interface Statement {
@@ -101,6 +103,8 @@ export interface Statement {
   algorithm?: number;
   sqlSecurity?: number;
   parameters: string[];
+  tables: string[];
+  isCte?: boolean;
 }
 
 export interface ConcreteStatement extends Statement {
@@ -124,6 +128,7 @@ export interface Token {
     | 'semicolon'
     | 'keyword'
     | 'parameter'
+    | 'table'
     | 'unknown';
   value: string;
   start: number;

src/index.ts

@@ -21,7 +21,7 @@ export function identify(query: string, options: IdentifyOptions = {}): Identify
     throw new Error(`Unknown dialect. Allowed values: ${DIALECTS.join(', ')}`);
   }
 
-  const result = parse(query, isStrict, dialect);
+  const result = parse(query, isStrict, dialect, options.identifyTables);
 
   return result.body.map((statement) => {
     const result: IdentifyResult = {
@@ -32,6 +32,7 @@ export function identify(query: string, options: IdentifyOptions = {}): Identify
       executionType: statement.executionType,
       // we want to sort the postgres params: $1 $2 $3, regardless of the order they appear
       parameters: dialect === 'psql' ? statement.parameters.sort() : statement.parameters,
+      tables: statement.tables || [],
     };
     return result;
   });

src/parser.ts

@@ -98,6 +98,10 @@ const statementsWithEnds = [
   'UNKNOWN',
 ];
 
+// keywords that come directly before a table name.
+// v1 - keeping it very simple.
+const PRE_TABLE_KEYWORDS = /^from$|^join$|^into$/i;
+
 const blockOpeners: Record<Dialect, string[]> = {
   generic: ['BEGIN', 'CASE'],
   psql: ['BEGIN', 'CASE', 'LOOP', 'IF'],
@@ -111,29 +115,36 @@ const blockOpeners: Record<Dialect, string[]> = {
 interface ParseOptions {
   isStrict: boolean;
   dialect: Dialect;
+  identifyTables: boolean;
 }
 
 function createInitialStatement(): Statement {
   return {
     start: -1,
     end: 0,
     parameters: [],
+    tables: [],
   };
 }
 
-function nextNonWhitespaceToken(state: State): Token {
+function nextNonWhitespaceToken(state: State, dialect: Dialect): Token {
   let token: Token;
   do {
     state = initState({ prevState: state });
-    token = scanToken(state);
+    token = scanToken(state, dialect);
   } while (token.type === 'whitespace');
   return token;
 }
 
 /**
  * Parser
  */
-export function parse(input: string, isStrict = true, dialect: Dialect = 'generic'): ParseResult {
+export function parse(
+  input: string,
+  isStrict = true,
+  dialect: Dialect = 'generic',
+  identifyTables = false,
+): ParseResult {
   const topLevelState = initState({ input });
   const topLevelStatement: ParseResult = {
     type: 'QUERY',
@@ -164,14 +175,13 @@ export function parse(input: string, isStrict = true, dialect: Dialect = 'generi
   while (prevState.position < topLevelState.end) {
     const tokenState = initState({ prevState });
     const token = scanToken(tokenState, dialect);
-    const nextToken = nextNonWhitespaceToken(tokenState);
+    const nextToken = nextNonWhitespaceToken(tokenState, dialect);
 
     if (!statementParser) {
       // ignore blank tokens before the start of a CTE / not part of a statement
       if (!cteState.isCte && ignoreOutsideBlankTokens.includes(token.type)) {
         topLevelStatement.tokens.push(token);
         prevState = tokenState;
-        continue;
       } else if (
         !cteState.isCte &&
         token.type === 'keyword' &&
@@ -181,7 +191,7 @@ export function parse(input: string, isStrict = true, dialect: Dialect = 'generi
         topLevelStatement.tokens.push(token);
         cteState.state = tokenState;
         prevState = tokenState;
-        continue;
+
         // If we're scanning in a CTE, handle someone putting a semicolon anywhere (after 'with',
         // after semicolon, etc.) along it to "early terminate".
       } else if (cteState.isCte && token.type === 'semicolon') {
@@ -193,12 +203,12 @@ export function parse(input: string, isStrict = true, dialect: Dialect = 'generi
           type: 'UNKNOWN',
           executionType: 'UNKNOWN',
           parameters: [],
+          tables: [],
         });
         cteState.isCte = false;
         cteState.asSeen = false;
         cteState.statementEnd = false;
         cteState.parens = 0;
-        continue;
       } else if (cteState.isCte && !cteState.statementEnd) {
         if (cteState.asSeen) {
           if (token.value === '(') {
@@ -215,14 +225,13 @@ export function parse(input: string, isStrict = true, dialect: Dialect = 'generi
 
         topLevelStatement.tokens.push(token);
         prevState = tokenState;
-        continue;
       } else if (cteState.isCte && cteState.statementEnd && token.value === ',') {
         cteState.asSeen = false;
         cteState.statementEnd = false;
 
         topLevelStatement.tokens.push(token);
         prevState = tokenState;
-        continue;
+
         // Ignore blank tokens after the end of the CTE till start of statement
       } else if (
         cteState.isCte &&
@@ -231,28 +240,32 @@ export function parse(input: string, isStrict = true, dialect: Dialect = 'generi
       ) {
         topLevelStatement.tokens.push(token);
         prevState = tokenState;
-        continue;
+      } else {
+        statementParser = createStatementParserByToken(token, nextToken, {
+          isStrict,
+          dialect,
+          identifyTables,
+        });
+        if (cteState.isCte) {
+          statementParser.getStatement().start = cteState.state.start;
+          statementParser.getStatement().isCte = true;
+          cteState.isCte = false;
+          cteState.asSeen = false;
+          cteState.statementEnd = false;
+        }
       }
+    } else {
+      statementParser.addToken(token, nextToken);
+      topLevelStatement.tokens.push(token);
+      prevState = tokenState;
 
-      statementParser = createStatementParserByToken(token, nextToken, { isStrict, dialect });
-      if (cteState.isCte) {
-        statementParser.getStatement().start = cteState.state.start;
-        cteState.isCte = false;
-        cteState.asSeen = false;
-        cteState.statementEnd = false;
+      const statement = statementParser.getStatement();
+      if (statement.endStatement) {
+        statement.end = token.end;
+        topLevelStatement.body.push(statement as ConcreteStatement);
+        statementParser = null;
       }
     }
-
-    statementParser.addToken(token, nextToken);
-    topLevelStatement.tokens.push(token);
-    prevState = tokenState;
-
-    const statement = statementParser.getStatement();
-    if (statement.endStatement) {
-      statement.end = token.end;
-      topLevelStatement.body.push(statement as ConcreteStatement);
-      statementParser = null;
-    }
   }
 
   // last statement without ending key
@@ -708,7 +721,7 @@ function createUnknownStatementParser(options: ParseOptions) {
 function stateMachineStatementParser(
   statement: Statement,
   steps: Step[],
-  { isStrict, dialect }: ParseOptions,
+  { isStrict, dialect, identifyTables }: ParseOptions,
 ): StatementParser {
   let currentStepIndex = 0;
   let prevToken: Token | undefined;
@@ -808,6 +821,18 @@ function stateMachineStatementParser(
         }
       }
 
+      if (
+        identifyTables &&
+        PRE_TABLE_KEYWORDS.exec(token.value) &&
+        !statement.isCte &&
+        statement.type?.match(/SELECT|INSERT/)
+      ) {
+        const tableValue = nextToken.value;
+        if (!statement.tables.includes(tableValue)) {
+          statement.tables.push(tableValue);
+        }
+      }
+
       if (
         token.type === 'parameter' &&
         (token.value === '?' || !statement.parameters.includes(token.value))

test/identifier/inner-statements.spec.ts

@@ -16,6 +16,7 @@ describe('identifier', () => {
           type: 'INSERT',
           executionType: 'MODIFICATION',
           parameters: [],
+          tables: [],
         },
       ];
 
@@ -34,6 +35,7 @@ describe('identifier', () => {
           type: 'INSERT',
           executionType: 'MODIFICATION',
           parameters: [],
+          tables: [],
         },
       ];
 
@@ -54,6 +56,7 @@ describe('identifier', () => {
           type: 'INSERT',
           executionType: 'MODIFICATION',
           parameters: [],
+          tables: [],
         },
       ];
 
@@ -75,6 +78,7 @@ describe('identifier', () => {
           type: 'INSERT',
           executionType: 'MODIFICATION',
           parameters: [],
+          tables: [],
         },
       ];