Repository: ariga/atlas
Pull request #427 (2022-01-17 09:13:56)

sql/postgres: minor fixes to schema/tables inspection


Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index a8ab32d2..5d298bc2 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -8,6 +8,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strconv"
 	"strings"
 
 	"ariga.io/atlas/sql/internal/sqlx"
@@ -456,12 +457,10 @@ func (i *inspect) addChecks(t *schema.Table, rows *sql.Rows) error {
 func (i *inspect) schemas(ctx context.Context, opts *schema.InspectRealmOption) ([]*schema.Schema, error) {
 	var (
 		args  []interface{}
-		query = schemasQuery + " WHERE schema_name "
+		query = schemasQuery
 	)
 	if opts != nil && len(opts.Schemas) > 0 {
-		query, args = inStrings(opts.Schemas, query, args)
-	} else {
-		query += "NOT IN ('information_schema', 'pg_catalog', 'pg_toast')"
+		query, args = inStrings(opts.Schemas, schemasQueryArgs, args)
 	}
 	rows, err := i.QueryContext(ctx, query, args...)
 	if err != nil {
@@ -485,8 +484,7 @@ func (i *inspect) schemas(ctx context.Context, opts *schema.InspectRealmOption)
 func (i *inspect) tableNames(ctx context.Context, schema string, opts *schema.InspectOptions) ([]string, error) {
 	query, args := tablesQuery, []interface{}{schema}
 	if opts != nil && len(opts.Tables) > 0 {
-		query += " AND table_name "
-		query, args = inStrings(opts.Tables, query, args)
+		query, args = inStrings(opts.Tables, tablesQueryArgs, args)
 	}
 	rows, err := i.QueryContext(ctx, query, args...)
 	if err != nil {
@@ -500,16 +498,25 @@ func (i *inspect) tableNames(ctx context.Context, schema string, opts *schema.In
 }
 
 func inStrings(s []string, query string, args []interface{}) (string, []interface{}) {
-	query += "IN ("
-	for i := range s {
-		args = append(args, s[i])
-		if i > 0 {
-			query += ", "
+	var b strings.Builder
+	switch len(s) {
+	case 1:
+		args = append(args, s[0])
+		b.WriteString("= $")
+		b.WriteString(strconv.Itoa(len(args)))
+	default:
+		b.WriteString("IN (")
+		for i := range s {
+			args = append(args, s[i])
+			if i > 0 {
+				b.WriteString(", ")
+			}
+			b.WriteByte('$')
+			b.WriteString(strconv.Itoa(len(args)))
 		}
-		query += fmt.Sprintf("$%d", len(args))
+		b.WriteByte(')')
 	}
-	query += ")"
-	return query, args
+	return fmt.Sprintf(query, b.String()), args
 }
 
 func defaultExpr(c *schema.Column, x string) schema.Expr {
@@ -684,11 +691,17 @@ const (
 	paramsQuery = `SELECT setting FROM pg_settings WHERE name IN ('lc_collate', 'lc_ctype', 'server_version_num') ORDER BY name`
 
 	// Query to list database schemas.
-	schemasQuery = "SELECT schema_name FROM information_schema.schemata"
+	schemasQuery = "SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'pg_toast') ORDER BY schema_name"
+
+	// 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 schema tables.
 	tablesQuery = "SELECT table_name FROM information_schema.tables WHERE table_type = 'BASE TABLE' AND table_schema = $1 ORDER BY table_name"
 
+	// Query to list specific schema tables.
+	tablesQueryArgs = "SELECT table_name FROM information_schema.tables WHERE table_type = 'BASE TABLE' AND table_schema = $1 AND table_name %s ORDER BY table_name"
+
 	// Query to list table information.
 	tableQuery = `
 SELECT


Test patch:
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 07f4bd89..ae95eb08 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -6,6 +6,7 @@ package postgres
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"ariga.io/atlas/sql/internal/sqltest"
@@ -309,7 +310,7 @@ func TestDriver_Realm(t *testing.T) {
 	mk.version("130000")
 	drv, err := Open(db)
 	require.NoError(t, err)
-	mk.ExpectQuery(sqltest.Escape(schemasQuery + " WHERE schema_name NOT IN ('information_schema', 'pg_catalog', 'pg_toast')")).
+	mk.ExpectQuery(sqltest.Escape(schemasQuery)).
 		WillReturnRows(sqltest.Rows(`
     schema_name
 --------------------
@@ -345,7 +346,7 @@ func TestDriver_Realm(t *testing.T) {
 		return r
 	}(), realm)
 
-	mk.ExpectQuery(sqltest.Escape(schemasQuery+" WHERE schema_name IN ($1, $2)")).
+	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "IN ($1, $2)"))).
 		WithArgs("test", "public").
 		WillReturnRows(sqltest.Rows(`
     schema_name
@@ -381,6 +382,37 @@ func TestDriver_Realm(t *testing.T) {
 		r.Schemas[1].Realm = r
 		return r
 	}(), realm)
+
+	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= $1"))).
+		WithArgs("test").
+		WillReturnRows(sqltest.Rows(`
+    schema_name
+--------------------
+ test
+`))
+	mk.tables("test")
+	realm, err = drv.InspectRealm(context.Background(), &schema.InspectRealmOption{Schemas: []string{"test"}})
+	require.NoError(t, err)
+	require.EqualValues(t, func() *schema.Realm {
+		r := &schema.Realm{
+			Schemas: []*schema.Schema{
+				{
+					Name: "test",
+				},
+			},
+			// Server default configuration.
+			Attrs: []schema.Attr{
+				&schema.Collation{
+					V: "en_US.utf8",
+				},
+				&CType{
+					V: "en_US.utf8",
+				},
+			},
+		}
+		r.Schemas[0].Realm = r
+		return r
+	}(), realm)
 }
 
 type mock struct {