feat: Optional switch for instrumentedsql (#2261)

Makefile

@@ -64,7 +64,7 @@ sweep: ## destroy the whole architecture; USE ONLY FOR DEVELOPMENT ACCOUNTS
 			else echo "Aborting..."; \
 		fi;
 
-test: ## run unit and integration tests
+test: test-client ## run unit and integration tests
 	go test -v -cover -timeout=30m ./...
 
 test-acceptance: ## run acceptance tests
@@ -76,6 +76,9 @@ test-integration: ## run SDK integration tests
 test-architecture: ## check architecture constraints between packages
 	go test ./pkg/architests/... -v
 
+test-client: ## runs test that checks sdk.Client without instrumentedsql
+	SF_TF_NO_INSTRUMENTED_SQL=1 go test ./pkg/sdk/internal/client/... -v
+
 build-local: ## build the binary locally
 	go build -o $(BASE_BINARY_NAME) .
 

README.md

@@ -14,6 +14,8 @@ This is a terraform provider for managing [Snowflake](https://www.snowflake.com/
   - [Getting started](#getting-started)
   - [SDK migration table](#sdk-migration-table)
   - [Getting Help](#getting-help)
+  - [Additional debug logs for `snowflake_grant_privileges_to_role` resource](#additional-debug-logs-for-snowflake_grant_privileges_to_role-resource)
+  - [Additional SQL Client configuration](#additional-sql-client-configuration)
   - [Contributing](#contributing)
 
 
@@ -138,6 +140,9 @@ Set environment variable `SF_TF_ADDITIONAL_DEBUG_LOGGING` to a non-empty value.
 2023/12/08 12:58:22.497078 sf-tf-additional-debug [DEBUG] Creating new client from db
 ```
 
+## Additional SQL Client configuration
+Currently underlying sql [gosnowflake](https://github.com/snowflakedb/gosnowflake) driver is wrapped with [instrumentedsql](https://github.com/luna-duclos/instrumentedsql). In order to use raw [gosnowflake](https://github.com/snowflakedb/gosnowflake) driver, set environment variable `SF_TF_NO_INSTRUMENTED_SQL` to a non-empty value.
+
 ## Contributing
 
 Cf. [Contributing](./CONTRIBUTING.md).

pkg/sdk/client.go

@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"fmt"
 	"log"
+	"os"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/luna-duclos/instrumentedsql"
@@ -13,6 +14,12 @@ import (
 	"github.com/snowflakedb/gosnowflake"
 )
 
+var instrumentedSQL bool
+
+func init() {
+	instrumentedSQL = os.Getenv("SF_TF_NO_INSTRUMENTED_SQL") == ""
+}
+
 type Client struct {
 	config         *gosnowflake.Config
 	db             *sqlx.DB
@@ -96,24 +103,30 @@ func NewClient(cfg *gosnowflake.Config) (*Client, error) {
 
 	var client *Client
 	// register the snowflake driver if it hasn't been registered yet
-	if !slices.Contains(sql.Drivers(), "snowflake-instrumented") {
-		logger := instrumentedsql.LoggerFunc(func(ctx context.Context, s string, kv ...interface{}) {
-			switch s {
-			case "sql-conn-query", "sql-conn-exec":
-				log.Printf("[DEBUG] %s: %v (%s)\n", s, kv, ctx.Value(snowflakeAccountLocatorContextKey))
-			default:
-				return
-			}
-		})
-		sql.Register("snowflake-instrumented", instrumentedsql.WrapDriver(gosnowflake.SnowflakeDriver{}, instrumentedsql.WithLogger(logger)))
+
+	driverName := "snowflake"
+	if instrumentedSQL {
+		if !slices.Contains(sql.Drivers(), "snowflake-instrumented") {
+			log.Println("[DEBUG] Registering snowflake-instrumented driver")
+			logger := instrumentedsql.LoggerFunc(func(ctx context.Context, s string, kv ...interface{}) {
+				switch s {
+				case "sql-conn-query", "sql-conn-exec":
+					log.Printf("[DEBUG] %s: %v (%s)\n", s, kv, ctx.Value(snowflakeAccountLocatorContextKey))
+				default:
+					return
+				}
+			})
+			sql.Register("snowflake-instrumented", instrumentedsql.WrapDriver(new(gosnowflake.SnowflakeDriver), instrumentedsql.WithLogger(logger)))
+		}
+		driverName = "snowflake-instrumented"
 	}
 
 	dsn, err := gosnowflake.DSN(cfg)
 	if err != nil {
 		return nil, err
 	}
 
-	db, err := sqlx.Connect("snowflake-instrumented", dsn)
+	db, err := sqlx.Connect(driverName, dsn)
 	if err != nil {
 		return nil, fmt.Errorf("open snowflake connection: %w", err)
 	}

pkg/sdk/client_integration_test.go

@@ -2,8 +2,11 @@ package sdk
 
 import (
 	"context"
+	"database/sql"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/stretchr/testify/require"
 )
 
@@ -13,13 +16,22 @@ func TestClient_NewClient(t *testing.T) {
 		_, err := NewClient(config)
 		require.NoError(t, err)
 	})
+
 	t.Run("uses env vars if values are missing", func(t *testing.T) {
 		cleanupEnvVars := setupEnvVars(t, "TEST_ACCOUNT", "TEST_USER", "abcd1234", "ACCOUNTADMIN", "")
 		t.Cleanup(cleanupEnvVars)
 		config := EnvConfig()
 		_, err := NewClient(config)
 		require.Error(t, err)
 	})
+
+	t.Run("registers snowflake-instrumented driver", func(t *testing.T) {
+		config := DefaultConfig()
+		_, err := NewClient(config)
+		require.NoError(t, err)
+
+		assert.ElementsMatch(t, sql.Drivers(), []string{"snowflake-instrumented", "snowflake"})
+	})
 }
 
 func TestClient_ping(t *testing.T) {

pkg/sdk/internal/client/client_test.go

@@ -0,0 +1,29 @@
+package client
+
+import (
+	"database/sql"
+	"os"
+	"testing"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestNewClientWithoutInstrumentedSQL checks if the client is initialized with the different driver implementation.
+// This is dependent on the SF_TF_NO_INSTRUMENTED_SQL env variable setting. That's why it was extracted to another file.
+// To run this test use: `make test-client` command.
+func TestNewClientWithoutInstrumentedSQL(t *testing.T) {
+	if os.Getenv("SF_TF_NO_INSTRUMENTED_SQL") == "" {
+		t.Skip("Skipping TestNewClientWithoutInstrumentedSQL, because SF_TF_NO_INSTRUMENTED_SQL is not set")
+	}
+
+	t.Run("registers snowflake-not-instrumented driver", func(t *testing.T) {
+		config := sdk.DefaultConfig()
+		_, err := sdk.NewClient(config)
+		require.NoError(t, err)
+
+		assert.NotContains(t, sql.Drivers(), "snowflake-instrumented")
+		assert.Contains(t, sql.Drivers(), "snowflake")
+	})
+}