Repository: ariga/atlas
Pull request #1861 (2023-07-16 03:20:49)

sql/sqlx: add TableResource and Quote helpers


Patch:
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index 6eda80e5..84e9dc91 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -261,9 +261,20 @@ func (b *Builder) Table(t *schema.Table) *Builder {
 	return b.mayQualify(t.Schema, t.Name)
 }
 
-// Table writes the table identifier to the builder, prefixed
+// TableResource writes the table's resource identifier to the builder, prefixed
 // with the schema name if exists.
-func (b *Builder) mayQualify(s *schema.Schema, ident string) *Builder {
+func (b *Builder) TableResource(t *schema.Table, r any) *Builder {
+	switch c := r.(type) {
+	case *schema.Column:
+		return b.mayQualify(t.Schema, t.Name, c.Name)
+	case *schema.Index:
+		return b.mayQualify(t.Schema, t.Name, c.Name)
+	default:
+		panic(fmt.Sprintf("unexpected table resource: %T", r))
+	}
+}
+
+func (b *Builder) mayQualify(s *schema.Schema, top string, children ...string) *Builder {
 	switch {
 	// Custom qualifier.
 	case b.Schema != nil:
@@ -277,7 +288,11 @@ func (b *Builder) mayQualify(s *schema.Schema, ident string) *Builder {
 		b.Ident(s.Name)
 		b.rewriteLastByte('.')
 	}
-	b.Ident(ident)
+	b.Ident(top)
+	for _, ident := range children {
+		b.rewriteLastByte('.')
+		b.Ident(ident)
+	}
 	return b
 }
 
@@ -334,6 +349,19 @@ func (b *Builder) MapComma(x any, f func(i int, b *Builder)) *Builder {
 	return b
 }
 
+// Quote wraps the given function with a single quote and a prefix
+func (b *Builder) Quote(prefix string, fn func(b *Builder)) *Builder {
+	b.WriteString(prefix)
+	b.WriteByte('\'')
+	fn(b)
+	if b.lastByte() != ' ' {
+		b.WriteByte('\'')
+	} else {
+		b.rewriteLastByte('\'')
+	}
+	return b
+}
+
 // MapIndent is like MapComma, but writes a new line before each element.
 func (b *Builder) MapIndent(x any, f func(i int, b *Builder)) *Builder {
 	return b.MapComma(x, func(i int, b *Builder) {
diff --git a/sql/postgres/diff.go b/sql/postgres/diff.go
index 083d6fc3..cb35d9ab 100644
--- a/sql/postgres/diff.go
+++ b/sql/postgres/diff.go
@@ -29,12 +29,20 @@ type diff struct{ conn }
 // SchemaAttrDiff returns a changeset for migrating schema attributes from one state to the other.
 func (*diff) SchemaAttrDiff(from, to *schema.Schema) []schema.Change {
 	var changes []schema.Change
-	if change := sqlx.CommentDiff(from.Attrs, to.Attrs); change != nil {
+	if change := sqlx.CommentDiff(skipDefaultComment(from), skipDefaultComment(to)); change != nil {
 		changes = append(changes, change)
 	}
 	return changes
 }
 
+func skipDefaultComment(s *schema.Schema) []schema.Attr {
+	attrs := s.Attrs
+	if c := (schema.Comment{}); sqlx.Has(attrs, &c) && c.Text == "standard public schema" && (s.Name == "" || s.Name == "public") {
+		attrs = schema.RemoveAttr[*schema.Comment](attrs)
+	}
+	return attrs
+}
+
 // SchemaObjectDiff returns a changeset for migrating schema objects from
 // one state to the other.
 func (*diff) SchemaObjectDiff(from, to *schema.Schema) ([]schema.Change, error) {


Test patch:
diff --git a/internal/integration/postgres_test.go b/internal/integration/postgres_test.go
index 259d6301..bea11297 100644
--- a/internal/integration/postgres_test.go
+++ b/internal/integration/postgres_test.go
@@ -991,9 +991,7 @@ CREATE TABLE "public"."users" ("id" integer NOT NULL);
 CREATE TABLE "other"."posts" ("id" integer NOT NULL);`, diff("test1?sslmode=disable", "test2?sslmode=disable"))
 
 		// 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';
--- Create "users" table
+		require.Equal(t, `-- Create "users" table
 CREATE TABLE "users" ("id" integer NOT NULL);
 -- Drop "posts" table
 DROP TABLE "posts";`, diff("test2?sslmode=disable&search_path=other", "test2?sslmode=disable&search_path=public"))
diff --git a/sql/internal/sqlx/sqlx_test.go b/sql/internal/sqlx/sqlx_test.go
index 96708802..6b51a988 100644
--- a/sql/internal/sqlx/sqlx_test.go
+++ b/sql/internal/sqlx/sqlx_test.go
@@ -79,9 +79,21 @@ func TestBuilder_Qualifier(t *testing.T) {
 	b.Reset()
 	b.P("CREATE TABLE").Table(schema.NewTable("users").SetSchema(schema.New("test")))
 	require.Equal(t, `CREATE TABLE "users"`, b.String())
-
 }
 
+func TestQuote(t *testing.T) {
+	var (
+		s = "s1"
+		b = &Builder{QuoteChar: '"', Schema: &s}
+	)
+	b.P("EXECUTE sp_rename").
+		P("@newname = N'c2'").Comma().
+		P("@objtype = N'COLUMN'").Comma()
+	b.P("@objname = ").Quote("N", func(b *Builder) {
+		b.TableResource(schema.NewTable("t1"), &schema.Column{Name: "c1"})
+	})
+	require.Equal(t, `EXECUTE sp_rename @newname = N'c2', @objtype = N'COLUMN', @objname = N'"s1"."t1"."c1"'`, b.String())
+}
 func TestMayWrap(t *testing.T) {
 	tests := []struct {
 		input   string
diff --git a/sql/postgres/diff_test.go b/sql/postgres/diff_test.go
index cb741052..cceda18f 100644
--- a/sql/postgres/diff_test.go
+++ b/sql/postgres/diff_test.go
@@ -522,6 +522,48 @@ func TestDiff_SchemaDiff(t *testing.T) {
 			},
 		}},
 	}, changes)
+
+	t.Run("DefaultComment", func(t *testing.T) {
+		from, to := schema.New("public").SetComment("standard public schema"), schema.New("public")
+		changes, err = drv.SchemaDiff(from, to)
+		require.NoError(t, err)
+		require.Empty(t, changes)
+
+		changes, err = drv.SchemaDiff(to, from)
+		require.NoError(t, err)
+		require.Empty(t, changes)
+
+		from.Name, to.Name = "", ""
+		changes, err = drv.SchemaDiff(from, to)
+		require.NoError(t, err)
+		require.Empty(t, changes)
+
+		from.Name, to.Name = "other", "other"
+		changes, err = drv.SchemaDiff(from, to)
+		require.NoError(t, err)
+		require.Len(t, changes, 1)
+		require.EqualValues(t, []schema.Change{
+			&schema.ModifySchema{S: to, Changes: schema.Changes{
+				&schema.ModifyAttr{
+					From: &schema.Comment{Text: "standard public schema"},
+					To:   &schema.Comment{Text: ""},
+				},
+			}},
+		}, changes)
+
+		from.Name, to.Name = "public", "public"
+		from.SetComment("custom public schema")
+		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: "custom public schema"},
+					To:   &schema.Comment{Text: ""},
+				},
+			}},
+		}, changes)
+	})
 }
 
 func TestDefaultDiff(t *testing.T) {