Skip to content

Commit cbecb2a

Browse files
tianzhouclaude
andauthored
fix: strip temp schema prefix from GRANT/REVOKE function signatures (#376) (#379)
* fix: strip temp schema prefix from GRANT/REVOKE function signatures (#376) When using an external database for plan generation, pg_get_function_identity_arguments() qualifies user-defined argument types with the temporary schema name (e.g., pgschema_tmp_20260326_161823_31f3dbda.my_input[]). This leaked into GRANT/REVOKE statements, causing apply to fail with "schema does not exist". Fix: normalize privilege ObjectName to strip the schema prefix from function/procedure argument types, consistent with how other schema objects are normalized. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * refactor: use splitTableColumns and add unit tests for normalizePrivilegeObjectName Address review feedback: - Use splitTableColumns helper for arg splitting to correctly handle nested parentheses in type modifiers (e.g., numeric(10, 2)) - Add table-driven unit tests covering edge cases: empty inputs, no-arg functions, array types, multi-arg, temp schema prefix, precision types Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 079bde8 commit cbecb2a

8 files changed

Lines changed: 223 additions & 0 deletions

File tree

ir/normalize.go

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,18 @@ func normalizeSchema(schema *Schema) {
5757
for _, typeObj := range schema.Types {
5858
normalizeType(typeObj)
5959
}
60+
61+
// Normalize privileges - strip schema qualifiers from function/procedure signatures
62+
for _, priv := range schema.Privileges {
63+
if priv.ObjectType == PrivilegeObjectTypeFunction || priv.ObjectType == PrivilegeObjectTypeProcedure {
64+
priv.ObjectName = normalizePrivilegeObjectName(priv.ObjectName, schema.Name)
65+
}
66+
}
67+
for _, revoked := range schema.RevokedDefaultPrivileges {
68+
if revoked.ObjectType == PrivilegeObjectTypeFunction || revoked.ObjectType == PrivilegeObjectTypeProcedure {
69+
revoked.ObjectName = normalizePrivilegeObjectName(revoked.ObjectName, schema.Name)
70+
}
71+
}
6072
}
6173

6274
// normalizeTable normalizes table-related objects
@@ -634,6 +646,51 @@ func stripSchemaPrefix(typeName, prefix string) string {
634646
return typeName
635647
}
636648

649+
// normalizePrivilegeObjectName strips schema qualifiers from argument types
650+
// in function/procedure signatures used in GRANT/REVOKE statements.
651+
// e.g., "f_test(p_items pgschema_tmp_20260326_161823_31f3dbda.my_input[])" → "f_test(p_items my_input[])"
652+
func normalizePrivilegeObjectName(objectName, schemaName string) string {
653+
if objectName == "" || schemaName == "" {
654+
return objectName
655+
}
656+
657+
// Find the argument list in the signature: name(args)
658+
parenOpen := strings.Index(objectName, "(")
659+
parenClose := strings.LastIndex(objectName, ")")
660+
if parenOpen < 0 || parenClose < 0 || parenClose <= parenOpen {
661+
return objectName
662+
}
663+
664+
funcName := objectName[:parenOpen]
665+
args := objectName[parenOpen+1 : parenClose]
666+
667+
if args == "" {
668+
return objectName
669+
}
670+
671+
prefix := schemaName + "."
672+
673+
// Use splitTableColumns to correctly handle nested parentheses in type modifiers
674+
// (e.g., numeric(10, 2)) — consistent with other normalization helpers in this file
675+
parts := splitTableColumns(args)
676+
changed := false
677+
for i, part := range parts {
678+
trimmed := strings.TrimSpace(part)
679+
if strings.Contains(trimmed, prefix) {
680+
parts[i] = strings.ReplaceAll(trimmed, prefix, "")
681+
changed = true
682+
} else {
683+
parts[i] = trimmed
684+
}
685+
}
686+
687+
if !changed {
688+
return objectName
689+
}
690+
691+
return funcName + "(" + strings.Join(parts, ", ") + ")"
692+
}
693+
637694
// normalizeTrigger normalizes trigger representation
638695
func normalizeTrigger(trigger *Trigger) {
639696
if trigger == nil {

ir/normalize_test.go

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,85 @@ func TestNormalizePolicyExpression(t *testing.T) {
390390
}
391391
}
392392

393+
func TestNormalizePrivilegeObjectName(t *testing.T) {
394+
tests := []struct {
395+
name string
396+
objectName string
397+
schemaName string
398+
expected string
399+
}{
400+
{
401+
name: "empty object name",
402+
objectName: "",
403+
schemaName: "public",
404+
expected: "",
405+
},
406+
{
407+
name: "empty schema name",
408+
objectName: "f_test(p_items my_input[])",
409+
schemaName: "",
410+
expected: "f_test(p_items my_input[])",
411+
},
412+
{
413+
name: "no args",
414+
objectName: "f_test()",
415+
schemaName: "public",
416+
expected: "f_test()",
417+
},
418+
{
419+
name: "no schema prefix in args",
420+
objectName: "f_test(p_items my_input[])",
421+
schemaName: "public",
422+
expected: "f_test(p_items my_input[])",
423+
},
424+
{
425+
name: "temp schema prefix in array type",
426+
objectName: "f_test(p_items pgschema_tmp_20260326_161823_31f3dbda.my_input[])",
427+
schemaName: "pgschema_tmp_20260326_161823_31f3dbda",
428+
expected: "f_test(p_items my_input[])",
429+
},
430+
{
431+
name: "public schema prefix",
432+
objectName: "f_test(p_items public.my_input[])",
433+
schemaName: "public",
434+
expected: "f_test(p_items my_input[])",
435+
},
436+
{
437+
name: "multiple args with schema prefix",
438+
objectName: "f_test(a public.my_type, b integer, c public.other_type[])",
439+
schemaName: "public",
440+
expected: "f_test(a my_type, b integer, c other_type[])",
441+
},
442+
{
443+
name: "no parentheses",
444+
objectName: "my_table",
445+
schemaName: "public",
446+
expected: "my_table",
447+
},
448+
{
449+
name: "builtin types unchanged",
450+
objectName: "calculate_total(quantity integer, unit_price numeric)",
451+
schemaName: "public",
452+
expected: "calculate_total(quantity integer, unit_price numeric)",
453+
},
454+
{
455+
name: "type with precision in parens",
456+
objectName: "f_test(a public.my_type, b numeric(10, 2))",
457+
schemaName: "public",
458+
expected: "f_test(a my_type, b numeric(10, 2))",
459+
},
460+
}
461+
462+
for _, tt := range tests {
463+
t.Run(tt.name, func(t *testing.T) {
464+
result := normalizePrivilegeObjectName(tt.objectName, tt.schemaName)
465+
if result != tt.expected {
466+
t.Errorf("normalizePrivilegeObjectName(%q, %q) = %q, want %q", tt.objectName, tt.schemaName, result, tt.expected)
467+
}
468+
})
469+
}
470+
}
471+
393472
func TestNormalizeCheckClause(t *testing.T) {
394473
tests := []struct {
395474
name string
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
REVOKE EXECUTE ON FUNCTION f_test(p_items my_input[]) FROM PUBLIC;
2+
3+
GRANT EXECUTE ON FUNCTION f_test(p_items my_input[]) TO appname_apiuser;
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
DO $$
2+
BEGIN
3+
IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'appname_apiuser') THEN
4+
CREATE ROLE appname_apiuser;
5+
END IF;
6+
END $$;
7+
8+
CREATE TYPE my_input AS (
9+
id uuid
10+
);
11+
12+
CREATE FUNCTION f_test(p_items my_input[])
13+
RETURNS integer
14+
LANGUAGE sql
15+
AS $$
16+
SELECT COALESCE(array_length(p_items, 1), 0);
17+
$$;
18+
19+
REVOKE ALL ON FUNCTION f_test(my_input[]) FROM PUBLIC;
20+
GRANT EXECUTE ON FUNCTION f_test(my_input[]) TO appname_apiuser;
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
DO $$
2+
BEGIN
3+
IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'appname_apiuser') THEN
4+
CREATE ROLE appname_apiuser;
5+
END IF;
6+
END $$;
7+
8+
CREATE TYPE my_input AS (
9+
id uuid
10+
);
11+
12+
CREATE FUNCTION f_test(p_items my_input[])
13+
RETURNS integer
14+
LANGUAGE sql
15+
AS $$
16+
SELECT COALESCE(array_length(p_items, 1), 0);
17+
$$;
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
{
2+
"version": "1.0.0",
3+
"pgschema_version": "1.8.0",
4+
"created_at": "1970-01-01T00:00:00Z",
5+
"source_fingerprint": {
6+
"hash": "27f3398a052dfc62c099d3054981ea0e5493b954bff2ffd763222a4fc21de01c"
7+
},
8+
"groups": [
9+
{
10+
"steps": [
11+
{
12+
"sql": "REVOKE EXECUTE ON FUNCTION f_test(p_items my_input[]) FROM PUBLIC;",
13+
"type": "revoked_default_privilege",
14+
"operation": "create",
15+
"path": "revoked_default.FUNCTION.f_test(p_items my_input[])"
16+
},
17+
{
18+
"sql": "GRANT EXECUTE ON FUNCTION f_test(p_items my_input[]) TO appname_apiuser;",
19+
"type": "privilege",
20+
"operation": "create",
21+
"path": "privileges.FUNCTION.f_test(p_items my_input[]).appname_apiuser"
22+
}
23+
]
24+
}
25+
]
26+
}
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
REVOKE EXECUTE ON FUNCTION f_test(p_items my_input[]) FROM PUBLIC;
2+
3+
GRANT EXECUTE ON FUNCTION f_test(p_items my_input[]) TO appname_apiuser;
Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
Plan: 2 to add.
2+
3+
Summary by type:
4+
privileges: 1 to add
5+
revoked default privileges: 1 to add
6+
7+
Privileges:
8+
+ appname_apiuser
9+
10+
Revoked default privileges:
11+
+ f_test(p_items my_input[])
12+
13+
DDL to be executed:
14+
--------------------------------------------------
15+
16+
REVOKE EXECUTE ON FUNCTION f_test(p_items my_input[]) FROM PUBLIC;
17+
18+
GRANT EXECUTE ON FUNCTION f_test(p_items my_input[]) TO appname_apiuser;

0 commit comments

Comments
 (0)