Repository: ariga/atlas
Pull request #2763 (2024-05-08 20:05:25)

sql/postgres: add nulls_first/last to sqlspec
#2762


Patch:
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index fb9355af..576c6d55 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -634,6 +634,29 @@ func convertPart(spec *sqlspec.IndexPart, part *schema.IndexPart) error {
 		}
 		part.Attrs = append(part.Attrs, &IndexOpClass{Name: name})
 	}
+	// The following attributes are mutually exclusive.
+	if nulls, ok := spec.Attr("nulls_last"); ok {
+		b, err := nulls.Bool()
+		if err != nil {
+			return err
+		}
+		if b {
+			at := sqlx.AttrOr(part.Attrs, &IndexColumnProperty{})
+			at.NullsLast = true
+			schema.ReplaceOrAppend(&part.Attrs, at)
+		}
+	}
+	if nulls, ok := spec.Attr("nulls_first"); ok {
+		b, err := nulls.Bool()
+		if err != nil {
+			return err
+		}
+		if b {
+			at := sqlx.AttrOr(part.Attrs, &IndexColumnProperty{})
+			at.NullsFirst = true
+			schema.ReplaceOrAppend(&part.Attrs, at)
+		}
+	}
 	return nil
 }
 
@@ -792,20 +815,26 @@ func indexPKSpec(idx *schema.Index, attrs []*schemahcl.Attr) []*schemahcl.Attr {
 }
 
 func partAttr(idx *schema.Index, part *schema.IndexPart, spec *sqlspec.IndexPart) error {
-	var op IndexOpClass
-	if !sqlx.Has(part.Attrs, &op) {
-		return nil
+	if op := (IndexOpClass{}); sqlx.Has(part.Attrs, &op) {
+		switch d, err := op.DefaultFor(idx, part); {
+		case err != nil:
+			return err
+		case d:
+		case len(op.Params) > 0:
+			spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.RawAttr("ops", op.String()))
+		case postgresop.HasClass(op.String()):
+			spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("ops", op.String()))
+		default:
+			spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.StringAttr("ops", op.String()))
+		}
 	}
-	switch d, err := op.DefaultFor(idx, part); {
-	case err != nil:
-		return err
-	case d:
-	case len(op.Params) > 0:
-		spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.RawAttr("ops", op.String()))
-	case postgresop.HasClass(op.String()):
-		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.VarAttr("ops", op.String()))
-	default:
-		spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.StringAttr("ops", op.String()))
+	if nulls := (IndexColumnProperty{}); sqlx.Has(part.Attrs, &nulls) {
+		switch {
+		case part.Desc && nulls.NullsLast:
+			spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.BoolAttr("nulls_last", true))
+		case !part.Desc && nulls.NullsFirst:
+			spec.Extra.Attrs = append(spec.Extra.Attrs, schemahcl.BoolAttr("nulls_first", true))
+		}
 	}
 	return nil
 }


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index ab343eb9..f4df7db6 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -978,6 +978,71 @@ schema "public" {
 	require.EqualValues(t, expected, string(buf))
 }
 
+func TestMarshalSpec_IndexNullsLastFirst(t *testing.T) {
+	s := schema.New("public").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(
+					schema.NewIntColumn("c", "int"),
+				).
+				AddIndexes(
+					// Default behavior.
+					schema.NewIndex("i1").
+						AddParts(schema.NewColumnPart(schema.NewColumn("c")).SetDesc(true).AddAttrs(&IndexColumnProperty{NullsFirst: true})),
+					schema.NewIndex("i2").
+						AddParts(schema.NewColumnPart(schema.NewColumn("c")).SetDesc(false).AddAttrs(&IndexColumnProperty{NullsLast: true})),
+					// Require emitting the NULLS FIRST/LAST attributes.
+					schema.NewIndex("i3").
+						AddParts(schema.NewColumnPart(schema.NewColumn("c")).SetDesc(true).AddAttrs(&IndexColumnProperty{NullsLast: true})),
+					schema.NewIndex("i4").
+						AddParts(schema.NewColumnPart(schema.NewColumn("c")).SetDesc(false).AddAttrs(&IndexColumnProperty{NullsFirst: true})),
+				),
+		)
+	buf, err := MarshalSpec(s, hclState)
+	require.NoError(t, err)
+	const expected = `table "users" {
+  schema = schema.public
+  column "c" {
+    null = false
+    type = int
+  }
+  index "i1" {
+    on {
+      desc   = true
+      column = column.c
+    }
+  }
+  index "i2" {
+    columns = [column.c]
+  }
+  index "i3" {
+    on {
+      desc       = true
+      column     = column.c
+      nulls_last = true
+    }
+  }
+  index "i4" {
+    on {
+      column      = column.c
+      nulls_first = true
+    }
+  }
+}
+schema "public" {
+}
+`
+	require.EqualValues(t, expected, string(buf))
+
+	var got schema.Schema
+	err = EvalHCLBytes([]byte(expected), &got, nil)
+	require.NoError(t, err)
+
+	buf, err = MarshalSpec(s, hclState)
+	require.NoError(t, err)
+	require.EqualValues(t, expected, string(buf))
+}
+
 func TestMarshalSpec_BRINIndex(t *testing.T) {
 	s := &schema.Schema{
 		Name: "test",