Repository: ariga/atlas (go)

Issue:
Missing schema comment (on Postgres)
When setting up as

```sql
CREATE SCHEMA "foo";
COMMENT ON SCHEMA "foo" IS 'bar';
```

I get from `atlas schema inspect`

```
schema "foo" {
}
schema "public" {
}
```

Am I wrong to expect something along the lines of

```
schema "foo" {
  comment = "bar"
}
schema "public" {
}
```

(using atlas version v0.12.1-17d97d4-canary)?

Pull request:
sql/postgres: support schema comments
Fixed #1819

Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index e77aae7f..fa095d3d 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -58,6 +58,9 @@ func Scan(r *schema.Realm, doc *ScanDoc, funcs *ScanFuncs) error {
 	byName := make(map[string]*schema.Schema)
 	for _, s := range doc.Schemas {
 		s1 := schema.New(s.Name)
+		if err := convertCommentFromSpec(s, &s1.Attrs); err != nil {
+			return err
+		}
 		r.AddSchemas(s1)
 		byName[s.Name] = s1
 	}
@@ -404,6 +407,7 @@ func FromSchema(s *schema.Schema, specT TableSpecFunc, specV ViewSpecFunc) (*Sch
 		}
 		views = append(views, view)
 	}
+	convertCommentFromSchema(s.Attrs, &spec.Extra.Attrs)
 	return &SchemaSpec{
 		Schema: spec,
 		Tables: tables,
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index 94bc8cfd..b43ebff4 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -27,9 +27,12 @@ var DefaultDiff schema.Differ = &sqlx.Diff{DiffDriver: &diff{}}
 type diff struct{ conn }
 
 // SchemaAttrDiff returns a changeset for migrating schema attributes from one state to the other.
-func (*diff) SchemaAttrDiff(_, _ *schema.Schema) []schema.Change {
-	// No special schema attribute diffing for PostgreSQL.
-	return nil
+func (*diff) SchemaAttrDiff(from, to *schema.Schema) []schema.Change {
+	var changes []schema.Change
+	if change := sqlx.CommentDiff(from.Attrs, to.Attrs); change != nil {
+		changes = append(changes, change)
+	}
+	return changes
 }
 
 // SchemaObjectDiff returns a changeset for migrating schema objects from
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 492f6891..31b477fc 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -601,13 +601,18 @@ func (i *inspect) schemas(ctx context.Context, opts *schema.InspectRealmOption)
 	defer rows.Close()
 	var schemas []*schema.Schema
 	for rows.Next() {
-		var name string
-		if err := rows.Scan(&name); err != nil {
+		var (
+			name    string
+			comment sql.NullString
+		)
+		if err := rows.Scan(&name, &comment); err != nil {
 			return nil, err
 		}
-		schemas = append(schemas, &schema.Schema{
-			Name: name,
-		})
+		s := schema.New(name)
+		if comment.Valid {
+			s.SetComment(comment.String)
+		}
+		schemas = append(schemas, s)
 	}
 	if err := rows.Close(); err != nil {
 		return nil, err
@@ -1121,10 +1126,29 @@ const (
 	paramsQuery = `SELECT setting FROM pg_settings WHERE name IN ('lc_collate', 'lc_ctype', 'server_version_num', 'crdb_version') ORDER BY name DESC`
 
 	// Query to list database schemas.
-	schemasQuery = "SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'pg_toast', 'crdb_internal', 'pg_extension') AND schema_name NOT LIKE 'pg_%temp_%' ORDER BY schema_name"
+	schemasQuery = `
+SELECT
+	nspname AS schema_name,
+	pg_catalog.obj_description(oid) AS comment
+FROM
+    pg_catalog.pg_namespace
+WHERE
+    nspname NOT IN ('information_schema', 'pg_catalog', 'pg_toast', 'crdb_internal', 'pg_extension')
+    AND nspname NOT LIKE 'pg_%temp_%'
+ORDER BY
+    nspname`
 
-	// Query to list specific database schemas.
-	schemasQueryArgs = "SELECT schema_name FROM information_schema.schemata WHERE schema_name %s ORDER BY schema_name"
+	// Query to list database schemas.
+	schemasQueryArgs = `
+SELECT
+	nspname AS schema_name,
+	pg_catalog.obj_description(oid) AS comment
+FROM
+    pg_catalog.pg_namespace
+WHERE
+    nspname %s
+ORDER BY
+    nspname`
 
 	// Query to list table information.
 	tablesQuery = `
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index a8ca38d6..e21b9026 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -147,6 +147,31 @@ func (s *state) topLevel(changes []schema.Change) ([]schema.Change, error) {
 				Reverse: s.Build("DROP SCHEMA").Ident(c.S.Name).P("CASCADE").String(),
 				Comment: fmt.Sprintf("Add new schema named %q", c.S.Name),
 			})
+			if cm := (schema.Comment{}); sqlx.Has(c.S.Attrs, &cm) {
+				s.append(s.schemaComment(c.S, cm.Text, ""))
+			}
+		case *schema.ModifySchema:
+			for i := range c.Changes {
+				switch change := c.Changes[i].(type) {
+				// Add schema attributes to an existing schema only if
+				// it is different from the default server configuration.
+				case *schema.AddAttr:
+					a, ok := change.A.(*schema.Comment)
+					if !ok {
+						return nil, fmt.Errorf("unexpected schema AddAttr: %T", change.A)
+					}
+					s.append(s.schemaComment(c.S, a.Text, ""))
+				case *schema.ModifyAttr:
+					to, ok1 := change.To.(*schema.Comment)
+					from, ok2 := change.From.(*schema.Comment)
+					if !ok1 || !ok2 {
+						return nil, fmt.Errorf("unexpected schema ModifyAttr: (%T, %T)", change.To, change.From)
+					}
+					s.append(s.schemaComment(c.S, to.Text, from.Text))
+				default:
+					return nil, fmt.Errorf("unsupported ModifySchema change: %T", change)
+				}
+			}
 		case *schema.DropSchema:
 			b := s.Build("DROP SCHEMA")
 			if sqlx.Has(c.Extra, &schema.IfExists{}) {
@@ -713,6 +738,15 @@ func (s *state) addComments(t *schema.Table) {
 	}
 }
 
+func (s *state) schemaComment(sc *schema.Schema, to, from string) *migrate.Change {
+	b := s.Build("COMMENT ON SCHEMA").Ident(sc.Name).P("IS")
+	return &migrate.Change{
+		Cmd:     b.Clone().P(quote(to)).String(),
+		Comment: fmt.Sprintf("set comment to schema: %q", sc.Name),
+		Reverse: b.Clone().P(quote(from)).String(),
+	}
+}
+
 func (s *state) tableComment(t *schema.Table, to, from string) *migrate.Change {
 	b := s.Build("COMMENT ON TABLE").Table(t).P("IS")
 	return &migrate.Change{


Test patch:
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index 5c4302a8..4094c0d6 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -81,8 +81,10 @@ func testAddDrop(t T) {
 	t.migrate(&schema.AddTable{T: petsT}, &schema.AddTable{T: usersT}, &schema.AddTable{T: postsT})
 	ensureNoChange(t, usersT, petsT, postsT)
 	t.migrate(&schema.DropTable{T: usersT}, &schema.DropTable{T: postsT}, &schema.DropTable{T: petsT})
-	// Ensure the realm is empty.
-	require.EqualValues(t, t.realm(), t.loadRealm())
+	// Ensure the realm has no tables.
+	for _, s := range t.loadRealm().Schemas {
+		require.Empty(t, s.Tables)
+	}
 }
 
 func testRelation(t T) {
diff --git a/internal/integration/postgres_test.go b/internal/integration/postgres_test.go
index 1948d924..6fc1d9c1 100644
--- a/internal/integration/postgres_test.go
+++ b/internal/integration/postgres_test.go
@@ -523,6 +523,7 @@ func TestPostgres_HCL_Realm(t *testing.T) {
 		require.NoError(t, err)
 		wa := string(hcl) + `
 schema "second" {
+  comment = "second schema"
 }
 `
 		t.applyRealmHcl(wa)
@@ -530,8 +531,10 @@ schema "second" {
 		require.NoError(t, err)
 		_, ok := realm.Schema("public")
 		require.True(t, ok)
-		_, ok = realm.Schema("second")
+		s2, ok := realm.Schema("second")
 		require.True(t, ok)
+		require.Len(t, s2.Attrs, 1)
+		require.Equal(t, "second schema", s2.Attrs[0].(*schema.Comment).Text)
 	})
 }
 
@@ -847,6 +850,7 @@ enum "status" {
 schema "other" {
 }
 schema "public" {
+  comment = "standard public schema"
 }
 `, string(out))
 		err = t.drv.(migrate.CleanChecker).CheckClean(context.Background(), nil)
@@ -958,7 +962,9 @@ func TestPostgres_SchemaDiff(t *testing.T) {
 
 		// Create schemas on test2 database.
 		hcl := `
-schema "public" {}
+schema "public" {
+  comment = "standard public schema"
+}
 table "users" {
 	schema = schema.public
 	column "id" { type = integer }
@@ -986,8 +992,11 @@ CREATE SCHEMA "other";
 CREATE TABLE "public"."users" ("id" integer NOT NULL);
 -- Create "posts" table
 CREATE TABLE "other"."posts" ("id" integer NOT NULL);`, diff("test1?sslmode=disable", "test2?sslmode=disable"))
-		// diff schemas
-		require.Equal(t, `-- Drop "posts" table
+
+		// Diffing schema should both tables and comments (from 'public' to 'other').
+		require.Equal(t, `-- Set comment to schema: ""
+COMMENT ON SCHEMA IS 'standard public schema';
+-- Drop "posts" table
 DROP TABLE "posts";
 -- Create "users" table
 CREATE TABLE "users" ("id" integer NOT NULL);`, diff("test2?sslmode=disable&search_path=other", "test2?sslmode=disable&search_path=public"))
@@ -1532,6 +1541,9 @@ func (t *pgTest) realm() *schema.Realm {
 		Schemas: []*schema.Schema{
 			{
 				Name: "public",
+				Attrs: []schema.Attr{
+					&schema.Comment{Text: "standard public schema"},
+				},
 			},
 		},
 		Attrs: []schema.Attr{
diff --git a/sql/postgres/diff_test.go b/sql/postgres/diff_test.go
index 14ac707a..cb741052 100644
--- a/sql/postgres/diff_test.go
+++ b/sql/postgres/diff_test.go
@@ -502,6 +502,26 @@ func TestDiff_SchemaDiff(t *testing.T) {
 		&schema.DropTable{T: from.Tables[1]},
 		&schema.AddTable{T: to.Tables[1]},
 	}, changes)
+
+	// Add comment.
+	from, to = schema.New("public"), schema.New("public").SetComment("comment")
+	changes, err = drv.SchemaDiff(from, to)
+	require.NoError(t, err)
+	require.EqualValues(t, []schema.Change{
+		&schema.ModifySchema{S: to, Changes: schema.Changes{&schema.AddAttr{A: &schema.Comment{Text: "comment"}}}},
+	}, changes)
+	// Modify comment.
+	from.SetComment("boring")
+	changes, err = drv.SchemaDiff(from, to)
+	require.NoError(t, err)
+	require.EqualValues(t, []schema.Change{
+		&schema.ModifySchema{S: to, Changes: schema.Changes{
+			&schema.ModifyAttr{
+				From: &schema.Comment{Text: "boring"},
+				To:   &schema.Comment{Text: "comment"},
+			},
+		}},
+	}, changes)
 }
 
 func TestDefaultDiff(t *testing.T) {
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index a0f87421..2ce1d7b9 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -340,9 +340,9 @@ users        | users_check1       | (((c2 + c1) + c3) > 10) | c3          | {2,1
 			mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= $1"))).
 				WithArgs("public").
 				WillReturnRows(sqltest.Rows(`
-    schema_name
---------------------
- public
+ schema_name | comment 
+-------------+---------
+ public      | nil
 `))
 			tt.before(mk)
 			s, err := drv.InspectSchema(context.Background(), "public", &schema.InspectOptions{
@@ -363,9 +363,9 @@ func TestDriver_InspectPartitionedTable(t *testing.T) {
 	require.NoError(t, err)
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= CURRENT_SCHEMA()"))).
 		WillReturnRows(sqltest.Rows(`
-   schema_name
---------------------
-public
+ schema_name | comment 
+-------------+---------
+ public      | nil
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1"))).
 		WithArgs("public").
@@ -443,9 +443,9 @@ func TestDriver_InspectCRDBSchema(t *testing.T) {
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= $1"))).
 		WithArgs("public").
 		WillReturnRows(sqltest.Rows(`
-schema_name
---------------------
-public
+ schema_name | comment 
+-------------+---------
+ public      | nil
 `))
 	mk.tableExists("public", "users", true)
 	mk.ExpectQuery(queryCRDBColumns).
@@ -510,9 +510,9 @@ func TestDriver_InspectSchema(t *testing.T) {
 	require.NoError(t, err)
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= CURRENT_SCHEMA()"))).
 		WillReturnRows(sqltest.Rows(`
-   schema_name
---------------------
-test
+ schema_name | comment 
+-------------+---------
+ test        | boring
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1"))).
 		WithArgs("test").
@@ -525,9 +525,7 @@ test
 	require.EqualValues(t, func() *schema.Schema {
 		r := &schema.Realm{
 			Schemas: []*schema.Schema{
-				{
-					Name: "test",
-				},
+				schema.New("test").SetComment("boring"),
 			},
 			// Server default configuration.
 			Attrs: []schema.Attr{
@@ -553,10 +551,10 @@ func TestDriver_Realm(t *testing.T) {
 	require.NoError(t, err)
 	mk.ExpectQuery(sqltest.Escape(schemasQuery)).
 		WillReturnRows(sqltest.Rows(`
-   schema_name
---------------------
-test
-public
+ schema_name | comment 
+-------------+---------
+ test        | nil
+ public      | nil
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1, $2"))).
 		WithArgs("test", "public").
@@ -595,10 +593,10 @@ public
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "IN ($1, $2)"))).
 		WithArgs("test", "public").
 		WillReturnRows(sqltest.Rows(`
-   schema_name
---------------------
-  test
-  public
+ schema_name | comment 
+-------------+---------
+ test        | nil
+ public      | nil
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1, $2"))).
 		WithArgs("test", "public").
@@ -638,9 +636,9 @@ public
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= $1"))).
 		WithArgs("test").
 		WillReturnRows(sqltest.Rows(`
- schema_name
---------------------
- test
+ schema_name | comment 
+-------------+---------
+ test        | nil
 `))
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(tablesQuery, "$1"))).
 		WithArgs("test").
@@ -680,10 +678,10 @@ func TestInspectMode_InspectRealm(t *testing.T) {
 	mk.version("130000")
 	mk.ExpectQuery(sqltest.Escape(schemasQuery)).
 		WillReturnRows(sqltest.Rows(`
-   schema_name
---------------------
-test
-public
+ schema_name | comment 
+-------------+---------
+ test        | nil
+ public      | nil
 `))
 	drv, err := Open(db)
 	m.ExpectQuery(sqltest.Escape(fmt.Sprintf(enumsQuery, "$1, $2"))).
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 04aa76fc..6bf90491 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -27,10 +27,12 @@ func TestPlanChanges(t *testing.T) {
 	}{
 		{
 			changes: []schema.Change{
-				&schema.AddSchema{S: schema.New("public")},
+				&schema.AddSchema{S: schema.New("public").SetComment("public schema")},
 				&schema.AddSchema{S: schema.New("test"), Extra: []schema.Clause{&schema.IfNotExists{}}},
 				&schema.DropSchema{S: schema.New("test"), Extra: []schema.Clause{&schema.IfExists{}}},
 				&schema.DropSchema{S: schema.New("test"), Extra: []schema.Clause{}},
+				&schema.ModifySchema{S: schema.New("modify1").SetComment("comment"), Changes: schema.Changes{&schema.AddAttr{A: &schema.Comment{Text: "comment"}}}},
+				&schema.ModifySchema{S: schema.New("modify2").SetComment("comment"), Changes: schema.Changes{&schema.ModifyAttr{From: &schema.Comment{Text: "old"}, To: &schema.Comment{Text: "new"}}}},
 			},
 			wantPlan: &migrate.Plan{
 				Reversible:    false,
@@ -40,6 +42,10 @@ func TestPlanChanges(t *testing.T) {
 						Cmd:     `CREATE SCHEMA IF NOT EXISTS "public"`,
 						Reverse: `DROP SCHEMA "public" CASCADE`,
 					},
+					{
+						Cmd:     `COMMENT ON SCHEMA "public" IS 'public schema'`,
+						Reverse: `COMMENT ON SCHEMA "public" IS ''`,
+					},
 					{
 						Cmd:     `CREATE SCHEMA IF NOT EXISTS "test"`,
 						Reverse: `DROP SCHEMA "test" CASCADE`,
@@ -50,6 +56,14 @@ func TestPlanChanges(t *testing.T) {
 					{
 						Cmd: `DROP SCHEMA "test" CASCADE`,
 					},
+					{
+						Cmd:     `COMMENT ON SCHEMA "modify1" IS 'comment'`,
+						Reverse: `COMMENT ON SCHEMA "modify1" IS ''`,
+					},
+					{
+						Cmd:     `COMMENT ON SCHEMA "modify2" IS 'new'`,
+						Reverse: `COMMENT ON SCHEMA "modify2" IS 'old'`,
+					},
 				},
 			},
 		},
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 3348053c..9e018df2 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -262,6 +262,48 @@ enum "account_type" {
 	require.EqualValues(t, exp, &s)
 }
 
+func TestMarshalSpec_Schema(t *testing.T) {
+	s := schema.New("public").SetComment("schema comment")
+	buf, err := MarshalHCL(s)
+	require.NoError(t, err)
+	require.Equal(t, `schema "public" {
+  comment = "schema comment"
+}
+`, string(buf))
+
+	r := schema.NewRealm(
+		schema.New("s1").SetComment("c1"),
+		schema.New("s2").SetComment("c2"),
+		schema.New("s3"),
+	)
+	buf, err = MarshalHCL(r)
+	require.NoError(t, err)
+	require.Equal(t, `schema "s1" {
+  comment = "c1"
+}
+schema "s2" {
+  comment = "c2"
+}
+schema "s3" {
+}
+`, string(buf))
+}
+
+func TestUnmarshalSpec_Schema(t *testing.T) {
+	var (
+		s schema.Schema
+		f = `
+schema "public" {
+  comment = "schema comment"
+}
+`
+	)
+	require.NoError(t, EvalHCLBytes([]byte(f), &s, nil))
+	require.Equal(t, "public", s.Name)
+	require.Len(t, s.Attrs, 1)
+	require.Equal(t, "schema comment", s.Attrs[0].(*schema.Comment).Text)
+}
+
 func TestMarshalViews(t *testing.T) {
 	s := schema.New("public").
 		AddTables(