Repository: ariga/atlas
Pull request #766 (2022-05-03 18:37:28)

sql/postgres: support inspecting partitioned tables


Patch:
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index 67e12e87..2b3824c6 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -386,35 +386,45 @@ func DefaultValue(c *schema.Column) (string, bool) {
 // Used by the different drivers to turn strings valid expressions.
 func MayWrap(s string) string {
 	n := len(s) - 1
-	if len(s) < 2 || s[0] != '(' || s[n] != ')' || !balanced(s[1:n]) {
+	if len(s) < 2 || s[0] != '(' || s[n] != ')' || !balanced(s) {
 		return "(" + s + ")"
 	}
 	return s
 }
 
-func balanced(s string) bool {
-	var d int
-	for i := 0; i < len(s) && d >= 0; i++ {
+func balanced(expr string) bool {
+	return ExprLastIndex(expr) == len(expr)-1
+}
+
+// ExprLastIndex scans the expression string (wrapped with parens)
+// and returns its last. e.g. "(a+1), c int ...".
+func ExprLastIndex(expr string) int {
+	var r, l int
+	for i := 0; i < len(expr); i++ {
 	Top:
-		switch s[i] {
+		switch expr[i] {
 		case '(':
-			d++
+			r++
 		case ')':
-			d--
+			l++
 		// String or identifier.
 		case '\'', '"', '`':
-			for j := i + 1; j < len(s); j++ {
-				switch s[j] {
+			for j := i + 1; j < len(expr); j++ {
+				switch expr[j] {
 				case '\\':
 					j++
-				case s[i]:
+				case expr[i]:
 					i = j
 					break Top
 				}
 			}
 			// Unexpected EOS.
-			return false
+			return -1
+		}
+		// Balanced parens.
+		if r == l {
+			return i
 		}
 	}
-	return d == 0
+	return -1
 }
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 2e175cc5..be523d5b 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -237,8 +237,15 @@ const (
 	IndexTypeGiST  = "GIST"
 )
 
-// List of "GENERATED" TYPES.
+// List of "GENERATED" types.
 const (
 	GeneratedTypeAlways    = "ALWAYS"
 	GeneratedTypeByDefault = "BY_DEFAULT" // BY DEFAULT.
 )
+
+// List of PARTITION KEY types.
+const (
+	PartitionTypeRange = "RANGE"
+	PartitionTypeList  = "LIST"
+	PartitionTypeHash  = "HASH"
+)
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 19460ff8..912190f8 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -76,6 +76,9 @@ func (i *inspect) inspectTables(ctx context.Context, r *schema.Realm, opts *sche
 		if err := i.indexes(ctx, s); err != nil {
 			return err
 		}
+		if err := i.partitions(s); err != nil {
+			return err
+		}
 		if err := i.fks(ctx, s); err != nil {
 			return err
 		}
@@ -107,8 +110,8 @@ func (i *inspect) tables(ctx context.Context, realm *schema.Realm, opts *schema.
 	}
 	defer rows.Close()
 	for rows.Next() {
-		var tSchema, name, comment sql.NullString
-		if err := rows.Scan(&tSchema, &name, &comment); err != nil {
+		var tSchema, name, comment, partattrs, partstart, partexprs sql.NullString
+		if err := rows.Scan(&tSchema, &name, &comment, &partattrs, &partstart, &partexprs); err != nil {
 			return fmt.Errorf("scan table information: %w", err)
 		}
 		if !sqlx.ValidString(tSchema) || !sqlx.ValidString(name) {
@@ -121,8 +124,13 @@ func (i *inspect) tables(ctx context.Context, realm *schema.Realm, opts *schema.
 		t := &schema.Table{Name: name.String}
 		s.AddTables(t)
 		if sqlx.ValidString(comment) {
-			t.Attrs = append(t.Attrs, &schema.Comment{
-				Text: comment.String,
+			t.SetComment(comment.String)
+		}
+		if sqlx.ValidString(partattrs) {
+			t.AddAttrs(&Partition{
+				start: partstart.String,
+				attrs: partattrs.String,
+				exprs: partexprs.String,
 			})
 		}
 	}
@@ -405,6 +413,56 @@ func (i *inspect) addIndexes(s *schema.Schema, rows *sql.Rows) error {
 	return nil
 }
 
+// partitions builds the partition each table in the schema.
+func (i *inspect) partitions(s *schema.Schema) error {
+	for _, t := range s.Tables {
+		var d Partition
+		if !sqlx.Has(t.Attrs, &d) {
+			continue
+		}
+		switch s := strings.ToLower(d.start); s {
+		case "r":
+			d.T = PartitionTypeRange
+		case "l":
+			d.T = PartitionTypeList
+		case "h":
+			d.T = PartitionTypeHash
+		default:
+			return fmt.Errorf("postgres: unexpected partition strategy %q", s)
+		}
+		idxs := strings.Split(strings.TrimSpace(d.attrs), " ")
+		if len(idxs) == 0 {
+			return fmt.Errorf("postgres: no columns/expressions were found in partition key for column %q", t.Name)
+		}
+		for i := range idxs {
+			switch idx, err := strconv.Atoi(idxs[i]); {
+			case err != nil:
+				return fmt.Errorf("postgres: faild parsing partition key index %q", idxs[i])
+			// An expression.
+			case idx == 0:
+				j := sqlx.ExprLastIndex(d.exprs)
+				if j == -1 {
+					return fmt.Errorf("postgres: no expression found in partition key: %q", d.exprs)
+				}
+				d.Parts = append(d.Parts, &PartitionPart{
+					X: &schema.RawExpr{X: d.exprs[:j+1]},
+				})
+				d.exprs = strings.TrimPrefix(d.exprs[j+1:], ", ")
+			// A column at index idx-1.
+			default:
+				if idx > len(t.Columns) {
+					return fmt.Errorf("postgres: unexpected column index %d", idx)
+				}
+				d.Parts = append(d.Parts, &PartitionPart{
+					C: t.Columns[idx-1],
+				})
+			}
+		}
+		schema.ReplaceOrAppend(&t.Attrs, &d)
+	}
+	return nil
+}
+
 // fks queries and appends the foreign keys of the given table.
 func (i *inspect) fks(ctx context.Context, s *schema.Schema) error {
 	rows, err := i.querySchema(ctx, fksQuery, s)
@@ -691,6 +749,28 @@ type (
 		schema.Attr
 		Columns []string
 	}
+
+	// Partition defines the spec of a partitioned table.
+	Partition struct {
+		schema.Attr
+		// T defines the type/strategy of the partition.
+		// Can be one of: RANGE, LIST, HASH.
+		T string
+		// Partition parts. The additional attributes
+		// on each part can be used to control collation.
+		Parts []*PartitionPart
+
+		// Internal info returned from pg_partitioned_table.
+		start, attrs, exprs string
+	}
+
+	// An PartitionPart represents an index part that
+	// can be either an expression or a column.
+	PartitionPart struct {
+		X     schema.Expr
+		C     *schema.Column
+		Attrs []schema.Attr
+	}
 )
 
 const (
@@ -708,10 +788,14 @@ const (
 SELECT
 	t1.table_schema,
 	t1.table_name,
-	pg_catalog.obj_description(t2.oid, 'pg_class') AS COMMENT
+	pg_catalog.obj_description(t2.oid, 'pg_class') AS comment,
+	t3.partattrs AS partition_attrs,
+	t3.partstrat AS partition_strategy,
+	pg_get_expr(t3.partexprs, t3.partrelid) AS partition_exprs
 FROM
 	INFORMATION_SCHEMA.TABLES AS t1
 	JOIN pg_catalog.pg_class AS t2 ON t2.oid = to_regclass(quote_ident(t1.table_schema) || '.' || quote_ident(t1.table_name))::oid
+	LEFT JOIN pg_catalog.pg_partitioned_table AS t3 ON t3.partrelid = t2.oid
 WHERE
 	t1.table_type = 'BASE TABLE'
 	AND t1.table_schema IN (%s)
@@ -722,10 +806,14 @@ ORDER BY
 SELECT
 	t1.table_schema,
 	t1.table_name,
-	pg_catalog.obj_description(t2.oid, 'pg_class') AS COMMENT
+	pg_catalog.obj_description(t2.oid, 'pg_class') AS comment,
+	t3.partattrs AS partition_attrs,
+	t3.partstrat AS partition_strategy,
+	pg_get_expr(t3.partexprs, t3.partrelid) AS partition_exprs
 FROM
 	INFORMATION_SCHEMA.TABLES AS t1
 	JOIN pg_catalog.pg_class AS t2 ON t2.oid = to_regclass(quote_ident(t1.table_schema) || '.' || quote_ident(t1.table_name))::oid
+	LEFT JOIN pg_catalog.pg_partitioned_table AS t3 ON t3.partrelid = t2.oid
 WHERE
 	t1.table_type = 'BASE TABLE'
 	AND t1.table_schema IN (%s)


Test patch:
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index b6769424..371b106b 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -300,6 +300,74 @@ users        | users_check1       | (((c2 + c1) + c3) > 10) | c3          | {2,1
 	}
 }
 
+func TestDriver_InspectPartitionedTable(t *testing.T) {
+	db, m, err := sqlmock.New()
+	require.NoError(t, err)
+	mk := mock{m}
+	mk.version("130000")
+	drv, err := Open(db)
+	require.NoError(t, err)
+	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= CURRENT_SCHEMA()"))).
+		WillReturnRows(sqltest.Rows(`
+   schema_name
+--------------------
+public
+`))
+	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))
+
+`))
+	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(columnsQuery, "$2, $3, $4"))).
+		WithArgs("public", "logs1", "logs2", "logs3").
+		WillReturnRows(sqltest.Rows(`
+table_name |column_name | data_type | is_nullable | column_default | character_maximum_length | numeric_precision | datetime_precision | numeric_scale | character_set_name | collation_name | udt_name | is_identity | identity_start | identity_increment | identity_generation | generation_expression | comment | typtype | oid
+-----------+------------+-----------+-------------+----------------+--------------------------+-------------------+--------------------+---------------+--------------------+----------------+----------+-------------+----------------+--------------------+---------------------+-----------------------+---------+---------+-----
+logs1      | c1         | integer   | NO          |                |                          |                32 |                    |             0 |                    |                | int4     | NO          |                |                    |                     |                       |         | b       |  23
+logs2      | c2         | integer   | NO          |                |                          |                32 |                    |             0 |                    |                | int4     | NO          |                |                    |                     |                       |         | b       |  23
+logs2      | c3         | integer   | NO          |                |                          |                32 |                    |             0 |                    |                | int4     | NO          |                |                    |                     |                       |         | b       |  23
+logs3      | c4         | integer   | NO          |                |                          |                32 |                    |             0 |                    |                | int4     | NO          |                |                    |                     |                       |         | b       |  23
+logs3      | c5         | integer   | NO          |                |                          |                32 |                    |             0 |                    |                | int4     | NO          |                |                    |                     |                       |         | b       |  23
+`))
+	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(indexesQuery, "$2, $3, $4"))).
+		WillReturnRows(sqlmock.NewRows([]string{"table_name", "index_name", "column_name", "primary", "unique", "constraint_type", "predicate", "expression"}))
+	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(fksQuery, "$2, $3, $4"))).
+		WillReturnRows(sqlmock.NewRows([]string{"constraint_name", "table_name", "column_name", "referenced_table_name", "referenced_column_name", "referenced_table_schema", "update_rule", "delete_rule"}))
+	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(checksQuery, "$2, $3, $4"))).
+		WillReturnRows(sqlmock.NewRows([]string{"table_name", "constraint_name", "expression", "column_name", "column_indexes"}))
+	s, err := drv.InspectSchema(context.Background(), "", &schema.InspectOptions{})
+	require.NoError(t, err)
+
+	t1, ok := s.Table("logs1")
+	require.True(t, ok)
+	require.Empty(t, t1.Attrs)
+
+	t2, ok := s.Table("logs2")
+	require.True(t, ok)
+	require.Len(t, t2.Attrs, 1)
+	key := t2.Attrs[0].(*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"}}}},
+	}, key.Parts)
+
+	t3, ok := s.Table("logs3")
+	require.True(t, ok)
+	require.Len(t, t3.Attrs, 1)
+	key = t3.Attrs[0].(*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"}}}},
+		{X: &schema.RawExpr{X: "(a + b)"}},
+		{X: &schema.RawExpr{X: "(a + (b * 2))"}},
+	}, key.Parts)
+}
+
 func TestDriver_InspectSchema(t *testing.T) {
 	db, m, err := sqlmock.New()
 	require.NoError(t, err)
@@ -315,7 +383,7 @@ test
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1"))).
 		WithArgs("test").
-		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment"}))
+		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment", "partition_attrs", "partition_strategy", "partition_exprs"}))
 	s, err := drv.InspectSchema(context.Background(), "", &schema.InspectOptions{})
 	require.NoError(t, err)
 	require.EqualValues(t, func() *schema.Schema {
@@ -356,7 +424,7 @@ public
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1, $2"))).
 		WithArgs("test", "public").
-		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment"}))
+		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment", "partition_attrs", "partition_strategy", "partition_exprs"}))
 	realm, err := drv.InspectRealm(context.Background(), &schema.InspectRealmOption{})
 	require.NoError(t, err)
 	require.EqualValues(t, func() *schema.Realm {
@@ -394,7 +462,7 @@ public
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1, $2"))).
 		WithArgs("test", "public").
-		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment"}))
+		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment", "partition_attrs", "partition_strategy", "partition_exprs"}))
 	realm, err = drv.InspectRealm(context.Background(), &schema.InspectRealmOption{Schemas: []string{"test", "public"}})
 	require.NoError(t, err)
 	require.EqualValues(t, func() *schema.Realm {
@@ -431,7 +499,7 @@ public
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1"))).
 		WithArgs("test").
-		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment"}))
+		WillReturnRows(sqlmock.NewRows([]string{"table_schema", "table_name", "comment", "partition_attrs", "partition_strategy", "partition_exprs"}))
 	realm, err = drv.InspectRealm(context.Background(), &schema.InspectRealmOption{Schemas: []string{"test"}})
 	require.NoError(t, err)
 	require.EqualValues(t, func() *schema.Realm {
@@ -513,9 +581,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"})
+	rows := sqlmock.NewRows([]string{"table_schema", "table_name", "table_comment", "partition_attrs", "partition_strategy", "partition_exprs"})
 	if exists {
-		rows.AddRow(schema, table, nil)
+		rows.AddRow(schema, table, nil, nil, nil, nil)
 	}
 	m.ExpectQuery(queryTables).
 		WithArgs(schema).