Repository: ariga/atlas
Pull request #404 (2022-01-13 13:15:35)

Add support for SQL comments
This PR adds support for adding comments to the HCL. Opposed to the previous implementation where there was an `Comment` field added the the `sqlspec` structs this is no longer the case since SQLite for example does not support comments.

The current implementation also makes sure that no `commen = ""` attributes are rendered to the marshaled HCL if there are no comments.


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index ccd53e74..09131f12 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -11,6 +11,7 @@ import (
 	"strings"
 
 	"ariga.io/atlas/schema/schemaspec"
+	"ariga.io/atlas/sql/internal/sqlx"
 	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlspec"
 )
@@ -106,6 +107,9 @@ func Table(spec *sqlspec.Table, parent *schema.Schema, convertColumn ConvertColu
 		}
 		tbl.Indexes = append(tbl.Indexes, i)
 	}
+	if err := convertCommentFromSpec(spec, &tbl.Attrs); err != nil {
+		return nil, err
+	}
 	return tbl, nil
 }
 
@@ -132,6 +136,9 @@ func Column(spec *sqlspec.Column, conv ConvertTypeFunc) (*schema.Column, error)
 		return nil, err
 	}
 	out.Type.Type = ct
+	if err := convertCommentFromSpec(spec, &out.Attrs); err != nil {
+		return nil, err
+	}
 	return out, err
 }
 
@@ -152,12 +159,16 @@ func Index(spec *sqlspec.Index, parent *schema.Table) (*schema.Index, error) {
 			C:     col,
 		})
 	}
-	return &schema.Index{
+	i := &schema.Index{
 		Name:   spec.Name,
 		Unique: spec.Unique,
 		Table:  parent,
 		Parts:  parts,
-	}, nil
+	}
+	if err := convertCommentFromSpec(spec, &i.Attrs); err != nil {
+		return nil, err
+	}
+	return i, nil
 }
 
 // PrimaryKey converts a sqlspec.PrimaryKey to a schema.Index.
@@ -314,6 +325,7 @@ func FromTable(t *schema.Table, colFn ColumnSpecFunc, pkFn PrimaryKeySpecFunc, i
 		}
 		spec.ForeignKeys = append(spec.ForeignKeys, f)
 	}
+	convertCommentFromSchema(t.Attrs, &spec.Extra.Attrs)
 	return spec, nil
 }
 
@@ -349,6 +361,7 @@ func FromColumn(col *schema.Column, columnTypeSpec ColumnTypeSpecFunc) (*sqlspec
 		}
 		spec.Default = lv
 	}
+	convertCommentFromSchema(col.Attrs, &spec.Extra.Attrs)
 	return spec, nil
 }
 
@@ -392,11 +405,13 @@ func FromIndex(idx *schema.Index) (*sqlspec.Index, error) {
 		}
 		c = append(c, colRef(p.C.Name, idx.Table.Name))
 	}
-	return &sqlspec.Index{
+	i := &sqlspec.Index{
 		Name:    idx.Name,
 		Unique:  idx.Unique,
 		Columns: c,
-	}, nil
+	}
+	convertCommentFromSchema(idx.Attrs, &i.Extra.Attrs)
+	return i, nil
 }
 
 // FromForeignKey converts schema.ForeignKey to sqlspec.ForeignKey
@@ -457,3 +472,25 @@ func colRef(cName string, tName string) *schemaspec.Ref {
 func schemaRef(n string) *schemaspec.Ref {
 	return &schemaspec.Ref{V: "$schema." + n}
 }
+
+// convertCommentFromSpec converts a spec comment attribute to a schema element attribute.
+func convertCommentFromSpec(spec interface {
+	Attr(string) (*schemaspec.Attr, bool)
+}, attrs *[]schema.Attr) error {
+	if c, ok := spec.Attr("comment"); ok {
+		s, err := c.String()
+		if err != nil {
+			return err
+		}
+		*attrs = append(*attrs, &schema.Comment{Text: s})
+	}
+	return nil
+}
+
+// convertCommentFromSchema converts a schema element comment attribute to a spec comment attribute.
+func convertCommentFromSchema(src []schema.Attr, trgt *[]*schemaspec.Attr) {
+	var c schema.Comment
+	if sqlx.Has(src, &c) {
+		*trgt = append(*trgt, StrAttr("comment", c.Text))
+	}
+}
diff --git a/sql/schema/dsl.go b/sql/schema/dsl.go
index 8ad026d5..737b6b50 100644
--- a/sql/schema/dsl.go
+++ b/sql/schema/dsl.go
@@ -452,6 +452,13 @@ func (i *Index) SetTable(t *Table) *Index {
 	return i
 }
 
+// SetComment sets or appends the Comment attribute
+// to the index with the given value.
+func (i *Index) SetComment(v string) *Index {
+	replaceOrAppend(&i.Attrs, &Comment{Text: v})
+	return i
+}
+
 // AddAttrs adds additional attributes to the index.
 func (i *Index) AddAttrs(attrs ...Attr) *Index {
 	i.Attrs = append(i.Attrs, attrs...)


Test patch:
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index 83cc4058..a4d12f01 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -16,6 +16,7 @@ schema "schema" {
 table "table" {
 	column "col" {
 		type = int
+		comment = "column comment"
 	}
 	column "age" {
 		type = int
@@ -32,6 +33,7 @@ table "table" {
 			table.table.column.col,
 			table.table.column.age,
 		]
+		comment = "index comment"
 	}
 	foreign_key "accounts" {
 		columns = [
@@ -42,6 +44,7 @@ table "table" {
 		]
 		on_delete = "SET NULL"
 	}
+	comment = "table comment"
 }
 
 table "accounts" {
@@ -72,6 +75,9 @@ table "accounts" {
 							T: TypeInt,
 						},
 					},
+					Attrs: []schema.Attr{
+						&schema.Comment{Text: "column comment"},
+					},
 				},
 				{
 					Name: "age",
@@ -91,6 +97,9 @@ table "accounts" {
 					},
 				},
 			},
+			Attrs: []schema.Attr{
+				&schema.Comment{Text: "table comment"},
+			},
 		},
 		{
 			Name:   "accounts",
@@ -123,6 +132,9 @@ table "accounts" {
 				{SeqNo: 0, C: exp.Tables[0].Columns[0]},
 				{SeqNo: 1, C: exp.Tables[0].Columns[1]},
 			},
+			Attrs: []schema.Attr{
+				&schema.Comment{Text: "index comment"},
+			},
 		},
 	}
 	exp.Tables[0].ForeignKeys = []*schema.ForeignKey{
@@ -272,6 +284,79 @@ schema "test" {
 	require.Equal(t, utf8mb4, b.Attrs)
 }
 
+func TestMarshalSpec_Comment(t *testing.T) {
+	s := &schema.Schema{
+		Name: "test",
+		Tables: []*schema.Table{
+			{
+				Name: "users",
+				Attrs: []schema.Attr{
+					&schema.Comment{Text: "table comment"},
+				},
+				Columns: []*schema.Column{
+					{
+						Name: "a",
+						Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}},
+						Attrs: []schema.Attr{
+							&schema.Comment{Text: "column comment"},
+						},
+					},
+				},
+			},
+			{
+				Name: "posts",
+				Columns: []*schema.Column{
+					{
+						Name: "a",
+						Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}},
+					},
+				},
+			},
+		},
+	}
+	s.Tables[0].Schema = s
+	s.Tables[1].Schema = s
+	s.Tables[0].Indexes = []*schema.Index{
+		{
+			Name:   "index",
+			Table:  s.Tables[0],
+			Unique: true,
+			Parts:  []*schema.IndexPart{{SeqNo: 0, C: s.Tables[0].Columns[0]}},
+			Attrs: []schema.Attr{
+				&schema.Comment{Text: "index comment"},
+			},
+		},
+	}
+	buf, err := MarshalSpec(s, hclState)
+	require.NoError(t, err)
+	// We expect a zero value comment to not bre present in the marshaled HCL.
+	const expected = `table "users" {
+  schema  = schema.test
+  comment = "table comment"
+  column "a" {
+    null    = false
+    type    = text
+    comment = "column comment"
+  }
+  index "index" {
+    unique  = true
+    columns = [table.users.column.a, ]
+    comment = "index comment"
+  }
+}
+table "posts" {
+  schema = schema.test
+  column "a" {
+    null = false
+    type = text
+  }
+}
+schema "test" {
+}
+`
+	require.EqualValues(t, expected, string(buf))
+}
+
 func TestTypes(t *testing.T) {
 	tests := []struct {
 		typeExpr  string
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index e8f4c87e..15dedbcd 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -17,6 +17,7 @@ schema "schema" {
 table "table" {
 	column "col" {
 		type = integer
+		comment = "column comment"
 	}
 	column "age" {
 		type = integer
@@ -36,6 +37,7 @@ table "table" {
 			table.table.column.col,
 			table.table.column.age,
 		]
+		comment = "index comment"
 	}
 	foreign_key "accounts" {
 		columns = [
@@ -46,6 +48,7 @@ table "table" {
 		]
 		on_delete = "SET NULL"
 	}
+	comment = "table comment"
 }
 
 table "accounts" {
@@ -75,6 +78,9 @@ table "accounts" {
 							T: "integer",
 						},
 					},
+					Attrs: []schema.Attr{
+						&schema.Comment{Text: "column comment"},
+					},
 				},
 				{
 					Name: "age",
@@ -102,6 +108,9 @@ table "accounts" {
 					},
 				},
 			},
+			Attrs: []schema.Attr{
+				&schema.Comment{Text: "table comment"},
+			},
 		},
 		{
 			Name:   "accounts",
@@ -134,6 +143,9 @@ table "accounts" {
 				{SeqNo: 0, C: exp.Tables[0].Columns[0]},
 				{SeqNo: 1, C: exp.Tables[0].Columns[1]},
 			},
+			Attrs: []schema.Attr{
+				&schema.Comment{Text: "index comment"},
+			},
 		},
 	}
 	exp.Tables[0].ForeignKeys = []*schema.ForeignKey{
diff --git a/sql/schema/dsl_test.go b/sql/schema/dsl_test.go
index 549d48cd..93eeff75 100644
--- a/sql/schema/dsl_test.go
+++ b/sql/schema/dsl_test.go
@@ -48,7 +48,9 @@ func TestSchema_AddTables(t *testing.T) {
 		SetPrimaryKey(schema.NewPrimaryKey(userColumns[0])).
 		SetComment("users table").
 		AddIndexes(
-			schema.NewUniqueIndex("unique_name").AddColumns(userColumns[2]),
+			schema.NewUniqueIndex("unique_name").
+				AddColumns(userColumns[2]).
+				SetComment("index comment"),
 		)
 	postColumns := []*schema.Column{
 		schema.NewIntColumn("id", "int"),
@@ -87,7 +89,12 @@ func TestSchema_AddTables(t *testing.T) {
 			users.PrimaryKey = &schema.Index{Unique: true, Parts: []*schema.IndexPart{{C: users.Columns[0]}}}
 			users.PrimaryKey.Table = users
 			users.Columns[0].Indexes = append(users.Columns[0].Indexes, users.PrimaryKey)
-			users.Indexes = append(users.Indexes, &schema.Index{Name: "unique_name", Unique: true, Parts: []*schema.IndexPart{{C: users.Columns[2]}}})
+			users.Indexes = append(users.Indexes, &schema.Index{
+				Name:   "unique_name",
+				Unique: true,
+				Parts:  []*schema.IndexPart{{C: users.Columns[2]}},
+				Attrs:  []schema.Attr{&schema.Comment{Text: "index comment"}},
+			})
 			users.Indexes[0].Table = users
 			users.Columns[2].Indexes = users.Indexes