From ce3db31098f2d5ec94d3866fd7ac24322b760e21 Mon Sep 17 00:00:00 2001
From: Piotr Fus <piotr.fus@snowflake.com>
Date: Fri, 22 Nov 2024 12:30:36 +0100
Subject: [PATCH] SNOW-1814531 Implement DriverContext interface (#1248)

---
 auth_test.go | 1 -
 connector.go | 2 +-
 driver.go    | 9 +++++++++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/auth_test.go b/auth_test.go
index ccc0f2cd7..c559a2407 100644
--- a/auth_test.go
+++ b/auth_test.go
@@ -979,7 +979,6 @@ func TestOktaRetryWithNewToken(t *testing.T) {
 }
 
 func TestContextPropagatedToAuthWhenUsingOpen(t *testing.T) {
-	t.Skip("dsnConnector loses context when opening using sql.Open")
 	db, err := sql.Open("snowflake", dsn)
 	assertNilF(t, err)
 	defer db.Close()
diff --git a/connector.go b/connector.go
index dbaf361dc..b25fa1139 100644
--- a/connector.go
+++ b/connector.go
@@ -20,7 +20,7 @@ type Connector struct {
 }
 
 // NewConnector creates a new connector with the given SnowflakeDriver and Config.
-func NewConnector(driver InternalSnowflakeDriver, config Config) Connector {
+func NewConnector(driver InternalSnowflakeDriver, config Config) driver.Connector {
 	return Connector{driver, config}
 }
 
diff --git a/driver.go b/driver.go
index 99c825ec2..7fee0bd2a 100644
--- a/driver.go
+++ b/driver.go
@@ -33,6 +33,15 @@ func (d SnowflakeDriver) Open(dsn string) (driver.Conn, error) {
 	return d.OpenWithConfig(ctx, *cfg)
 }
 
+// OpenConnector creates a new connector with parsed DSN.
+func (d SnowflakeDriver) OpenConnector(dsn string) (driver.Connector, error) {
+	cfg, err := ParseDSN(dsn)
+	if err != nil {
+		return Connector{}, err
+	}
+	return NewConnector(d, *cfg), nil
+}
+
 // OpenWithConfig creates a new connection with the given Config.
 func (d SnowflakeDriver) OpenWithConfig(ctx context.Context, config Config) (driver.Conn, error) {
 	if err := config.Validate(); err != nil {