fix: support routine privilege diffs

internal/migration_acceptance_tests/privilege_cases_test.go

@@ -156,6 +156,39 @@ var privilegeAcceptanceTestCases = []acceptanceTestCase{
 			diff.MigrationHazardTypeAuthzUpdate,
 		},
 	},
+	{
+		name:  "Revoke implicit PUBLIC EXECUTE and grant function EXECUTE to role",
+		roles: []string{"app_user"},
+		oldSchemaDDL: []string{
+			`
+				CREATE SCHEMA welcome;
+				CREATE FUNCTION welcome.available_deposit_ordinal() RETURNS integer
+					LANGUAGE SQL
+					RETURN 1;
+				CREATE FUNCTION welcome.on_deposit(deposit_id uuid) RETURNS void
+					LANGUAGE SQL
+					AS $$ SELECT $$;
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+				CREATE SCHEMA welcome;
+				CREATE FUNCTION welcome.available_deposit_ordinal() RETURNS integer
+					LANGUAGE SQL
+					RETURN 1;
+				REVOKE EXECUTE ON FUNCTION welcome.available_deposit_ordinal() FROM PUBLIC;
+				GRANT EXECUTE ON FUNCTION welcome.available_deposit_ordinal() TO app_user;
+
+				CREATE FUNCTION welcome.on_deposit(deposit_id uuid) RETURNS void
+					LANGUAGE SQL
+					AS $$ SELECT $$;
+				REVOKE EXECUTE ON FUNCTION welcome.on_deposit(uuid) FROM PUBLIC;
+			`,
+		},
+		expectedHazardTypes: []diff.MigrationHazardType{
+			diff.MigrationHazardTypeAuthzUpdate,
+		},
+	},
 	{
 		name:  "Grant on partitioned parent table",
 		roles: []string{"app_user"},

internal/queries/queries.sql

@@ -266,7 +266,19 @@ SELECT
     pg_catalog.pg_get_function_identity_arguments(
         pg_proc.oid
     ) AS func_identity_arguments,
-    pg_catalog.pg_get_functiondef(pg_proc.oid) AS func_def
+    pg_catalog.pg_get_functiondef(pg_proc.oid) AS func_def,
+    ARRAY(
+        SELECT json_build_object(
+            'grantee', COALESCE(grantee_role.rolname, ''),
+            'privilege', acl.privilege_type,
+            'is_grantable', acl.is_grantable
+        )::TEXT
+        FROM ACLEXPLODE(COALESCE(pg_proc.proacl, ACLDEFAULT('f', pg_proc.proowner))) AS acl
+        LEFT JOIN pg_catalog.pg_roles AS grantee_role
+            ON acl.grantee = grantee_role.oid
+        WHERE acl.grantee != pg_proc.proowner OR acl.grantee = 0
+        ORDER BY COALESCE(grantee_role.rolname, ''), acl.privilege_type
+    ) AS privileges
 FROM pg_catalog.pg_proc
 INNER JOIN
     pg_catalog.pg_namespace AS proc_namespace

internal/schema/schema.go

@@ -86,11 +86,17 @@ func (s Schema) Normalize() Schema {
 	var normFunctions []Function
 	for _, function := range sortSchemaObjectsByName(s.Functions) {
 		function.DependsOnFunctions = sortSchemaObjectsByName(function.DependsOnFunctions)
+		function.Privileges = sortSchemaObjectsByName(function.Privileges)
 		normFunctions = append(normFunctions, function)
 	}
 	s.Functions = normFunctions
 
-	s.Procedures = sortSchemaObjectsByName(s.Procedures)
+	var normProcedures []Procedure
+	for _, procedure := range sortSchemaObjectsByName(s.Procedures) {
+		procedure.Privileges = sortSchemaObjectsByName(procedure.Privileges)
+		normProcedures = append(normProcedures, procedure)
+	}
+	s.Procedures = normProcedures
 	s.Triggers = sortSchemaObjectsByName(s.Triggers)
 
 	var normViews []View
@@ -242,24 +248,27 @@ func (t Table) IsPartition() bool {
 	return t.ParentTable != nil
 }
 
-// TablePrivilege represents a privilege granted on a table
-type TablePrivilege struct {
+// Privilege represents a privilege granted on a schema object.
+type Privilege struct {
 	// Grantee is the role that has the privilege. Empty string means PUBLIC.
 	Grantee string
-	// Privilege is the type of privilege (SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER)
+	// Privilege is the type of privilege (SELECT, EXECUTE, etc.)
 	Privilege string
 	// IsGrantable indicates if the grantee can grant this privilege to others (WITH GRANT OPTION)
 	IsGrantable bool
 }
 
-func (p TablePrivilege) GetName() string {
+func (p Privilege) GetName() string {
 	grantee := p.Grantee
 	if grantee == "" {
 		grantee = "PUBLIC"
 	}
 	return fmt.Sprintf("%s:%s", grantee, p.Privilege)
 }
 
+// TablePrivilege represents a privilege granted on a table.
+type TablePrivilege = Privilege
+
 type ColumnIdentityType string
 
 const (
@@ -449,14 +458,16 @@ type Function struct {
 	// can track the dependencies of the function (or not)
 	Language           string
 	DependsOnFunctions []SchemaQualifiedName
+	Privileges         []Privilege
 }
 
 type Procedure struct {
 	SchemaQualifiedName
 	// Def is the statement required to completely (re)create
 	// the procedure, as returned by `pg_get_functiondef`. It is a CREATE OR REPLACE
 	// statement.
-	Def string
+	Def        string
+	Privileges []Privilege
 }
 
 var (
@@ -1319,12 +1330,17 @@ func (s *schemaFetcher) buildFunction(ctx context.Context, rawFunction queries.G
 	if err != nil {
 		return Function{}, fmt.Errorf("fetchDependsOnFunctions(%s): %w", rawFunction.Oid, err)
 	}
+	privileges, err := parseJSONPrivileges(rawFunction.Privileges)
+	if err != nil {
+		return Function{}, fmt.Errorf("parseJSONPrivileges(%s): %w", rawFunction.Oid, err)
+	}
 
 	return Function{
 		SchemaQualifiedName: buildProcName(rawFunction.FuncName, rawFunction.FuncIdentityArguments, rawFunction.FuncSchemaName),
 		FunctionDef:         rawFunction.FuncDef,
 		Language:            rawFunction.FuncLang,
 		DependsOnFunctions:  dependsOnFunctions,
+		Privileges:          privileges,
 	}, nil
 }
 
@@ -1353,9 +1369,14 @@ func (s *schemaFetcher) fetchProcedures(ctx context.Context) ([]Procedure, error
 
 	var procedures []Procedure
 	for _, rawProcedure := range rawProcedures {
+		privileges, err := parseJSONPrivileges(rawProcedure.Privileges)
+		if err != nil {
+			return nil, fmt.Errorf("parseJSONPrivileges(%s): %w", rawProcedure.Oid, err)
+		}
 		p := Procedure{
 			SchemaQualifiedName: buildProcName(rawProcedure.FuncName, rawProcedure.FuncIdentityArguments, rawProcedure.FuncSchemaName),
 			Def:                 rawProcedure.FuncDef,
+			Privileges:          privileges,
 		}
 		procedures = append(procedures, p)
 	}
@@ -1583,6 +1604,26 @@ func parseJSONTableDependencies(vals []string) ([]TableDependency, error) {
 	return out, nil
 }
 
+func parseJSONPrivileges(vals []string) ([]Privilege, error) {
+	var out []Privilege
+	for _, v := range vals {
+		var p struct {
+			Grantee     string `json:"grantee"`
+			Privilege   string `json:"privilege"`
+			IsGrantable bool   `json:"is_grantable"`
+		}
+		if err := json.Unmarshal([]byte(v), &p); err != nil {
+			return nil, fmt.Errorf("json.Unmarshal(%q, Privilege): %w", string(v), err)
+		}
+		out = append(out, Privilege{
+			Grantee:     p.Grantee,
+			Privilege:   p.Privilege,
+			IsGrantable: p.IsGrantable,
+		})
+	}
+	return out, nil
+}
+
 // buildProcName is used to build the schema qualified name for a proc (function, procedure), i.e., anything
 // identified by a name AND its arguments.
 func buildProcName(name, identityArguments, schemaName string) SchemaQualifiedName {

internal/schema/schema_test.go

@@ -34,6 +34,7 @@ var (
 		EscapedName: `"C"`,
 		SchemaName:  "pg_catalog",
 	}
+	defaultExecutePrivileges = []Privilege{{Privilege: "EXECUTE"}}
 
 	testCases = []*testCase{
 		// Exclude materialized views from the test for now because Postgres 14-15 fully qualify column names while Postgres
@@ -239,7 +240,7 @@ var (
 			GRANT SELECT ON schema_2.foo TO some_role_1;
 			GRANT INSERT ON schema_2.foo TO some_role_2 WITH GRANT OPTION;
 		`},
-			expectedHash: "4c2174e2cac3956b",
+			expectedHash: "228b09af8d579433",
 			expectedSchema: Schema{
 				NamedSchemas: []NamedSchema{
 					{Name: "public"},
@@ -466,26 +467,31 @@ var (
 							{EscapedName: "\"add\"(a integer, b integer)", SchemaName: "schema_filtered_1"},
 							{EscapedName: "\"increment\"(i integer)", SchemaName: "schema_1"},
 						},
+						Privileges: defaultExecutePrivileges,
 					},
 					{
 						SchemaQualifiedName: SchemaQualifiedName{EscapedName: "\"increment\"(i integer)", SchemaName: "schema_1"},
 						FunctionDef:         "CREATE OR REPLACE FUNCTION schema_1.increment(i integer)\n RETURNS integer\n LANGUAGE plpgsql\nAS $function$\n\t\t\t\t\tBEGIN\n\t\t\t\t\t\t\tRETURN i + 1;\n\t\t\t\t\tEND;\n\t\t\t$function$\n",
 						Language:            "plpgsql",
+						Privileges:          defaultExecutePrivileges,
 					},
 					{
 						SchemaQualifiedName: SchemaQualifiedName{EscapedName: "\"increment_version\"()", SchemaName: "public"},
 						FunctionDef:         "CREATE OR REPLACE FUNCTION public.increment_version()\n RETURNS trigger\n LANGUAGE plpgsql\nAS $function$\n\t\t\t\tBEGIN\n\t\t\t\t\tNEW.version = OLD.version + 1;\n\t\t\t\t\tRETURN NEW;\n\t\t\t\tEND;\n\t\t\t$function$\n",
 						Language:            "plpgsql",
+						Privileges:          defaultExecutePrivileges,
 					},
 				},
 				Procedures: []Procedure{
 					{
 						SchemaQualifiedName: SchemaQualifiedName{SchemaName: "public", EscapedName: "\"some_plpgsql_procedure\"(IN foobar numeric)"},
 						Def:                 "CREATE OR REPLACE PROCEDURE public.some_plpgsql_procedure(IN foobar numeric)\n LANGUAGE plpgsql\nAS $procedure$\n\t\t\t\tBEGIN\n\t\t\t\t\tRAISE NOTICE 'some notice';\n\t\t\t\tEND\n\t\t\t\t$procedure$\n",
+						Privileges:          defaultExecutePrivileges,
 					},
 					{
 						SchemaQualifiedName: SchemaQualifiedName{SchemaName: "schema_2", EscapedName: "\"some_insert_procedure\"(IN a integer, IN b integer)"},
 						Def:                 "CREATE OR REPLACE PROCEDURE schema_2.some_insert_procedure(IN a integer, IN b integer)\n LANGUAGE sql\nBEGIN ATOMIC\n INSERT INTO schema_2.foo DEFAULT VALUES;\nEND\n",
+						Privileges:          defaultExecutePrivileges,
 					},
 				},
 				Triggers: []Trigger{
@@ -591,7 +597,7 @@ var (
 			ALTER TABLE foo_fk_1 ADD CONSTRAINT foo_fk_1_fk FOREIGN KEY (author, content) REFERENCES foo_1 (author, content)
 				NOT VALID;
 		`},
-			expectedHash: "32c5a9c52dcfb15e",
+			expectedHash: "85d661d829e05ae1",
 			expectedSchema: Schema{
 				NamedSchemas: []NamedSchema{
 					{Name: "public"},
@@ -881,6 +887,7 @@ var (
 						SchemaQualifiedName: SchemaQualifiedName{EscapedName: "\"increment_version\"()", SchemaName: "public"},
 						FunctionDef:         "CREATE OR REPLACE FUNCTION public.increment_version()\n RETURNS trigger\n LANGUAGE plpgsql\nAS $function$\n\t\t\t\tBEGIN\n\t\t\t\t\tNEW.version = OLD.version + 1;\n\t\t\t\t\tRETURN NEW;\n\t\t\t\tEND;\n\t\t\t$function$\n",
 						Language:            "plpgsql",
+						Privileges:          defaultExecutePrivileges,
 					},
 				},
 				Triggers: []Trigger{

pkg/diff/function_sql_vertex_generator.go

@@ -30,12 +30,21 @@ func (f *functionSQLVertexGenerator) Add(function schema.Function) ([]Statement,
 				"created/altered before this statement.",
 		})
 	}
-	return []Statement{{
+	stmts := []Statement{{
 		DDL:         function.FunctionDef,
 		Timeout:     statementTimeoutDefault,
 		LockTimeout: lockTimeoutDefault,
 		Hazards:     hazards,
-	}}, nil
+	}}
+	privilegeStmts, err := exactRoutinePrivilegeStatements(
+		newFunctionPrivilegeSQLVertexGenerator(function.SchemaQualifiedName),
+		function.Privileges,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("generating function privilege statements: %w", err)
+	}
+	stmts = append(stmts, privilegeStmts...)
+	return stmts, nil
 }
 
 func (f *functionSQLVertexGenerator) Delete(function schema.Function) ([]Statement, error) {
@@ -60,10 +69,33 @@ func (f *functionSQLVertexGenerator) Delete(function schema.Function) ([]Stateme
 func (f *functionSQLVertexGenerator) Alter(diff functionDiff) ([]Statement, error) {
 	// We are assuming the function has been normalized, i.e., we don't have to worry DependsOnFunctions ordering
 	// causing a false positive diff detected.
-	if cmp.Equal(diff.old, diff.new) {
-		return nil, nil
+	oldWithoutPrivileges := diff.old
+	oldWithoutPrivileges.Privileges = nil
+	newWithoutPrivileges := diff.new
+	newWithoutPrivileges.Privileges = nil
+
+	var stmts []Statement
+	if !cmp.Equal(oldWithoutPrivileges, newWithoutPrivileges) {
+		addStmts, err := f.Add(diff.new)
+		if err != nil {
+			return nil, err
+		}
+		stmts = append(stmts, addStmts...)
+	} else {
+		privilegesPartialGraph, err := generatePartialGraph(
+			newFunctionPrivilegeSQLVertexGenerator(diff.new.SchemaQualifiedName),
+			diff.privilegesDiff,
+		)
+		if err != nil {
+			return nil, fmt.Errorf("resolving function privilege sql: %w", err)
+		}
+		privilegeStmts, err := graphStatements(privilegesPartialGraph)
+		if err != nil {
+			return nil, fmt.Errorf("ordering function privilege sql: %w", err)
+		}
+		stmts = append(stmts, privilegeStmts...)
 	}
-	return f.Add(diff.new)
+	return stmts, nil
 }
 
 func canFunctionDependenciesBeTracked(function schema.Function) bool {

pkg/diff/plan_generator.go

@@ -279,23 +279,38 @@ func schemaFromTempDb(ctx context.Context, db *tempdb.Database, plan *planOption
 	return schema.GetSchema(ctx, db.ConnPool, append(plan.getSchemaOpts, db.ExcludeMetadataOptions...)...)
 }
 
-// clearTablePrivileges returns a copy of the schema with all table privileges cleared.
+// clearPrivileges returns a copy of the schema with all privileges cleared.
 // This is used during plan validation because privilege statements are skipped (roles don't exist in temp DB).
-func clearTablePrivileges(s schema.Schema) schema.Schema {
+func clearPrivileges(s schema.Schema) schema.Schema {
 	tables := make([]schema.Table, len(s.Tables))
 	for i, t := range s.Tables {
 		t.Privileges = nil
 		tables[i] = t
 	}
 	s.Tables = tables
+
+	functions := make([]schema.Function, len(s.Functions))
+	for i, f := range s.Functions {
+		f.Privileges = nil
+		functions[i] = f
+	}
+	s.Functions = functions
+
+	procedures := make([]schema.Procedure, len(s.Procedures))
+	for i, p := range s.Procedures {
+		p.Privileges = nil
+		procedures[i] = p
+	}
+	s.Procedures = procedures
+
 	return s
 }
 
 func assertMigratedSchemaMatchesTarget(migratedSchema, targetSchema schema.Schema, planOptions *planOptions) error {
 	// Clear privileges from both schemas since privilege statements are skipped during validation
 	// (roles don't exist in temp DB). We make copies to avoid modifying the original schemas.
-	migratedSchema = clearTablePrivileges(migratedSchema)
-	targetSchema = clearTablePrivileges(targetSchema)
+	migratedSchema = clearPrivileges(migratedSchema)
+	targetSchema = clearPrivileges(targetSchema)
 
 	toTargetSchemaStmts, err := generateMigrationStatements(migratedSchema, targetSchema, planOptions)
 	if err != nil {