refactor: support return stmt sql body in create function in pg

Author: taozhi8833998

pegjs/postgresql.pegjs

@@ -698,19 +698,20 @@ create_func_opt
       value,
     }
   }
-
+  / return_stmt
+  
 create_function_stmt
   = a:KW_CREATE __
   or:(KW_OR __ KW_REPLACE)? __
   t:'FUNCTION'i __
-  c:table_name __ LPAREN __ args:alter_func_args? __ RPAREN __
+  c:proc_func_name __ LPAREN __ args:alter_func_args? __ RPAREN __
   r:func_returns? __
   fo:create_func_opt* __ SEMICOLON? __ {
     /*
       export type create_function_stmt_t = {
         type: 'create';
         replace?: string;
-        name: { schema?: string; name: string };
+        name: proc_func_name;
         args?: alter_func_args;
         returns?: func_returns;
         keyword: 'function';
@@ -725,7 +726,7 @@ create_function_stmt
           args: args || [],
           type: 'create',
           replace: or && 'or replace',
-          name: { schema: c.db, name: c.table },
+          name: c,
           returns: r,
           keyword: t && t.toLowerCase(),
           options: fo || [],
@@ -5281,6 +5282,7 @@ KW_DELETE   = "DELETE"i     !ident_start
 KW_INSERT   = "INSERT"i     !ident_start
 KW_RECURSIVE= "RECURSIVE"   !ident_start { return 'RECURSIVE'; }
 KW_REPLACE  = "REPLACE"i    !ident_start
+KW_RETURN = 'RETURN'i  !ident_start { return 'RETURN' }
 KW_RETURNING  = "RETURNING"i    !ident_start { return 'RETURNING' }
 KW_RENAME   = "RENAME"i     !ident_start
 KW_IGNORE   = "IGNORE"i     !ident_start
@@ -5445,7 +5447,6 @@ KW_VAR_PRE_DOLLAR = '$'
 KW_VAR_PRE_DOLLAR_DOUBLE = '$$'
 KW_VAR_PRE
   = KW_VAR__PRE_AT_AT / KW_VAR__PRE_AT / KW_VAR_PRE_DOLLAR / KW_VAR_PRE_DOLLAR
-KW_RETURN = 'return'i
 KW_ASSIGN = ':='
 KW_DOUBLE_COLON = '::'
 KW_ASSIGIN_EQUAL = '='
@@ -5575,7 +5576,6 @@ assign_stmt
     };
   }
 
-
 return_stmt
   = KW_RETURN __ e:proc_expr {
     // => { type: 'return'; expr: proc_expr; }

src/create.js

@@ -336,14 +336,16 @@ function createFunctionOptionToSQL(stmt) {
       return [toUpper(type), stmt.symbol, unionToSQL(stmt.declare), toUpper(stmt.begin), multipleToSQL(stmt.expr), toUpper(stmt.end), stmt.symbol].filter(hasVal).join(' ')
     case 'set':
       return [toUpper(type), stmt.parameter, toUpper(stmt.value && stmt.value.prefix), stmt.value && stmt.value.expr.map(exprToSQL).join(', ')].filter(hasVal).join(' ')
+    case 'return':
+      return [toUpper(type), exprToSQL(stmt.expr)].filter(hasVal).join(' ')
     default:
       return exprToSQL(stmt)
   }
 }
 function createFunctionToSQL(stmt) {
   const { type, replace, keyword, name, args, returns, options, last } = stmt
   const sql = [toUpper(type), toUpper(replace), toUpper(keyword)]
-  const functionName = [identifierToSql(name.schema), name.name].filter(hasVal).join('.')
+  const functionName = [literalToSQL(name.schema), name.name.map(literalToSQL).join('.')].filter(hasVal).join('.')
   const argsSQL = args.map(alterArgsToSQL).filter(hasVal).join(', ')
   sql.push(`${functionName}(${argsSQL})`, createFunctionReturnsToSQL(returns), options.map(createFunctionOptionToSQL).join(' '), last)
   return sql.filter(hasVal).join(' ')

test/postgres.spec.js

@@ -1131,7 +1131,7 @@ describe('Postgres', () => {
             ELSE $1 || ', ' || $2
           END
         $_$;`,
-        `CREATE FUNCTION "public"._group_concat(TEXT, TEXT) RETURNS TEXT LANGUAGE sql IMMUTABLE AS $_$ SELECT CASE WHEN $2 IS NULL THEN $1 WHEN $1 IS NULL THEN $2 ELSE $1 || ', ' || $2 END $_$`
+        `CREATE FUNCTION public._group_concat(TEXT, TEXT) RETURNS TEXT LANGUAGE sql IMMUTABLE AS $_$ SELECT CASE WHEN $2 IS NULL THEN $1 WHEN $1 IS NULL THEN $2 ELSE $1 || ', ' || $2 END $_$`
       ]
     },
     {
@@ -1146,7 +1146,7 @@ describe('Postgres', () => {
           AND store_id = $2
           AND NOT inventory_in_stock(inventory_id);
         $_$;`,
-        `CREATE FUNCTION "public".film_not_in_stock(p_film_id INTEGER DEFAULT 1, p_store_id INTEGER = 1, OUT p_film_count INTEGER) RETURNS SETOF INTEGER LANGUAGE sql AS $_$ SELECT inventory_id FROM "inventory" WHERE film_id = $1 AND store_id = $2 AND NOT inventory_in_stock(inventory_id) $_$`
+        `CREATE FUNCTION public.film_not_in_stock(p_film_id INTEGER DEFAULT 1, p_store_id INTEGER = 1, OUT p_film_count INTEGER) RETURNS SETOF INTEGER LANGUAGE sql AS $_$ SELECT inventory_id FROM "inventory" WHERE film_id = $1 AND store_id = $2 AND NOT inventory_in_stock(inventory_id) $_$`
       ]
     },
     {
@@ -1182,7 +1182,7 @@ describe('Postgres', () => {
     {
       title: 'create function with if else stmt',
       sql: [
-        `CREATE FUNCTION public.inventory_in_stock(p_inventory_id integer) RETURNS boolean
+        `CREATE FUNCTION "public".inventory_in_stock(p_inventory_id integer) RETURNS boolean
         LANGUAGE plpgsql
         AS $$
           DECLARE
@@ -1220,7 +1220,7 @@ describe('Postgres', () => {
     {
       title: 'create function without args',
       sql: [
-        `CREATE FUNCTION public.last_updated() RETURNS trigger
+        `CREATE FUNCTION "public".last_updated() RETURNS trigger
         LANGUAGE plpgsql
         AS $$
           BEGIN
@@ -1230,6 +1230,17 @@ describe('Postgres', () => {
         'CREATE FUNCTION "public".last_updated() RETURNS "trigger" LANGUAGE plpgsql AS $$ BEGIN NEW.last_update = CURRENT_TIMESTAMP ; RETURN NEW END $$'
       ]
     },
+    {
+      title: 'create function with sql body',
+      sql: [
+        `CREATE FUNCTION add(a integer, b integer) RETURNS integer
+          LANGUAGE SQL
+          IMMUTABLE
+          RETURNS NULL ON NULL INPUT
+          RETURN a + b;`,
+        'CREATE FUNCTION add(a INTEGER, b INTEGER) RETURNS INTEGER LANGUAGE SQL IMMUTABLE RETURNS NULL ON NULL INPUT RETURN a + b'
+      ]
+    },
     {
       title: 'create aggregate',
       sql: [