feat(postgres): add function variables for postgres (#11277)

sequelize · Aug 9, 2019 · ff97d93 · ff97d93
1 parent b445ec0
commit ff97d93
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 3 deletions.
diff --git a/lib/dialects/postgres/query-generator.js b/lib/dialects/postgres/query-generator.js
@@ -608,11 +608,11 @@ class PostgresQueryGenerator extends AbstractQueryGenerator {
     if (!functionName || !returnType || !language || !body) throw new Error('createFunction missing some parameters. Did you pass functionName, returnType, language and body?');
 
     const paramList = this.expandFunctionParamList(params);
+    const variableList = options && options.variables ? this.expandFunctionVariableList(options.variables) : '';
     const expandedOptionsArray = this.expandOptions(optionsArray);
-
     const statement = options && options.force ? 'CREATE OR REPLACE FUNCTION' : 'CREATE FUNCTION';
 
-    return `${statement} ${functionName}(${paramList}) RETURNS ${returnType} AS $func$ BEGIN ${body} END; $func$ language '${language}'${expandedOptionsArray};`;
+    return `${statement} ${functionName}(${paramList}) RETURNS ${returnType} AS $func$ ${variableList} BEGIN ${body} END; $func$ language '${language}'${expandedOptionsArray};`;
   }
 
   dropFunction(functionName, params) {
@@ -667,6 +667,25 @@ class PostgresQueryGenerator extends AbstractQueryGenerator {
     return paramList.join(', ');
   }
 
+  expandFunctionVariableList(variables) {
+    if (!Array.isArray(variables)) {
+      throw new Error('expandFunctionVariableList: function variables must be an array');
+    }
+    const variableDefinitions = [];
+    variables.forEach(variable => {
+      if (!variable.name || !variable.type) {
+        throw new Error('function variable must have a name and type');
+      }
+      let variableDefinition = `DECLARE ${variable.name} ${variable.type}`;
+      if (variable.default) {
+        variableDefinition += ` := ${variable.default}`;
+      }
+      variableDefinition += ';';
+      variableDefinitions.push(variableDefinition);
+    });
+    return variableDefinitions.join(' ');
+  }
+
   expandOptions(options) {
     return options === undefined || _.isEmpty(options) ?
       '' : options.join(' ');

diff --git a/lib/query-interface.js b/lib/query-interface.js
@@ -1231,7 +1231,14 @@ class QueryInterface {
    *   [
    *     'IMMUTABLE',
    *     'LEAKPROOF'
-   *   ]
+   *   ],
+   *   {
+   *    variables:
+   *      [
+   *        {type: 'integer', name: 'myVar', default: 100}
+   *      ],
+   *      force: true
+   *   };
    * );
    *
    * @param {string}  functionName  Name of SQL function to create
@@ -1242,6 +1249,7 @@ class QueryInterface {
    * @param {Array}   optionsArray  Extra-options for creation
    * @param {Object}  [options]     query options
    * @param {boolean} options.force If force is true, any existing functions with the same parameters will be replaced. For postgres, this means using `CREATE OR REPLACE FUNCTION` instead of `CREATE FUNCTION`. Default is false
+   * @param {Array<Object>}   options.variables List of declared variables. Each variable should be an object with string fields `type` and `name`, and optionally having a `default` field as well.
    *
    * @returns {Promise}
    */

diff --git a/test/integration/dialects/postgres/query-interface.test.js b/test/integration/dialects/postgres/query-interface.test.js
@@ -173,6 +173,34 @@ if (dialect.match(/^postgres/)) {
             expect(res[0].my_func).to.be.eql('second');
           });
       });
+
+      it('produces an error when options.variables is missing expected parameters', function() {
+        const body = 'return 1;';
+        expect(() => {
+          const options = { variables: 100 };
+          return this.queryInterface.createFunction('test_func', [], 'integer', 'plpgsql', body, [], options);
+        }).to.throw(/expandFunctionVariableList: function variables must be an array/);
+
+        expect(() => {
+          const options = { variables: [{ name: 'myVar' }] };
+          return this.queryInterface.createFunction('test_func', [], 'integer', 'plpgsql', body, [], options);
+        }).to.throw(/function variable must have a name and type/);
+
+        expect(() => {
+          const options = { variables: [{ type: 'integer' }] };
+          return this.queryInterface.createFunction('test_func', [], 'integer', 'plpgsql', body, [], options);
+        }).to.throw(/function variable must have a name and type/);
+      });
+
+      it('uses declared variables', function() {
+        const body = 'RETURN myVar + 1;';
+        const options = { variables: [{ type: 'integer', name: 'myVar', default: 100 }] };
+        return this.queryInterface.createFunction('add_one', [], 'integer', 'plpgsql', body, [], options)
+          .then(() => this.sequelize.query('select add_one();', { type: this.sequelize.QueryTypes.SELECT }))
+          .then(res => {
+            expect(res[0].add_one).to.be.eql(101);
+          });
+      });
     });
 
     describe('dropFunction', () => {