Fix connector_test.go and gofmt lint

Signed-off-by: Matthew Kim <11141331+mattdeekay@users.noreply.github.com>
databricks · Aug 4, 2023 · 86b6ec4 · 86b6ec4
1 parent 4be8ace
commit 86b6ec4
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 72 deletions.
diff --git a/connector_test.go b/connector_test.go
@@ -36,24 +36,33 @@ func TestNewConnector(t *testing.T) {
 			WithSessionParams(sessionParams),
 			WithRetries(10, 3*time.Second, 60*time.Second),
 			WithTransport(roundTripper),
+			WithCloudFetch(true),
+			WithMaxDownloadThreads(15),
 		)
+		expectedCloudFetchConfig := config.CloudFetchConfig{
+			UseCloudFetch:      true,
+			MaxDownloadThreads: 15,
+			MaxFilesInMemory:   10,
+			MinTimeToExpiry:    0 * time.Second,
+		}
 		expectedUserConfig := config.UserConfig{
-			Host:           host,
-			Port:           port,
-			Protocol:       "https",
-			AccessToken:    accessToken,
-			Authenticator:  &pat.PATAuth{AccessToken: accessToken},
-			HTTPPath:       "/" + httpPath,
-			MaxRows:        maxRows,
-			QueryTimeout:   timeout,
-			Catalog:        catalog,
-			Schema:         schema,
-			UserAgentEntry: userAgentEntry,
-			SessionParams:  sessionParams,
-			RetryMax:       10,
-			RetryWaitMin:   3 * time.Second,
-			RetryWaitMax:   60 * time.Second,
-			Transport:      roundTripper,
+			Host:             host,
+			Port:             port,
+			Protocol:         "https",
+			AccessToken:      accessToken,
+			Authenticator:    &pat.PATAuth{AccessToken: accessToken},
+			HTTPPath:         "/" + httpPath,
+			MaxRows:          maxRows,
+			QueryTimeout:     timeout,
+			Catalog:          catalog,
+			Schema:           schema,
+			UserAgentEntry:   userAgentEntry,
+			SessionParams:    sessionParams,
+			RetryMax:         10,
+			RetryWaitMin:     3 * time.Second,
+			RetryWaitMax:     60 * time.Second,
+			Transport:        roundTripper,
+			CloudFetchConfig: expectedCloudFetchConfig,
 		}
 		expectedCfg := config.WithDefaults()
 		expectedCfg.DriverVersion = DriverVersion
@@ -75,18 +84,25 @@ func TestNewConnector(t *testing.T) {
 			WithAccessToken(accessToken),
 			WithHTTPPath(httpPath),
 		)
+		expectedCloudFetchConfig := config.CloudFetchConfig{
+			UseCloudFetch:      false,
+			MaxDownloadThreads: 10,
+			MaxFilesInMemory:   10,
+			MinTimeToExpiry:    0 * time.Second,
+		}
 		expectedUserConfig := config.UserConfig{
-			Host:          host,
-			Port:          port,
-			Protocol:      "https",
-			AccessToken:   accessToken,
-			Authenticator: &pat.PATAuth{AccessToken: accessToken},
-			HTTPPath:      "/" + httpPath,
-			MaxRows:       maxRows,
-			SessionParams: sessionParams,
-			RetryMax:      4,
-			RetryWaitMin:  1 * time.Second,
-			RetryWaitMax:  30 * time.Second,
+			Host:             host,
+			Port:             port,
+			Protocol:         "https",
+			AccessToken:      accessToken,
+			Authenticator:    &pat.PATAuth{AccessToken: accessToken},
+			HTTPPath:         "/" + httpPath,
+			MaxRows:          maxRows,
+			SessionParams:    sessionParams,
+			RetryMax:         4,
+			RetryWaitMin:     1 * time.Second,
+			RetryWaitMax:     30 * time.Second,
+			CloudFetchConfig: expectedCloudFetchConfig,
 		}
 		expectedCfg := config.WithDefaults()
 		expectedCfg.UserConfig = expectedUserConfig
@@ -109,18 +125,25 @@ func TestNewConnector(t *testing.T) {
 			WithHTTPPath(httpPath),
 			WithRetries(-1, 0, 0),
 		)
+		expectedCloudFetchConfig := config.CloudFetchConfig{
+			UseCloudFetch:      false,
+			MaxDownloadThreads: 10,
+			MaxFilesInMemory:   10,
+			MinTimeToExpiry:    0 * time.Second,
+		}
 		expectedUserConfig := config.UserConfig{
-			Host:          host,
-			Port:          port,
-			Protocol:      "https",
-			AccessToken:   accessToken,
-			Authenticator: &pat.PATAuth{AccessToken: accessToken},
-			HTTPPath:      "/" + httpPath,
-			MaxRows:       maxRows,
-			SessionParams: sessionParams,
-			RetryMax:      -1,
-			RetryWaitMin:  0,
-			RetryWaitMax:  0,
+			Host:             host,
+			Port:             port,
+			Protocol:         "https",
+			AccessToken:      accessToken,
+			Authenticator:    &pat.PATAuth{AccessToken: accessToken},
+			HTTPPath:         "/" + httpPath,
+			MaxRows:          maxRows,
+			SessionParams:    sessionParams,
+			RetryMax:         -1,
+			RetryWaitMin:     0,
+			RetryWaitMax:     0,
+			CloudFetchConfig: expectedCloudFetchConfig,
 		}
 		expectedCfg := config.WithDefaults()
 		expectedCfg.DriverVersion = DriverVersion

diff --git a/internal/config/config.go b/internal/config/config.go
@@ -82,23 +82,23 @@ func (c *Config) DeepCopy() *Config {
 
 // UserConfig is the set of configurations exposed to users
 type UserConfig struct {
-	Protocol       string
-	Host           string // from databricks UI
-	Port           int    // from databricks UI
-	HTTPPath       string // from databricks UI
-	Catalog        string
-	Schema         string
-	Authenticator  auth.Authenticator
-	AccessToken    string        // from databricks UI
-	MaxRows        int           // max rows per page
-	QueryTimeout   time.Duration // Timeout passed to server for query processing
-	UserAgentEntry string
-	Location       *time.Location
-	SessionParams  map[string]string
-	RetryWaitMin   time.Duration
-	RetryWaitMax   time.Duration
-	RetryMax       int
-	Transport      http.RoundTripper
+	Protocol          string
+	Host              string // from databricks UI
+	Port              int    // from databricks UI
+	HTTPPath          string // from databricks UI
+	Catalog           string
+	Schema            string
+	Authenticator     auth.Authenticator
+	AccessToken       string        // from databricks UI
+	MaxRows           int           // max rows per page
+	QueryTimeout      time.Duration // Timeout passed to server for query processing
+	UserAgentEntry    string
+	Location          *time.Location
+	SessionParams     map[string]string
+	RetryWaitMin      time.Duration
+	RetryWaitMax      time.Duration
+	RetryMax          int
+	Transport         http.RoundTripper
 	UseLz4Compression bool
 	CloudFetchConfig
 }
@@ -123,23 +123,23 @@ func (ucfg UserConfig) DeepCopy() UserConfig {
 	}
 
 	return UserConfig{
-		Protocol:       ucfg.Protocol,
-		Host:           ucfg.Host,
-		Port:           ucfg.Port,
-		HTTPPath:       ucfg.HTTPPath,
-		Catalog:        ucfg.Catalog,
-		Schema:         ucfg.Schema,
-		Authenticator:  ucfg.Authenticator,
-		AccessToken:    ucfg.AccessToken,
-		MaxRows:        ucfg.MaxRows,
-		QueryTimeout:   ucfg.QueryTimeout,
-		UserAgentEntry: ucfg.UserAgentEntry,
-		Location:       loccp,
-		SessionParams:  sessionParams,
-		RetryWaitMin:   ucfg.RetryWaitMin,
-		RetryWaitMax:   ucfg.RetryWaitMax,
-		RetryMax:       ucfg.RetryMax,
-		Transport:      ucfg.Transport,
+		Protocol:          ucfg.Protocol,
+		Host:              ucfg.Host,
+		Port:              ucfg.Port,
+		HTTPPath:          ucfg.HTTPPath,
+		Catalog:           ucfg.Catalog,
+		Schema:            ucfg.Schema,
+		Authenticator:     ucfg.Authenticator,
+		AccessToken:       ucfg.AccessToken,
+		MaxRows:           ucfg.MaxRows,
+		QueryTimeout:      ucfg.QueryTimeout,
+		UserAgentEntry:    ucfg.UserAgentEntry,
+		Location:          loccp,
+		SessionParams:     sessionParams,
+		RetryWaitMin:      ucfg.RetryWaitMin,
+		RetryWaitMax:      ucfg.RetryWaitMax,
+		RetryMax:          ucfg.RetryMax,
+		Transport:         ucfg.Transport,
 		UseLz4Compression: ucfg.UseLz4Compression,
 		CloudFetchConfig:  ucfg.CloudFetchConfig,
 	}