Repository: ariga/atlas
Pull request #2323 (2023-11-28 09:27:24)

sql/postgres: add object ids to inspection


Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index f147ea35..99b2427e 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -159,8 +159,11 @@ func (i *inspect) tables(ctx context.Context, realm *schema.Realm, opts *schema.
 	}
 	defer rows.Close()
 	for rows.Next() {
-		var tSchema, name, comment, partattrs, partstart, partexprs sql.NullString
-		if err := rows.Scan(&tSchema, &name, &comment, &partattrs, &partstart, &partexprs); err != nil {
+		var (
+			oid                                                     sql.NullInt64
+			tSchema, name, comment, partattrs, partstart, partexprs sql.NullString
+		)
+		if err := rows.Scan(&oid, &tSchema, &name, &comment, &partattrs, &partstart, &partexprs); err != nil {
 			return fmt.Errorf("scan table information: %w", err)
 		}
 		if !sqlx.ValidString(tSchema) || !sqlx.ValidString(name) {
@@ -172,6 +175,9 @@ func (i *inspect) tables(ctx context.Context, realm *schema.Realm, opts *schema.
 		}
 		t := schema.NewTable(name.String)
 		s.AddTables(t)
+		if oid.Valid {
+			t.AddAttrs(&OID{V: oid.Int64})
+		}
 		if sqlx.ValidString(comment) {
 			t.SetComment(comment.String)
 		}
@@ -776,6 +782,12 @@ type (
 		T string // e.g., void, any, cstring, etc.
 	}
 
+	// OID is the object identifier as defined in the Postgres catalog.
+	OID struct {
+		schema.Attr
+		V int64
+	}
+
 	// enumType represents an enum type. It serves aa intermediate representation of a Postgres enum type,
 	// to temporary save TypeID and TypeName of an enum column until the enum values can be extracted.
 	enumType struct {
@@ -1275,6 +1287,7 @@ ORDER BY
 	// 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,
@@ -1297,6 +1310,7 @@ ORDER BY
 `
 	tablesQueryArgs = `
 SELECT
+	t3.oid,
 	t1.table_schema,
 	t1.table_name,
 	pg_catalog.obj_description(t3.oid, 'pg_class') AS comment,
diff --git a/sql/schema/schema.go b/sql/schema/schema.go
index ad96a70d..0115da06 100644
--- a/sql/schema/schema.go
+++ b/sql/schema/schema.go
@@ -39,7 +39,8 @@ type (
 		Indexes     []*Index
 		PrimaryKey  *Index
 		ForeignKeys []*ForeignKey
-		Attrs       []Attr // Attrs, constraints and options.
+		Attrs       []Attr   // Attrs, constraints and options.
+		Deps        []Object // Objects this table depends on.
 	}
 
 	// A View represents a view definition.
@@ -49,8 +50,8 @@ type (
 		Schema  *Schema
 		Columns []*Column
 		Attrs   []Attr   // Attrs and options.
-		Deps    []Object // Tables and views used in view definition.
 		Indexes []*Index // Indexes on materialized view.
+		Deps    []Object // Objects this view depends on.
 	}
 
 	// A Column represents a column definition.
@@ -114,9 +115,10 @@ type (
 		Schema *Schema
 		Args   []*FuncArg
 		Ret    Type
-		Body   string // Function body only.
-		Lang   string // Language (e.g. SQL, PL/pgSQL, etc.).
-		Attrs  []Attr // Extra driver specific attributes.
+		Body   string   // Function body only.
+		Lang   string   // Language (e.g. SQL, PL/pgSQL, etc.).
+		Attrs  []Attr   // Extra driver specific attributes.
+		Deps   []Object // Objects this function depends on.
 	}
 
 	// Proc represents a procedure definition.
@@ -124,9 +126,10 @@ type (
 		Name   string
 		Schema *Schema
 		Args   []*FuncArg
-		Body   string // Function body only.
-		Lang   string // Language (e.g. SQL, PL/pgSQL, etc.).
-		Attrs  []Attr // Extra driver specific attributes.
+		Body   string   // Function body only.
+		Lang   string   // Language (e.g. SQL, PL/pgSQL, etc.).
+		Attrs  []Attr   // Extra driver specific attributes.
+		Deps   []Object // Objects this function depends on.
 	}
 
 	// A FuncArg represents a single function argument.
@@ -531,6 +534,8 @@ const (
 // objects.
 func (*Table) obj()    {}
 func (*View) obj()     {}
+func (*Func) obj()     {}
+func (*Proc) obj()     {}
 func (*EnumType) obj() {}
 
 // expressions.


Test patch:
diff --git a/internal/integration/postgres_test.go b/internal/integration/postgres_test.go
index 9f1ba4c5..0c329765 100644
--- a/internal/integration/postgres_test.go
+++ b/internal/integration/postgres_test.go
@@ -1134,8 +1134,11 @@ create table atlas_types_sanity
 		require.Len(t, realm.Schemas, 1)
 		ts, ok := realm.Schemas[0].Table(n)
 		require.True(t, ok)
+		require.Len(t, ts.Attrs, 1)
+		require.IsType(t, &postgres.OID{}, ts.Attrs[0])
 		expected := schema.Table{
 			Name:   n,
+			Attrs:  ts.Attrs,
 			Schema: realm.Schemas[0],
 			Columns: []*schema.Column{
 				{
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 00d039ff..e0ad3672 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -370,11 +370,11 @@ func TestDriver_InspectPartitionedTable(t *testing.T) {
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1"))).
 		WithArgs("public").
 		WillReturnRows(sqltest.Rows(`
- table_schema | table_name  | comment | partition_attrs | partition_strategy |                  partition_exprs                   
---------------+-------------+---------+-----------------+--------------------+----------------------------------------------------
- public       | logs1       |         |                 |                    | 
- public       | logs2       |         | 1               | r                  | 
- public       | logs3       |         | 2 0 0           | l                  | (a + b), (a + (b * 2))
+ 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))
 
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "$2, $3, $4"))).
@@ -402,12 +402,12 @@ logs3      | c5         | integer   | integer   | NO          |                |
 
 	t1, ok := s.Table("logs1")
 	require.True(t, ok)
-	require.Empty(t, t1.Attrs)
+	require.Equal(t, []schema.Attr{&OID{V: 112}}, t1.Attrs)
 
 	t2, ok := s.Table("logs2")
 	require.True(t, ok)
-	require.Len(t, t2.Attrs, 1)
-	key := t2.Attrs[0].(*Partition)
+	require.Len(t, t2.Attrs, 2)
+	key := t2.Attrs[1].(*Partition)
 	require.Equal(t, PartitionTypeRange, key.T)
 	require.Equal(t, []*PartitionPart{
 		{C: &schema.Column{Name: "c2", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}}},
@@ -415,8 +415,8 @@ logs3      | c5         | integer   | integer   | NO          |                |
 
 	t3, ok := s.Table("logs3")
 	require.True(t, ok)
-	require.Len(t, t3.Attrs, 1)
-	key = t3.Attrs[0].(*Partition)
+	require.Len(t, t3.Attrs, 2)
+	key = t3.Attrs[1].(*Partition)
 	require.Equal(t, PartitionTypeList, key.T)
 	require.Equal(t, []*PartitionPart{
 		{C: &schema.Column{Name: "c5", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}}},
@@ -694,9 +694,9 @@ func (m mock) version(version string) {
 }
 
 func (m mock) tableExists(schema, table string, exists bool) {
-	rows := sqlmock.NewRows([]string{"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"})
 	if exists {
-		rows.AddRow(schema, table, nil, nil, nil, nil)
+		rows.AddRow(nil, schema, table, nil, nil, nil, nil)
 	}
 	m.ExpectQuery(queryTables).
 		WithArgs(schema).