Fix issue where schema extraction is broken for remote services. Closes #1497

db/db_client/db_client.go

@@ -3,9 +3,7 @@ package db_client
 import (
 	"context"
 	"database/sql"
-	"fmt"
 	"log"
-	"strings"
 	"sync"
 
 	"github.com/jackc/pgx/v4"
@@ -184,13 +182,13 @@ func (c *DbClient) RefreshConnectionAndSearchPaths(ctx context.Context) *steampi
 	return res
 }
 
-func (c *DbClient) GetSchemaFromDB(ctx context.Context, schemas []string) (*schema.Metadata, error) {
+func (c *DbClient) GetSchemaFromDB(ctx context.Context) (*schema.Metadata, error) {
 	utils.LogTime("db_client.GetSchemaFromDB start")
 	defer utils.LogTime("db_client.GetSchemaFromDB end")
 	connection, err := c.dbClient.Conn(ctx)
 	utils.FailOnError(err)
 
-	query := c.buildSchemasQuery(schemas)
+	query := c.buildSchemasQuery()
 
 	tablesResult, err := connection.QueryContext(ctx, query)
 	if err != nil {
@@ -212,14 +210,15 @@ func (c *DbClient) GetSchemaFromDB(ctx context.Context, schemas []string) (*sche
 	return metadata, nil
 }
 
-func (c *DbClient) buildSchemasQuery(schemas []string) string {
-	schemasClause := ""
-	if len(schemas) > 0 {
-		schemasClause = fmt.Sprintf(`
-    cols.table_schema in ('%s')
-OR`, strings.Join(schemas, "','"))
-	}
-	query := fmt.Sprintf(`
+func (c *DbClient) buildSchemasQuery() string {
+	query := `
+WITH distinct_schema AS (
+	SELECT DISTINCT(foreign_table_schema) 
+	FROM 
+		information_schema.foreign_tables 
+	WHERE 
+		foreign_table_schema <> 'steampipe_command'
+)
 SELECT
     table_name,
     column_name,
@@ -235,12 +234,12 @@ LEFT JOIN
     pg_catalog.pg_namespace nsp ON nsp.nspname = cols.table_schema
 LEFT JOIN
     pg_catalog.pg_class c ON c.relname = cols.table_name AND c.relnamespace = nsp.oid
-WHERE %s
+WHERE
+	cols.table_schema in (select * from distinct_schema)
+	OR
     LEFT(cols.table_schema,8) = 'pg_temp_'
-ORDER BY
-    cols.table_schema, cols.table_name, cols.column_name;
 
-`, schemasClause)
+`
 	return query
 }
 

db/db_common/client.go

@@ -34,7 +34,7 @@ type Client interface {
 
 	SetEnsureSessionDataFunc(EnsureSessionStateCallback)
 	RefreshSessions(ctx context.Context) *AcquireSessionResult
-	GetSchemaFromDB(context.Context, []string) (*schema.Metadata, error)
+	GetSchemaFromDB(context.Context) (*schema.Metadata, error)
 	// remote client will have empty implementation
 	RefreshConnectionAndSearchPaths(context.Context) *steampipeconfig.RefreshConnectionResult
 	LoadForeignSchemaNames(context.Context) error

db/db_local/local_db_client.go

@@ -154,8 +154,90 @@ func (c *LocalDbClient) ContructSearchPath(ctx context.Context, requiredSearchPa
 	return c.client.ContructSearchPath(ctx, requiredSearchPath, searchPathPrefix, currentSearchPath)
 }
 
-func (c *LocalDbClient) GetSchemaFromDB(ctx context.Context, schemas []string) (*schema.Metadata, error) {
-	return c.client.GetSchemaFromDB(ctx, schemas)
+func (c *LocalDbClient) GetSchemaFromDB(ctx context.Context) (*schema.Metadata, error) {
+	// build a ConnectionSchemaMap object to identify the schemas to load
+	// (pass nil for connection state - this forces NewConnectionSchemaMap to load it)
+	connectionSchemaMap, err := steampipeconfig.NewConnectionSchemaMap()
+	if err != nil {
+		return nil, err
+	}
+	// get the unique schema - we use this to limit the schemas we load from the database
+	schemas := connectionSchemaMap.UniqueSchemas()
+	query := c.buildSchemasQuery(schemas)
+
+	acquireSessionResult := c.AcquireSession(ctx)
+	if acquireSessionResult.Error != nil {
+		acquireSessionResult.Session.Close(false)
+		return nil, err
+	}
+
+	tablesResult, err := acquireSessionResult.Session.Connection.QueryContext(ctx, query)
+	if err != nil {
+		return nil, err
+	}
+
+	metadata, err := db_common.BuildSchemaMetadata(tablesResult)
+	if err != nil {
+		acquireSessionResult.Session.Close(false)
+		return nil, err
+	}
+	acquireSessionResult.Session.Close(false)
+
+	c.populateSchemaMetadata(metadata, connectionSchemaMap)
+
+	searchPath, err := c.GetCurrentSearchPath(ctx)
+	if err != nil {
+		return nil, err
+	}
+	metadata.SearchPath = searchPath
+
+	return metadata, nil
+}
+
+func (c *LocalDbClient) populateSchemaMetadata(schemaMetadata *schema.Metadata, connectionSchemaMap steampipeconfig.ConnectionSchemaMap) {
+	// we now need to add in all other schemas which have the same schemas as those we have loaded
+	for loadedSchema, otherSchemas := range connectionSchemaMap {
+		// all 'otherSchema's have the same schema as loadedSchema
+		exemplarSchema, ok := schemaMetadata.Schemas[loadedSchema]
+		if !ok {
+			// should can happen in the case of a dynamic plugin with no tables - use empty schema
+			exemplarSchema = make(map[string]schema.TableSchema)
+		}
+
+		for _, s := range otherSchemas {
+			schemaMetadata.Schemas[s] = exemplarSchema
+		}
+	}
+}
+
+func (c *LocalDbClient) buildSchemasQuery(hintSchemas []string) string {
+	for idx, s := range hintSchemas {
+		hintSchemas[idx] = fmt.Sprintf("'%s'", s)
+	}
+	schemaClause := strings.Join(hintSchemas, ",")
+	query := fmt.Sprintf(`
+SELECT
+    table_name,
+    column_name,
+    column_default,
+    is_nullable,
+    data_type,
+    table_schema,
+    (COALESCE(pg_catalog.col_description(c.oid, cols.ordinal_position :: int),'')) as column_comment,
+    (COALESCE(pg_catalog.obj_description(c.oid),'')) as table_comment
+FROM
+    information_schema.columns cols
+LEFT JOIN
+    pg_catalog.pg_namespace nsp ON nsp.nspname = cols.table_schema
+LEFT JOIN
+    pg_catalog.pg_class c ON c.relname = cols.table_name AND c.relnamespace = nsp.oid
+WHERE
+	cols.table_schema in (%s)
+	OR
+    LEFT(cols.table_schema,8) = 'pg_temp_'
+
+`, schemaClause)
+	return query
 }
 
 func (c *LocalDbClient) LoadForeignSchemaNames(ctx context.Context) error {

interactive/interactive_client.go

@@ -25,7 +25,6 @@ import (
 	"github.com/turbot/steampipe/query/queryresult"
 	"github.com/turbot/steampipe/schema"
 	"github.com/turbot/steampipe/statushooks"
-	"github.com/turbot/steampipe/steampipeconfig"
 	"github.com/turbot/steampipe/steampipeconfig/modconfig"
 	"github.com/turbot/steampipe/utils"
 	"github.com/turbot/steampipe/version"
@@ -155,23 +154,14 @@ func (c *InteractiveClient) LoadSchema() error {
 	utils.LogTime("db_client.LoadSchema start")
 	defer utils.LogTime("db_client.LoadSchema end")
 
-	// build a ConnectionSchemaMap object to identify the schemas to load
-	// (pass nil for connection state - this forces NewConnectionSchemaMap to load it)
-	connectionSchemaMap, err := steampipeconfig.NewConnectionSchemaMap()
-	if err != nil {
-		return err
-	}
-	// get the unique schema - we use this to limit the schemas we load from the database
-	schemas := connectionSchemaMap.UniqueSchemas()
 	// load these schemas
 	// in a background context, since we are not running in a context - but GetSchemaFromDB needs one
-	metadata, err := c.client().GetSchemaFromDB(context.Background(), schemas)
+	metadata, err := c.client().GetSchemaFromDB(context.Background())
 	if err != nil {
 		return err
 	}
 
-	c.populateSchemaMetadata(metadata, connectionSchemaMap)
-
+	c.schemaMetadata = metadata
 	return nil
 }
 
@@ -589,21 +579,3 @@ func (c *InteractiveClient) addControlSuggestion(control *modconfig.Control, con
 	}
 	return prompt.Suggest{Text: controlName, Output: controlName, Description: description}
 }
-
-func (c *InteractiveClient) populateSchemaMetadata(schemaMetadata *schema.Metadata, connectionSchemaMap steampipeconfig.ConnectionSchemaMap) error {
-	// we now need to add in all other schemas which have the same schemas as those we have loaded
-	for loadedSchema, otherSchemas := range connectionSchemaMap {
-		// all 'otherSchema's have the same schema as loadedSchema
-		exemplarSchema, ok := schemaMetadata.Schemas[loadedSchema]
-		if !ok {
-			// should can happen in the case of a dynamic plugin with no tables - use empty schema
-			exemplarSchema = make(map[string]schema.TableSchema)
-		}
-
-		for _, s := range otherSchemas {
-			schemaMetadata.Schemas[s] = exemplarSchema
-		}
-	}
-	c.schemaMetadata = schemaMetadata
-	return nil
-}