Repository: ariga/atlas
Pull request #431 (2022-01-17 16:07:40)

sql:  add support for querying attached schema


Patch:
diff --git a/sql/mysql/inspect.go b/sql/mysql/inspect.go
index 4db18e0a..974467f6 100644
--- a/sql/mysql/inspect.go
+++ b/sql/mysql/inspect.go
@@ -46,11 +46,10 @@ func (i *inspect) InspectRealm(ctx context.Context, opts *schema.InspectRealmOpt
 	return realm, nil
 }
 
-// InspectSchema returns schema descriptions of all tables in the given schema.
+// InspectSchema returns schema descriptions of the tables in the given schema.
+// If the schema name is empty, the result will be the attached schema.
 func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.InspectOptions) (*schema.Schema, error) {
-	schemas, err := i.schemas(ctx, &schema.InspectRealmOption{
-		Schemas: []string{name},
-	})
+	schemas, err := i.schemas(ctx, &schema.InspectRealmOption{Schemas: []string{name}})
 	if err != nil {
 		return nil, err
 	}
@@ -59,11 +58,11 @@ func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.I
 			Err: fmt.Errorf("mysql: schema %q was not found", name),
 		}
 	}
-	names, err := i.tableNames(ctx, name, opts)
+	s := schemas[0]
+	names, err := i.tableNames(ctx, s.Name, opts)
 	if err != nil {
 		return nil, err
 	}
-	s := schemas[0]
 	for _, name := range names {
 		t, err := i.inspectTable(ctx, name, &schema.InspectTableOptions{Schema: s.Name}, s)
 		if err != nil {
@@ -111,13 +110,19 @@ func (i *inspect) schemas(ctx context.Context, opts *schema.InspectRealmOption)
 		args  []interface{}
 		query = schemasQuery
 	)
-	if opts != nil && len(opts.Schemas) > 0 {
-		query += " WHERE `SCHEMA_NAME` IN (" + strings.Repeat("?, ", len(opts.Schemas)-1) + "?)"
-		for _, s := range opts.Schemas {
-			args = append(args, s)
+	if opts != nil {
+		switch n := len(opts.Schemas); {
+		case n == 1 && opts.Schemas[0] == "":
+			query = fmt.Sprintf(schemasQueryArgs, "= SCHEMA()")
+		case n == 1 && opts.Schemas[0] != "":
+			query = fmt.Sprintf(schemasQueryArgs, "= ?")
+			args = append(args, opts.Schemas[0])
+		case n > 0:
+			query = fmt.Sprintf(schemasQueryArgs, "IN ("+strings.Repeat("?, ", len(opts.Schemas)-1)+"?)")
+			for _, s := range opts.Schemas {
+				args = append(args, s)
+			}
 		}
-	} else {
-		query += " WHERE `SCHEMA_NAME` NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')"
 	}
 	rows, err := i.QueryContext(ctx, query, args...)
 	if err != nil {
@@ -583,7 +588,10 @@ const (
 	variablesQuery = "SELECT @@version, @@collation_server, @@character_set_server"
 
 	// Query to list database schemas.
-	schemasQuery = "SELECT `SCHEMA_NAME`, `DEFAULT_CHARACTER_SET_NAME`, `DEFAULT_COLLATION_NAME` from `INFORMATION_SCHEMA`.`SCHEMATA`"
+	schemasQuery = "SELECT `SCHEMA_NAME`, `DEFAULT_CHARACTER_SET_NAME`, `DEFAULT_COLLATION_NAME` from `INFORMATION_SCHEMA`.`SCHEMATA` WHERE `SCHEMA_NAME` NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys') ORDER BY `SCHEMA_NAME`"
+
+	// Query to list specific database schemas.
+	schemasQueryArgs = "SELECT `SCHEMA_NAME`, `DEFAULT_CHARACTER_SET_NAME`, `DEFAULT_COLLATION_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` = ? ORDER BY `TABLE_NAME`"
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 5d298bc2..9b4ef947 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -45,24 +45,31 @@ func (i *inspect) InspectRealm(ctx context.Context, opts *schema.InspectRealmOpt
 	return realm, nil
 }
 
-// InspectSchema returns schema descriptions of all tables in the given schema.
-func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.InspectOptions) (*schema.Schema, error) {
-	schemas, err := i.schemas(ctx, &schema.InspectRealmOption{
-		Schemas: []string{name},
-	})
-	if err != nil {
-		return nil, err
-	}
-	if len(schemas) == 0 {
-		return nil, &schema.NotExistError{
-			Err: fmt.Errorf("postgres: schema %q was not found", name),
+// InspectSchema returns schema descriptions of the tables in the given schema.
+// If the schema name is empty, the result will be the attached schema.
+func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.InspectOptions) (s *schema.Schema, err error) {
+	var schemas []*schema.Schema
+	switch name {
+	case "":
+		if err := i.QueryRowContext(ctx, "SELECT CURRENT_SCHEMA()").Scan(&name); err != nil {
+			return nil, fmt.Errorf("postgres: failed querying attached schema: %w", err)
+		}
+		schemas = append(schemas, &schema.Schema{Name: name})
+	default:
+		if schemas, err = i.schemas(ctx, &schema.InspectRealmOption{Schemas: []string{name}}); err != nil {
+			return nil, err
+		}
+		if len(schemas) == 0 {
+			return nil, &schema.NotExistError{
+				Err: fmt.Errorf("postgres: schema %q was not found", name),
+			}
 		}
 	}
 	names, err := i.tableNames(ctx, name, opts)
 	if err != nil {
 		return nil, err
 	}
-	s := schemas[0]
+	s = schemas[0]
 	for _, name := range names {
 		t, err := i.inspectTable(ctx, name, &schema.InspectTableOptions{Schema: s.Name}, s)
 		if err != nil {
diff --git a/sql/schema/inspect.go b/sql/schema/inspect.go
index f3913c67..aa13f4d0 100644
--- a/sql/schema/inspect.go
+++ b/sql/schema/inspect.go
@@ -58,8 +58,9 @@ type (
 	// Inspector is the interface implemented by the different database
 	// drivers for inspecting multiple tables.
 	Inspector interface {
-		// InspectSchema returns the schema description by its name. A NotExistError
-		// error is returned if the schema does not exists in the database.
+		// InspectSchema returns the schema description by its name. An empty name means the
+		// "attached schema" (e.g. SCHEMA() in MySQL or CURRENT_SCHEMA() in PostgreSQL).
+		// A NotExistError error is returned if the schema does not exists in the database.
 		InspectSchema(ctx context.Context, name string, opts *InspectOptions) (*Schema, error)
 
 		// InspectTable returns the table description by its name. A NotExistError
diff --git a/sql/sqlite/inspect.go b/sql/sqlite/inspect.go
index f7732fb8..80dc9c4d 100644
--- a/sql/sqlite/inspect.go
+++ b/sql/sqlite/inspect.go
@@ -51,8 +51,12 @@ func (i *inspect) InspectRealm(ctx context.Context, opts *schema.InspectRealmOpt
 	return realm, nil
 }
 
-// InspectSchema returns schema descriptions of all tables in the given schema.
+// InspectSchema returns schema descriptions of the tables in the given schema.
+// If the schema name is empty, the "main" database is used.
 func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.InspectOptions) (*schema.Schema, error) {
+	if name == "" {
+		name = mainFile
+	}
 	schemas, err := i.databases(ctx, &schema.InspectRealmOption{
 		Schemas: []string{name},
 	})
@@ -84,10 +88,10 @@ func (i *inspect) InspectSchema(ctx context.Context, name string, opts *schema.I
 
 // InspectTable returns the schema description of the given table.
 func (i *inspect) InspectTable(ctx context.Context, name string, opts *schema.InspectTableOptions) (*schema.Table, error) {
-	if opts != nil && opts.Schema != "main" {
+	if opts != nil && opts.Schema != mainFile {
 		return nil, fmt.Errorf("sqlite: querying attached database is not supported. got: %q", opts.Schema)
 	}
-	s, err := i.InspectSchema(ctx, "main", &schema.InspectOptions{
+	s, err := i.InspectSchema(ctx, mainFile, &schema.InspectOptions{
 		Tables: []string{name},
 	})
 	if err != nil {
@@ -648,6 +652,8 @@ func fillChecks(t *schema.Table) error {
 }
 
 const (
+	// Name of main database file.
+	mainFile = "main"
 	// Query to list attached database files.
 	databasesQuery = "SELECT `name`, `file` FROM pragma_database_list()"
 	// Query to list database tables.


Test patch:
diff --git a/sql/mysql/inspect_test.go b/sql/mysql/inspect_test.go
index 74ef79cd..e4b7a727 100644
--- a/sql/mysql/inspect_test.go
+++ b/sql/mysql/inspect_test.go
@@ -6,6 +6,7 @@ package mysql
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"ariga.io/atlas/sql/internal/sqltest"
@@ -625,15 +626,16 @@ func TestDriver_InspectTable(t *testing.T) {
 func TestDriver_InspectSchema(t *testing.T) {
 	tests := []struct {
 		name   string
+		schema string
 		opts   *schema.InspectOptions
 		before func(mock)
 		expect func(*require.Assertions, *schema.Schema, error)
 	}{
 		{
-			name: "no tables",
+			name: "attached schema",
 			before: func(m mock) {
 				m.version("5.7.23")
-				m.ExpectQuery(sqltest.Escape(schemasQuery + " WHERE `SCHEMA_NAME` IN (?)")).
+				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= SCHEMA()"))).
 					WillReturnRows(sqltest.Rows(`
 +-------------+----------------------------+------------------------+
 | SCHEMA_NAME | DEFAULT_CHARACTER_SET_NAME | DEFAULT_COLLATION_NAME |
@@ -671,10 +673,12 @@ func TestDriver_InspectSchema(t *testing.T) {
 			},
 		},
 		{
-			name: "multi table",
+			name:   "multi table",
+			schema: "public",
 			before: func(m mock) {
 				m.version("8.0.13")
-				m.ExpectQuery(sqltest.Escape(schemasQuery + " WHERE `SCHEMA_NAME` IN (?)")).
+				m.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= ?"))).
+					WithArgs("public").
 					WillReturnRows(sqltest.Rows(`
 +-------------+----------------------------+------------------------+
 | SCHEMA_NAME | DEFAULT_CHARACTER_SET_NAME | DEFAULT_COLLATION_NAME |
@@ -767,7 +771,7 @@ func TestDriver_InspectSchema(t *testing.T) {
 			tt.before(mock{m})
 			drv, err := Open(db)
 			require.NoError(t, err)
-			tables, err := drv.InspectSchema(context.Background(), "public", tt.opts)
+			tables, err := drv.InspectSchema(context.Background(), tt.schema, tt.opts)
 			tt.expect(require.New(t), tables, err)
 		})
 	}
@@ -778,7 +782,7 @@ func TestDriver_Realm(t *testing.T) {
 	require.NoError(t, err)
 	mk := mock{m}
 	mk.version("8.0.13")
-	mk.ExpectQuery(sqltest.Escape(schemasQuery + " WHERE `SCHEMA_NAME` NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys')")).
+	mk.ExpectQuery(sqltest.Escape(schemasQuery)).
 		WillReturnRows(sqltest.Rows(`
 +-------------+----------------------------+------------------------+
 | SCHEMA_NAME | DEFAULT_CHARACTER_SET_NAME | DEFAULT_COLLATION_NAME |
@@ -816,7 +820,7 @@ func TestDriver_Realm(t *testing.T) {
 		return r
 	}(), realm)
 
-	mk.ExpectQuery(sqltest.Escape(schemasQuery+" WHERE `SCHEMA_NAME` IN (?, ?)")).
+	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "IN (?, ?)"))).
 		WithArgs("test", "public").
 		WillReturnRows(sqltest.Rows(`
 +-------------+----------------------------+------------------------+
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index ae95eb08..641f99be 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -303,6 +303,44 @@ func TestDriver_InspectTable(t *testing.T) {
 	}
 }
 
+func TestDriver_InspectSchema(t *testing.T) {
+	db, m, err := sqlmock.New()
+	require.NoError(t, err)
+	mk := mock{m}
+	mk.version("130000")
+	drv, err := Open(db)
+	require.NoError(t, err)
+	mk.ExpectQuery(sqltest.Escape("SELECT CURRENT_SCHEMA()")).
+		WillReturnRows(sqltest.Rows(`
+    schema_name
+--------------------
+ test
+`))
+	mk.tables("test")
+	s, err := drv.InspectSchema(context.Background(), "", &schema.InspectOptions{})
+	require.NoError(t, err)
+	require.EqualValues(t, func() *schema.Schema {
+		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.Schemas[0]
+	}(), s)
+}
+
 func TestDriver_Realm(t *testing.T) {
 	db, m, err := sqlmock.New()
 	require.NoError(t, err)