From e8b8f7b350916d2b2bc2425d4fb90863f876bd8d Mon Sep 17 00:00:00 2001
From: Rene Saarsoo <nene@triin.net>
Date: Wed, 22 Mar 2023 12:25:28 +0200
Subject: [PATCH] Support custom key funcion in custom param types

---
 src/lexer/Tokenizer.ts        | 14 ++++++++------
 src/lexer/TokenizerOptions.ts | 12 ++++++++++--
 src/lexer/regexFactory.ts     |  9 +--------
 test/options/paramTypes.ts    | 27 ++++++++++++++++++++++++++-
 4 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/src/lexer/Tokenizer.ts b/src/lexer/Tokenizer.ts
index ec6b1bacf0..f1a8e6f3a5 100644
--- a/src/lexer/Tokenizer.ts
+++ b/src/lexer/Tokenizer.ts
@@ -2,7 +2,7 @@ import { Token, TokenType } from './token.js';
 import * as regex from './regexFactory.js';
 import { ParamTypes, TokenizerOptions } from './TokenizerOptions.js';
 import TokenizerEngine, { TokenRule } from './TokenizerEngine.js';
-import { escapeRegExp } from './regexUtil.js';
+import { escapeRegExp, patternToRegex } from './regexUtil.js';
 import { equalizeWhitespace, Optional } from '../utils.js';
 import { NestedComment } from './NestedComment.js';
 
@@ -227,11 +227,13 @@ export default class Tokenizer {
         type: TokenType.POSITIONAL_PARAMETER,
         regex: paramTypes.positional ? /[?]/y : undefined,
       },
-      {
-        type: TokenType.CUSTOM_PARAMETER,
-        regex: regex.customParameter(paramTypes.custom),
-        key: v => v,
-      },
+      ...paramTypes.custom.map(
+        (customParam): TokenRule => ({
+          type: TokenType.CUSTOM_PARAMETER,
+          regex: patternToRegex(customParam.regex),
+          key: customParam.key ?? (v => v),
+        })
+      ),
     ]);
   }
 
diff --git a/src/lexer/TokenizerOptions.ts b/src/lexer/TokenizerOptions.ts
index 06609ef8d4..beca4575e5 100644
--- a/src/lexer/TokenizerOptions.ts
+++ b/src/lexer/TokenizerOptions.ts
@@ -40,8 +40,16 @@ export interface ParamTypes {
   // Prefixes for quoted parameter placeholders to support, e.g. :"name"
   // The type of quotes will depend on `identifierTypes` option.
   quoted?: (':' | '@' | '$')[];
-  // Array of regular expressions
-  custom?: RegexPattern[];
+  // Custom parameter type definitions
+  custom?: CustomParameter[];
+}
+
+export interface CustomParameter {
+  // Regex pattern for matching the parameter
+  regex: string;
+  // Takes the matched parameter string and returns the name of the parameter
+  // For example we might match "{foo}" and the name would be "foo".
+  key?: (text: string) => string;
 }
 
 export interface TokenizerOptions {
diff --git a/src/lexer/regexFactory.ts b/src/lexer/regexFactory.ts
index 122d5dd288..a466e0ce57 100644
--- a/src/lexer/regexFactory.ts
+++ b/src/lexer/regexFactory.ts
@@ -1,6 +1,6 @@
 import { sortByLengthDesc } from '../utils.js';
 
-import { IdentChars, QuoteType, RegexPattern, VariableType } from './TokenizerOptions.js';
+import { IdentChars, QuoteType, VariableType } from './TokenizerOptions.js';
 import { escapeRegExp, patternToRegex, prefixesPattern, withDashes } from './regexUtil.js';
 
 /**
@@ -66,13 +66,6 @@ export const parameter = (paramTypes: string[], pattern: string): RegExp | undef
   return patternToRegex(`(?:${typesRegex})(?:${pattern})`);
 };
 
-export const customParameter = (patterns: RegexPattern[]): RegExp | undefined => {
-  if (!patterns.length) {
-    return undefined;
-  }
-  return patternToRegex(patterns.map(p => p.regex).join('|'));
-};
-
 const buildQStringPatterns = () => {
   const specialDelimiterMap = {
     '<': '>',
diff --git a/test/options/paramTypes.ts b/test/options/paramTypes.ts
index 72e5317715..3b1d466d77 100644
--- a/test/options/paramTypes.ts
+++ b/test/options/paramTypes.ts
@@ -57,7 +57,7 @@ export default function supportsParamTypes(format: FormatFn) {
   // - it likely works when the other paramTypes tests work
   // - it's the config that's least likely to be actually used in practice.
 
-  describe('when paramTypes.custom=[regex]', () => {
+  describe('when paramTypes.custom=[...]', () => {
     it('replaces %blah% numbered placeholders with param values', () => {
       const result = format('SELECT %1%, %2%, %3%;', {
         paramTypes: { custom: [{ regex: '%[0-9]+%' }] },
@@ -70,5 +70,30 @@ export default function supportsParamTypes(format: FormatFn) {
           third;
       `);
     });
+
+    it('supports custom function for extracting parameter name', () => {
+      const result = format('SELECT %1%, %2%, %3%;', {
+        paramTypes: { custom: [{ regex: '%[0-9]+%', key: v => v.slice(1, -1) }] },
+        params: { '1': 'first', '2': 'second', '3': 'third' },
+      });
+      expect(result).toBe(dedent`
+        SELECT
+          first,
+          second,
+          third;
+      `);
+    });
+
+    it('supports multiple custom param types', () => {
+      const result = format('SELECT %1%, {2};', {
+        paramTypes: { custom: [{ regex: '%[0-9]+%' }, { regex: String.raw`\{[0-9]\}` }] },
+        params: { '%1%': 'first', '{2}': 'second' },
+      });
+      expect(result).toBe(dedent`
+        SELECT
+          first,
+          second;
+      `);
+    });
   });
 }