replace Timeout with ConnectTimeout/SocketTimeout/HttpReadTimeout

beltran · Oct 22, 2021 · d164d8c · d164d8c
1 parent 9701375
commit d164d8c
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 25 deletions.
diff --git a/hive.go b/hive.go
@@ -59,7 +59,9 @@ type ConnectConfiguration struct {
 	TLSConfig            *tls.Config
 	ZookeeperNamespace   string
 	Database             string
-	Timeout              time.Duration
+	ConnectTimeout       time.Duration
+	SocketTimeout        time.Duration
+	HttpReadTimeout      time.Duration
 	DialContext          DialContextFunc
 }
 
@@ -187,33 +189,33 @@ func innerConnect(ctx context.Context, host string, port int, auth string,
 	addr := fmt.Sprintf("%s:%d", host, port)
 	if configuration.DialContext != nil {
 		var netConn net.Conn
-		netConn, err = dial(ctx, addr, configuration.DialContext, configuration.Timeout)
+		netConn, err = dial(ctx, addr, configuration.DialContext, configuration.ConnectTimeout)
 		if err != nil {
 			return
 		}
 		if configuration.TLSConfig != nil {
 			socket = thrift.NewTSSLSocketFromConnConf(netConn, &thrift.TConfiguration{
-				ConnectTimeout: configuration.Timeout,
-				SocketTimeout:  configuration.Timeout,
+				ConnectTimeout: configuration.ConnectTimeout,
+				SocketTimeout:  configuration.SocketTimeout,
 				TLSConfig:      configuration.TLSConfig,
 			})
 		} else {
 			socket = thrift.NewTSocketFromConnConf(netConn, &thrift.TConfiguration{
-				ConnectTimeout: configuration.Timeout,
-				SocketTimeout:  configuration.Timeout,
+				ConnectTimeout: configuration.ConnectTimeout,
+				SocketTimeout:  configuration.SocketTimeout,
 			})
 		}
 	} else {
 		if configuration.TLSConfig != nil {
 			socket, err = thrift.NewTSSLSocketConf(addr, &thrift.TConfiguration{
-				ConnectTimeout: configuration.Timeout,
-				SocketTimeout:  configuration.Timeout,
+				ConnectTimeout: configuration.ConnectTimeout,
+				SocketTimeout:  configuration.SocketTimeout,
 				TLSConfig:      configuration.TLSConfig,
 			})
 		} else {
 			socket, err = thrift.NewTSocketConf(addr, &thrift.TConfiguration{
-				ConnectTimeout: configuration.Timeout,
-				SocketTimeout:  configuration.Timeout,
+				ConnectTimeout: configuration.ConnectTimeout,
+				SocketTimeout:  configuration.SocketTimeout,
 			})
 		}
 		if err != nil {
@@ -366,7 +368,7 @@ func innerConnect(ctx context.Context, host string, port int, auth string,
 func getHTTPClient(configuration *ConnectConfiguration) (httpClient *http.Client, protocol string, err error) {
 	if configuration.TLSConfig != nil {
 		httpClient = &http.Client{
-			Timeout: configuration.Timeout,
+			Timeout: configuration.HttpReadTimeout,
 			Transport: &http.Transport{
 				TLSClientConfig: configuration.TLSConfig,
 				DialContext:     configuration.DialContext,
@@ -375,7 +377,7 @@ func getHTTPClient(configuration *ConnectConfiguration) (httpClient *http.Client
 		protocol = "https"
 	} else {
 		httpClient = &http.Client{
-			Timeout: configuration.Timeout,
+			Timeout: configuration.HttpReadTimeout,
 			Transport: &http.Transport{
 				DialContext: configuration.DialContext,
 			},

diff --git a/hive_test.go b/hive_test.go
@@ -1,3 +1,4 @@
+//go:build all || integration
 // +build all integration
 
 package gohive
@@ -585,15 +586,15 @@ func NoopDialContext(ctx context.Context, network string, addr string) (net.Conn
 }
 
 func sleepContext(ctx context.Context, delay time.Duration) {
-    select {
-    case <-ctx.Done():
-    case <-time.After(delay):
-    }
+	select {
+	case <-ctx.Done():
+	case <-time.After(delay):
+	}
 }
 
 func SleepDialContext(ctx context.Context, network string, addr string) (net.Conn, error) {
 	var d net.Dialer
-	sleepContext(ctx, 10 * time.Second)
+	sleepContext(ctx, 10*time.Second)
 	return d.DialContext(ctx, network, addr)
 }
 
@@ -627,7 +628,7 @@ func TestSimpleSelectWithDialFunctionAndTimeout(t *testing.T) {
 	configuration.TransportMode = getTransport()
 	configuration.Service = "hive"
 	configuration.FetchSize = 1000
-	configuration.Timeout = time.Hour
+	configuration.ConnectTimeout = time.Hour
 
 	connection, cursor := makeConnectionWithConnectConfiguration(t, configuration)
 	createTable(t, cursor)
@@ -648,7 +649,7 @@ func TestSimpleSelectWithDialFunctionAndTimeout(t *testing.T) {
 
 func TestSimpleSelectWithTimeout(t *testing.T) {
 	configuration := NewConnectConfiguration()
-	configuration.Timeout = time.Hour
+	configuration.ConnectTimeout = time.Hour
 	configuration.TransportMode = getTransport()
 	configuration.Service = "hive"
 	configuration.FetchSize = 1000
@@ -676,7 +677,7 @@ func TestConnectTimeoutWithDialFn(t *testing.T) {
 	configuration := NewConnectConfiguration()
 	configuration.Service = "hive"
 	configuration.TransportMode = mode
-	configuration.Timeout = 3 * time.Second
+	configuration.ConnectTimeout = 3 * time.Second
 	configuration.DialContext = SleepDialContext
 
 	if ssl {
@@ -701,10 +702,10 @@ func TestConnectTimeoutWithDialFn(t *testing.T) {
 	if !strings.Contains(errConn.Error(), "timeout") {
 		t.Fatalf("Expected a timeout error, but received: %+v", errConn)
 	}
-	if elapsed <= 2 * time.Second {
+	if elapsed <= 2*time.Second {
 		t.Fatalf("Timed out too fast: %v", elapsed)
 	}
-	if elapsed >= 4 * time.Second {
+	if elapsed >= 4*time.Second {
 		t.Fatalf("Timed out too slow: %v", elapsed)
 	}
 }
@@ -715,7 +716,7 @@ func TestConnectTimeout(t *testing.T) {
 	configuration := NewConnectConfiguration()
 	configuration.Service = "hive"
 	configuration.TransportMode = mode
-	configuration.Timeout = 3 * time.Second
+	configuration.ConnectTimeout = 3 * time.Second
 
 	if ssl {
 		tlsConfig, err := getTlsConfiguration("client.cer.pem", "client.cer.key")
@@ -739,10 +740,10 @@ func TestConnectTimeout(t *testing.T) {
 	if !strings.Contains(errConn.Error(), "timeout") {
 		t.Fatalf("Expected a timeout error, but received: %+v", errConn)
 	}
-	if elapsed <= 2 * time.Second {
+	if elapsed <= 2*time.Second {
 		t.Fatalf("Timed out too fast: %v", elapsed)
 	}
-	if elapsed >= 4 * time.Second {
+	if elapsed >= 4*time.Second {
 		t.Fatalf("Timed out too slow: %v", elapsed)
 	}
 }