From 3b8431d32384db0955675067c8c84cbbd28f9f32 Mon Sep 17 00:00:00 2001 From: tianzhou Date: Wed, 22 Oct 2025 13:37:41 +0800 Subject: [PATCH] fix: table function --- internal/diff/function.go | 15 ++++++++++++--- internal/diff/procedure.go | 15 ++++++++++++--- ir/queries/queries.sql | 14 ++++++++++++-- ir/queries/queries.sql.go | 14 ++++++++++++-- .../diff/create_function/add_function/diff.sql | 7 +++++++ .../diff/create_function/add_function/new.sql | 7 ++++++- .../diff/create_function/add_function/plan.json | 6 ++++++ .../diff/create_function/add_function/plan.sql | 7 +++++++ .../diff/create_function/add_function/plan.txt | 12 ++++++++++-- 9 files changed, 84 insertions(+), 13 deletions(-) diff --git a/internal/diff/function.go b/internal/diff/function.go index b352ab6c..7f855fc8 100644 --- a/internal/diff/function.go +++ b/internal/diff/function.go @@ -135,10 +135,19 @@ func generateFunctionSQL(function *ir.Function, targetSchema string) string { stmt.WriteString("\nSTRICT") } - // Add the function body with proper dollar quoting + // Add the function body if function.Definition != "" { - tag := generateDollarQuoteTag(function.Definition) - stmt.WriteString(fmt.Sprintf("\nAS %s%s%s;", tag, function.Definition, tag)) + // Check if this uses RETURN clause syntax (PG14+) + // pg_get_function_sqlbody returns "RETURN expression" which should not be wrapped + // Use case-insensitive comparison to handle all variations + trimmedDef := strings.TrimSpace(function.Definition) + if len(trimmedDef) >= 7 && strings.EqualFold(trimmedDef[:7], "RETURN ") { + stmt.WriteString(fmt.Sprintf("\n%s;", trimmedDef)) + } else { + // Traditional AS $$ ... $$ syntax + tag := generateDollarQuoteTag(function.Definition) + stmt.WriteString(fmt.Sprintf("\nAS %s%s%s;", tag, function.Definition, tag)) + } } else { stmt.WriteString("\nAS $$$$;") } diff --git a/internal/diff/procedure.go b/internal/diff/procedure.go index 24a46740..477ebd3d 100644 --- a/internal/diff/procedure.go +++ b/internal/diff/procedure.go @@ -157,10 +157,19 @@ func generateProcedureSQL(procedure *ir.Procedure, targetSchema string) string { // Note: Procedures don't have SECURITY DEFINER/INVOKER in PostgreSQL // This is a function-only feature - // Add the procedure body with proper dollar quoting + // Add the procedure body if procedure.Definition != "" { - tag := generateProcedureDollarQuoteTag(procedure.Definition) - stmt.WriteString(fmt.Sprintf("\nAS %s%s%s;", tag, procedure.Definition, tag)) + // Check if this uses RETURN clause syntax (PG14+) + // pg_get_function_sqlbody returns "RETURN expression" which should not be wrapped + // Use case-insensitive comparison to handle all variations + trimmedDef := strings.TrimSpace(procedure.Definition) + if len(trimmedDef) >= 7 && strings.EqualFold(trimmedDef[:7], "RETURN ") { + stmt.WriteString(fmt.Sprintf("\n%s;", trimmedDef)) + } else { + // Traditional AS $$ ... $$ syntax + tag := generateProcedureDollarQuoteTag(procedure.Definition) + stmt.WriteString(fmt.Sprintf("\nAS %s%s%s;", tag, procedure.Definition, tag)) + } } else { stmt.WriteString("\nAS $$$$;") } diff --git a/ir/queries/queries.sql b/ir/queries/queries.sql index 99b21942..61c03a6e 100644 --- a/ir/queries/queries.sql +++ b/ir/queries/queries.sql @@ -741,7 +741,12 @@ ORDER BY s.schemaname, s.sequencename; SELECT r.routine_schema, r.routine_name, - CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END AS routine_definition, + -- Use pg_get_function_sqlbody for RETURN clause syntax (PG14+) + -- Fall back to prosrc for traditional AS $$ ... $$ syntax + COALESCE( + pg_get_function_sqlbody(p.oid), + CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END + ) AS routine_definition, r.routine_type, COALESCE(pg_get_function_result(p.oid), r.data_type) AS data_type, r.external_language, @@ -774,7 +779,12 @@ ORDER BY r.routine_schema, r.routine_name; SELECT r.routine_schema, r.routine_name, - CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END AS routine_definition, + -- Use pg_get_function_sqlbody for RETURN clause syntax (PG14+) + -- Fall back to prosrc for traditional AS $$ ... $$ syntax + COALESCE( + pg_get_function_sqlbody(p.oid), + CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END + ) AS routine_definition, r.routine_type, r.external_language, COALESCE(desc_proc.description, '') AS procedure_comment, diff --git a/ir/queries/queries.sql.go b/ir/queries/queries.sql.go index f792c08a..2967fe0c 100644 --- a/ir/queries/queries.sql.go +++ b/ir/queries/queries.sql.go @@ -1185,7 +1185,12 @@ const getFunctionsForSchema = `-- name: GetFunctionsForSchema :many SELECT r.routine_schema, r.routine_name, - CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END AS routine_definition, + -- Use pg_get_function_sqlbody for RETURN clause syntax (PG14+) + -- Fall back to prosrc for traditional AS $$ ... $$ syntax + COALESCE( + pg_get_function_sqlbody(p.oid), + CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END + ) AS routine_definition, r.routine_type, COALESCE(pg_get_function_result(p.oid), r.data_type) AS data_type, r.external_language, @@ -1633,7 +1638,12 @@ const getProceduresForSchema = `-- name: GetProceduresForSchema :many SELECT r.routine_schema, r.routine_name, - CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END AS routine_definition, + -- Use pg_get_function_sqlbody for RETURN clause syntax (PG14+) + -- Fall back to prosrc for traditional AS $$ ... $$ syntax + COALESCE( + pg_get_function_sqlbody(p.oid), + CASE WHEN p.prosrc ~ E'\n$' THEN p.prosrc ELSE p.prosrc || E'\n' END + ) AS routine_definition, r.routine_type, r.external_language, COALESCE(desc_proc.description, '') AS procedure_comment, diff --git a/testdata/diff/create_function/add_function/diff.sql b/testdata/diff/create_function/add_function/diff.sql index fd556b37..a4e635ba 100644 --- a/testdata/diff/create_function/add_function/diff.sql +++ b/testdata/diff/create_function/add_function/diff.sql @@ -1,3 +1,10 @@ +CREATE OR REPLACE FUNCTION days_since_special_date() +RETURNS SETOF timestamp with time zone +LANGUAGE sql +SECURITY INVOKER +STABLE +RETURN generate_series((date_trunc('day'::text, '2025-01-01 00:00:00'::timestamp without time zone))::timestamp with time zone, date_trunc('day'::text, now()), '1 day'::interval); + CREATE OR REPLACE FUNCTION process_order( order_id integer, discount_percent numeric DEFAULT 0, diff --git a/testdata/diff/create_function/add_function/new.sql b/testdata/diff/create_function/add_function/new.sql index 460749fc..a884a33f 100644 --- a/testdata/diff/create_function/add_function/new.sql +++ b/testdata/diff/create_function/add_function/new.sql @@ -22,4 +22,9 @@ BEGIN SELECT amount INTO total FROM orders WHERE id = order_id; RETURN total - (total * discount_percent / 100); END; -$$; \ No newline at end of file +$$; + +-- Table function with RETURN clause (bug report test case) +CREATE FUNCTION days_since_special_date() RETURNS SETOF timestamptz + LANGUAGE sql STABLE PARALLEL SAFE + RETURN generate_series(date_trunc('day', '2025-01-01'::timestamp), date_trunc('day', NOW()), '1 day'::interval); \ No newline at end of file diff --git a/testdata/diff/create_function/add_function/plan.json b/testdata/diff/create_function/add_function/plan.json index 831f7fd5..0b22a107 100644 --- a/testdata/diff/create_function/add_function/plan.json +++ b/testdata/diff/create_function/add_function/plan.json @@ -8,6 +8,12 @@ "groups": [ { "steps": [ + { + "sql": "CREATE OR REPLACE FUNCTION days_since_special_date()\nRETURNS SETOF timestamp with time zone\nLANGUAGE sql\nSECURITY INVOKER\nSTABLE\nRETURN generate_series((date_trunc('day'::text, '2025-01-01 00:00:00'::timestamp without time zone))::timestamp with time zone, date_trunc('day'::text, now()), '1 day'::interval);", + "type": "function", + "operation": "create", + "path": "public.days_since_special_date" + }, { "sql": "CREATE OR REPLACE FUNCTION process_order(\n order_id integer,\n discount_percent numeric DEFAULT 0,\n priority_level integer DEFAULT 1,\n note varchar DEFAULT '',\n status text DEFAULT 'pending',\n apply_tax boolean DEFAULT true,\n is_priority boolean DEFAULT false\n)\nRETURNS numeric\nLANGUAGE plpgsql\nSECURITY DEFINER\nVOLATILE\nSTRICT\nAS $$\nDECLARE\n total numeric;\nBEGIN\n SELECT amount INTO total FROM orders WHERE id = order_id;\n RETURN total - (total * discount_percent / 100);\nEND;\n$$;", "type": "function", diff --git a/testdata/diff/create_function/add_function/plan.sql b/testdata/diff/create_function/add_function/plan.sql index fd556b37..a4e635ba 100644 --- a/testdata/diff/create_function/add_function/plan.sql +++ b/testdata/diff/create_function/add_function/plan.sql @@ -1,3 +1,10 @@ +CREATE OR REPLACE FUNCTION days_since_special_date() +RETURNS SETOF timestamp with time zone +LANGUAGE sql +SECURITY INVOKER +STABLE +RETURN generate_series((date_trunc('day'::text, '2025-01-01 00:00:00'::timestamp without time zone))::timestamp with time zone, date_trunc('day'::text, now()), '1 day'::interval); + CREATE OR REPLACE FUNCTION process_order( order_id integer, discount_percent numeric DEFAULT 0, diff --git a/testdata/diff/create_function/add_function/plan.txt b/testdata/diff/create_function/add_function/plan.txt index 3b64a999..cae2ecd3 100644 --- a/testdata/diff/create_function/add_function/plan.txt +++ b/testdata/diff/create_function/add_function/plan.txt @@ -1,14 +1,22 @@ -Plan: 1 to add. +Plan: 2 to add. Summary by type: - functions: 1 to add + functions: 2 to add Functions: + + days_since_special_date + process_order DDL to be executed: -------------------------------------------------- +CREATE OR REPLACE FUNCTION days_since_special_date() +RETURNS SETOF timestamp with time zone +LANGUAGE sql +SECURITY INVOKER +STABLE +RETURN generate_series((date_trunc('day'::text, '2025-01-01 00:00:00'::timestamp without time zone))::timestamp with time zone, date_trunc('day'::text, now()), '1 day'::interval); + CREATE OR REPLACE FUNCTION process_order( order_id integer, discount_percent numeric DEFAULT 0,