Merge pull request #36 from weissleb/configure-timeout

add dml and ddl timeout overrides

README.md

@@ -642,6 +642,58 @@ Sample Output:
 2020/01/20 15:28:35 context deadline exceeded
 ```
 
+### Overriding Athena Service Limits for Query Timeout
+This library assumes default [Athena service limits](https://docs.aws.amazon.com/athena/latest/ug/service-limits.html) for DDL and DML query timeouts, as can be found in `athenadriver/go/constants.go`.
+If you've increased your service limits, for example via the [Athena Service Quotas](https://console.aws.amazon.com/servicequotas/home/services/athena/quotas) console,
+you can override them on your `Config`.
+
+Here's the same example found at [Query Cancellation](#query-cancellation), but with an *increased* query timeout.
+
+```go
+package main
+
+import (
+	"context"
+	"database/sql"
+	"log"
+	"time"
+	drv "github.com/uber/athenadriver/go"
+)
+
+func main() {
+	// 1. Set AWS Credential in Driver Config.
+	conf, _ := drv.NewDefaultConfig("s3://myqueryresults/",
+		"us-east-2", "DummyAccessID", "DummySecretAccessKey")
+	
+	// 2. Override the DML query timeout to 60 minutes (3600 seconds). 
+	serviceLimitOverride := drv.NewServiceLimitOverride()
+	serviceLimitOverride.SetDMLQueryTimeout(3600)
+	conf.SetServiceLimitOverride(*serviceLimitOverride)
+
+	// 3. Open Connection.
+	dsn := conf.Stringify()
+	db, _ := sql.Open(drv.DriverName, dsn)
+	
+	// 4. Run the query.
+	rows, err := db.QueryContext(context.Background(), "select count(*) from sampledb.elb_logs")
+	if err != nil {
+		log.Fatal(err)
+		return
+	}
+	defer rows.Close()
+
+	var requestTimestamp string
+	var url string
+	for rows.Next() {
+		if err := rows.Scan(&requestTimestamp, &url); err != nil {
+			log.Fatal(err)
+		}
+		println(requestTimestamp + "," + url)
+	}
+}
+```
+
+
 ### Missing Value Handling 
 
 It is common to have missing values in S3 file, or Athena DB. When this happens, you can specify if you want to use

go/config.go

@@ -421,3 +421,17 @@ func (c *Config) SetAWSProfile(profile string) {
 func (c *Config) GetAWSProfile() string {
 	return c.values.Get("AWSProfile")
 }
+
+// SetServiceLimitOverride is to set values from a ServiceLimitOverride
+func (c *Config) SetServiceLimitOverride(serviceLimitOverride ServiceLimitOverride) {
+	for k, v := range serviceLimitOverride.GetAsStringMap() {
+		c.values.Set(k, v)
+	}
+}
+
+// GetServiceLimitOverride is to get the ServiceLimitOverride manually set by a user
+func (c *Config) GetServiceLimitOverride() *ServiceLimitOverride {
+	serviceLimitOverride := NewServiceLimitOverride()
+	serviceLimitOverride.SetFromValues(c.values)
+	return serviceLimitOverride
+}

go/config_test.go

@@ -286,3 +286,29 @@ func TestConfig_SetAWSProfile(t *testing.T) {
 	testConf.SetAWSProfile("development")
 	assert.Equal(t, testConf.GetAWSProfile(), "development")
 }
+
+func TestConfig_SetServiceLimitOverride(t *testing.T) {
+	var s3bucket string = "s3://query-results-henry-wu-us-east-2/"
+
+	testConf := NewNoOpsConfig()
+	_ = testConf.SetOutputBucket(s3bucket)
+	serviceLimitOverride := NewServiceLimitOverride()
+	ddlQueryTimeout := 1000 * 60 // 1000 minutes
+	_ = serviceLimitOverride.SetDDLQueryTimeout(ddlQueryTimeout)
+	testConf.SetServiceLimitOverride(*serviceLimitOverride)
+	testServiceLimitOverride := testConf.GetServiceLimitOverride()
+	assert.Equal(t, ddlQueryTimeout, testServiceLimitOverride.GetDDLQueryTimeout())
+
+	expected := "s3://query-results-henry-wu-us-east-2?DDLQueryTimeout=60000&DMLQueryTimeout=0&WGRemoteCreation=true&db=default&missingAsEmptyString=true&region=us-east-1"
+	assert.Equal(t, expected, testConf.Stringify())
+
+	dmlQueryTimeout := 60 * 60 // 60 minutes
+	_ = serviceLimitOverride.SetDMLQueryTimeout(dmlQueryTimeout)
+	testConf.SetServiceLimitOverride(*serviceLimitOverride)
+	testServiceLimitOverride = testConf.GetServiceLimitOverride()
+	assert.Equal(t, ddlQueryTimeout, testServiceLimitOverride.GetDDLQueryTimeout())
+	assert.Equal(t, dmlQueryTimeout, testServiceLimitOverride.GetDMLQueryTimeout())
+
+	expected = "s3://query-results-henry-wu-us-east-2?DDLQueryTimeout=60000&DMLQueryTimeout=3600&WGRemoteCreation=true&db=default&missingAsEmptyString=true&region=us-east-1"
+	assert.Equal(t, expected, testConf.Stringify())
+}

go/connection.go

@@ -432,7 +432,7 @@ WAITING_FOR_RESULT:
 			obs.Log(ErrorLevel, "query canceled", zap.String("queryID", queryID))
 			return nil, ctx.Err()
 		case <-time.After(PoolInterval * time.Second):
-			if isQueryTimeOut(startOfStartQueryExecution, *statusResp.QueryExecution.StatementType) {
+			if isQueryTimeOut(startOfStartQueryExecution, *statusResp.QueryExecution.StatementType, c.connector.config.GetServiceLimitOverride()) {
 				obs.Log(ErrorLevel, "Query timeout failure",
 					zap.String("workgroup", wg.Name),
 					zap.String("queryID", queryID),

go/errors.go

@@ -22,6 +22,7 @@ package athenadriver
 
 import (
 	"errors"
+	"fmt"
 )
 
 // Various errors the driver might return. Can change between driver versions.
@@ -41,4 +42,5 @@ var (
 	ErrAthenaNilAPI                 = errors.New("athenaAPI must not be nil")
 	ErrTestMockGeneric              = errors.New("some_mock_error_for_test")
 	ErrTestMockFailedByAthena       = errors.New("the reason why Athena failed the query")
+	ErrServiceLimitOverride         = errors.New(fmt.Sprintf("service limit override must be greater than %d", PoolInterval))
 )

go/servicelimitoverride.go

@@ -0,0 +1,87 @@
+// Copyright (c) 2020 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package athenadriver
+
+import (
+	"fmt"
+	"net/url"
+	"strconv"
+)
+
+// ServiceLimitOverride allows users to override service limits, hardcoded in constants.go.
+// This assumes the service limits have been raised in the AWS account.
+// https://docs.aws.amazon.com/athena/latest/ug/service-limits.html
+type ServiceLimitOverride struct {
+	ddlQueryTimeout int
+	dmlQueryTimeout int
+}
+
+// NewServiceLimitOverride is to create an empty ServiceLimitOverride.
+// Values can be set using setters.
+func NewServiceLimitOverride() *ServiceLimitOverride {
+	return &ServiceLimitOverride{}
+}
+
+// SetDDLQueryTimeout is to set the DDLQueryTimeout override.
+func (c *ServiceLimitOverride) SetDDLQueryTimeout(seconds int) error {
+	if seconds < PoolInterval {
+		return ErrServiceLimitOverride
+	}
+	c.ddlQueryTimeout = seconds
+	return nil
+}
+
+// GetDDLQueryTimeout is to get the DDLQueryTimeout override.
+func (c *ServiceLimitOverride) GetDDLQueryTimeout() int {
+	return c.ddlQueryTimeout
+}
+
+// SetDMLQueryTimeout is to set the DMLQueryTimeout override.
+func (c *ServiceLimitOverride) SetDMLQueryTimeout(seconds int) error {
+	if seconds < PoolInterval {
+		return ErrServiceLimitOverride
+	}
+	c.dmlQueryTimeout = seconds
+	return nil
+}
+
+// GetDMLQueryTimeout is to get the DMLQueryTimeout override.
+func (c *ServiceLimitOverride) GetDMLQueryTimeout() int {
+	return c.dmlQueryTimeout
+}
+
+// GetAsStringMap is to get the ServiceLimitOverride as a map of strings
+// and aids in setting url.Values in Config
+func (c *ServiceLimitOverride) GetAsStringMap() map[string]string {
+	res := map[string]string{}
+	res["DDLQueryTimeout"] = fmt.Sprintf("%d", c.ddlQueryTimeout)
+	res["DMLQueryTimeout"] = fmt.Sprintf("%d", c.dmlQueryTimeout)
+	return res
+}
+
+// SetFromValues is to set ServiceLimitOverride properties from a url.Values
+// which might be a list of override and other ignored values from a dsn
+func (c *ServiceLimitOverride) SetFromValues(kvp url.Values) {
+	ddlQueryTimeout, _ := strconv.Atoi(kvp.Get("DDLQueryTimeout"))
+	_ = c.SetDDLQueryTimeout(ddlQueryTimeout)
+	dmlQueryTimeout, _ := strconv.Atoi(kvp.Get("DMLQueryTimeout"))
+	_ = c.SetDMLQueryTimeout(dmlQueryTimeout)
+}

go/servicelimitoverride_test.go

@@ -0,0 +1,37 @@
+package athenadriver
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+// Tests for ServiceLimitOverride.
+func TestNewServiceLimitOverride(t *testing.T) {
+	testConf := NewServiceLimitOverride()
+	assert.Zero(t, testConf.GetDDLQueryTimeout())
+	assert.Zero(t, testConf.GetDMLQueryTimeout())
+
+	ddlQueryTimeout := 30 * 60 // seconds
+	dmlQueryTimeout := 60 * 60 // seconds
+	testConf.SetDDLQueryTimeout(ddlQueryTimeout)
+	assert.Equal(t, ddlQueryTimeout, testConf.GetDDLQueryTimeout()) // seconds
+
+	testConf.SetDMLQueryTimeout(dmlQueryTimeout)
+	assert.Equal(t, dmlQueryTimeout, testConf.GetDMLQueryTimeout()) // seconds
+
+	ddlQueryTimeout = 0
+	dmlQueryTimeout = 0
+	err := testConf.SetDDLQueryTimeout(ddlQueryTimeout)
+	assert.NotNil(t, err)
+
+	err = testConf.SetDMLQueryTimeout(dmlQueryTimeout)
+	assert.NotNil(t, err)
+
+	ddlQueryTimeout = -1
+	dmlQueryTimeout = -1
+	err = testConf.SetDDLQueryTimeout(ddlQueryTimeout)
+	assert.NotNil(t, err)
+
+	err = testConf.SetDMLQueryTimeout(dmlQueryTimeout)
+	assert.NotNil(t, err)
+}

go/utils.go

@@ -665,22 +665,32 @@ func valueToNamedValue(args []driver.Value) []driver.NamedValue {
 	return nameValues
 }
 
-func isQueryTimeOut(startOfStartQueryExecution time.Time, queryType string) bool {
+func isQueryTimeOut(startOfStartQueryExecution time.Time, queryType string, serviceLimitOverride *ServiceLimitOverride) bool {
+	ddlQueryTimeout := DDLQueryTimeout
+	dmlQueryTimeout := DMLQueryTimeout
+	if serviceLimitOverride != nil {
+		if serviceLimitOverride.GetDDLQueryTimeout() > 0 {
+			ddlQueryTimeout = serviceLimitOverride.GetDDLQueryTimeout()
+		}
+		if serviceLimitOverride.GetDMLQueryTimeout() > 0 {
+			dmlQueryTimeout = serviceLimitOverride.GetDMLQueryTimeout()
+		}
+	}
 	switch queryType {
 	case "DDL":
 		return time.Since(startOfStartQueryExecution) >
-			DDLQueryTimeout*time.Second
+			time.Duration(ddlQueryTimeout)*time.Second
 	case "DML":
 		return time.Since(startOfStartQueryExecution) >
-			DMLQueryTimeout*time.Second
+			time.Duration(dmlQueryTimeout)*time.Second
 	case "UTILITY":
 		return time.Since(startOfStartQueryExecution) >
-			DMLQueryTimeout*time.Second
+			time.Duration(dmlQueryTimeout)*time.Second
 	case "TIMEOUT_NOW":
 		return true
 	default:
 		return time.Since(startOfStartQueryExecution) >
-			DDLQueryTimeout*time.Second
+			time.Duration(ddlQueryTimeout)*time.Second
 	}
 }
 

go/utils_test.go

@@ -206,14 +206,22 @@ func TestValueToNamedValue(t *testing.T) {
 }
 
 func TestIsQueryTimeOut(t *testing.T) {
-	assert.False(t, isQueryTimeOut(time.Now(), athena.StatementTypeDdl))
-	assert.False(t, isQueryTimeOut(time.Now(), athena.StatementTypeDml))
-	assert.False(t, isQueryTimeOut(time.Now(), athena.StatementTypeUtility))
+	assert.False(t, isQueryTimeOut(time.Now(), athena.StatementTypeDdl, nil))
+	assert.False(t, isQueryTimeOut(time.Now(), athena.StatementTypeDml, nil))
+	assert.False(t, isQueryTimeOut(time.Now(), athena.StatementTypeUtility, nil))
 	now := time.Now()
 	OneHourAgo := now.Add(-3600 * time.Second)
-	assert.True(t, isQueryTimeOut(OneHourAgo, athena.StatementTypeDml))
-	assert.False(t, isQueryTimeOut(OneHourAgo, athena.StatementTypeDdl))
-	assert.False(t, isQueryTimeOut(OneHourAgo, "UNKNOWN"))
+	assert.True(t, isQueryTimeOut(OneHourAgo, athena.StatementTypeDml, nil))
+	assert.False(t, isQueryTimeOut(OneHourAgo, athena.StatementTypeDdl, nil))
+	assert.False(t, isQueryTimeOut(OneHourAgo, "UNKNOWN", nil))
+
+	testConf := NewServiceLimitOverride()
+	testConf.SetDMLQueryTimeout(65 * 60) // 65 minutes
+	assert.False(t, isQueryTimeOut(OneHourAgo, athena.StatementTypeDml, testConf))
+
+	testConf.SetDDLQueryTimeout(30 * 60) // 30 minutes
+	assert.True(t, isQueryTimeOut(OneHourAgo, athena.StatementTypeDdl, testConf))
+	assert.True(t, isQueryTimeOut(OneHourAgo, "UNKNOWN", testConf))
 }
 
 func TestEscapeBytesBackslash(t *testing.T) {