You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by jo...@apache.org on 2024/10/02 21:11:16 UTC

(arrow-adbc) branch main updated: fix(go/adbc/driver/snowflake): call GetObjects with null catalog at catalog depth (#2194)

This is an automated email from the ASF dual-hosted git repository.

joellubi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git


The following commit(s) were added to refs/heads/main by this push:
     new 0cb51b8d4 fix(go/adbc/driver/snowflake): call GetObjects with null catalog at catalog depth (#2194)
0cb51b8d4 is described below

commit 0cb51b8d499148b5b79ffce59947ccf35422fde9
Author: davidhcoe <13...@users.noreply.github.com>
AuthorDate: Wed Oct 2 17:11:12 2024 -0400

    fix(go/adbc/driver/snowflake): call GetObjects with null catalog at catalog depth (#2194)
    
    This PR partially addresses
    https://github.com/apache/arrow-adbc/issues/2171 by calling SHOW TERSE
    DATABASES if GetObjects is called with the catalog depth and a null
    catalog is passed.
    
    It does not address other parts of
    https://github.com/apache/arrow-adbc/issues/2171 where a catalog can be
    null at other depths.
    
    ---------
    
    Co-authored-by: David Coe <co...@umich.edu>
---
 go/adbc/driver/snowflake/connection.go             | 48 +++++++++++++++-------
 go/adbc/driver/snowflake/driver_test.go            |  6 +++
 .../queries/get_objects_terse_catalogs.sql         | 24 +++++++++++
 3 files changed, 64 insertions(+), 14 deletions(-)

diff --git a/go/adbc/driver/snowflake/connection.go b/go/adbc/driver/snowflake/connection.go
index e6ec69938..a8361a365 100644
--- a/go/adbc/driver/snowflake/connection.go
+++ b/go/adbc/driver/snowflake/connection.go
@@ -41,10 +41,11 @@ const (
 	defaultStatementQueueSize  = 200
 	defaultPrefetchConcurrency = 10
 
-	queryTemplateGetObjectsAll       = "get_objects_all.sql"
-	queryTemplateGetObjectsCatalogs  = "get_objects_catalogs.sql"
-	queryTemplateGetObjectsDbSchemas = "get_objects_dbschemas.sql"
-	queryTemplateGetObjectsTables    = "get_objects_tables.sql"
+	queryTemplateGetObjectsAll           = "get_objects_all.sql"
+	queryTemplateGetObjectsCatalogs      = "get_objects_catalogs.sql"
+	queryTemplateGetObjectsDbSchemas     = "get_objects_dbschemas.sql"
+	queryTemplateGetObjectsTables        = "get_objects_tables.sql"
+	queryTemplateGetObjectsTerseCatalogs = "get_objects_terse_catalogs.sql"
 )
 
 //go:embed queries/*
@@ -74,7 +75,7 @@ type connectionImpl struct {
 
 func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth, catalog *string, dbSchema *string, tableName *string, columnName *string, tableType []string) (array.RecordReader, error) {
 	var (
-		pkQueryID, fkQueryID, uniqueQueryID string
+		pkQueryID, fkQueryID, uniqueQueryID, terseDbQueryID string
 	)
 
 	conn, err := c.sqldb.Conn(ctx)
@@ -83,11 +84,28 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
 	}
 	defer conn.Close()
 
-	gConstraints, gContraintsCtx := errgroup.WithContext(ctx)
+	gQueryIDs, gQueryIDsCtx := errgroup.WithContext(ctx)
 	queryFile := queryTemplateGetObjectsAll
 	switch depth {
 	case adbc.ObjectDepthCatalogs:
-		queryFile = queryTemplateGetObjectsCatalogs
+		if catalog == nil {
+			queryFile = queryTemplateGetObjectsTerseCatalogs
+			// if the catalog is null, show the terse databases
+			// which doesn't require a database context
+			gQueryIDs.Go(func() error {
+				return conn.Raw(func(driverConn any) error {
+					rows, err := driverConn.(driver.QueryerContext).QueryContext(gQueryIDsCtx, "SHOW TERSE DATABASES", nil)
+					if err != nil {
+						return err
+					}
+
+					terseDbQueryID = rows.(gosnowflake.SnowflakeRows).GetQueryID()
+					return rows.Close()
+				})
+			})
+		} else {
+			queryFile = queryTemplateGetObjectsCatalogs
+		}
 	case adbc.ObjectDepthDBSchemas:
 		queryFile = queryTemplateGetObjectsDbSchemas
 	case adbc.ObjectDepthTables:
@@ -96,9 +114,9 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
 	default:
 		// Detailed constraint info not available in information_schema
 		// Need to dispatch SHOW queries and use conn.Raw to extract the queryID for reuse in GetObjects query
-		gConstraints.Go(func() error {
+		gQueryIDs.Go(func() error {
 			return conn.Raw(func(driverConn any) error {
-				rows, err := driverConn.(driver.QueryerContext).QueryContext(gContraintsCtx, "SHOW PRIMARY KEYS", nil)
+				rows, err := driverConn.(driver.QueryerContext).QueryContext(gQueryIDsCtx, "SHOW PRIMARY KEYS", nil)
 				if err != nil {
 					return err
 				}
@@ -108,9 +126,9 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
 			})
 		})
 
-		gConstraints.Go(func() error {
+		gQueryIDs.Go(func() error {
 			return conn.Raw(func(driverConn any) error {
-				rows, err := driverConn.(driver.QueryerContext).QueryContext(gContraintsCtx, "SHOW IMPORTED KEYS", nil)
+				rows, err := driverConn.(driver.QueryerContext).QueryContext(gQueryIDsCtx, "SHOW IMPORTED KEYS", nil)
 				if err != nil {
 					return err
 				}
@@ -120,9 +138,9 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
 			})
 		})
 
-		gConstraints.Go(func() error {
+		gQueryIDs.Go(func() error {
 			return conn.Raw(func(driverConn any) error {
-				rows, err := driverConn.(driver.QueryerContext).QueryContext(gContraintsCtx, "SHOW UNIQUE KEYS", nil)
+				rows, err := driverConn.(driver.QueryerContext).QueryContext(gQueryIDsCtx, "SHOW UNIQUE KEYS", nil)
 				if err != nil {
 					return err
 				}
@@ -145,7 +163,7 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
 	}
 
 	// Need constraint subqueries to complete before we can query GetObjects
-	if err := gConstraints.Wait(); err != nil {
+	if err := gQueryIDs.Wait(); err != nil {
 		return nil, err
 	}
 
@@ -157,9 +175,11 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
 		driverbase.PatternToNamedArg("COLUMN", columnName),
 
 		// QueryIDs for constraint data if depth is tables or deeper
+		// or if the depth is catalog and catalog is null
 		sql.Named("PK_QUERY_ID", pkQueryID),
 		sql.Named("FK_QUERY_ID", fkQueryID),
 		sql.Named("UNIQUE_QUERY_ID", uniqueQueryID),
+		sql.Named("SHOW_DB_QUERY_ID", terseDbQueryID),
 	}
 
 	// the connection that is used is not the same connection context where the database may have been set
diff --git a/go/adbc/driver/snowflake/driver_test.go b/go/adbc/driver/snowflake/driver_test.go
index d4a12a6cd..895015ffd 100644
--- a/go/adbc/driver/snowflake/driver_test.go
+++ b/go/adbc/driver/snowflake/driver_test.go
@@ -2158,3 +2158,9 @@ func (suite *SnowflakeTests) TestGetSetClientConfigFile() {
 	suite.NoError(err)
 	suite.True(file == result)
 }
+
+func (suite *SnowflakeTests) TestGetObjectsWithNilCatalog() {
+	// this test demonstrates calling GetObjects with the catalog depth and a nil catalog
+	_, err := suite.cnxn.GetObjects(suite.ctx, adbc.ObjectDepthCatalogs, nil, nil, nil, nil, nil)
+	suite.NoError(err)
+}
diff --git a/go/adbc/driver/snowflake/queries/get_objects_terse_catalogs.sql b/go/adbc/driver/snowflake/queries/get_objects_terse_catalogs.sql
new file mode 100644
index 000000000..be56795e0
--- /dev/null
+++ b/go/adbc/driver/snowflake/queries/get_objects_terse_catalogs.sql
@@ -0,0 +1,24 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements.  See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership.  The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License.  You may obtain a copy of the License at
+--
+--   http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied.  See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- SHOW TERSE DATABASES would run before this
+SELECT
+    {
+    'catalog_name': "name",
+    'catalog_db_schemas': null
+    } get_objects
+FROM TABLE(RESULT_SCAN(:SHOW_DB_QUERY_ID))