Support AWS Managed Output Buckets in AthenaDriver Configuration

README.md

@@ -168,9 +168,13 @@ import (
 )
 
 func main() {
-	// Step 1. Set AWS Credential in Driver Config.
+	// Step 1. Set AWS Credential in Driver Config (default)
 	conf, _ := drv.NewDefaultConfig("s3://myqueryresults/",
 		"us-east-2", "DummyAccessID", "DummySecretAccessKey")
+
+	// (optional) - if using aws managed output buckets
+	// conf, _ := drv.NewDefaultConfig("", "us-east-2", "DummyAccessID", "DummySecretAccessKey", true)
+
 	// Step 2. Open Connection.
 	db, _ := sql.Open(drv.DriverName, conf.Stringify())
 	// Step 3. Query and print results

examples/auth.go

@@ -114,16 +114,32 @@ func useAWSCLIConfigForAuthProfileByManualSetup(profile string) {
 	os.Unsetenv("AWS_SDK_LOAD_CONFIG")
 }
 
+// To use AthenaDriver Config for authentication with AWS managed bucket enabled
+func useAthenaDriverConfigForAuthWithManagedBucket() {
+	conf, err := drv.NewDefaultConfig(secret.OutputBucketDev, secret.Region,
+		secret.AccessID, secret.SecretAccessKey, true)
+	if err != nil {
+		return
+	}
+
+	db, _ := sql.Open(drv.DriverName, conf.Stringify())
+	var i int
+	_ = db.QueryRow("SELECT 789").Scan(&i)
+	println("with AthenaDriver Config With Managed Bucket:", i)
+}
+
 func main() {
 	useAthenaDriverConfigForAuth()
 	useAWSCLIConfigForAuth()
 	useAWSCLIConfigForAuthProfileByEnv("henry")
 	useAWSCLIConfigForAuthProfileByManualSetup("profile-development")
+	useAthenaDriverConfigForAuthWithManagedBucket()
 }
 
 /*
 Sample Output:
 with AthenaDriver Config: 123
 with AWS CLI Config: 456
 with AWS CLI Config With Profile: 789
+with AthenaDriver Config With Managed Bucket: 789
 */

go/config.go

@@ -23,8 +23,8 @@ package athenadriver
 import (
 	"net/url"
 	"regexp"
-	"strings"
 	"strconv"
+	"strings"
 	"time"
 )
 
@@ -63,12 +63,23 @@ var (
 
 // NewDefaultConfig is to new a Config with some default values.
 func NewDefaultConfig(outputBucket string, region string, accessID string,
-	secretAccessKey string) (*Config, error) {
+	secretAccessKey string, managedBucket ...bool) (*Config, error) {
 	conf := NewNoOpsConfig()
-	err := conf.SetOutputBucket(outputBucket)
-	if err != nil {
-		return nil, err
+	var err error
+
+	if len(managedBucket) == 0 {
+		managedBucket = []bool{false}
+	}
+
+	conf.SetManagedBucket(managedBucket[0])
+
+	if !conf.GetManagedBucket() {
+		err = conf.SetOutputBucket(outputBucket)
+		if err != nil {
+			return nil, err
+		}
 	}
+
 	err = conf.SetRegion(region)
 	if err != nil {
 		return nil, err
@@ -293,6 +304,20 @@ func (c *Config) GetOutputBucket() string {
 	return c.dsn.Scheme + "://" + c.dsn.Host + "/" + c.dsn.Path
 }
 
+// GetManagedBucket is getter of ManagedBucket.
+func (c *Config) GetManagedBucket() bool {
+	return c.values.Get("managedBucket") == "true"
+}
+
+// SetManagedBucket is a setter of ManagedBucket.
+func (c *Config) SetManagedBucket(b bool) {
+	if b {
+		c.values.Set("managedBucket", "true")
+	} else {
+		c.values.Set("managedBucket", "false")
+	}
+}
+
 // GetWorkgroup is getter of Workgroup.
 func (c *Config) GetWorkgroup() Workgroup {
 	tagString := c.values.Get("tag")

go/connection.go

@@ -335,10 +335,11 @@ func (c *Connection) QueryContext(ctx context.Context, query string, namedArgs [
 		return nil, ErrInvalidQuery
 	}
 	wg := c.connector.config.GetWorkgroup()
+	var athenaWG *athena.WorkGroup
 	if wg.Name == "" {
 		wg.Name = DefaultWGName
 	} else if wg.Name != DefaultWGName {
-		athenaWG, err := getWG(ctx, c.athenaAPI, wg.Name)
+		athenaWG, err = getWG(ctx, c.athenaAPI, wg.Name)
 		if err != nil {
 			obs.Scope().Counter(DriverName + ".failure.querycontext.getwg").Inc(1)
 			obs.Log(WarnLevel, "Didn't find workgroup "+wg.Name+" due to: "+err.Error())
@@ -409,17 +410,26 @@ func (c *Connection) QueryContext(ctx context.Context, query string, namedArgs [
 	if err != nil {
 		return nil, err
 	}
-	resp, err := c.athenaAPI.StartQueryExecution(&athena.StartQueryExecutionInput{
+
+	// Build StartQueryExecutionInput
+	startInput := &athena.StartQueryExecutionInput{
 		QueryString:         aws.String(queryWithPlaceholders),
 		ExecutionParameters: executionParams,
 		QueryExecutionContext: &athena.QueryExecutionContext{
 			Database: aws.String(c.connector.config.GetDB()),
 		},
-		ResultConfiguration: &athena.ResultConfiguration{
-			OutputLocation: aws.String(c.connector.config.GetOutputBucket()),
-		},
 		WorkGroup: aws.String(wg.Name),
-	})
+	}
+
+	// Only set ResultConfiguration if awsmanaged bucked is not enabled
+	// This avoids conflicts with workgroups that have awsmanaged bucked enabled
+	if !c.connector.config.GetManagedBucket() {
+		startInput.ResultConfiguration = &athena.ResultConfiguration{
+			OutputLocation: aws.String(c.connector.config.GetOutputBucket()),
+		}
+	}
+
+	resp, err := c.athenaAPI.StartQueryExecution(startInput)
 	if err != nil {
 		if pseudoCommand == PCGetQID {
 			if reqerr, ok := err.(awserr.RequestFailure); ok {