Repository: ariga/atlas
Pull request #2929 (2024-07-08 08:00:09)

sql/postgres: expose postgres policies
Already approved internally. 


Patch:
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index 6e89fd28..4a253ac9 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -63,6 +63,11 @@ func (d *diff) TableAttrDiff(from, to *schema.Table) ([]schema.Change, error) {
 	if err := d.partitionChanged(from, to); err != nil {
 		return nil, err
 	}
+	change, err := d.tableAttrDiff(from, to)
+	if err != nil {
+		return nil, err
+	}
+	changes = append(changes, change...)
 	return append(changes, sqlx.CheckDiff(from, to, func(c1, c2 *schema.Check) bool {
 		return sqlx.Has(c1.Attrs, &NoInherit{}) == sqlx.Has(c2.Attrs, &NoInherit{})
 	})...), nil
diff --git a/sql/postgres/driver_oss.go b/sql/postgres/driver_oss.go
index 62d0e780..1b756104 100644
--- a/sql/postgres/driver_oss.go
+++ b/sql/postgres/driver_oss.go
@@ -30,6 +30,34 @@ var (
 	}
 )
 
+// newTable creates a new table.
+func (*inspect) newTable(_ context.Context, name, _ string) *schema.Table {
+	return schema.NewTable(name) // Default implementation.
+}
+
+func tableAttrsSpec(*schema.Table, *sqlspec.Table) {
+	// unimplemented.
+}
+
+func convertTableAttrs(*sqlspec.Table, *schema.Table) error {
+	return nil // unimplemented.
+}
+
+// tableAttrDiff allows extending table attributes diffing with build-specific logic.
+func (*diff) tableAttrDiff(_, _ *schema.Table) ([]schema.Change, error) {
+	return nil, nil // unimplemented.
+}
+
+// addTableAttrs allows extending table attributes creation with build-specific logic.
+func (*state) addTableAttrs(_ *schema.AddTable) {
+	// unimplemented.
+}
+
+// alterTableAttr allows extending table attributes alteration with build-specific logic.
+func (s *state) alterTableAttr(*sqlx.Builder, *schema.ModifyAttr) {
+	// unimplemented.
+}
+
 func realmObjectsSpec(*doc, *schema.Realm) error {
 	return nil // unimplemented.
 }
@@ -50,7 +78,7 @@ func (*inspect) inspectTypes(context.Context, *schema.Realm, *schema.InspectOpti
 	return nil // unimplemented.
 }
 
-func (*inspect) inspectSequences(context.Context, *schema.Realm, *schema.InspectOptions) error {
+func (*inspect) inspectObjects(context.Context, *schema.Realm, *schema.InspectOptions) error {
 	return nil // unimplemented.
 }
 
@@ -242,6 +270,13 @@ func convertSequences(_ []*sqlspec.Table, seqs []*sqlspec.Sequence, _ *schema.Re
 	return nil
 }
 
+func convertPolicies(_ []*sqlspec.Table, ps []*policy, _ *schema.Realm) error {
+	if len(ps) > 0 {
+		return fmt.Errorf("postgres: policies are not supported by this version. Use: https://atlasgo.io/getting-started")
+	}
+	return nil
+}
+
 func convertExtensions(exs []*extension, _ *schema.Realm) error {
 	if len(exs) > 0 {
 		return fmt.Errorf("postgres: extensions are not supported by this version. Use: https://atlasgo.io/getting-started")
@@ -381,3 +416,59 @@ func (*state) detachCycles(changes []schema.Change) ([]schema.Change, error) {
 func excludeSpec(*sqlspec.Table, *sqlspec.Index, *schema.Index, *Constraint) error {
 	return nil // unimplemented.
 }
+
+const (
+	// Query to list tables information.
+	// Note, 'attrs' are not supported in this version.
+	tablesQuery = `
+SELECT
+	t3.oid,
+	t1.table_schema,
+	t1.table_name,
+	pg_catalog.obj_description(t3.oid, 'pg_class') AS comment,
+	t4.partattrs AS partition_attrs,
+	t4.partstrat AS partition_strategy,
+	pg_get_expr(t4.partexprs, t4.partrelid) AS partition_exprs,
+	'{}' AS attrs
+FROM
+	INFORMATION_SCHEMA.TABLES AS t1
+	JOIN pg_catalog.pg_namespace AS t2 ON t2.nspname = t1.table_schema
+	JOIN pg_catalog.pg_class AS t3 ON t3.relnamespace = t2.oid AND t3.relname = t1.table_name
+	LEFT JOIN pg_catalog.pg_partitioned_table AS t4 ON t4.partrelid = t3.oid
+	LEFT JOIN pg_depend AS t5 ON t5.classid = 'pg_catalog.pg_class'::regclass::oid AND t5.objid = t3.oid AND t5.deptype = 'e'
+WHERE
+	t1.table_type = 'BASE TABLE'
+	AND NOT COALESCE(t3.relispartition, false)
+	AND t1.table_schema IN (%s)
+	AND t5.objid IS NULL
+ORDER BY
+	t1.table_schema, t1.table_name
+`
+	// Query to list tables by their names.
+	// Note, 'attrs' are not supported in this version.
+	tablesQueryArgs = `
+SELECT
+	t3.oid,
+	t1.table_schema,
+	t1.table_name,
+	pg_catalog.obj_description(t3.oid, 'pg_class') AS comment,
+	t4.partattrs AS partition_attrs,
+	t4.partstrat AS partition_strategy,
+	pg_get_expr(t4.partexprs, t4.partrelid) AS partition_exprs,
+	'{}' AS attrs
+FROM
+	INFORMATION_SCHEMA.TABLES AS t1
+	JOIN pg_catalog.pg_namespace AS t2 ON t2.nspname = t1.table_schema
+	JOIN pg_catalog.pg_class AS t3 ON t3.relnamespace = t2.oid AND t3.relname = t1.table_name
+	LEFT JOIN pg_catalog.pg_partitioned_table AS t4 ON t4.partrelid = t3.oid
+	LEFT JOIN pg_depend AS t5 ON t5.classid = 'pg_catalog.pg_class'::regclass::oid AND t5.objid = t3.oid AND t5.deptype = 'e'
+WHERE
+	t1.table_type = 'BASE TABLE'
+	AND NOT COALESCE(t3.relispartition, false)
+	AND t1.table_schema IN (%s)
+	AND t1.table_name IN (%s)
+	AND t5.objid IS NULL
+ORDER BY
+	t1.table_schema, t1.table_name
+`
+)
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 996549c1..8a63d3c4 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -72,7 +72,7 @@ func (i *inspect) InspectRealm(ctx context.Context, opts *schema.InspectRealmOpt
 			}
 		}
 		if mode.Is(schema.InspectObjects) {
-			if err := i.inspectSequences(ctx, r, nil); err != nil {
+			if err := i.inspectObjects(ctx, r, nil); err != nil {
 				return nil, err
 			}
 			if err := i.inspectRealmObjects(ctx, r, nil); err != nil {
@@ -168,7 +168,7 @@ func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.I
 		}
 	}
 	if mode.Is(schema.InspectObjects) {
-		if err := i.inspectSequences(ctx, r, opts); err != nil {
+		if err := i.inspectObjects(ctx, r, opts); err != nil {
 			return nil, err
 		}
 	}
@@ -232,10 +232,10 @@ func (i *inspect) tables(ctx context.Context, realm *schema.Realm, opts *schema.
 	defer rows.Close()
 	for rows.Next() {
 		var (
-			oid                                                     sql.NullInt64
-			tSchema, name, comment, partattrs, partstart, partexprs sql.NullString
+			oid                                                            sql.NullInt64
+			tSchema, name, comment, partattrs, partstart, partexprs, extra sql.NullString
 		)
-		if err := rows.Scan(&oid, &tSchema, &name, &comment, &partattrs, &partstart, &partexprs); err != nil {
+		if err := rows.Scan(&oid, &tSchema, &name, &comment, &partattrs, &partstart, &partexprs, &extra); err != nil {
 			return fmt.Errorf("scan table information: %w", err)
 		}
 		if !sqlx.ValidString(tSchema) || !sqlx.ValidString(name) {
@@ -245,7 +245,7 @@ func (i *inspect) tables(ctx context.Context, realm *schema.Realm, opts *schema.
 		if !ok {
 			return fmt.Errorf("schema %q was not found in realm", tSchema.String)
 		}
-		t := schema.NewTable(name.String)
+		t := i.newTable(ctx, name.String, extra.String)
 		s.AddTables(t)
 		if oid.Valid {
 			t.AddAttrs(&OID{V: oid.Int64})
@@ -1517,54 +1517,6 @@ WHERE
 ORDER BY
     nspname`
 
-	// Query to list table information.
-	tablesQuery = `
-SELECT
-	t3.oid,
-	t1.table_schema,
-	t1.table_name,
-	pg_catalog.obj_description(t3.oid, 'pg_class') AS comment,
-	t4.partattrs AS partition_attrs,
-	t4.partstrat AS partition_strategy,
-	pg_get_expr(t4.partexprs, t4.partrelid) AS partition_exprs
-FROM
-	INFORMATION_SCHEMA.TABLES AS t1
-	JOIN pg_catalog.pg_namespace AS t2 ON t2.nspname = t1.table_schema
-	JOIN pg_catalog.pg_class AS t3 ON t3.relnamespace = t2.oid AND t3.relname = t1.table_name
-	LEFT JOIN pg_catalog.pg_partitioned_table AS t4 ON t4.partrelid = t3.oid
-	LEFT JOIN pg_depend AS t5 ON t5.classid = 'pg_catalog.pg_class'::regclass::oid AND t5.objid = t3.oid AND t5.deptype = 'e'
-WHERE
-	t1.table_type = 'BASE TABLE'
-	AND NOT COALESCE(t3.relispartition, false)
-	AND t1.table_schema IN (%s)
-	AND t5.objid IS NULL
-ORDER BY
-	t1.table_schema, t1.table_name
-`
-	tablesQueryArgs = `
-SELECT
-	t3.oid,
-	t1.table_schema,
-	t1.table_name,
-	pg_catalog.obj_description(t3.oid, 'pg_class') AS comment,
-	t4.partattrs AS partition_attrs,
-	t4.partstrat AS partition_strategy,
-	pg_get_expr(t4.partexprs, t4.partrelid) AS partition_exprs
-FROM
-	INFORMATION_SCHEMA.TABLES AS t1
-	JOIN pg_catalog.pg_namespace AS t2 ON t2.nspname = t1.table_schema
-	JOIN pg_catalog.pg_class AS t3 ON t3.relnamespace = t2.oid AND t3.relname = t1.table_name
-	LEFT JOIN pg_catalog.pg_partitioned_table AS t4 ON t4.partrelid = t3.oid
-	LEFT JOIN pg_depend AS t5 ON t5.classid = 'pg_catalog.pg_class'::regclass::oid AND t5.objid = t3.oid AND t5.deptype = 'e'
-WHERE
-	t1.table_type = 'BASE TABLE'
-	AND NOT COALESCE(t3.relispartition, false)
-	AND t1.table_schema IN (%s)
-	AND t1.table_name IN (%s)
-	AND t5.objid IS NULL
-ORDER BY
-	t1.table_schema, t1.table_name
-`
 	// Query to list table columns.
 	columnsQuery = `
 SELECT
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 3858ea22..ee994085 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -287,6 +287,7 @@ func (s *state) addTable(add *schema.AddTable) error {
 		}
 	}
 	s.addComments(add, add.T)
+	s.addTableAttrs(add)
 	return nil
 }
 
@@ -339,11 +340,23 @@ func (s *state) modifyTable(modify *schema.ModifyTable) error {
 	)
 	for _, change := range skipAutoChanges(modify.Changes) {
 		switch change := change.(type) {
-		case *schema.AddAttr, *schema.ModifyAttr:
+		case *schema.ModifyAttr:
+			if _, ok := change.From.(*schema.Comment); !ok {
+				alter = append(alter, change)
+				continue
+			}
 			from, to, err := commentChange(change)
 			if err != nil {
 				return err
 			}
+			// Comments are not part of the ALTER command.
+			changes = append(changes, s.tableComment(modify, modify.T, to, from))
+		case *schema.AddAttr:
+			from, to, err := commentChange(change)
+			if err != nil {
+				return err
+			}
+			// Comments are not part of the ALTER command.
 			changes = append(changes, s.tableComment(modify, modify.T, to, from))
 		case *schema.DropAttr:
 			return fmt.Errorf("unsupported change type: %T", change)
@@ -613,6 +626,12 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 					From: change.To,
 					To:   change.From,
 				})
+			case *schema.ModifyAttr:
+				s.alterTableAttr(b, change)
+				reverse = append(reverse, &schema.ModifyAttr{
+					From: change.To,
+					To:   change.From,
+				})
 			}
 			return nil
 		})
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 3a460fe5..32430c7d 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -34,6 +34,7 @@ type (
 		Procs         []*sqlspec.Func     `spec:"procedure"`
 		Aggregates    []*aggregate        `spec:"aggregate"`
 		Triggers      []*sqlspec.Trigger  `spec:"trigger"`
+		Policies      []*policy           `spec:"policy"`
 		EventTriggers []*eventTrigger     `spec:"event_trigger"`
 		Extensions    []*extension        `spec:"extension"`
 		Schemas       []*sqlspec.Schema   `spec:"schema"`
@@ -105,6 +106,15 @@ type (
 		// are appended after the function arguments.
 		schemahcl.DefaultExtension
 	}
+
+	// Policy defines row-level security policy for a table.
+	// See: https://www.postgresql.org/docs/current/view-pg-policies.html.
+	policy struct {
+		Name string         `spec:",name"`
+		On   *schemahcl.Ref `spec:"on"`
+		// Rest of the attributes are appended after the table reference.
+		schemahcl.DefaultExtension
+	}
 )
 
 // merge merges the doc d1 into d.
@@ -121,6 +131,7 @@ func (d *doc) merge(d1 *doc) {
 	d.Sequences = append(d.Sequences, d1.Sequences...)
 	d.Extensions = append(d.Extensions, d1.Extensions...)
 	d.Triggers = append(d.Triggers, d1.Triggers...)
+	d.Policies = append(d.Policies, d1.Policies...)
 	d.EventTriggers = append(d.EventTriggers, d1.EventTriggers...)
 	d.Materialized = append(d.Materialized, d1.Materialized...)
 }
@@ -188,6 +199,7 @@ func (a *aggregate) SchemaRef() *schemahcl.Ref { return a.Schema }
 func init() {
 	schemahcl.Register("enum", &enum{})
 	schemahcl.Register("domain", &domain{})
+	schemahcl.Register("policy", &policy{})
 	schemahcl.Register("composite", &composite{})
 	schemahcl.Register("aggregate", &aggregate{})
 	schemahcl.Register("extension", &extension{})
@@ -214,6 +226,9 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		if err := convertSequences(d.Tables, d.Sequences, v); err != nil {
 			return err
 		}
+		if err := convertPolicies(d.Tables, d.Policies, v); err != nil {
+			return err
+		}
 		if err := convertExtensions(d.Extensions, v); err != nil {
 			return err
 		}
@@ -244,6 +259,9 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		if err := convertSequences(d.Tables, d.Sequences, r); err != nil {
 			return err
 		}
+		if err := convertPolicies(d.Tables, d.Policies, r); err != nil {
+			return err
+		}
 		// Extensions are skipped in schema scope.
 		if err := normalizeRealm(r); err != nil {
 			return err
@@ -376,6 +394,9 @@ func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, er
 	if err := convertPartition(spec.Extra, t); err != nil {
 		return nil, err
 	}
+	if err := convertTableAttrs(spec, t); err != nil {
+		return nil, err
+	}
 	return t, nil
 }
 
@@ -782,6 +803,7 @@ func tableSpec(t *schema.Table) (*sqlspec.Table, error) {
 	if p := (Partition{}); sqlx.Has(t.Attrs, &p) {
 		spec.Extra.Children = append(spec.Extra.Children, fromPartition(p))
 	}
+	tableAttrsSpec(t, spec)
 	return spec, nil
 }
 


Test patch:
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index d15a60b4..7b192145 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -378,11 +378,11 @@ func TestDriver_InspectPartitionedTable(t *testing.T) {
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1"))).
 		WithArgs("public").
 		WillReturnRows(sqltest.Rows(`
- oid   | table_schema | table_name  | comment | partition_attrs | partition_strategy |                  partition_exprs
--------+--------------+-------------+---------+-----------------+--------------------+----------------------------------------------------
- 112  | public       | logs1       |         |                 |                     | 
- 113  | public       | logs2       |         | 1               | r                   | 
- 114  | public       | logs3       |         | 2 0 0           | l                   | (a + b), (a + (b * 2))
+ oid   | table_schema | table_name  | comment | partition_attrs | partition_strategy |                  partition_exprs                   |                  extra                   
+-------+--------------+-------------+---------+-----------------+--------------------+----------------------------------------------------+----------------------------------------------------
+ 112  | public       | logs1       |         |                 |                     |                                                    |                                                    
+ 113  | public       | logs2       |         | 1               | r                   |                                                    |                                                    
+ 114  | public       | logs3       |         | 2 0 0           | l                   | (a + b), (a + (b * 2))                             |                              
 
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "$2, $3, $4"))).
@@ -732,9 +732,9 @@ func (m mock) version(version string) {
 }
 
 func (m mock) tableExists(schema, table string, exists bool) {
-	rows := sqlmock.NewRows([]string{"oid", "table_schema", "table_name", "table_comment", "partition_attrs", "partition_strategy", "partition_exprs"})
+	rows := sqlmock.NewRows([]string{"oid", "table_schema", "table_name", "table_comment", "partition_attrs", "partition_strategy", "partition_exprs", "attrs"})
 	if exists {
-		rows.AddRow(nil, schema, table, nil, nil, nil, nil)
+		rows.AddRow(nil, schema, table, nil, nil, nil, nil, nil)
 	}
 	m.ExpectQuery(queryTables).
 		WithArgs(schema).