Repository: ariga/atlas
Pull request #769 (2022-05-04 10:42:19)

sql/postgres: support unmarshaling partitions


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 195f4b48..75e8d8f3 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -165,7 +165,7 @@ func Index(spec *sqlspec.Index, parent *schema.Table, partFns ...func(*sqlspec.I
 		return nil, fmt.Errorf(`multiple definitions for index %q, use "columns" or "on"`, spec.Name)
 	case n > 0:
 		for i, c := range spec.Columns {
-			c, err := column(parent, c)
+			c, err := ColumnByRef(parent, c)
 			if err != nil {
 				return nil, err
 			}
@@ -185,7 +185,7 @@ func Index(spec *sqlspec.Index, parent *schema.Table, partFns ...func(*sqlspec.I
 			case p.Expr != "":
 				part.X = &schema.RawExpr{X: p.Expr}
 			case p.Column != nil:
-				c, err := column(parent, p.Column)
+				c, err := ColumnByRef(parent, p.Column)
 				if err != nil {
 					return nil, err
 				}
@@ -223,7 +223,7 @@ func Check(spec *sqlspec.Check) (*schema.Check, error) {
 func PrimaryKey(spec *sqlspec.PrimaryKey, parent *schema.Table) (*schema.Index, error) {
 	parts := make([]*schema.IndexPart, 0, len(spec.Columns))
 	for seqno, c := range spec.Columns {
-		c, err := column(parent, c)
+		c, err := ColumnByRef(parent, c)
 		if err != nil {
 			return nil, nil
 		}
@@ -254,7 +254,7 @@ func linkForeignKeys(tbl *schema.Table, sch *schema.Schema, table *sqlspec.Table
 			return fmt.Errorf("sqlspec: number of referencing and referenced columns do not match for foreign-key %q", fk.Symbol)
 		}
 		for _, ref := range spec.Columns {
-			c, err := column(tbl, ref)
+			c, err := ColumnByRef(tbl, ref)
 			if err != nil {
 				return err
 			}
@@ -264,7 +264,7 @@ func linkForeignKeys(tbl *schema.Table, sch *schema.Schema, table *sqlspec.Table
 			t, c, err := externalRef(ref, sch)
 			if isLocalRef(ref) {
 				t = fk.Table
-				c, err = column(fk.Table, ref)
+				c, err = ColumnByRef(fk.Table, ref)
 			}
 			if err != nil {
 				return err
@@ -544,7 +544,8 @@ func SchemaName(ref *schemaspec.Ref) (string, error) {
 	return parts[1], nil
 }
 
-func column(t *schema.Table, ref *schemaspec.Ref) (*schema.Column, error) {
+// ColumnByRef returns a column from the table by its reference.
+func ColumnByRef(t *schema.Table, ref *schemaspec.Ref) (*schema.Column, error) {
 	s := strings.Split(ref.V, "$column.")
 	if len(s) != 2 {
 		return nil, fmt.Errorf("specutil: failed to extract column name from %q", ref)
@@ -561,7 +562,7 @@ func externalRef(ref *schemaspec.Ref, sch *schema.Schema) (*schema.Table, *schem
 	if err != nil {
 		return nil, nil, err
 	}
-	c, err := column(tbl, ref)
+	c, err := ColumnByRef(tbl, ref)
 	if err != nil {
 		return nil, nil, err
 	}
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 38f9beb0..727e2f45 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -106,6 +106,7 @@ var (
 	hclState = schemahcl.New(
 		schemahcl.WithTypes(TypeRegistry.Specs()),
 		schemahcl.WithScopedEnums("table.index.type", IndexTypeBTree, IndexTypeHash, IndexTypeGIN, IndexTypeGiST),
+		schemahcl.WithScopedEnums("table.partition.type", PartitionTypeRange, PartitionTypeList, PartitionTypeHash),
 		schemahcl.WithScopedEnums("table.column.identity.generated", GeneratedTypeAlways, GeneratedTypeByDefault),
 		schemahcl.WithScopedEnums("table.column.as.type", "STORED"),
 		schemahcl.WithScopedEnums("table.foreign_key.on_update", specutil.ReferenceVars...),
@@ -127,7 +128,70 @@ var (
 // ForeignKeySpecs into ForeignKeys, as the target tables do not necessarily exist in the schema
 // at this point. Instead, the linking is done by the convertSchema function.
 func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, error) {
-	return specutil.Table(spec, parent, convertColumn, specutil.PrimaryKey, convertIndex, specutil.Check)
+	t, err := specutil.Table(spec, parent, convertColumn, specutil.PrimaryKey, convertIndex, specutil.Check)
+	if err != nil {
+		return nil, err
+	}
+	if err := convertPartition(spec.Extra, t); err != nil {
+		return nil, err
+	}
+	return t, nil
+}
+
+// convertPartition converts and appends the partition block into the table attributes if exists.
+func convertPartition(spec schemaspec.Resource, table *schema.Table) error {
+	r, ok := spec.Resource("partition")
+	if !ok {
+		return nil
+	}
+	var p struct {
+		Type    string            `spec:"type"`
+		Columns []*schemaspec.Ref `spec:"columns"`
+		Parts   []*struct {
+			Expr   string          `spec:"expr"`
+			Column *schemaspec.Ref `spec:"column"`
+		} `spec:"by"`
+	}
+	if err := r.As(&p); err != nil {
+		return fmt.Errorf("parsing %s.partition: %w", table.Name, err)
+	}
+	if p.Type == "" {
+		return fmt.Errorf("missing attribute %s.partition.type", table.Name)
+	}
+	key := &Partition{T: p.Type}
+	switch n, m := len(p.Columns), len(p.Parts); {
+	case n == 0 && m == 0:
+		return fmt.Errorf("missing columns or expressions for %s.partition", table.Name)
+	case n > 0 && m > 0:
+		return fmt.Errorf(`multiple definitions for %s.partition, use "columns" or "by"`, table.Name)
+	case n > 0:
+		for _, r := range p.Columns {
+			c, err := specutil.ColumnByRef(table, r)
+			if err != nil {
+				return err
+			}
+			key.Parts = append(key.Parts, &PartitionPart{C: c})
+		}
+	case m > 0:
+		for i, p := range p.Parts {
+			switch {
+			case p.Column == nil && p.Expr == "":
+				return fmt.Errorf("missing column or expression for %s.partition.by at position %d", table.Name, i)
+			case p.Column != nil && p.Expr != "":
+				return fmt.Errorf("multiple definitions for  %s.partition.by at position %d", table.Name, i)
+			case p.Column != nil:
+				c, err := specutil.ColumnByRef(table, p.Column)
+				if err != nil {
+					return err
+				}
+				key.Parts = append(key.Parts, &PartitionPart{C: c})
+			case p.Expr != "":
+				key.Parts = append(key.Parts, &PartitionPart{X: &schema.RawExpr{X: p.Expr}})
+			}
+		}
+	}
+	table.AddAttrs(key)
+	return nil
 }
 
 // convertColumn converts a sqlspec.Column into a schema.Column.


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index cf3674f4..60130fa4 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -285,6 +285,105 @@ table "t" {
 	})
 }
 
+func TestUnmarshalSpec_Partitioned(t *testing.T) {
+	t.Run("Columns", func(t *testing.T) {
+		var (
+			s = &schema.Schema{}
+			f = `
+schema "test" {}
+table "logs" {
+	schema = schema.test
+	column "name" {
+		type = text
+	}
+	partition {
+		type = HASH
+		columns = [
+			column.name
+		]
+	}
+}
+`
+		)
+		err := UnmarshalHCL([]byte(f), s)
+		require.NoError(t, err)
+		c := schema.NewStringColumn("name", "text")
+		expected := schema.New("test").
+			AddTables(schema.NewTable("logs").AddColumns(c).AddAttrs(&Partition{T: PartitionTypeHash, Parts: []*PartitionPart{{C: c}}}))
+		require.Equal(t, expected, s)
+	})
+
+	t.Run("Parts", func(t *testing.T) {
+		var (
+			s = &schema.Schema{}
+			f = `
+schema "test" {}
+table "logs" {
+	schema = schema.test
+	column "name" {
+		type = text
+	}
+	partition {
+		type = RANGE
+		by {
+			column = column.name
+		}
+		by {
+			expr = "lower(name)"
+		}
+	}
+}
+`
+		)
+		err := UnmarshalHCL([]byte(f), s)
+		require.NoError(t, err)
+		c := schema.NewStringColumn("name", "text")
+		expected := schema.New("test").
+			AddTables(schema.NewTable("logs").AddColumns(c).AddAttrs(&Partition{T: PartitionTypeRange, Parts: []*PartitionPart{{C: c}, {X: &schema.RawExpr{X: "lower(name)"}}}}))
+		require.Equal(t, expected, s)
+	})
+
+	t.Run("Invalid", func(t *testing.T) {
+		err := UnmarshalHCL([]byte(`
+			schema "test" {}
+			table "logs" {
+				schema = schema.test
+				column "name" { type = text }
+				partition {
+					columns = [column.name]
+				}
+			}
+		`), &schema.Schema{})
+		require.EqualError(t, err, "missing attribute logs.partition.type")
+
+		err = UnmarshalHCL([]byte(`
+			schema "test" {}
+			table "logs" {
+				schema = schema.test
+				column "name" { type = text }
+				partition {
+					type = HASH
+				}
+			}
+		`), &schema.Schema{})
+		require.EqualError(t, err, `missing columns or expressions for logs.partition`)
+
+		err = UnmarshalHCL([]byte(`
+			schema "test" {}
+			table "logs" {
+				schema = schema.test
+				column "name" { type = text }
+				partition {
+					type = HASH
+					columns = [column.name]
+					by { column = column.name }
+				}
+			}
+		`), &schema.Schema{})
+		require.EqualError(t, err, `multiple definitions for logs.partition, use "columns" or "by"`)
+	})
+}
+
 func TestMarshalSpec_IndexPredicate(t *testing.T) {
 	s := &schema.Schema{
 		Name: "test",