Skip to content

Commit

Permalink
New operation: create_constraint and support unique constraints wit…
Browse files Browse the repository at this point in the history
…h multiple columns (#459)

This PR adds a new operation named `create_constraint`. Previously, we
only supported adding constraints to columns. This operation lets us
define table level constraints including multiple columns.

The operation supports `unique` constraints for now. In a follow-up PR I
am adding `foreign_key` and `check` constraints as well.

### Unique

```json
{
  "name": "44_add_table_unique_constraint",
  "operations": [
    {
      "create_constraint": {
        "type": "unique",
        "table": "tickets",
        "name": "unique_zip_name",
        "columns": [
          "sellers_name",
          "sellers_zip"
        ],
        "up": {
          "sellers_name": "sellers_name",
          "sellers_zip": "sellers_zip"
        },
        "down": {
          "sellers_name": "sellers_name",
          "sellers_zip": "sellers_zip"
        }
      }
    }
  ]
}
```

### Related

Created from #411

---------

Co-authored-by: Andrew Farries <[email protected]>
Co-authored-by: Ryan Slade <[email protected]>
  • Loading branch information
3 people authored Nov 13, 2024
1 parent 0e34f78 commit a5fb72e
Show file tree
Hide file tree
Showing 10 changed files with 688 additions and 0 deletions.
36 changes: 36 additions & 0 deletions docs/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
* [Add unique constraint](#add-unique-constraint)
* [Create index](#create-index)
* [Create table](#create-table)
* [Create constraint](#create-constraint)
* [Drop column](#drop-column)
* [Drop constraint](#drop-constraint)
* [Drop index](#drop-index)
Expand Down Expand Up @@ -687,6 +688,7 @@ See the [examples](../examples) directory for examples of each kind of operation
* [Add unique constraint](#add-unique-constraint)
* [Create index](#create-index)
* [Create table](#create-table)
* [Create constraint](#create-constraint)
* [Drop column](#drop-column)
* [Drop constraint](#drop-constraint)
* [Drop index](#drop-index)
Expand Down Expand Up @@ -1037,6 +1039,40 @@ Example **create table** migrations:
* [25_add_table_with_check_constraint.json](../examples/25_add_table_with_check_constraint.json)
* [28_different_defaults.json](../examples/28_different_defaults.json)

### Create constraint

A create constraint operation adds a new constraint to an existing table.

Only `UNIQUE` constraints are supported.

Required fields: `name`, `table`, `type`, `up`, `down`.

**create constraint** operations have this structure:

```json
{
"create_constraint": {
"table": "name of table",
"name": "my_unique_constraint",
"columns": ["col1", "col2"],
"type": "unique"
"up": {
"col1": "col1 || random()",
"col2": "col2 || random()"
},
"down": {
"col1": "col1",
"col2": "col2"
}
}
}
```

Example **create constraint** migrations:

* [44_add_table_unique_constraint.json](../examples/44_add_table_unique_constraint.json)


### Drop column

A drop column operation drops a column from an existing table.
Expand Down
2 changes: 2 additions & 0 deletions examples/.ledger
Original file line number Diff line number Diff line change
Expand Up @@ -40,3 +40,5 @@
40_create_enum_type.json
41_add_enum_column.json
42_create_unique_index.json
43_create_tickets_table.json
44_add_table_unique_constraint.json
25 changes: 25 additions & 0 deletions examples/43_create_tickets_table.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
{
"name": "43_create_tickets_table",
"operations": [
{
"create_table": {
"name": "tickets",
"columns": [
{
"name": "ticket_id",
"type": "serial",
"pk": true
},
{
"name": "sellers_name",
"type": "varchar(255)"
},
{
"name": "sellers_zip",
"type": "integer"
}
]
}
}
]
}
24 changes: 24 additions & 0 deletions examples/44_add_table_unique_constraint.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
{
"name": "44_add_table_unique_constraint",
"operations": [
{
"create_constraint": {
"type": "unique",
"table": "tickets",
"name": "unique_zip_name",
"columns": [
"sellers_name",
"sellers_zip"
],
"up": {
"sellers_name": "sellers_name",
"sellers_zip": "sellers_zip"
},
"down": {
"sellers_name": "sellers_name",
"sellers_zip": "sellers_zip"
}
}
}
]
}
9 changes: 9 additions & 0 deletions pkg/migrations/errors.go
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,15 @@ func (e ColumnDoesNotExistError) Error() string {
return fmt.Sprintf("column %q does not exist on table %q", e.Name, e.Table)
}

type ColumnMigrationMissingError struct {
Table string
Name string
}

func (e ColumnMigrationMissingError) Error() string {
return fmt.Sprintf("migration for column %q in %q is missing", e.Name, e.Table)
}

type ColumnIsNotNullableError struct {
Table string
Name string
Expand Down
7 changes: 7 additions & 0 deletions pkg/migrations/op_common.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ const (
OpNameDropConstraint OpName = "drop_constraint"
OpNameSetReplicaIdentity OpName = "set_replica_identity"
OpRawSQLName OpName = "sql"
OpCreateConstraintName OpName = "create_constraint"

// Internal operation types used by `alter_column`
OpNameRenameColumn OpName = "rename_column"
Expand Down Expand Up @@ -124,6 +125,9 @@ func (v *Operations) UnmarshalJSON(data []byte) error {
case OpRawSQLName:
item = &OpRawSQL{}

case OpCreateConstraintName:
item = &OpCreateConstraint{}

default:
return fmt.Errorf("unknown migration type: %v", opName)
}
Expand Down Expand Up @@ -210,6 +214,9 @@ func OperationName(op Operation) OpName {
case *OpRawSQL:
return OpRawSQLName

case *OpCreateConstraint:
return OpCreateConstraintName

}

panic(fmt.Errorf("unknown operation for %T", op))
Expand Down
216 changes: 216 additions & 0 deletions pkg/migrations/op_create_constraint.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,216 @@
// SPDX-License-Identifier: Apache-2.0

package migrations

import (
"context"
"fmt"
"strings"

"github.com/lib/pq"

"github.com/xataio/pgroll/pkg/db"
"github.com/xataio/pgroll/pkg/schema"
)

var _ Operation = (*OpCreateConstraint)(nil)

func (o *OpCreateConstraint) Start(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, s *schema.Schema, cbs ...CallbackFn) (*schema.Table, error) {
var err error
var table *schema.Table
for _, col := range o.Columns {
if table, err = o.duplicateColumnBeforeStart(ctx, conn, latestSchema, tr, col, s); err != nil {
return nil, err
}
}

switch o.Type { //nolint:gocritic // more cases will be added
case OpCreateConstraintTypeUnique:
return table, o.addUniqueIndex(ctx, conn)
}

return table, nil
}

func (o *OpCreateConstraint) duplicateColumnBeforeStart(ctx context.Context, conn db.DB, latestSchema string, tr SQLTransformer, colName string, s *schema.Schema) (*schema.Table, error) {
table := s.GetTable(o.Table)
column := table.GetColumn(colName)

d := NewColumnDuplicator(conn, table, column)
if err := d.Duplicate(ctx); err != nil {
return nil, fmt.Errorf("failed to duplicate column for new constraint: %w", err)
}

upSQL, ok := o.Up[colName]
if !ok {
return nil, fmt.Errorf("up migration is missing for column %s", colName)
}
physicalColumnName := TemporaryName(colName)
err := createTrigger(ctx, conn, tr, triggerConfig{
Name: TriggerName(o.Table, colName),
Direction: TriggerDirectionUp,
Columns: table.Columns,
SchemaName: s.Name,
LatestSchema: latestSchema,
TableName: o.Table,
PhysicalColumn: physicalColumnName,
SQL: upSQL,
})
if err != nil {
return nil, fmt.Errorf("failed to create up trigger: %w", err)
}

table.AddColumn(colName, schema.Column{
Name: physicalColumnName,
})

downSQL, ok := o.Down[colName]
if !ok {
return nil, fmt.Errorf("down migration is missing for column %s", colName)
}
err = createTrigger(ctx, conn, tr, triggerConfig{
Name: TriggerName(o.Table, physicalColumnName),
Direction: TriggerDirectionDown,
Columns: table.Columns,
LatestSchema: latestSchema,
SchemaName: s.Name,
TableName: o.Table,
PhysicalColumn: colName,
SQL: downSQL,
})
if err != nil {
return nil, fmt.Errorf("failed to create down trigger: %w", err)
}
return table, nil
}

func (o *OpCreateConstraint) Complete(ctx context.Context, conn db.DB, tr SQLTransformer, s *schema.Schema) error {
switch o.Type { //nolint:gocritic // more cases will be added
case OpCreateConstraintTypeUnique:
uniqueOp := &OpSetUnique{
Table: o.Table,
Name: o.Name,
}
err := uniqueOp.Complete(ctx, conn, tr, s)
if err != nil {
return err
}
}

// remove old columns
_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s %s",
pq.QuoteIdentifier(o.Table),
dropMultipleColumns(quoteColumnNames(o.Columns)),
))
if err != nil {
return err
}

// rename new columns to old name
table := s.GetTable(o.Table)
for _, col := range o.Columns {
column := table.GetColumn(col)
if err := RenameDuplicatedColumn(ctx, conn, table, column); err != nil {
return err
}
}

return o.removeTriggers(ctx, conn)
}

func (o *OpCreateConstraint) Rollback(ctx context.Context, conn db.DB, tr SQLTransformer, s *schema.Schema) error {
_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s %s",
pq.QuoteIdentifier(o.Table),
dropMultipleColumns(quotedTemporaryNames(o.Columns)),
))
if err != nil {
return err
}

return o.removeTriggers(ctx, conn)
}

func (o *OpCreateConstraint) removeTriggers(ctx context.Context, conn db.DB) error {
dropFuncs := make([]string, len(o.Columns)*2)
for i, j := 0, 0; i < len(o.Columns); i, j = i+1, j+2 {
dropFuncs[j] = pq.QuoteIdentifier(TriggerFunctionName(o.Table, o.Columns[i]))
dropFuncs[j+1] = pq.QuoteIdentifier(TriggerFunctionName(o.Table, TemporaryName(o.Columns[i])))
}
_, err := conn.ExecContext(ctx, fmt.Sprintf("DROP FUNCTION IF EXISTS %s CASCADE",
strings.Join(dropFuncs, ", "),
))
return err
}

func dropMultipleColumns(columns []string) string {
for i, col := range columns {
columns[i] = "DROP COLUMN IF EXISTS " + col
}
return strings.Join(columns, ", ")
}

func (o *OpCreateConstraint) Validate(ctx context.Context, s *schema.Schema) error {
table := s.GetTable(o.Table)
if table == nil {
return TableDoesNotExistError{Name: o.Table}
}

if err := ValidateIdentifierLength(o.Name); err != nil {
return err
}

if table.ConstraintExists(o.Name) {
return ConstraintAlreadyExistsError{
Table: o.Table,
Constraint: o.Name,
}
}

for _, col := range o.Columns {
if table.GetColumn(col) == nil {
return ColumnDoesNotExistError{
Table: o.Table,
Name: col,
}
}
if _, ok := o.Up[col]; !ok {
return ColumnMigrationMissingError{
Table: o.Table,
Name: col,
}
}
if _, ok := o.Down[col]; !ok {
return ColumnMigrationMissingError{
Table: o.Table,
Name: col,
}
}
}

switch o.Type { //nolint:gocritic // more cases will be added
case OpCreateConstraintTypeUnique:
if len(o.Columns) == 0 {
return FieldRequiredError{Name: "columns"}
}
}

return nil
}

func (o *OpCreateConstraint) addUniqueIndex(ctx context.Context, conn db.DB) error {
_, err := conn.ExecContext(ctx, fmt.Sprintf("CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS %s ON %s (%s)",
pq.QuoteIdentifier(o.Name),
pq.QuoteIdentifier(o.Table),
strings.Join(quotedTemporaryNames(o.Columns), ", "),
))

return err
}

func quotedTemporaryNames(columns []string) []string {
names := make([]string, len(columns))
for i, col := range columns {
names[i] = pq.QuoteIdentifier(TemporaryName(col))
}
return names
}
Loading

0 comments on commit a5fb72e

Please sign in to comment.