fix thread safety issue (#74)

Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with
[ReviewStack](https://reviewstack.dev/databricks/databricks-sql-go/pull/74).
* __->__ #74

fix thread safety issue

This fixes a thread safety (and structural issue) with connector.
Because connectors have a long lifetime (potentially process-lifetime),
and are used by multiple threads it's not safe to share the same thrift
client connection.

This changes the connector so that a new thrift client is created on
each connection. The go DB library will perform pooling of connections,
so we don't need to worry about it here.

Signed-off-by: Carl Verge <[email protected]>
Co-authored-by: Carl Verge <[email protected]>

Author: carl-db

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## 0.1.x (Unreleased)
 
+- Fix thread safety issue in connector
+
 ## 0.2.0 (2022-11-18)
 
 - Support for DirectResults

connector.go

@@ -15,8 +15,7 @@ import (
 )
 
 type connector struct {
-	cfg    *config.Config
-	client cli_service.TCLIService
+	cfg *config.Config
 }
 
 func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
@@ -29,8 +28,13 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		schemaName = cli_service.TIdentifierPtr(cli_service.TIdentifier(c.cfg.Schema))
 	}
 
+	tclient, err := client.InitThriftClient(c.cfg)
+	if err != nil {
+		return nil, wrapErr(err, "error initializing thrift client")
+	}
+
 	// we need to ensure that open session will eventually end
-	session, err := c.client.OpenSession(ctx, &cli_service.TOpenSessionReq{
+	session, err := tclient.OpenSession(ctx, &cli_service.TOpenSessionReq{
 		ClientProtocol: c.cfg.ThriftProtocolVersion,
 		Configuration:  make(map[string]string),
 		InitialNamespace: &cli_service.TNamespace{
@@ -47,7 +51,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	conn := &conn{
 		id:      client.SprintGuid(session.SessionHandle.GetSessionId().GUID),
 		cfg:     c.cfg,
-		client:  c.client,
+		client:  tclient,
 		session: session,
 	}
 	log := logger.WithContext(conn.id, driverctx.CorrelationIdFromContext(ctx), "")
@@ -82,13 +86,8 @@ func NewConnector(options ...connOption) (driver.Connector, error) {
 	for _, opt := range options {
 		opt(cfg)
 	}
-	// validate config?
-	tclient, err := client.InitThriftClient(cfg)
-	if err != nil {
-		return nil, wrapErr(err, "error initializing thrift client")
-	}
 
-	return &connector{cfg, tclient}, nil
+	return &connector{cfg: cfg}, nil
 }
 
 // WithServerHostname sets up the server hostname. Mandatory.

connector_test.go

@@ -1,38 +1,14 @@
 package dbsql
 
 import (
-	"context"
 	"testing"
 	"time"
 
-	"github.com/databricks/databricks-sql-go/internal/cli_service"
-	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func TestConnector_Connect(t *testing.T) {
-	t.Run("Connect returns err when thrift client initialization fails", func(t *testing.T) {
-		cfg := config.WithDefaults()
-		var openSessionCount int
-		openSession := func(ctx context.Context, req *cli_service.TOpenSessionReq) (r *cli_service.TOpenSessionResp, err error) {
-			openSessionCount++
-			return getTestSession(), nil
-		}
-		testClient := &client.TestClient{
-			FnOpenSession: openSession,
-		}
-		testConnector := connector{
-			cfg:    cfg,
-			client: testClient,
-		}
-		conn, err := testConnector.Connect(context.Background())
-		assert.NotNil(t, conn)
-		assert.NoError(t, err)
-	})
-}
-
 func TestNewConnector(t *testing.T) {
 	t.Run("Connector initialized with functional options should have all options set", func(t *testing.T) {
 		host := "databricks-host"

driver.go

@@ -5,7 +5,6 @@ import (
 	"database/sql"
 	"database/sql/driver"
 
-	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
 	_ "github.com/databricks/databricks-sql-go/logger"
 )
@@ -36,11 +35,8 @@ func (d *databricksDriver) OpenConnector(dsn string) (driver.Connector, error) {
 		return nil, err
 	}
 	cfg.UserConfig = ucfg
-	tclient, err := client.InitThriftClient(cfg)
-	if err != nil {
-		return nil, wrapErr(err, "error initializing thrift client")
-	}
-	return &connector{cfg, tclient}, nil
+
+	return &connector{cfg: cfg}, nil
 }
 
 var _ driver.Driver = (*databricksDriver)(nil)