fix numeric and custom tokens

Author: not-night-but

src/defines.ts

@@ -81,9 +81,8 @@ export interface ParamTypes {
   numbered?: Array<"?" | ":" | "$">,
   named?: Array<":" | "@" | "$">,
   quoted?: Array<":" | "@" | "$">,
-  // regex is for identifying that it is a param, key is how the token is translated to an object value for the formatter,
-  // may not be necessary here, we shal see
-  custom?: Array<{regex: string, key?: (text: string) => string }>
+  // regex for identifying that it is a param
+  custom?: Array<string>
 }
 
 export interface IdentifyOptions {

src/index.ts

@@ -22,6 +22,7 @@ export function identify(query: string, options: IdentifyOptions = {}): Identify
   }
 
   const result = parse(query, isStrict, dialect, options.identifyTables, options.paramTypes);
+  const sort = dialect === 'psql' && !options.paramTypes;
 
   return result.body.map((statement) => {
     const result: IdentifyResult = {
@@ -31,7 +32,7 @@ export function identify(query: string, options: IdentifyOptions = {}): Identify
       type: statement.type,
       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,
+      parameters: sort ? statement.parameters.sort() : statement.parameters,
       tables: statement.tables || [],
     };
     return result;

src/tokenizer.ts

@@ -254,85 +254,85 @@ function scanString(state: State, endToken: Char): Token {
 }
 
 function getCustomParam(state: State, paramTypes: ParamTypes): string | null | undefined {
-  const matches = paramTypes?.custom?.map(({ regex }) => {
-    const reg = new RegExp(`(?:${regex})`, 'u');
-    return reg.exec(state.input);
+  const matches = paramTypes?.custom?.map((regex) => {
+    const reg = new RegExp(`^(?:${regex})`, 'u');
+    return reg.exec(state.input.slice(state.start));
   }).filter((value) => !!value)[0];
 
   return matches ? matches[0] : null;
 }
 
-function scanParameter(state: State, dialect: Dialect, paramTypes?: ParamTypes): Token {
-  // user has defined wanted param types, so we only evaluate them
-  if (paramTypes) {
-    const curCh: any = state.input[0];
-    let nextChar = peek(state);
-    let matched = false
+function scanCustomParameter(state: State, dialect: Dialect, paramTypes: ParamTypes): Token {
 
-    // this could be a named parameter that just starts with a number (ugh)
-    if (paramTypes.numbered && paramTypes.numbered.length && paramTypes.numbered.includes(curCh)) {
-      const maybeNumbers = state.input.slice(1, state.input.length);
-      if (nextChar !== null && !isNaN(Number(nextChar)) && /^\d+$/.test(maybeNumbers)) {
-        do {
-          nextChar = read(state);
-        } while (nextChar !== null && !isNaN(Number(nextChar)) && !isWhitespace(nextChar));
-
-        if (nextChar !== null) unread(state);
-        matched = true;
-      }
-    } 
-    
-    if (!matched && paramTypes.named && paramTypes.named.length && paramTypes.named.includes(curCh)) {
-      if (!isQuotedIdentifier(nextChar, dialect)) {
-        while (isAlphaNumeric(peek(state))) read(state);
-        matched = true;
-      }
-    } 
-    
-    if (!matched && paramTypes.quoted && paramTypes.quoted.length && paramTypes.quoted.includes(curCh)) {
-      if (isQuotedIdentifier(nextChar, dialect)) {
-        const endChars = new Map<string, string>([
-          ['"', '"'],
-          ['[', ']'],
-          ['`', '`']
-        ]);
-        const quoteChar = read(state) as string;
-        const end = endChars.get(quoteChar);
-        // end when we reach the end quote
-        while ((isAlphaNumeric(peek(state)) || peek(state) === ' ') && peek(state) != end) read(state);
-
-        // read the end quote
-        read(state);
-
-        matched = true;
-      }
-    } 
-    
-    if (!matched && paramTypes.custom && paramTypes.custom.length) {
-      const custom = getCustomParam(state, paramTypes);
+  const curCh: any = state.input[state.start];
+  let nextChar = peek(state);
+  let matched = false
 
-      if (custom) {
-        read(state, custom.length);
-        matched = true;
-      }
-    } 
-    
-    if (!matched && curCh !== '?' && nextChar !== null) { // not positional, panic
-      return {
-        type: 'parameter',
-        value: 'unknown',
-        start: state.start,
-        end: state.end
-      }
-    }
+  if (paramTypes.numbered && paramTypes.numbered.length && paramTypes.numbered.includes(curCh)) {
+    const endIndex = state.input.slice(state.start).split('').findIndex((val) => isWhitespace(val));
+    const maybeNumbers = state.input.slice(state.start + 1, endIndex > 0 ? state.start + endIndex : state.end + 1);
+    if (nextChar !== null && !isNaN(Number(nextChar)) && /^\d+$/.test(maybeNumbers)) {
+      let nextChar: Char = null;
+      do {
+        nextChar = read(state);
+      } while (nextChar !== null && !isNaN(Number(nextChar)) && !isWhitespace(nextChar));
 
-    const value = state.input.slice(state.start, state.position + 1);
+      if (nextChar !== null) unread(state);
+      matched = true;
+    }
+  } 
+  
+  if (!matched && paramTypes.named && paramTypes.named.length && paramTypes.named.includes(curCh)) {
+    if (!isQuotedIdentifier(nextChar, dialect)) {
+      while (isAlphaNumeric(peek(state))) read(state);
+      matched = true;
+    }
+  } 
+  
+  if (!matched && paramTypes.quoted && paramTypes.quoted.length && paramTypes.quoted.includes(curCh)) {
+    if (isQuotedIdentifier(nextChar, dialect)) {
+      const quoteChar = read(state) as string;
+      // end when we reach the end quote
+      while ((isAlphaNumeric(peek(state)) || peek(state) === ' ') && peek(state) != ENDTOKENS[quoteChar]) read(state);
+
+      // read the end quote
+      read(state);
+
+      matched = true;
+    }
+  } 
+  
+  if (!matched && paramTypes.custom && paramTypes.custom.length) {
+    const custom = getCustomParam(state, paramTypes);
+
+    if (custom) {
+      read(state, custom.length);
+      matched = true;
+    }
+  } 
+  
+  if (!matched && !paramTypes.positional) { // not positional, panic
     return {
       type: 'parameter',
-      value,
+      value: 'unknown',
       start: state.start,
-      end: state.start + value.length - 1,
-    };
+      end: state.end
+    }
+  }
+
+  const value = state.input.slice(state.start, state.position + 1);
+  return {
+    type: 'parameter',
+    value,
+    start: state.start,
+    end: state.start + value.length - 1,
+  };
+}
+
+function scanParameter(state: State, dialect: Dialect, paramTypes?: ParamTypes): Token {
+  // user has defined wanted param types, so we only evaluate them
+  if (paramTypes) {
+    return scanCustomParameter(state, dialect, paramTypes);
   }
 
   if (['mysql', 'generic', 'sqlite'].includes(dialect)) {
@@ -495,17 +495,17 @@ function isString(ch: Char, dialect: Dialect): boolean {
 }
 
 function isCustomParam(state: State, paramTypes: ParamTypes): boolean | undefined {
-  return paramTypes?.custom?.some(({ regex }) => {
-    const reg = new RegExp(`(?:${regex})`, 'uy');
-    return reg.test(state.input);
+  return paramTypes?.custom?.some((regex) => {
+    const reg = new RegExp(`^(?:${regex})`, 'uy');
+    return reg.test(state.input.slice(state.start));
   })
 }
 
 function isParameter(ch: Char, state: State, dialect: Dialect, paramTypes?: ParamTypes): boolean {
   if (paramTypes && ch !== null) {
     const curCh: any = ch;
     const nextChar = peek(state);
-    if (paramTypes.positional && ch === '?' && nextChar === null) return true;
+    if (paramTypes.positional && ch === '?' && (nextChar === null || isWhitespace(nextChar))) return true;
 
     if (paramTypes.numbered && paramTypes.numbered.length && paramTypes.numbered.includes(curCh)) {
       if (nextChar !== null && !isNaN(Number(nextChar))) {

test/index.spec.ts

@@ -1,5 +1,6 @@
 import { Dialect, getExecutionType, identify } from '../src/index';
 import { expect } from 'chai';
+import { ParamTypes } from '../src/defines';
 
 describe('identify', () => {
   it('should throw error for invalid dialect', () => {
@@ -22,6 +23,29 @@ describe('identify', () => {
     ]);
   });
 
+  it('should identify custom parameters', () => {
+    const paramTypes: ParamTypes = {
+      positional: true,
+      numbered: ['$'],
+      named: [':'],
+      quoted: [':'],
+      custom: [ '\\{[a-zA-Z0-9_]+\\}' ]
+    };
+    const query = `SELECT * FROM foo WHERE bar = ? AND baz = $1 AND fizz = :fizzz  AND buzz = :"buzz buzz" AND foo2 = {fooo}`;
+
+    expect(identify(query, { dialect: 'psql', paramTypes })).to.eql([
+      {
+        start: 0,
+        end: 104,
+        text: query,
+        type: 'SELECT',
+        executionType: 'LISTING',
+        parameters: ['?', '$1', ':fizzz', ':"buzz buzz"', '{fooo}'],
+        tables: []
+      }
+    ])
+  })
+
   it('should identify tables in simple for basic cases', () => {
     expect(
       identify('SELECT * FROM foo JOIN bar ON foo.id = bar.id', { identifyTables: true }),

test/tokenizer/index.spec.ts

@@ -562,7 +562,7 @@ describe('scan', () => {
 
         it('should allow custom parameters for all dialects', () => {
           const paramTypes: ParamTypes = {
-            custom: [{ regex: '\\{[a-zA-Z0-9_]+\\}' }]
+            custom: [ '\\{[a-zA-Z0-9_]+\\}' ]
           };
 
           const expected = {
@@ -585,6 +585,37 @@ describe('scan', () => {
             quoted: [':'],
             custom: []
           };
+
+          const expected = [
+            {
+              type: 'parameter',
+              value: '?',
+              start: 0,
+              end: 0
+            },
+            {
+              type: 'parameter',
+              value: ':123',
+              start: 0,
+              end: 3
+            },
+            {
+              type: 'parameter',
+              value: ':123hello',
+              start: 0,
+              end: 8
+            },
+            {
+              type: 'parameter',
+              value: ':"named param"',
+              start: 0,
+              end: 13
+            }
+          ];
+
+          expected.forEach((expected) => {
+            expect(scanToken(initState(expected.value), 'mssql', paramTypes)).to.eql(expected);
+          })
         })
       });
     });