Repository: ariga/atlas
Pull request #770 (2022-05-04 12:11:30)

sql/postgres: support marashaling partitioned tables


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 75e8d8f3..3ef08f66 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -346,7 +346,7 @@ func FromTable(t *schema.Table, colFn ColumnSpecFunc, pkFn PrimaryKeySpecFunc, i
 func FromPrimaryKey(s *schema.Index) (*sqlspec.PrimaryKey, error) {
 	c := make([]*schemaspec.Ref, 0, len(s.Parts))
 	for _, v := range s.Parts {
-		c = append(c, colRef(v.C.Name))
+		c = append(c, ColumnRef(v.C.Name))
 	}
 	return &sqlspec.PrimaryKey{
 		Columns: c,
@@ -469,7 +469,7 @@ func FromIndex(idx *schema.Index, partFns ...func(*schema.IndexPart, *sqlspec.In
 		case p.C != nil && p.X != nil:
 			return nil, fmt.Errorf("multiple key part definitions for index %q", idx.Name)
 		case p.C != nil:
-			part.Column = colRef(p.C.Name)
+			part.Column = ColumnRef(p.C.Name)
 		case p.X != nil:
 			x, ok := p.X.(*schema.RawExpr)
 			if !ok {
@@ -491,7 +491,7 @@ func columnsOnly(idx *schema.Index) ([]*schemaspec.Ref, bool) {
 		if p.C == nil || p.Desc {
 			return nil, false
 		}
-		parts[i] = colRef(p.C.Name)
+		parts[i] = ColumnRef(p.C.Name)
 	}
 	return parts, true
 }
@@ -500,11 +500,11 @@ func columnsOnly(idx *schema.Index) ([]*schemaspec.Ref, bool) {
 func FromForeignKey(s *schema.ForeignKey) (*sqlspec.ForeignKey, error) {
 	c := make([]*schemaspec.Ref, 0, len(s.Columns))
 	for _, v := range s.Columns {
-		c = append(c, colRef(v.Name))
+		c = append(c, ColumnRef(v.Name))
 	}
 	r := make([]*schemaspec.Ref, 0, len(s.RefColumns))
 	for _, v := range s.RefColumns {
-		ref := colRef(v.Name)
+		ref := ColumnRef(v.Name)
 		if s.Table != s.RefTable {
 			ref = externalColRef(v.Name, s.RefTable.Name)
 		}
@@ -622,7 +622,8 @@ func isLocalRef(r *schemaspec.Ref) bool {
 	return strings.HasPrefix(r.V, "$column")
 }
 
-func colRef(cName string) *schemaspec.Ref {
+// ColumnRef returns the reference of a column by its name.
+func ColumnRef(cName string) *schemaspec.Ref {
 	return &schemaspec.Ref{V: "$column." + cName}
 }
 
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index 292ecbd5..ea955488 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -63,6 +63,14 @@ func VarAttr(k, v string) *schemaspec.Attr {
 	}
 }
 
+// RefAttr is a helper method for constructing *schemaspec.Attr instances that contain a reference.
+func RefAttr(k string, r *schemaspec.Ref) *schemaspec.Attr {
+	return &schemaspec.Attr{
+		K: k,
+		V: r,
+	}
+}
+
 // ListAttr is a helper method for constructing *schemaspec.Attr instances that contain list values.
 func ListAttr(k string, litValues ...string) *schemaspec.Attr {
 	lv := &schemaspec.ListValue{}
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 727e2f45..e7e574b6 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -194,6 +194,41 @@ func convertPartition(spec schemaspec.Resource, table *schema.Table) error {
 	return nil
 }
 
+// fromPartition returns the resource spec for representing the partition block.
+func fromPartition(p Partition) *schemaspec.Resource {
+	key := &schemaspec.Resource{
+		Type: "partition",
+		Attrs: []*schemaspec.Attr{
+			specutil.VarAttr("type", strings.ToUpper(specutil.Var(p.T))),
+		},
+	}
+	columns, ok := func() (*schemaspec.ListValue, bool) {
+		parts := make([]schemaspec.Value, 0, len(p.Parts))
+		for _, p := range p.Parts {
+			if p.C == nil {
+				return nil, false
+			}
+			parts = append(parts, specutil.ColumnRef(p.C.Name))
+		}
+		return &schemaspec.ListValue{V: parts}, true
+	}()
+	if ok {
+		key.Attrs = append(key.Attrs, &schemaspec.Attr{K: "columns", V: columns})
+		return key
+	}
+	for _, p := range p.Parts {
+		part := &schemaspec.Resource{Type: "by"}
+		switch {
+		case p.C != nil:
+			part.Attrs = append(part.Attrs, specutil.RefAttr("column", specutil.ColumnRef(p.C.Name)))
+		case p.X != nil:
+			part.Attrs = append(part.Attrs, specutil.StrAttr("expr", p.X.(*schema.RawExpr).X))
+		}
+		key.Children = append(key.Children, part)
+	}
+	return key
+}
+
 // convertColumn converts a sqlspec.Column into a schema.Column.
 func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error) {
 	if err := fixDefaultQuotes(spec.Default); err != nil {
@@ -380,15 +415,22 @@ func schemaSpec(schem *schema.Schema) (*doc, error) {
 }
 
 // tableSpec converts from a concrete Postgres sqlspec.Table to a schema.Table.
-func tableSpec(tab *schema.Table) (*sqlspec.Table, error) {
-	return specutil.FromTable(
-		tab,
+func tableSpec(table *schema.Table) (*sqlspec.Table, error) {
+	spec, err := specutil.FromTable(
+		table,
 		columnSpec,
 		specutil.FromPrimaryKey,
 		indexSpec,
 		specutil.FromForeignKey,
 		specutil.FromCheck,
 	)
+	if err != nil {
+		return nil, err
+	}
+	if p := (Partition{}); sqlx.Has(table.Attrs, &p) {
+		spec.Extra.Children = append(spec.Extra.Children, fromPartition(p))
+	}
+	return spec, nil
 }
 
 func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 60130fa4..56a1fbbb 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -384,6 +384,57 @@ table "logs" {
 	})
 }
 
+func TestMarshalSpec_Partitioned(t *testing.T) {
+	t.Run("Columns", func(t *testing.T) {
+		c := schema.NewStringColumn("name", "text")
+		s := schema.New("test").
+			AddTables(schema.NewTable("logs").AddColumns(c).AddAttrs(&Partition{T: PartitionTypeHash, Parts: []*PartitionPart{{C: c}}}))
+		buf, err := MarshalHCL(s)
+		require.NoError(t, err)
+		require.Equal(t, `table "logs" {
+  schema = schema.test
+  column "name" {
+    null = false
+    type = text
+  }
+  partition {
+    type    = HASH
+    columns = [column.name]
+  }
+}
+schema "test" {
+}
+`, string(buf))
+	})
+
+	t.Run("Parts", func(t *testing.T) {
+		c := schema.NewStringColumn("name", "text")
+		s := schema.New("test").
+			AddTables(schema.NewTable("logs").AddColumns(c).AddAttrs(&Partition{T: PartitionTypeHash, Parts: []*PartitionPart{{C: c}, {X: &schema.RawExpr{X: "lower(name)"}}}}))
+		buf, err := MarshalHCL(s)
+		require.NoError(t, err)
+		require.Equal(t, `table "logs" {
+  schema = schema.test
+  column "name" {
+    null = false
+    type = text
+  }
+  partition {
+    type = HASH
+    by {
+      column = column.name
+    }
+    by {
+      expr = "lower(name)"
+    }
+  }
+}
+schema "test" {
+}
+`, string(buf))
+	})
+}
+
 func TestMarshalSpec_IndexPredicate(t *testing.T) {
 	s := &schema.Schema{
 		Name: "test",