fix: preserve schema qualifiers in function/procedure bodies (#354) (#358)

Previously, schema qualifiers (e.g., public.test) were unconditionally
stripped from function/procedure bodies during normalization. This broke
functions with restrictive search_path settings (SET search_path = '',
pg_temp, etc.) where qualified references are required.

Changes:
- Move schema stripping from normalization time to comparison time
- Make stripSchemaQualifications skip dollar-quoted blocks so function
  bodies applied to embedded postgres preserve their qualifiers
- Export StripSchemaPrefixFromBody for use by the diff package
- Fix dollar-quote regex to match PostgreSQL grammar (no $1$ false positives)
- Add unit tests for splitDollarQuotedSegments

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: tianzhou

internal/diff/function.go

@@ -371,7 +371,7 @@ func functionsEqualExceptAttributes(old, new *ir.Function) bool {
 	if old.Name != new.Name {
 		return false
 	}
-	if old.Definition != new.Definition {
+	if !definitionsEqualIgnoringSchema(old.Definition, new.Definition, old.Schema) {
 		return false
 	}
 	if old.ReturnType != new.ReturnType {
@@ -409,7 +409,7 @@ func functionsEqual(old, new *ir.Function) bool {
 	if old.Name != new.Name {
 		return false
 	}
-	if old.Definition != new.Definition {
+	if !definitionsEqualIgnoringSchema(old.Definition, new.Definition, old.Schema) {
 		return false
 	}
 	if old.ReturnType != new.ReturnType {
@@ -458,7 +458,7 @@ func functionsEqualExceptComment(old, new *ir.Function) bool {
 	if old.Name != new.Name {
 		return false
 	}
-	if old.Definition != new.Definition {
+	if !definitionsEqualIgnoringSchema(old.Definition, new.Definition, old.Schema) {
 		return false
 	}
 	if old.ReturnType != new.ReturnType {
@@ -522,6 +522,15 @@ func functionRequiresRecreate(old, new *ir.Function) bool {
 	return false
 }
 
+// definitionsEqualIgnoringSchema compares two function/procedure definitions,
+// stripping the given schema qualifier from both before comparing. This allows
+// definitions that differ only in schema qualification (e.g., "public.test" vs "test")
+// to be treated as equal, while preserving the original qualifiers in the IR for
+// correct DDL generation. (Issue #354)
+func definitionsEqualIgnoringSchema(a, b, schema string) bool {
+	return ir.StripSchemaPrefixFromBody(a, schema) == ir.StripSchemaPrefixFromBody(b, schema)
+}
+
 // filterNonTableParameters filters out TABLE mode parameters
 // TABLE parameters are output columns in RETURNS TABLE() and shouldn't be compared as input parameters
 func filterNonTableParameters(params []*ir.Parameter) []*ir.Parameter {

internal/diff/procedure.go

@@ -260,7 +260,7 @@ func proceduresEqual(old, new *ir.Procedure) bool {
 	if old.Name != new.Name {
 		return false
 	}
-	if old.Definition != new.Definition {
+	if !definitionsEqualIgnoringSchema(old.Definition, new.Definition, old.Schema) {
 		return false
 	}
 	if old.Language != new.Language {
@@ -296,7 +296,7 @@ func proceduresEqualExceptComment(old, new *ir.Procedure) bool {
 	if old.Name != new.Name {
 		return false
 	}
-	if old.Definition != new.Definition {
+	if !definitionsEqualIgnoringSchema(old.Definition, new.Definition, old.Schema) {
 		return false
 	}
 	if old.Language != new.Language {

internal/postgres/desired_state.go

@@ -96,6 +96,77 @@ func stripSchemaQualifications(sql string, schemaName string) string {
 		return sql
 	}
 
+	// Split SQL into dollar-quoted and non-dollar-quoted segments.
+	// Schema qualifiers inside function/procedure bodies (dollar-quoted blocks)
+	// must be preserved — the user may need them when search_path doesn't include
+	// the function's schema (e.g., SET search_path = ''). (Issue #354)
+	segments := splitDollarQuotedSegments(sql)
+	var result strings.Builder
+	result.Grow(len(sql))
+	for _, seg := range segments {
+		if seg.quoted {
+			// Preserve dollar-quoted content as-is
+			result.WriteString(seg.text)
+		} else {
+			result.WriteString(stripSchemaQualificationsFromText(seg.text, schemaName))
+		}
+	}
+	return result.String()
+}
+
+// dollarQuotedSegment represents a segment of SQL text, either inside or outside a dollar-quoted block.
+type dollarQuotedSegment struct {
+	text   string
+	quoted bool // true if this segment is inside dollar quotes (including the delimiters)
+}
+
+// splitDollarQuotedSegments splits SQL text into segments that are either inside or outside
+// dollar-quoted blocks ($$...$$, $tag$...$tag$, etc.). This allows callers to process
+// only the non-quoted parts while preserving function/procedure bodies verbatim.
+// dollarQuoteRe matches PostgreSQL dollar-quote tags: $$ or $identifier$ where the
+// identifier must start with a letter or underscore (not a digit). This avoids
+// false positives on $1, $2 etc. parameter references.
+var dollarQuoteRe = regexp.MustCompile(`\$(?:[a-zA-Z_][a-zA-Z0-9_]*)?\$`)
+
+func splitDollarQuotedSegments(sql string) []dollarQuotedSegment {
+	var segments []dollarQuotedSegment
+
+	pos := 0
+	for pos < len(sql) {
+		// Find the next dollar-quote opening tag
+		loc := dollarQuoteRe.FindStringIndex(sql[pos:])
+		if loc == nil {
+			// No more dollar quotes — rest is unquoted
+			segments = append(segments, dollarQuotedSegment{text: sql[pos:], quoted: false})
+			break
+		}
+
+		openStart := pos + loc[0]
+		openEnd := pos + loc[1]
+		tag := sql[openStart:openEnd]
+
+		// Add the unquoted segment before this tag
+		if openStart > pos {
+			segments = append(segments, dollarQuotedSegment{text: sql[pos:openStart], quoted: false})
+		}
+
+		// Find the matching closing tag
+		closeIdx := strings.Index(sql[openEnd:], tag)
+		if closeIdx == -1 {
+			// No closing tag — treat rest as quoted (unterminated)
+			segments = append(segments, dollarQuotedSegment{text: sql[openStart:], quoted: true})
+			pos = len(sql)
+		} else {
+			closeEnd := openEnd + closeIdx + len(tag)
+			segments = append(segments, dollarQuotedSegment{text: sql[openStart:closeEnd], quoted: true})
+			pos = closeEnd
+		}
+	}
+	return segments
+}
+
+// stripSchemaQualificationsFromText performs the actual schema qualification stripping on a text segment.
+func stripSchemaQualificationsFromText(text string, schemaName string) string {
 	// Escape the schema name for use in regex
 	escapedSchema := regexp.QuoteMeta(schemaName)
 
@@ -133,7 +204,7 @@ func stripSchemaQualifications(sql string, schemaName string) string {
 	pattern4 := fmt.Sprintf(`(?:^|[^"])%s\.([a-zA-Z_][a-zA-Z0-9_$]*)`, escapedSchema)
 	re4 := regexp.MustCompile(pattern4)
 
-	result := sql
+	result := text
 	// Apply in order: quoted schema first to avoid double-matching
 	result = re1.ReplaceAllString(result, "$1")
 	result = re2.ReplaceAllString(result, "$1")

internal/postgres/desired_state_test.go

@@ -1,9 +1,82 @@
 package postgres
 
 import (
+	"reflect"
 	"testing"
 )
 
+func TestSplitDollarQuotedSegments(t *testing.T) {
+	tests := []struct {
+		name     string
+		sql      string
+		expected []dollarQuotedSegment
+	}{
+		{
+			name:     "no dollar quotes",
+			sql:      "SELECT 1 FROM public.users;",
+			expected: []dollarQuotedSegment{{text: "SELECT 1 FROM public.users;", quoted: false}},
+		},
+		{
+			name: "simple dollar-quoted body",
+			sql:  "CREATE FUNCTION f() AS $$body$$ LANGUAGE sql;",
+			expected: []dollarQuotedSegment{
+				{text: "CREATE FUNCTION f() AS ", quoted: false},
+				{text: "$$body$$", quoted: true},
+				{text: " LANGUAGE sql;", quoted: false},
+			},
+		},
+		{
+			name: "named dollar-quote tag",
+			sql:  "AS $func$body$func$ LANGUAGE sql;",
+			expected: []dollarQuotedSegment{
+				{text: "AS ", quoted: false},
+				{text: "$func$body$func$", quoted: true},
+				{text: " LANGUAGE sql;", quoted: false},
+			},
+		},
+		{
+			name: "parameter references not treated as dollar quotes",
+			sql:  "SELECT $1 + $2 FROM t;",
+			expected: []dollarQuotedSegment{
+				{text: "SELECT $1 + $2 FROM t;", quoted: false},
+			},
+		},
+		{
+			name: "multiple dollar-quoted blocks",
+			sql:  "AS $$body1$$; AS $f$body2$f$;",
+			expected: []dollarQuotedSegment{
+				{text: "AS ", quoted: false},
+				{text: "$$body1$$", quoted: true},
+				{text: "; AS ", quoted: false},
+				{text: "$f$body2$f$", quoted: true},
+				{text: ";", quoted: false},
+			},
+		},
+		{
+			name: "unterminated dollar quote",
+			sql:  "AS $$body without end",
+			expected: []dollarQuotedSegment{
+				{text: "AS ", quoted: false},
+				{text: "$$body without end", quoted: true},
+			},
+		},
+		{
+			name:     "empty input",
+			sql:      "",
+			expected: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := splitDollarQuotedSegments(tt.sql)
+			if !reflect.DeepEqual(result, tt.expected) {
+				t.Errorf("splitDollarQuotedSegments(%q)\n  got:  %+v\n  want: %+v", tt.sql, result, tt.expected)
+			}
+		})
+	}
+}
+
 func TestReplaceSchemaInSearchPath(t *testing.T) {
 	tests := []struct {
 		name         string

ir/normalize.go

@@ -280,7 +280,7 @@ func normalizeView(view *View) {
 
 	// Strip same-schema qualifiers from view definition for consistent comparison.
 	// This uses the same logic as function/procedure body normalization.
-	view.Definition = stripSchemaPrefixFromBody(view.Definition, view.Schema)
+	view.Definition = StripSchemaPrefixFromBody(view.Definition, view.Schema)
 
 	// Normalize triggers on the view (e.g., INSTEAD OF triggers)
 	for _, trigger := range view.Triggers {
@@ -321,8 +321,10 @@ func normalizeFunction(function *Function) {
 	}
 	// Normalize function body to handle whitespace differences
 	function.Definition = normalizeFunctionDefinition(function.Definition)
-	// Strip current schema qualifier from function body for consistent unqualified output
-	function.Definition = stripSchemaPrefixFromBody(function.Definition, function.Schema)
+	// Note: We intentionally do NOT strip schema qualifiers from function bodies here.
+	// Functions may have SET search_path that excludes their own schema, making
+	// qualified references (e.g., public.test) necessary. Stripping is done at
+	// comparison time in the diff package instead. (Issue #354)
 }
 
 // normalizeFunctionDefinition normalizes function body whitespace
@@ -344,10 +346,10 @@ func normalizeFunctionDefinition(def string) string {
 	return strings.Join(normalized, "\n")
 }
 
-// stripSchemaPrefixFromBody removes the current schema qualifier from identifiers
+// StripSchemaPrefixFromBody removes the current schema qualifier from identifiers
 // in a function or procedure body. For example, "public.users" becomes "users".
 // It skips single-quoted string literals to avoid modifying string constants.
-func stripSchemaPrefixFromBody(body, schema string) string {
+func StripSchemaPrefixFromBody(body, schema string) string {
 	if body == "" || schema == "" {
 		return body
 	}
@@ -445,8 +447,8 @@ func normalizeProcedure(procedure *Procedure) {
 		}
 	}
 
-	// Strip current schema qualifier from procedure body for consistent unqualified output
-	procedure.Definition = stripSchemaPrefixFromBody(procedure.Definition, procedure.Schema)
+	// Note: We intentionally do NOT strip schema qualifiers from procedure bodies here.
+	// Same rationale as functions — see normalizeFunction. (Issue #354)
 }
 
 // normalizeFunctionReturnType normalizes function return types, especially TABLE types

ir/normalize_test.go

@@ -93,9 +93,9 @@ func TestStripSchemaPrefixFromBody(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			result := stripSchemaPrefixFromBody(tt.body, tt.schema)
+			result := StripSchemaPrefixFromBody(tt.body, tt.schema)
 			if result != tt.expected {
-				t.Errorf("stripSchemaPrefixFromBody(%q, %q) = %q, want %q", tt.body, tt.schema, result, tt.expected)
+				t.Errorf("StripSchemaPrefixFromBody(%q, %q) = %q, want %q", tt.body, tt.schema, result, tt.expected)
 			}
 		})
 	}

testdata/diff/create_function/issue_354_empty_search_path/diff.sql

@@ -7,6 +7,6 @@ VOLATILE
 SET search_path = ''
 AS $$
 BEGIN
-  INSERT INTO test (title) VALUES (p_title);
+  INSERT INTO public.test (title) VALUES (p_title);
 END;
 $$;

testdata/diff/create_function/issue_354_empty_search_path/plan.json

@@ -9,7 +9,7 @@
     {
       "steps": [
         {
-          "sql": "CREATE OR REPLACE FUNCTION create_hello(\n    p_title text\n)\nRETURNS void\nLANGUAGE plpgsql\nVOLATILE\nSET search_path = ''\nAS $$\nBEGIN\n  INSERT INTO test (title) VALUES (p_title);\nEND;\n$$;",
+          "sql": "CREATE OR REPLACE FUNCTION create_hello(\n    p_title text\n)\nRETURNS void\nLANGUAGE plpgsql\nVOLATILE\nSET search_path = ''\nAS $$\nBEGIN\n  INSERT INTO public.test (title) VALUES (p_title);\nEND;\n$$;",
           "type": "function",
           "operation": "create",
           "path": "public.create_hello"

testdata/diff/create_function/issue_354_empty_search_path/plan.sql

@@ -7,6 +7,6 @@ VOLATILE
 SET search_path = ''
 AS $$
 BEGIN
-  INSERT INTO test (title) VALUES (p_title);
+  INSERT INTO public.test (title) VALUES (p_title);
 END;
 $$;

testdata/diff/create_function/issue_354_empty_search_path/plan.txt

@@ -18,6 +18,6 @@ VOLATILE
 SET search_path = ''
 AS $$
 BEGIN
-  INSERT INTO test (title) VALUES (p_title);
+  INSERT INTO public.test (title) VALUES (p_title);
 END;
 $$;