fix: order a table after the functions its policies call

internal/migration_acceptance_tests/policy_cases_test.go

@@ -736,6 +736,33 @@ var policyAcceptanceTestCases = []acceptanceTestCase{
 			diff.MigrationHazardTypeAuthzUpdate,
 		},
 	},
+	{
+		name: "Create policy whose expression calls a function created in the same plan",
+		oldSchemaDDL: []string{
+			``,
+		},
+		newSchemaDDL: []string{
+			`
+                CREATE FUNCTION current_user_id() RETURNS uuid
+                    LANGUAGE sql STABLE
+                    AS $$ SELECT '00000000-0000-0000-0000-000000000000'::uuid $$;
+                CREATE TABLE foobar(id uuid);
+                CREATE POLICY foobar_policy ON foobar
+                    AS PERMISSIVE
+                    FOR SELECT
+                    TO PUBLIC
+                    USING (id = current_user_id());
+			`,
+		},
+		// PostgreSQL resolves the function the USING expression calls at CREATE time, so the function
+		// must come first.
+		expectedPlanDDL: []string{
+			"CREATE OR REPLACE FUNCTION public.current_user_id()\n RETURNS uuid\n LANGUAGE sql\n STABLE\nAS $function$ SELECT '00000000-0000-0000-0000-000000000000'::uuid $function$\n",
+			"CREATE TABLE \"public\".\"foobar\" (\n\t\"id\" uuid\n)",
+			"CREATE POLICY \"foobar_policy\" ON \"public\".\"foobar\"\n\tAS PERMISSIVE\n\tFOR SELECT\n\tTO PUBLIC\n\tUSING ((id = current_user_id()))",
+		},
+		// No hazards: the table is new, so the policy statements' hazards are stripped.
+	},
 }
 
 func TestPolicyCases(t *testing.T) {

internal/queries/queries.sql

@@ -447,6 +447,7 @@ WITH roles AS (
 )
 
 SELECT
+    pol.oid,
     pol.polname::TEXT AS policy_name,
     table_c.relname::TEXT AS owning_table_name,
     table_namespace.nspname::TEXT AS owning_table_schema_name,

internal/schema/schema.go

@@ -129,6 +129,7 @@ func normalizeTable(t Table) Table {
 		p.Columns = sortByKey(p.Columns, func(s string) string {
 			return s
 		})
+		p.DependsOnFunctions = sortSchemaObjectsByName(p.DependsOnFunctions)
 		normPolicies = append(normPolicies, p)
 	}
 	t.Policies = normPolicies
@@ -496,6 +497,10 @@ type Policy struct {
 	UsingExpression string
 	// Columns are the columns that the policy applies to.
 	Columns []string
+	// DependsOnFunctions is the list of functions referenced by the policy's USING and WITH CHECK
+	// expressions. PostgreSQL resolves those references at CREATE time, so the policy must be
+	// created after the functions exist and dropped before they are dropped.
+	DependsOnFunctions []SchemaQualifiedName
 }
 
 func (p Policy) GetName() string {
@@ -1389,15 +1394,20 @@ func (s *schemaFetcher) fetchPolicies(ctx context.Context) ([]policyAndTable, er
 
 	var policies []policyAndTable
 	for _, rp := range rawPolicies {
+		dependsOnFunctions, err := s.fetchDependsOnFunctions(ctx, "pg_policy", rp.Oid)
+		if err != nil {
+			return nil, fmt.Errorf("fetchDependsOnFunctions(%s): %w", rp.Oid, err)
+		}
 		policies = append(policies, policyAndTable{
 			policy: Policy{
-				EscapedName:     EscapeIdentifier(rp.PolicyName),
-				IsPermissive:    rp.IsPermissive,
-				AppliesTo:       rp.AppliesTo,
-				Cmd:             PolicyCmd(rp.Cmd),
-				CheckExpression: rp.CheckExpression,
-				UsingExpression: rp.UsingExpression,
-				Columns:         rp.ColumnNames,
+				EscapedName:        EscapeIdentifier(rp.PolicyName),
+				IsPermissive:       rp.IsPermissive,
+				AppliesTo:          rp.AppliesTo,
+				Cmd:                PolicyCmd(rp.Cmd),
+				CheckExpression:    rp.CheckExpression,
+				UsingExpression:    rp.UsingExpression,
+				Columns:            rp.ColumnNames,
+				DependsOnFunctions: dependsOnFunctions,
 			},
 			table: buildNameFromUnescaped(rp.OwningTableName, rp.OwningTableSchemaName),
 		})

internal/schema/schema_test.go

@@ -239,7 +239,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: "eff2271f4336fe47",
 			expectedSchema: Schema{
 				NamedSchemas: []NamedSchema{
 					{Name: "public"},

pkg/diff/sql_generator.go

@@ -1175,6 +1175,7 @@ func (t *tableSQLVertexGenerator) GetAddAlterDependencies(table, _ schema.Table)
 			mustRun(t.GetSQLVertexId(table, diffTypeAddAlter)).after(buildTableVertexId(*table.ParentTable, diffTypeAddAlter)),
 		)
 	}
+	deps = append(deps, consumerPolicyFunctionDependencies(table)...)
 	return deps, nil
 }
 
@@ -1243,9 +1244,37 @@ func (t *tableSQLVertexGenerator) GetDeleteDependencies(table schema.Table) ([]d
 			mustRun(t.GetSQLVertexId(table, diffTypeDelete)).after(buildTableVertexId(*table.ParentTable, diffTypeDelete)),
 		)
 	}
+	deps = append(deps, policyFunctionDeleteDependencies(table)...)
 	return deps, nil
 }
 
+// consumerPolicyFunctionDependencies orders a table's add/alter after every function its policies
+// call. A table's policies are created in the same vertex as the table — inside the CREATE TABLE
+// statement list for a new table, appended by the table's alter otherwise — and PostgreSQL resolves
+// the functions named in a policy expression at CREATE POLICY time.
+func consumerPolicyFunctionDependencies(table schema.Table) []dependency {
+	var deps []dependency
+	for _, policy := range table.Policies {
+		for _, depFunction := range policy.DependsOnFunctions {
+			deps = append(deps, mustRun(buildTableVertexId(table.SchemaQualifiedName, diffTypeAddAlter)).after(buildFunctionVertexId(depFunction, diffTypeAddAlter)))
+		}
+	}
+	return deps
+}
+
+// policyFunctionDeleteDependencies orders a table's drop before the drop of every function its
+// policies call: a policy lives and dies with its table, so the function may only be dropped after
+// the policy is gone.
+func policyFunctionDeleteDependencies(table schema.Table) []dependency {
+	var deps []dependency
+	for _, policy := range table.Policies {
+		for _, depFunction := range policy.DependsOnFunctions {
+			deps = append(deps, mustRun(buildTableVertexId(table.SchemaQualifiedName, diffTypeDelete)).before(buildFunctionVertexId(depFunction, diffTypeDelete)))
+		}
+	}
+	return deps
+}
+
 type columnSQLVertexGenerator struct {
 	tableName schema.SchemaQualifiedName
 }