sql/sqlspec: allow labeling primary keys (#3351)

sql/internal/specutil/convert.go

@@ -522,6 +522,7 @@ func PrimaryKey(spec *sqlspec.PrimaryKey, parent *schema.Table) (*schema.Index,
 		})
 	}
 	pk := &schema.Index{
+		Name:  spec.Name,
 		Table: parent,
 		Parts: parts,
 	}

sql/sqlspec/sqlspec.go

@@ -63,6 +63,7 @@ type (
 
 	// PrimaryKey holds a specification for the primary key of a table.
 	PrimaryKey struct {
+		Name    string           `spec:",name"` // Optional name.
 		Parts   []*IndexPart     `spec:"on"`
 		Columns []*schemahcl.Ref `spec:"columns"`
 		schemahcl.DefaultExtension

sql/sqlspec/sqlspec_test.go

@@ -7,6 +7,8 @@ package sqlspec
 import (
 	"testing"
 
+	"ariga.io/atlas/schemahcl"
+
 	"github.com/stretchr/testify/require"
 )
 
@@ -55,3 +57,44 @@ SELECT
 		require.Equal(t, tt.expected, MightHeredoc(tt.input))
 	}
 }
+
+func TestMarshalPrimaryKey(t *testing.T) {
+	spec := &Table{
+		Name: "users",
+		Columns: []*Column{
+			{Name: "id", Type: &schemahcl.Type{T: "text"}},
+		},
+		PrimaryKey: &PrimaryKey{
+			Columns: []*schemahcl.Ref{
+				{V: "$column.id"},
+			},
+		},
+	}
+	buf, err := schemahcl.Marshal.MarshalSpec(spec)
+	require.NoError(t, err)
+	require.Equal(t, `table "users" {
+  column "id" {
+    null = false
+    type = sql("text")
+  }
+  primary_key {
+    columns = [column.id]
+  }
+}
+`, string(buf))
+
+	// Include primary key on marshaling.
+	spec.PrimaryKey.Name = "custom_name"
+	buf, err = schemahcl.Marshal.MarshalSpec(spec)
+	require.NoError(t, err)
+	require.Equal(t, `table "users" {
+  column "id" {
+    null = false
+    type = sql("text")
+  }
+  primary_key "custom_name" {
+    columns = [column.id]
+  }
+}
+`, string(buf))
+}