fix: rewrite SET search_path in function definitions for temp schema (#335)

cmd/plan/plan.go

@@ -515,6 +515,7 @@ func normalizeSchemaNames(irData *ir.IR, fromSchema, toSchema string) {
 			fn.Schema = toSchema
 			fn.ReturnType = replaceString(fn.ReturnType)
 			fn.Definition = replaceString(fn.Definition)
+			fn.SearchPath = replaceString(fn.SearchPath)
 			for _, param := range fn.Parameters {
 				param.DataType = replaceString(param.DataType)
 			}

internal/postgres/desired_state.go

@@ -174,6 +174,91 @@ func stripSchemaQualifications(sql string, schemaName string) string {
 	return result
 }
 
+// replaceSchemaInSearchPath replaces the target schema name in SET search_path clauses
+// within function/procedure definitions.
+//
+// Purpose:
+// When functions or procedures have SET search_path = public, pg_temp (or similar),
+// PostgreSQL uses that search_path during function body validation (for SQL-language functions)
+// and execution. When applying to a temporary schema, we need to replace the target schema
+// in these clauses so that table references in function bodies can be resolved.
+//
+// Example (when target schema is "public" and temp schema is "pgschema_tmp_xxx"):
+//
+//	SET search_path = public, pg_temp  ->  SET search_path = "pgschema_tmp_xxx", pg_temp
+//	SET search_path TO public          ->  SET search_path TO "pgschema_tmp_xxx"
+//
+// This handles both = and TO syntax, quoted and unquoted schema names (case-insensitive),
+// and preserves other schemas in the comma-separated list.
+//
+// Limitations:
+//   - Like stripSchemaQualifications and replaceSchemaInDefaultPrivileges, this function
+//     operates on the raw SQL string without dollar-quote awareness. A SET search_path
+//     inside a $$-quoted function body (e.g., dynamic SQL) would also be rewritten. In
+//     practice this is not an issue because such usage is extremely rare, and the round-trip
+//     through database inspection and normalizeSchemaNames restores the original schema name.
+//   - When targetSchema is "public", replacing it removes "public" from the function's
+//     search_path. If the function body references extension objects installed in "public"
+//     (e.g., citext), they may not be found. Most extension objects (uuid, jsonb, etc.) live
+//     in pg_catalog which is always searched, so this is rarely an issue in practice.
+func replaceSchemaInSearchPath(sql string, targetSchema, tempSchema string) string {
+	if targetSchema == "" || tempSchema == "" {
+		return sql
+	}
+
+	replacement := fmt.Sprintf(`"%s"`, tempSchema)
+
+	// Pattern: SET search_path = ... or SET search_path TO ...
+	// We match the entire SET search_path clause and replace the target schema within it.
+	searchPathPattern := regexp.MustCompile(`(?i)(SET\s+search_path\s*(?:=|TO)\s*)([^\n;]+)`)
+
+	// Pattern to detect trailing function body start in the captured value.
+	// When SET search_path and the body are on the same line, the value regex captures both.
+	// Handles both AS $$ (dollar-quoted) and BEGIN ATOMIC (SQL-standard, PG14+) syntax.
+	bodyStartPattern := regexp.MustCompile(`(?i)\s+(?:AS\s|BEGIN\s+ATOMIC\b)`)
+
+	return searchPathPattern.ReplaceAllStringFunc(sql, func(match string) string {
+		loc := searchPathPattern.FindStringSubmatchIndex(match)
+		if loc == nil {
+			return match
+		}
+		prefix := match[loc[2]:loc[3]]
+		value := match[loc[4]:loc[5]]
+
+		// Separate the search_path value from any trailing function body start
+		suffix := ""
+		if asLoc := bodyStartPattern.FindStringIndex(value); asLoc != nil {
+			suffix = value[asLoc[0]:]
+			value = value[:asLoc[0]]
+		}
+
+		// Tokenize the comma-separated search_path list and replace matching schemas.
+		// This avoids regex pitfalls with quoted identifiers (e.g., "PUBLIC" should not
+		// be matched by a case-insensitive unquoted pattern for "public").
+		tokens := strings.Split(value, ",")
+		for i, token := range tokens {
+			trimmed := strings.TrimSpace(token)
+			if strings.HasPrefix(trimmed, `"`) && strings.HasSuffix(trimmed, `"`) {
+				// Quoted identifier: case-sensitive exact match.
+				// "public" matches targetSchema "public", but "PUBLIC" does not
+				// (in PostgreSQL, quoted identifiers preserve case).
+				inner := trimmed[1 : len(trimmed)-1]
+				if inner == targetSchema {
+					tokens[i] = strings.Replace(token, trimmed, replacement, 1)
+				}
+			} else {
+				// Unquoted identifier: case-insensitive match.
+				// PostgreSQL folds unquoted identifiers to lowercase.
+				if strings.EqualFold(trimmed, targetSchema) {
+					tokens[i] = strings.Replace(token, trimmed, replacement, 1)
+				}
+			}
+		}
+
+		return prefix + strings.Join(tokens, ",") + suffix
+	})
+}
+
 // replaceSchemaInDefaultPrivileges replaces schema names in ALTER DEFAULT PRIVILEGES statements.
 // This is needed because stripSchemaQualifications only handles "schema.object" patterns,
 // not "IN SCHEMA <schema>" clauses used by ALTER DEFAULT PRIVILEGES.

internal/postgres/desired_state_test.go

@@ -0,0 +1,123 @@
+package postgres
+
+import (
+	"testing"
+)
+
+func TestReplaceSchemaInSearchPath(t *testing.T) {
+	tests := []struct {
+		name         string
+		sql          string
+		targetSchema string
+		tempSchema   string
+		expected     string
+	}{
+		{
+			name:         "unquoted with equals",
+			sql:          "SET search_path = public, pg_temp",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_20260302_000000_abcd1234",
+			expected:     `SET search_path = "pgschema_tmp_20260302_000000_abcd1234", pg_temp`,
+		},
+		{
+			name:         "unquoted with TO",
+			sql:          "SET search_path TO public",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_20260302_000000_abcd1234",
+			expected:     `SET search_path TO "pgschema_tmp_20260302_000000_abcd1234"`,
+		},
+		{
+			name:         "quoted target schema",
+			sql:          `SET search_path = "public", pg_temp`,
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_20260302_000000_abcd1234",
+			expected:     `SET search_path = "pgschema_tmp_20260302_000000_abcd1234", pg_temp`,
+		},
+		{
+			name:         "case insensitive schema match",
+			sql:          "SET search_path = PUBLIC, pg_temp",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_20260302_000000_abcd1234",
+			expected:     `SET search_path = "pgschema_tmp_20260302_000000_abcd1234", pg_temp`,
+		},
+		{
+			name:         "mixed case schema",
+			sql:          "SET search_path = Public, pg_temp",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_20260302_000000_abcd1234",
+			expected:     `SET search_path = "pgschema_tmp_20260302_000000_abcd1234", pg_temp`,
+		},
+		{
+			name:         "schema not in search_path is no-op",
+			sql:          "SET search_path = pg_catalog, pg_temp",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_20260302_000000_abcd1234",
+			expected:     "SET search_path = pg_catalog, pg_temp",
+		},
+		{
+			name:         "multiple functions in same SQL",
+			sql:          "CREATE FUNCTION f1() RETURNS void LANGUAGE sql SET search_path = public AS $$ SELECT 1; $$;\nCREATE FUNCTION f2() RETURNS void LANGUAGE sql SET search_path = public, pg_temp AS $$ SELECT 2; $$;",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_xxx",
+			expected:     "CREATE FUNCTION f1() RETURNS void LANGUAGE sql SET search_path = \"pgschema_tmp_xxx\" AS $$ SELECT 1; $$;\nCREATE FUNCTION f2() RETURNS void LANGUAGE sql SET search_path = \"pgschema_tmp_xxx\", pg_temp AS $$ SELECT 2; $$;",
+		},
+		{
+			name:         "empty target schema returns unchanged",
+			sql:          "SET search_path = public, pg_temp",
+			targetSchema: "",
+			tempSchema:   "pgschema_tmp_xxx",
+			expected:     "SET search_path = public, pg_temp",
+		},
+		{
+			name:         "empty temp schema returns unchanged",
+			sql:          "SET search_path = public, pg_temp",
+			targetSchema: "public",
+			tempSchema:   "",
+			expected:     "SET search_path = public, pg_temp",
+		},
+		{
+			name:         "no search_path in SQL is no-op",
+			sql:          "CREATE TABLE foo (id int);",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_xxx",
+			expected:     "CREATE TABLE foo (id int);",
+		},
+		{
+			name:         "non-public target schema",
+			sql:          "SET search_path = myschema, public",
+			targetSchema: "myschema",
+			tempSchema:   "pgschema_tmp_xxx",
+			expected:     `SET search_path = "pgschema_tmp_xxx", public`,
+		},
+		{
+			name:         "does not match partial schema names",
+			sql:          "SET search_path = public_data, pg_temp",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_xxx",
+			expected:     "SET search_path = public_data, pg_temp",
+		},
+		{
+			name:         "does not replace quoted schema with different case",
+			sql:          `SET search_path = "PUBLIC", pg_temp`,
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_xxx",
+			expected:     `SET search_path = "PUBLIC", pg_temp`,
+		},
+		{
+			name:         "single-line BEGIN ATOMIC function",
+			sql:          "CREATE FUNCTION f1() RETURNS int LANGUAGE sql SET search_path = public BEGIN ATOMIC SELECT 1; END;",
+			targetSchema: "public",
+			tempSchema:   "pgschema_tmp_xxx",
+			expected:     `CREATE FUNCTION f1() RETURNS int LANGUAGE sql SET search_path = "pgschema_tmp_xxx" BEGIN ATOMIC SELECT 1; END;`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := replaceSchemaInSearchPath(tt.sql, tt.targetSchema, tt.tempSchema)
+			if result != tt.expected {
+				t.Errorf("replaceSchemaInSearchPath() =\n%s\nwant:\n%s", result, tt.expected)
+			}
+		})
+	}
+}

internal/postgres/embedded.go

@@ -215,6 +215,11 @@ func (ep *EmbeddedPostgres) ApplySchema(ctx context.Context, schema string, sql
 	// These use "IN SCHEMA <schema>" syntax which isn't handled by stripSchemaQualifications
 	schemaAgnosticSQL = replaceSchemaInDefaultPrivileges(schemaAgnosticSQL, schema, ep.tempSchema)
 
+	// Replace schema names in SET search_path clauses within function/procedure definitions
+	// SQL-language functions are validated at creation time using the function's own search_path,
+	// so we need to rewrite it to point to the temporary schema (issue #335)
+	schemaAgnosticSQL = replaceSchemaInSearchPath(schemaAgnosticSQL, schema, ep.tempSchema)
+
 	// Execute the SQL directly
 	// Note: Desired state SQL should never contain operations like CREATE INDEX CONCURRENTLY
 	// that cannot run in transactions. Those are migration details, not state declarations.

internal/postgres/external.go

@@ -121,6 +121,11 @@ func (ed *ExternalDatabase) ApplySchema(ctx context.Context, schema string, sql
 	// These use "IN SCHEMA <schema>" syntax which isn't handled by stripSchemaQualifications
 	schemaAgnosticSQL = replaceSchemaInDefaultPrivileges(schemaAgnosticSQL, schema, ed.tempSchema)
 
+	// Replace schema names in SET search_path clauses within function/procedure definitions
+	// SQL-language functions are validated at creation time using the function's own search_path,
+	// so we need to rewrite it to point to the temporary schema (issue #335)
+	schemaAgnosticSQL = replaceSchemaInSearchPath(schemaAgnosticSQL, schema, ed.tempSchema)
+
 	// Execute the SQL directly
 	// Note: Desired state SQL should never contain operations like CREATE INDEX CONCURRENTLY
 	// that cannot run in transactions. Those are migration details, not state declarations.

testdata/diff/create_function/issue_335_search_path_rewrite/diff.sql

@@ -0,0 +1,28 @@
+CREATE TABLE IF NOT EXISTS person_accounts (
+    id uuid DEFAULT gen_random_uuid(),
+    first_name text,
+    last_name text,
+    email_address text NOT NULL,
+    created_at timestamptz DEFAULT now() NOT NULL,
+    modified_at timestamptz DEFAULT now() NOT NULL,
+    CONSTRAINT person_accounts_pkey PRIMARY KEY (id),
+    CONSTRAINT person_accounts_email_address_key UNIQUE (email_address)
+);
+
+CREATE OR REPLACE FUNCTION auth_lookup_account_by_email(
+    input_email text
+)
+RETURNS text
+LANGUAGE sql
+STABLE
+SECURITY DEFINER
+SET search_path = public, pg_temp
+AS $$
+    SELECT
+        pa.id::text AS person_account_id
+    FROM person_accounts pa
+    WHERE lower(pa.email_address) = lower(trim(input_email))
+    LIMIT 1;
+$$;
+
+REVOKE EXECUTE ON FUNCTION auth_lookup_account_by_email(input_email text) FROM PUBLIC;

testdata/diff/create_function/issue_335_search_path_rewrite/new.sql

@@ -0,0 +1,31 @@
+CREATE TABLE person_accounts (
+    id uuid DEFAULT gen_random_uuid(),
+    first_name text,
+    last_name text,
+    email_address text NOT NULL,
+    created_at timestamptz DEFAULT now() NOT NULL,
+    modified_at timestamptz DEFAULT now() NOT NULL,
+    CONSTRAINT person_accounts_pkey PRIMARY KEY (id),
+    CONSTRAINT person_accounts_email_address_key UNIQUE (email_address)
+);
+
+-- SQL-language function with SET search_path = public, pg_temp
+-- that references a table. PostgreSQL validates SQL function bodies at
+-- creation time using the function's own search_path, not the session's.
+-- This reproduces issue #335 where the function's search_path isn't
+-- rewritten to point to the temporary schema.
+CREATE OR REPLACE FUNCTION auth_lookup_account_by_email(input_email text)
+RETURNS text
+LANGUAGE sql
+STABLE
+SECURITY DEFINER
+SET search_path = public, pg_temp
+AS $$
+    SELECT
+        pa.id::text AS person_account_id
+    FROM person_accounts pa
+    WHERE lower(pa.email_address) = lower(trim(input_email))
+    LIMIT 1;
+$$;
+
+REVOKE ALL ON FUNCTION auth_lookup_account_by_email(text) FROM PUBLIC;

testdata/diff/create_function/issue_335_search_path_rewrite/old.sql

@@ -0,0 +1 @@
+-- Empty schema (no functions or tables)

testdata/diff/create_function/issue_335_search_path_rewrite/plan.json

@@ -0,0 +1,32 @@
+{
+  "version": "1.0.0",
+  "pgschema_version": "1.7.3",
+  "created_at": "1970-01-01T00:00:00Z",
+  "source_fingerprint": {
+    "hash": "965b1131737c955e24c7f827c55bd78e4cb49a75adfd04229e0ba297376f5085"
+  },
+  "groups": [
+    {
+      "steps": [
+        {
+          "sql": "CREATE TABLE IF NOT EXISTS person_accounts (\n    id uuid DEFAULT gen_random_uuid(),\n    first_name text,\n    last_name text,\n    email_address text NOT NULL,\n    created_at timestamptz DEFAULT now() NOT NULL,\n    modified_at timestamptz DEFAULT now() NOT NULL,\n    CONSTRAINT person_accounts_pkey PRIMARY KEY (id),\n    CONSTRAINT person_accounts_email_address_key UNIQUE (email_address)\n);",
+          "type": "table",
+          "operation": "create",
+          "path": "public.person_accounts"
+        },
+        {
+          "sql": "CREATE OR REPLACE FUNCTION auth_lookup_account_by_email(\n    input_email text\n)\nRETURNS text\nLANGUAGE sql\nSTABLE\nSECURITY DEFINER\nSET search_path = public, pg_temp\nAS $$\n    SELECT\n        pa.id::text AS person_account_id\n    FROM person_accounts pa\n    WHERE lower(pa.email_address) = lower(trim(input_email))\n    LIMIT 1;\n$$;",
+          "type": "function",
+          "operation": "create",
+          "path": "public.auth_lookup_account_by_email"
+        },
+        {
+          "sql": "REVOKE EXECUTE ON FUNCTION auth_lookup_account_by_email(input_email text) FROM PUBLIC;",
+          "type": "revoked_default_privilege",
+          "operation": "create",
+          "path": "revoked_default.FUNCTION.auth_lookup_account_by_email(input_email text)"
+        }
+      ]
+    }
+  ]
+}

testdata/diff/create_function/issue_335_search_path_rewrite/plan.sql

@@ -0,0 +1,28 @@
+CREATE TABLE IF NOT EXISTS person_accounts (
+    id uuid DEFAULT gen_random_uuid(),
+    first_name text,
+    last_name text,
+    email_address text NOT NULL,
+    created_at timestamptz DEFAULT now() NOT NULL,
+    modified_at timestamptz DEFAULT now() NOT NULL,
+    CONSTRAINT person_accounts_pkey PRIMARY KEY (id),
+    CONSTRAINT person_accounts_email_address_key UNIQUE (email_address)
+);
+
+CREATE OR REPLACE FUNCTION auth_lookup_account_by_email(
+    input_email text
+)
+RETURNS text
+LANGUAGE sql
+STABLE
+SECURITY DEFINER
+SET search_path = public, pg_temp
+AS $$
+    SELECT
+        pa.id::text AS person_account_id
+    FROM person_accounts pa
+    WHERE lower(pa.email_address) = lower(trim(input_email))
+    LIMIT 1;
+$$;
+
+REVOKE EXECUTE ON FUNCTION auth_lookup_account_by_email(input_email text) FROM PUBLIC;