Repository: ariga/atlas
Pull request #767 (2022-05-04 07:35:28)

sql/postgres: add support for diffing partitioned tables


Patch:
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index 18fc96d6..382e6aa9 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -6,6 +6,7 @@ package postgres
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"reflect"
 	"strconv"
@@ -31,6 +32,9 @@ func (d *diff) TableAttrDiff(from, to *schema.Table) ([]schema.Change, error) {
 	if change := sqlx.CommentDiff(from.Attrs, to.Attrs); change != nil {
 		changes = append(changes, change)
 	}
+	if err := d.partitionChanged(from, to); err != nil {
+		return nil, err
+	}
 	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
@@ -99,6 +103,30 @@ func (*diff) generatedChanged(from, to *schema.Column) (bool, error) {
 	}
 }
 
+// partitionChanged checks and returns an error if the partition key of a table was changed.
+func (*diff) partitionChanged(from, to *schema.Table) error {
+	var fromP, toP Partition
+	switch fromHas, toHas := sqlx.Has(from.Attrs, &fromP), sqlx.Has(to.Attrs, &toP); {
+	case fromHas && !toHas:
+		return fmt.Errorf("partition key cannot be dropped from %q (drop and add is required)", from.Name)
+	case !fromHas && toHas:
+		return fmt.Errorf("partition key cannot be added to %q (drop and add is required)", to.Name)
+	case fromHas && toHas:
+		s1, err := formatPartition(fromP)
+		if err != nil {
+			return err
+		}
+		s2, err := formatPartition(toP)
+		if err != nil {
+			return err
+		}
+		if s1 != s2 {
+			return fmt.Errorf("partition key of table %q cannot be changed from %s to %s (drop and add is required)", to.Name, s1, s2)
+		}
+	}
+	return nil
+}
+
 // IsGeneratedIndexName reports if the index name was generated by the database.
 func (d *diff) IsGeneratedIndexName(t *schema.Table, idx *schema.Index) bool {
 	names := make([]string, len(idx.Parts))
@@ -353,6 +381,32 @@ func identity(attrs []schema.Attr) (*Identity, bool) {
 	return i, true
 }
 
+// formatPartition returns the string representation of the
+// partition key according to the PostgreSQL format/grammar.
+func formatPartition(p Partition) (string, error) {
+	b := Build("PARTITION BY")
+	switch t := strings.ToUpper(p.T); t {
+	case PartitionTypeRange, PartitionTypeList, PartitionTypeHash:
+		b.P(t)
+	default:
+		return "", fmt.Errorf("unknown partition type: %q", t)
+	}
+	if len(p.Parts) == 0 {
+		return "", errors.New("missing parts for partition key")
+	}
+	b.Wrap(func(b *sqlx.Builder) {
+		b.MapComma(p.Parts, func(i int, b *sqlx.Builder) {
+			switch k := p.Parts[i]; {
+			case k.C != nil:
+				b.Ident(k.C.Name)
+			case k.X != nil:
+				b.P(sqlx.MayWrap(k.X.(*schema.RawExpr).X))
+			}
+		})
+	})
+	return b.String(), nil
+}
+
 func trimCast(s string) string {
 	i := strings.LastIndex(s, "::")
 	if i == -1 {


Test patch:
diff --git a/sql/postgres/diff_test.go b/sql/postgres/diff_test.go
index c7909a3b..17a27e79 100644
--- a/sql/postgres/diff_test.go
+++ b/sql/postgres/diff_test.go
@@ -69,6 +69,54 @@ func TestDiff_TableDiff(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "drop partition key",
+			from: schema.NewTable("logs").
+				AddAttrs(&Partition{
+					T:     PartitionTypeRange,
+					Parts: []*PartitionPart{{C: schema.NewColumn("c")}},
+				}),
+			to:      schema.NewTable("logs"),
+			wantErr: true,
+		},
+		{
+			name: "add partition key",
+			from: schema.NewTable("logs"),
+			to: schema.NewTable("logs").
+				AddAttrs(&Partition{
+					T:     PartitionTypeRange,
+					Parts: []*PartitionPart{{C: schema.NewColumn("c")}},
+				}),
+			wantErr: true,
+		},
+		{
+			name: "change partition key column",
+			from: schema.NewTable("logs").
+				AddAttrs(&Partition{
+					T:     PartitionTypeRange,
+					Parts: []*PartitionPart{{C: schema.NewColumn("c")}},
+				}),
+			to: schema.NewTable("logs").
+				AddAttrs(&Partition{
+					T:     PartitionTypeRange,
+					Parts: []*PartitionPart{{C: schema.NewColumn("d")}},
+				}),
+			wantErr: true,
+		},
+		{
+			name: "change partition key type",
+			from: schema.NewTable("logs").
+				AddAttrs(&Partition{
+					T:     PartitionTypeRange,
+					Parts: []*PartitionPart{{C: schema.NewColumn("c")}},
+				}),
+			to: schema.NewTable("logs").
+				AddAttrs(&Partition{
+					T:     PartitionTypeHash,
+					Parts: []*PartitionPart{{C: schema.NewColumn("c")}},
+				}),
+			wantErr: true,
+		},
 		{
 			name: "add check",
 			from: &schema.Table{Name: "t1", Schema: &schema.Schema{Name: "public"}},
@@ -309,7 +357,7 @@ func TestDiff_TableDiff(t *testing.T) {
 		require.NoError(t, err)
 		t.Run(tt.name, func(t *testing.T) {
 			changes, err := drv.TableDiff(tt.from, tt.to)
-			require.Equal(t, tt.wantErr, err != nil)
+			require.Equalf(t, tt.wantErr, err != nil, "got: %v", err)
 			require.EqualValues(t, tt.wantChanges, changes)
 		})
 	}