Merge pull request #137 from supabase/feat/func-def-update

w3b6x9 · web-flow · commit 003a26ce6201 · 2021-07-26T12:40:35.000-04:00
feat: enable function definition update
diff --git a/src/lib/PostgresMetaFunctions.ts b/src/lib/PostgresMetaFunctions.ts
@@ -3,6 +3,18 @@ import { DEFAULT_SYSTEM_SCHEMAS } from './constants'
 import { functionsSql } from './sql'
 import { PostgresMetaResult, PostgresFunction } from './types'
 
+type FunctionInputs = {
+  name: string
+  definition: string
+  args?: string[]
+  behavior?: 'IMMUTABLE' | 'STABLE' | 'VOLATILE'
+  config_params?: { [key: string]: string }
+  schema?: string
+  language?: string
+  return_type?: string
+  security_definer?: boolean
+}
+
 export default class PostgresMetaFunctions {
   query: (sql: string) => Promise<PostgresMetaResult<any>>
 
@@ -53,31 +65,7 @@ export default class PostgresMetaFunctions {
         return { data: data[0], error }
       }
     } else if (name && schema && args) {
-      const sql = `${enrichedFunctionsSql} JOIN pg_proc AS p ON id = p.oid WHERE schema = ${literal(
-        schema
-      )} AND name = ${literal(name)} AND p.proargtypes::text = ${
-        args.length
-          ? `(
-            SELECT STRING_AGG(type_oid::text, ' ') FROM (
-              SELECT (
-                split_args.arr[
-                  array_length(
-                    split_args.arr,
-                    1
-                  )
-                ]::regtype::oid
-              ) AS type_oid FROM (
-                SELECT STRING_TO_ARRAY(
-                  UNNEST(
-                    ARRAY[${args.map(literal)}]
-                  ),
-                  ' '
-                ) AS arr
-              ) AS split_args
-            ) args
-      );`
-          : literal('')
-      }`
+      const sql = this.generateRetrieveFunctionSql({ name, schema, args })
       const { data, error } = await this.query(sql)
       if (error) {
         return { data, error }
@@ -106,32 +94,18 @@ export default class PostgresMetaFunctions {
     behavior = 'VOLATILE',
     security_definer = false,
     config_params = {},
-  }: {
-    name: string
-    schema?: string
-    args?: string[]
-    definition: string
-    return_type?: string
-    language?: string
-    behavior?: 'IMMUTABLE' | 'STABLE' | 'VOLATILE'
-    security_definer?: boolean
-    config_params?: { [key: string]: string }
-  }): Promise<PostgresMetaResult<PostgresFunction>> {
-    const sql = `
-      CREATE FUNCTION ${ident(schema)}.${ident(name)}(${args.join(', ')})
-      RETURNS ${return_type}
-      AS ${literal(definition)}
-      LANGUAGE ${language}
-      ${behavior}
-      ${security_definer ? 'SECURITY DEFINER' : 'SECURITY INVOKER'}
-      ${Object.entries(config_params)
-        .map(
-          ([param, value]) =>
-            `SET ${param} ${value[0] === 'FROM CURRENT' ? 'FROM CURRENT' : 'TO ' + value}`
-        )
-        .join('\n')}
-      RETURNS NULL ON NULL INPUT;
-    `
+  }: FunctionInputs): Promise<PostgresMetaResult<PostgresFunction>> {
+    const sql = this.generateCreateFunctionSql({
+      name,
+      schema,
+      args,
+      definition,
+      return_type,
+      language,
+      behavior,
+      security_definer,
+      config_params,
+    })
     const { error } = await this.query(sql)
     if (error) {
       return { data: null, error }
@@ -144,36 +118,78 @@ export default class PostgresMetaFunctions {
     {
       name,
       schema,
+      definition,
     }: {
       name?: string
       schema?: string
+      definition?: string
     }
   ): Promise<PostgresMetaResult<PostgresFunction>> {
-    const { data: old, error: retrieveError } = await this.retrieve({ id })
-    if (retrieveError) {
-      return { data: null, error: retrieveError }
+    const { data: currentFunc, error } = await this.retrieve({ id })
+    if (error) {
+      return { data: null, error }
     }
 
+    const updateDefinitionSql = typeof definition === 'string' ? this.generateCreateFunctionSql(
+      { ...currentFunc!, definition },
+      { replace: true }
+    ) : ''
+
+    const retrieveFunctionSql = this.generateRetrieveFunctionSql(
+      {
+        schema: currentFunc!.schema,
+        name: currentFunc!.name,
+        args: currentFunc!.argument_types.split(', '),
+      },
+      { terminateCommand: false }
+    )
+
     const updateNameSql =
-      name && name !== old!.name
-        ? `ALTER FUNCTION ${ident(old!.schema)}.${ident(old!.name)}(${
-            old!.argument_types
+      name && name !== currentFunc!.name
+        ? `ALTER FUNCTION ${ident(currentFunc!.schema)}.${ident(currentFunc!.name)}(${
+            currentFunc!.argument_types
           }) RENAME TO ${ident(name)};`
         : ''
 
     const updateSchemaSql =
-      schema && schema !== old!.schema
-        ? `ALTER FUNCTION ${ident(old!.schema)}.${ident(name || old!.name)}(${
-            old!.argument_types
+      schema && schema !== currentFunc!.schema
+        ? `ALTER FUNCTION ${ident(currentFunc!.schema)}.${ident(name || currentFunc!.name)}(${
+            currentFunc!.argument_types
           })  SET SCHEMA ${ident(schema)};`
         : ''
 
-    const sql = `BEGIN; ${updateNameSql} ${updateSchemaSql} COMMIT;`
+    const sql = `
+      DO LANGUAGE plpgsql $$
+      DECLARE
+        function record;
+      BEGIN
+        IF ${typeof definition === 'string' ? 'TRUE' : 'FALSE'} THEN
+          ${updateDefinitionSql}
 
-    const { error } = await this.query(sql)
-    if (error) {
-      return { data: null, error }
+          ${retrieveFunctionSql} INTO function;
+
+          IF function.id != ${id} THEN
+            RAISE EXCEPTION 'Cannot find function "${currentFunc!.schema}"."${currentFunc!.name}"(${
+      currentFunc!.argument_types
+    })';
+          END IF;
+        END IF;
+
+        ${updateNameSql}
+
+        ${updateSchemaSql}
+      END;
+      $$;
+    `
+
+    {
+      const { error } = await this.query(sql)
+
+      if (error) {
+        return { data: null, error }
+      }
     }
+
     return await this.retrieve({ id })
   }
 
@@ -196,6 +212,84 @@ export default class PostgresMetaFunctions {
     }
     return { data: func!, error: null }
   }
+
+  private generateCreateFunctionSql(
+    {
+      name,
+      schema,
+      args,
+      argument_types,
+      definition,
+      return_type,
+      language,
+      behavior,
+      security_definer,
+      config_params,
+    }: Partial<Omit<FunctionInputs, 'config_params'> & PostgresFunction>,
+    { replace = false, terminateCommand = true } = {}
+  ): string {
+    return `
+      CREATE ${replace ? 'OR REPLACE' : ''} FUNCTION ${ident(schema!)}.${ident(name!)}(${
+      argument_types || args?.join(', ') || ''
+    })
+      RETURNS ${return_type}
+      AS ${literal(definition)}
+      LANGUAGE ${language}
+      ${behavior}
+      CALLED ON NULL INPUT
+      ${security_definer ? 'SECURITY DEFINER' : 'SECURITY INVOKER'}
+      ${
+        config_params
+          ? Object.entries(config_params)
+              .map(
+                ([param, value]) =>
+                  `SET ${param} ${value[0] === 'FROM CURRENT' ? 'FROM CURRENT' : 'TO ' + value}`
+              )
+              .join('\n')
+          : ''
+      }
+      ${terminateCommand ? ';' : ''}
+    `
+  }
+
+  private generateRetrieveFunctionSql(
+    {
+      schema,
+      name,
+      args,
+    }: {
+      schema: string
+      name: string
+      args: string[]
+    },
+    { terminateCommand = true } = {}
+  ): string {
+    return `${enrichedFunctionsSql} JOIN pg_proc AS p ON id = p.oid WHERE schema = ${literal(
+      schema
+    )} AND name = ${literal(name)} AND p.proargtypes::text = ${
+      args.length
+        ? `(
+          SELECT STRING_AGG(type_oid::text, ' ') FROM (
+            SELECT (
+              split_args.arr[
+                array_length(
+                  split_args.arr,
+                  1
+                )
+              ]::regtype::oid
+            ) AS type_oid FROM (
+              SELECT STRING_TO_ARRAY(
+                UNNEST(
+                  ARRAY[${args.map(literal)}]
+                ),
+                ' '
+              ) AS arr
+            ) AS split_args
+          ) args
+    ) ${terminateCommand ? ';' : ''}`
+        : literal('')
+    }`
+  }
 }
 
 const enrichedFunctionsSql = `
diff --git a/src/lib/sql/functions.sql b/src/lib/sql/functions.sql
@@ -3,10 +3,14 @@ SELECT
   n.nspname AS schema,
   p.proname AS name,
   l.lanname AS language,
+  CASE
+    WHEN l.lanname = 'internal' THEN ''
+    ELSE p.prosrc
+  END AS definition,
   CASE
     WHEN l.lanname = 'internal' THEN p.prosrc
     ELSE pg_get_functiondef(p.oid)
-  END AS definition,
+  END AS complete_statement,
   pg_get_function_arguments(p.oid) AS argument_types,
   t.typname AS return_type,
   CASE
diff --git a/src/lib/types.ts b/src/lib/types.ts
@@ -76,6 +76,7 @@ const postgresFunctionSchema = Type.Object({
   name: Type.String(),
   language: Type.String(),
   definition: Type.String(),
+  complete_statement: Type.String(),
   argument_types: Type.String(),
   return_type: Type.String(),
   behavior: Type.Union([
diff --git a/test/integration/index.spec.js b/test/integration/index.spec.js
@@ -215,6 +215,17 @@ describe('/functions', () => {
     assert.strictEqual(newFunc.schema, 'public')
     assert.strictEqual(newFunc.argument_types, 'a smallint, b smallint')
     assert.strictEqual(newFunc.language, 'sql')
+    assert.strictEqual(newFunc.definition, 'select a + b')
+    assert.strictEqual(
+      newFunc.complete_statement,
+      'CREATE OR REPLACE FUNCTION public.test_func(a smallint, b smallint)\n' +
+      ' RETURNS integer\n' +
+      ' LANGUAGE sql\n' +
+      ' STABLE SECURITY DEFINER\n' +
+      " SET search_path TO 'hooks', 'auth'\n" +
+      " SET role TO 'postgres'\n" +
+      'AS $function$select a + b$function$\n'
+    )
     assert.strictEqual(newFunc.return_type, 'int4')
     assert.strictEqual(newFunc.behavior, 'STABLE')
     assert.strictEqual(newFunc.security_definer, true)
@@ -225,12 +236,24 @@ describe('/functions', () => {
     const updates = {
       name: 'test_func_renamed',
       schema: 'test_schema',
+      definition: 'select b - a'
     }
 
     let { data: updated } = await axios.patch(`${URL}/functions/${func.id}`, updates)
     assert.strictEqual(updated.id, func.id)
     assert.strictEqual(updated.name, 'test_func_renamed')
     assert.strictEqual(updated.schema, 'test_schema')
+    assert.strictEqual(updated.definition, 'select b - a')
+    assert.strictEqual(
+      updated.complete_statement,
+      'CREATE OR REPLACE FUNCTION test_schema.test_func_renamed(a smallint, b smallint)\n' +
+      ' RETURNS integer\n' +
+      ' LANGUAGE sql\n' +
+      ' STABLE SECURITY DEFINER\n' +
+      " SET search_path TO 'hooks', 'auth'\n" +
+      " SET role TO 'postgres'\n" +
+      'AS $function$select b - a$function$\n'
+    )
   })
   it('DELETE', async () => {
     await axios.delete(`${URL}/functions/${func.id}`)
