You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by li...@apache.org on 2024/09/23 00:37:31 UTC
(arrow-adbc) branch main updated: fix(go/adbc/driver/snowflake): fix setting database and schema context after initial connection (#2169)
This is an automated email from the ASF dual-hosted git repository.
lidavidm 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 776726f42 fix(go/adbc/driver/snowflake): fix setting database and schema context after initial connection (#2169)
776726f42 is described below
commit 776726f42091caa4c29c8c44c9623c5da56ed38d
Author: davidhcoe <13...@users.noreply.github.com>
AuthorDate: Sun Sep 22 20:37:26 2024 -0400
fix(go/adbc/driver/snowflake): fix setting database and schema context after initial connection (#2169)
This PR removes the parameterization for calls to USE DATABASE and USE
SCHEMA, which isn't supported for Snowflake DDL. It also modifies the
GetObjects call to set the context to the database and schema, since the
GetObjects call uses a sql.DB connection context instead of the
snowflakeConn connection context.
Fixes https://github.com/apache/arrow-adbc/issues/2167
Fixes https://github.com/apache/arrow-adbc/issues/2168
---------
Co-authored-by: David Coe <co...@umich.edu>
---
csharp/src/Apache.Arrow.Adbc/AdbcConnection11.cs | 1 +
.../test/Drivers/Interop/Snowflake/DriverTests.cs | 36 ++++++++++++++++++++++
go/adbc/driver/snowflake/connection.go | 26 ++++++++++++++--
go/adbc/driver/snowflake/driver_test.go | 33 ++++++++++++++++++++
4 files changed, 94 insertions(+), 2 deletions(-)
diff --git a/csharp/src/Apache.Arrow.Adbc/AdbcConnection11.cs b/csharp/src/Apache.Arrow.Adbc/AdbcConnection11.cs
index e4f79cbbc..5ea433706 100644
--- a/csharp/src/Apache.Arrow.Adbc/AdbcConnection11.cs
+++ b/csharp/src/Apache.Arrow.Adbc/AdbcConnection11.cs
@@ -32,6 +32,7 @@ namespace Apache.Arrow.Adbc
, IAsyncDisposable
#endif
{
+
~AdbcConnection11() => Dispose(false);
/// <summary>
diff --git a/csharp/test/Drivers/Interop/Snowflake/DriverTests.cs b/csharp/test/Drivers/Interop/Snowflake/DriverTests.cs
index 2069baa9b..272e1e35a 100644
--- a/csharp/test/Drivers/Interop/Snowflake/DriverTests.cs
+++ b/csharp/test/Drivers/Interop/Snowflake/DriverTests.cs
@@ -125,6 +125,42 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.Interop.Snowflake
return expectedRole == stringArray.GetString(0);
}
+ [SkippableFact, Order(1)]
+ public void CanSetDatabase()
+ {
+ Skip.If(string.IsNullOrEmpty(_testConfiguration.Metadata.Catalog));
+
+ // connect without the parameter and ensure we get the DATABASE successfully
+ Dictionary<string, string> parameters = new Dictionary<string, string>();
+ Dictionary<string, string> options = new Dictionary<string, string>();
+
+ using AdbcDriver localSnowflakeDriver = SnowflakeTestingUtils.GetSnowflakeAdbcDriver(_testConfiguration, out parameters);
+ parameters.Remove(SnowflakeParameters.DATABASE);
+ using AdbcDatabase localDatabase = localSnowflakeDriver.Open(parameters);
+ using AdbcConnection localConnection = localDatabase.Connect(options);
+
+ localConnection.SetOption(AdbcOptions.Connection.CurrentCatalog, _testConfiguration.Metadata.Catalog);
+
+ Assert.True(CurrentDatabaseIsExpectedCatalog(localConnection, _testConfiguration.Metadata.Catalog));
+
+ localConnection.GetObjects(AdbcConnection.GetObjectsDepth.All, _testConfiguration.Metadata.Catalog, _testConfiguration.Metadata.Schema, _testConfiguration.Metadata.Table, _tableTypes, null);
+ }
+
+ private bool CurrentDatabaseIsExpectedCatalog(AdbcConnection cn, string expectedCatalog)
+ {
+ using AdbcStatement statement = cn.CreateStatement();
+ statement.SqlQuery = "SELECT CURRENT_DATABASE() as CURRENT_DATABASE;"; // GetOption doesn't exist in 1.0, only 1.1
+
+ QueryResult queryResult = statement.ExecuteQuery();
+ using RecordBatch? recordBatch = queryResult.Stream?.ReadNextRecordBatchAsync().Result;
+ Assert.True(recordBatch != null);
+
+ StringArray stringArray = (StringArray)recordBatch.Column("CURRENT_DATABASE");
+ Assert.True(stringArray.Length > 0);
+
+ return expectedCatalog == stringArray.GetString(0);
+ }
+
/// <summary>
/// Validates if the driver can connect to a live server and
/// parse the results.
diff --git a/go/adbc/driver/snowflake/connection.go b/go/adbc/driver/snowflake/connection.go
index 65b3b9c98..e6ec69938 100644
--- a/go/adbc/driver/snowflake/connection.go
+++ b/go/adbc/driver/snowflake/connection.go
@@ -162,6 +162,24 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
sql.Named("UNIQUE_QUERY_ID", uniqueQueryID),
}
+ // the connection that is used is not the same connection context where the database may have been set
+ // if the caller called SetCurrentCatalog() so need to ensure the database context is appropriate
+ if !isNilOrEmpty(catalog) {
+ _, e := conn.ExecContext(context.Background(), fmt.Sprintf("USE DATABASE %s;", quoteTblName(*catalog)), nil)
+ if e != nil {
+ return nil, errToAdbcErr(adbc.StatusIO, e)
+ }
+ }
+
+ // the connection that is used is not the same connection context where the schema may have been set
+ // if the caller called SetCurrentDbSchema() so need to ensure the schema context is appropriate
+ if !isNilOrEmpty(dbSchema) {
+ _, e2 := conn.ExecContext(context.Background(), fmt.Sprintf("USE SCHEMA %s;", quoteTblName(*dbSchema)), nil)
+ if e2 != nil {
+ return nil, errToAdbcErr(adbc.StatusIO, e2)
+ }
+ }
+
query := bldr.String()
rows, err := conn.QueryContext(ctx, query, args...)
if err != nil {
@@ -214,6 +232,10 @@ func (c *connectionImpl) GetObjects(ctx context.Context, depth adbc.ObjectDepth,
}
}
+func isNilOrEmpty(str *string) bool {
+ return str == nil || *str == ""
+}
+
// PrepareDriverInfo implements driverbase.DriverInfoPreparer.
func (c *connectionImpl) PrepareDriverInfo(ctx context.Context, infoCodes []adbc.InfoCode) error {
if err := c.ConnectionImplBase.DriverInfo.RegisterInfoCode(adbc.InfoVendorSql, true); err != nil {
@@ -239,13 +261,13 @@ func (c *connectionImpl) GetCurrentDbSchema() (string, error) {
// SetCurrentCatalog implements driverbase.CurrentNamespacer.
func (c *connectionImpl) SetCurrentCatalog(value string) error {
- _, err := c.cn.ExecContext(context.Background(), "USE DATABASE ?", []driver.NamedValue{{Value: value}})
+ _, err := c.cn.ExecContext(context.Background(), fmt.Sprintf("USE DATABASE %s;", quoteTblName(value)), nil)
return err
}
// SetCurrentDbSchema implements driverbase.CurrentNamespacer.
func (c *connectionImpl) SetCurrentDbSchema(value string) error {
- _, err := c.cn.ExecContext(context.Background(), "USE SCHEMA ?", []driver.NamedValue{{Value: value}})
+ _, err := c.cn.ExecContext(context.Background(), fmt.Sprintf("USE SCHEMA %s;", quoteTblName(value)), nil)
return err
}
diff --git a/go/adbc/driver/snowflake/driver_test.go b/go/adbc/driver/snowflake/driver_test.go
index 0448384d6..e56ce179a 100644
--- a/go/adbc/driver/snowflake/driver_test.go
+++ b/go/adbc/driver/snowflake/driver_test.go
@@ -2111,3 +2111,36 @@ func TestIngestCancelContext(t *testing.T) {
require.Equal(t, "", buf.String())
})
}
+
+func (suite *SnowflakeTests) TestChangeDatabaseAndGetObjects() {
+ // this test demonstrates:
+ // 1. changing the database context
+ // 2. being able to call GetObjects after changing the database context
+ // (this uses a different connection context but still executes successfully)
+
+ uri, ok := os.LookupEnv("SNOWFLAKE_URI")
+ if !ok {
+ suite.T().Skip("Cannot find the `SNOWFLAKE_URI` value")
+ }
+
+ newCatalog, ok := os.LookupEnv("SNOWFLAKE_NEW_CATALOG")
+ if !ok {
+ suite.T().Skip("Cannot find the `SNOWFLAKE_NEW_CATALOG` value")
+ }
+
+ getObjectsTable, ok := os.LookupEnv("SNOWFLAKE_TABLE_GETOBJECTS")
+ if !ok {
+ suite.T().Skip("Cannot find the `SNOWFLAKE_TABLE_GETOBJECTS` value")
+ }
+
+ cfg, err := gosnowflake.ParseDSN(uri)
+ suite.NoError(err)
+
+ cnxnopt, ok := suite.cnxn.(adbc.PostInitOptions)
+ suite.True(ok)
+ err = cnxnopt.SetOption(adbc.OptionKeyCurrentCatalog, newCatalog)
+ suite.NoError(err)
+
+ _, err2 := suite.cnxn.GetObjects(suite.ctx, adbc.ObjectDepthAll, &newCatalog, &cfg.Schema, &getObjectsTable, nil, nil)
+ suite.NoError(err2)
+}