Add support for nulls_not_distinct option in create_constraint operations

Author: kvch

docs/operations/create_constraint.mdx

@@ -18,6 +18,7 @@ Required fields: `name`, `table`, `type`, `up`, `down`.
     "type": "unique"| "check" | "foreign_key",
     "check": "SQL expression for CHECK constraint",
     "no_inherit": "true|false",
+    "nulls_not_distinct": "true|false",
     "references": {
       "name": "name of foreign key reference",
       "table": "name of referenced table",

docs/operations/create_index.mdx

@@ -27,6 +27,7 @@ description: A create index operation creates a new index on a set of columns.
     "predicate": "conditional expression for defining a partial index",
     "storage_parameters": "comma-separated list of storage parameters",
     "unique": true | false,
+    "nulls_not_distinct": true | false,
     "method": "btree"
   }
 }
@@ -35,6 +36,7 @@ description: A create index operation creates a new index on a set of columns.
 * The field `method` can be `btree`, `hash`, `gist`, `spgist`, `gin`, `brin`.
 * You can also specify storage parameters for the index in `storage_parameters`.
 * To create a unique index set `unique` to `true`.
+* If the index is unique, you can configure `nulls_not_distinct`.
 
 ## Examples
 

examples/44_add_table_unique_constraint.json

@@ -10,6 +10,7 @@
           "sellers_name",
           "sellers_zip"
         ],
+        "nulls_not_distinct": false,
         "up": {
           "sellers_name": "sellers_name",
           "sellers_zip": "sellers_zip"

internal/jsonschema/testdata/create-constraint-3-invalid-check-nulls.txtar

@@ -0,0 +1,30 @@
+This is an invalid 'create_constraint' migration.
+Check constraints cannot configure nulls_not_distinct.
+
+-- create_constraint.json --
+{
+  "name": "migration_name",
+  "operations": [
+    {
+      "create_constraint": {
+        "name": "my_invalid_check",
+        "table": "my_table",
+        "type": "check",
+        "check": "my_column > 5",
+        "nulls_not_distinct": true,
+        "columns": [
+          "my_column"
+        ],
+        "up": {
+          "my_column": "my_column"
+        },
+        "down": {
+          "my_column": "my_column"
+        }
+      }
+    }
+  ]
+}
+
+-- valid --
+false

pkg/migrations/duplicate.go

@@ -136,7 +136,7 @@ func (d *Duplicator) Duplicate(ctx context.Context) error {
 			continue
 		}
 		if duplicatedMember, constraintColumns := d.stmtBuilder.allConstraintColumns(uc.Columns, colNames...); duplicatedMember {
-			if err := createUniqueIndexConcurrently(ctx, d.conn, "", DuplicationName(uc.Name), d.stmtBuilder.table.Name, constraintColumns); err != nil {
+			if err := createUniqueIndexConcurrently(ctx, d.conn, "", DuplicationName(uc.Name), d.stmtBuilder.table.Name, constraintColumns, uc.NullsNotDistinct); err != nil {
 				return err
 			}
 		}

pkg/migrations/duplicate_test.go

@@ -231,7 +231,7 @@ func TestCreateIndexConcurrentlySqlGeneration(t *testing.T) {
 		},
 	} {
 		t.Run(name, func(t *testing.T) {
-			stmt := getCreateUniqueIndexConcurrentlySQL(testCases.indexName, testCases.schemaName, testCases.tableName, testCases.columns)
+			stmt := getCreateUniqueIndexConcurrentlySQL(testCases.indexName, testCases.schemaName, testCases.tableName, testCases.columns, false)
 			assert.Equal(t, testCases.expectedStmt, stmt)
 		})
 	}

pkg/migrations/index.go

@@ -12,15 +12,15 @@ import (
 	"github.com/xataio/pgroll/pkg/db"
 )
 
-func createUniqueIndexConcurrently(ctx context.Context, conn db.DB, schemaName string, indexName string, tableName string, columnNames []string) error {
+func createUniqueIndexConcurrently(ctx context.Context, conn db.DB, schemaName string, indexName string, tableName string, columnNames []string, nullsNotDistinct bool) error {
 	quotedQualifiedIndexName := pq.QuoteIdentifier(indexName)
 	if schemaName != "" {
 		quotedQualifiedIndexName = fmt.Sprintf("%s.%s", pq.QuoteIdentifier(schemaName), pq.QuoteIdentifier(indexName))
 	}
 	for retryCount := 0; retryCount < 5; retryCount++ {
 		// Add a unique index to the new column
 		// Indexes are created in the same schema with the table automatically. Instead of the qualified one, just pass the index name.
-		createIndexSQL := getCreateUniqueIndexConcurrentlySQL(indexName, schemaName, tableName, columnNames)
+		createIndexSQL := getCreateUniqueIndexConcurrentlySQL(indexName, schemaName, tableName, columnNames, nullsNotDistinct)
 		if _, err := conn.ExecContext(ctx, createIndexSQL); err != nil {
 			return fmt.Errorf("failed to add unique index %q: %w", indexName, err)
 		}
@@ -64,7 +64,7 @@ func createUniqueIndexConcurrently(ctx context.Context, conn db.DB, schemaName s
 	return fmt.Errorf("failed to create unique index %q", indexName)
 }
 
-func getCreateUniqueIndexConcurrentlySQL(indexName string, schemaName string, tableName string, columnNames []string) string {
+func getCreateUniqueIndexConcurrentlySQL(indexName string, schemaName string, tableName string, columnNames []string, nullsNotDistinct bool) string {
 	// create unique index concurrently
 	qualifiedTableName := pq.QuoteIdentifier(tableName)
 	if schemaName != "" {
@@ -77,6 +77,9 @@ func getCreateUniqueIndexConcurrentlySQL(indexName string, schemaName string, ta
 		qualifiedTableName,
 		strings.Join(quoteColumnNames(columnNames), ", "),
 	)
+	if nullsNotDistinct {
+		indexQuery += " NULLS NOT DISTINCT"
+	}
 
 	return indexQuery
 }

pkg/migrations/op_add_column.go

@@ -62,7 +62,7 @@ func (o *OpAddColumn) Start(ctx context.Context, conn db.DB, latestSchema string
 	}
 
 	if o.Column.Unique {
-		if err := createUniqueIndexConcurrently(ctx, conn, s.Name, UniqueIndexName(o.Column.Name), table.Name, []string{TemporaryName(o.Column.Name)}); err != nil {
+		if err := createUniqueIndexConcurrently(ctx, conn, s.Name, UniqueIndexName(o.Column.Name), table.Name, []string{TemporaryName(o.Column.Name)}, false); err != nil {
 			return nil, fmt.Errorf("failed to add unique index: %w", err)
 		}
 	}

pkg/migrations/op_create_constraint.go

@@ -82,7 +82,7 @@ func (o *OpCreateConstraint) Start(ctx context.Context, conn db.DB, latestSchema
 
 	switch o.Type {
 	case OpCreateConstraintTypeUnique:
-		return table, createUniqueIndexConcurrently(ctx, conn, s.Name, o.Name, table.Name, temporaryNames(o.Columns))
+		return table, createUniqueIndexConcurrently(ctx, conn, s.Name, o.Name, table.Name, temporaryNames(o.Columns), o.NullsNotDistinct)
 	case OpCreateConstraintTypeCheck:
 		return table, o.addCheckConstraint(ctx, conn, table.Name)
 	case OpCreateConstraintTypeForeignKey:
@@ -279,6 +279,10 @@ func (o *OpCreateConstraint) Validate(ctx context.Context, s *schema.Schema) err
 		}
 	}
 
+	if o.Type != OpCreateConstraintTypeUnique && o.NullsNotDistinct {
+		return fmt.Errorf("nulls_not_distinct can only be true for unique constraints")
+	}
+
 	return nil
 }
 

pkg/migrations/op_create_index.go

@@ -61,6 +61,10 @@ func (o *OpCreateIndex) Start(ctx context.Context, conn db.DB, latestSchema stri
 	}
 	stmt += fmt.Sprintf(" (%s)", strings.Join(colSQLs, ", "))
 
+	if o.Unique && o.NullsNotDistinct {
+		stmt += " NULLS NOT DISTINCT"
+	}
+
 	if o.StorageParameters != "" {
 		stmt += fmt.Sprintf(" WITH (%s)", o.StorageParameters)
 	}