Support for Interpolated Variables from Connections Config (#1064)

* [#744] Support for Interpolated Variables * Update extension.ts * Apply suggestions from code review Co-authored-by: John Murray <johnm@georgejames.com> Co-authored-by: John Murray <johnm@georgejames.com>
mtxr · Dec 22, 2022 · 6bb426d · 6bb426d
1 parent 777b021
commit 6bb426d
Show file tree

Hide file tree

Showing 6 changed files with 62 additions and 29 deletions.
diff --git a/packages/extension/package.json b/packages/extension/package.json
@@ -544,7 +544,7 @@
                 },
                 "sqltools.queryParams.regex": {
                     "type": "string",
-                    "default": "\\$[\\d]+|\\$\\[[\\d\\w]+\\]",
+                    "default": "\\$[\\d]+|\\$\\[([\\d\\w]+)\\]",
                     "description": "RegEx used to identify query parameters.",
                     "required": true
                 },
@@ -721,6 +721,15 @@
                                         "default": null
                                     }
                                 }
+                            },
+                            "variables": {
+                                "type": [
+                                    "object",
+                                    "null"
+                                ],
+                                "default": null,
+                                "description": "Connection variables in a key/value pair format. Use this property with `sqltools.queryParams.enableReplace` to replace the variables without prompting.",
+                                "properties": {}
                             }
                         }
                     }

diff --git a/packages/plugins/connection-manager/extension.ts b/packages/plugins/connection-manager/extension.ts
@@ -212,32 +212,42 @@ export class ConnectionManagerPlugin implements IExtensionPlugin {
     }
   }
 
-  private replaceParams = async (query: string) => {
+  private replaceParams = async (query: string, conn: IConnection) => {
     if (!Config['queryParams.enableReplace']) return query;
 
-    const params = getQueryParameters(query, Config['queryParams.regex']);
+    const regex = Config['queryParams.regex']
+    const params = getQueryParameters(query, regex);
     if (params.length > 0) {
-      await new Promise<void>((resolve, reject) => {
-        const ib = window.createInputBox();
-        ib.step = 1;
-        ib.totalSteps = params.length;
-        ib.ignoreFocusOut = true;
-        ib.title = `Value for '${params[ib.step - 1].param}' in '${params[ib.step - 1].string}'`;
-        ib.prompt = 'Remember to escape values if needed.'
-        ib.onDidAccept(() => {
-          const r = new RegExp(params[ib.step - 1].param.replace(/([\$\[\]])/g, '\\$1'), 'g');
-          query = query.replace(r, ib.value);
-          ib.step++;
-          if (ib.step > ib.totalSteps) {
-            ib.hide();
-            return resolve();
-          }
-          ib.value = '';
-          ib.title = `Value for '${params[ib.step - 1].param}' in '${params[ib.step - 1].string}'`;
+      const connVariables = conn.variables || {}
+      const connParams = params.filter(p => p.varName && Object.keys(connVariables).includes(p.varName))
+      for (const connParam of connParams) {
+        const r = new RegExp(connParam.param.replace(/([\$\[\]])/g, '\\$1'), 'g');
+        query = query.replace(r, connVariables[connParam.varName]);
+      }
+      const promptParams = params.filter(p => connParams.indexOf(p) === -1)
+      if (promptParams.length > 0) {
+        await new Promise<void>((resolve, reject) => {
+          const ib = window.createInputBox();
+          ib.step = 1;
+          ib.totalSteps = promptParams.length;
+          ib.ignoreFocusOut = true;
+          ib.title = `Value for '${promptParams[ib.step - 1].param}' in '${promptParams[ib.step - 1].string}'`;
+          ib.prompt = 'Remember to escape values if needed.'
+          ib.onDidAccept(() => {
+            const r = new RegExp(promptParams[ib.step - 1].param.replace(/([\$\[\]])/g, '\\$1'), 'g');
+            query = query.replace(r, ib.value);
+            ib.step++;
+            if (ib.step > ib.totalSteps) {
+              ib.hide();
+              return resolve();
+            }
+            ib.value = '';
+            ib.title = `Value for '${promptParams[ib.step - 1].param}' in '${promptParams[ib.step - 1].string}'`;
+          });
+          ib.onDidHide(() => ib.step >= ib.totalSteps && ib.value.trim() ? resolve() : reject(new Error('Didn\'t fill all params. Cancelling...')));
+          ib.show();
         });
-        ib.onDidHide(() => ib.step >= ib.totalSteps && ib.value.trim() ? resolve() : reject(new Error('Didn\'t fill all params. Cancelling...')));
-        ib.show();
-      });
+      }
     }
 
     return query;
@@ -266,8 +276,9 @@ export class ConnectionManagerPlugin implements IExtensionPlugin {
         await this._connect();
       }
 
-      query = await this.replaceParams(query);
       const conn = await this.explorer.getActive()
+      query = await this.replaceParams(query, conn);
+
       const view = await this._openResultsWebview(conn && conn.id, opt.requestId);
       const payload = await this._runConnectionCommandWithArgs('query', query, { ...opt, requestId: view.requestId });
       this.updateViewResults(view, payload);

diff --git a/packages/types/index.d.ts b/packages/types/index.d.ts
@@ -191,8 +191,14 @@ export interface IConnection<DriverOptions = any> {
     connected?: string;
     disconnected?: string;
   };
-
-
+  /**
+   * Connection variables. Use this property with `sqltools.queryParams.enableReplace` to replace the variables without prompting.
+   * @type {object}
+   * @memberof IConnection
+   */
+  variables?: {
+    [key: string]: string
+  }
 
   // WONT BE INCLUDED IN SETTINGS
   /**

diff --git a/packages/util/query/index.ts b/packages/util/query/index.ts
@@ -60,7 +60,7 @@ export function getQueryParameters(query: string, regexStr: string) {
 
   const regex = new RegExp(regexStr, 'g');
 
-  const paramsMap: { [k: string]: { param: string; string: string; }} = {};
+  const paramsMap: { [k: string]: { param: string; string: string; varName?: string }} = {};
 
   let match;
   while ((match = regex.exec(query)) !== null) {
@@ -70,6 +70,10 @@ export function getQueryParameters(query: string, regexStr: string) {
         param: match[0],
         string: `...${queryPart}...`,
       };
+
+      if (match[1]) {
+        paramsMap[match[0]].varName = match[1]
+      }
     }
   }
   return Object.values(paramsMap);

diff --git a/test/docker/sqlite/1.create-some-stuff.sql b/test/docker/sqlite/1.create-some-stuff.sql
@@ -17,7 +17,7 @@ CREATE TABLE contacts_metadata (
   FOREIGN KEY(contact_id) REFERENCES contact(id)
 );
 
-CREATE VIEW contacts_view
+CREATE VIEW contacts_view_$[env]
 AS
 SELECT
 *

diff --git a/test/test-vscode-sqltools.code-workspace b/test/test-vscode-sqltools.code-workspace
@@ -75,7 +75,10 @@
         "name": "SQLite",
         "driver": "SQLite",
         "database": "./docker/sqlite/test_db.db",
-        "connectionTimeout": 15
+        "connectionTimeout": 15,
+        "variables": {
+          "env": "dev"
+        }
       },
       {
         "askForPassword": false,