Merge pull request #11 from muroon/result-mode

Result Mode

Author: muroon

README.md

@@ -42,6 +42,20 @@ If a method must be supplied to satisfy a standard library interface but is unsu
 the driver will **panic** indicating so. If there are new offerings in Athena and/or
 helpful additions, feel free to PR.
 
+## Result Mode
+
+go-athena has the following modes to get the result of the query.
+
+- API (default)
+- DL
+- GZIP DL
+
+Note
+
+- DL and GZIP DL Mode are used only in the Select statement.
+  - Other statements automatically use API mode under DL or GZIP DL Mode.
+- Detailed explanation is described [here](doc/result_mode.md).
+- [Usages of Result Mode](doc/result_mode.md#usages).
 
 ## Testing
 

conn.go

@@ -4,10 +4,15 @@ import (
 	"context"
 	"database/sql/driver"
 	"errors"
+	"fmt"
 	"regexp"
+	"strings"
 	"time"
 
+	uuid "github.com/satori/go.uuid"
+
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/athena"
 	"github.com/aws/aws-sdk-go/service/athena/athenaiface"
 )
@@ -19,6 +24,11 @@ type conn struct {
 	workgroup      string
 
 	pollFrequency time.Duration
+
+	resultMode ResultMode
+	session    *session.Session
+	timeout    uint
+	catalog    string
 }
 
 func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
@@ -40,6 +50,38 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 }
 
 func (c *conn) runQuery(ctx context.Context, query string) (driver.Rows, error) {
+	// result mode
+	isSelect := isSelectQuery(query)
+	resultMode := c.resultMode
+	if rmode, ok := getResultMode(ctx); ok {
+		resultMode = rmode
+	}
+	if !isSelect {
+		resultMode = ResultModeAPI
+	}
+
+	// timeout
+	timeout := c.timeout
+	if to, ok := getTimeout(ctx); ok {
+		timeout = to
+	}
+
+	// catalog
+	catalog := c.catalog
+	if cat, ok := getCatalog(ctx); ok {
+		catalog = cat
+	}
+
+	// mode ctas
+	var ctasTable string
+	var afterDownload func() error
+	if isSelect && resultMode == ResultModeGzipDL {
+		// Create AS Select
+		ctasTable = fmt.Sprintf("tmp_ctas_%v", strings.Replace(uuid.NewV4().String(), "-", "", -1))
+		query = fmt.Sprintf("CREATE TABLE %s WITH (format='TEXTFILE') AS %s", ctasTable, query)
+		afterDownload = c.dropCTASTable(ctx, ctasTable)
+	}
+
 	queryID, err := c.startQuery(query)
 	if err != nil {
 		return nil, err
@@ -50,12 +92,33 @@ func (c *conn) runQuery(ctx context.Context, query string) (driver.Rows, error)
 	}
 
 	return newRows(rowsConfig{
-		Athena:     c.athena,
-		QueryID:    queryID,
-		SkipHeader: !isDDLQuery(query),
+		Athena:         c.athena,
+		QueryID:        queryID,
+		SkipHeader:     !isDDLQuery(query),
+		ResultMode:     resultMode,
+		Session:        c.session,
+		OutputLocation: c.OutputLocation,
+		Timeout:        timeout,
+		AfterDownload:  afterDownload,
+		CTASTable:      ctasTable,
+		DB:             c.db,
+		Catalog:        catalog,
 	})
 }
 
+func (c *conn) dropCTASTable(ctx context.Context, table string) func() error {
+	return func() error {
+		query := fmt.Sprintf("DROP TABLE %s", table)
+
+		queryID, err := c.startQuery(query)
+		if err != nil {
+			return err
+		}
+
+		return c.waitOnQuery(ctx, queryID)
+	}
+}
+
 // startQuery starts an Athena query and returns its ID.
 func (c *conn) startQuery(query string) (string, error) {
 	resp, err := c.athena.StartQueryExecution(&athena.StartQueryExecutionInput{
@@ -146,3 +209,11 @@ var ddlQueryRegex = regexp.MustCompile(`(?i)^(ALTER|CREATE|DESCRIBE|DROP|MSCK|SH
 func isDDLQuery(query string) bool {
 	return ddlQueryRegex.Match([]byte(query))
 }
+
+func isSelectQuery(query string) bool {
+	return regexp.MustCompile(`(?i)^SELECT`).Match([]byte(query))
+}
+
+func isCTASQuery(query string) bool {
+	return regexp.MustCompile(`(?i)^CREATE.+AS\s+SELECT`).Match([]byte(query))
+}

context.go

@@ -0,0 +1,72 @@
+package athena
+
+import "context"
+
+const contextPrefix string = "go-athena"
+
+/*
+ * Result Mode
+ */
+
+const resultModeContextKey string = "result_mode_key"
+
+// ResultModeContextKey context key of setting result mode
+var ResultModeContextKey string = contextPrefix + resultModeContextKey
+
+// SetAPIMode set APIMode to ResultMode from context
+func SetAPIMode(ctx context.Context) context.Context {
+	return context.WithValue(ctx, ResultModeContextKey, ResultModeAPI)
+}
+
+// SetDLMode set DownloadMode to ResultMode from context
+func SetDLMode(ctx context.Context) context.Context {
+	return context.WithValue(ctx, ResultModeContextKey, ResultModeDL)
+}
+
+// SetGzipDLMode set CTASMode to ResultMode from context
+func SetGzipDLMode(ctx context.Context) context.Context {
+	return context.WithValue(ctx, ResultModeContextKey, ResultModeGzipDL)
+}
+
+func getResultMode(ctx context.Context) (ResultMode, bool) {
+	val, ok := ctx.Value(ResultModeContextKey).(ResultMode)
+	return val, ok
+}
+
+/*
+ * timeout
+ */
+
+const timeoutContextKey string = "timeout_key"
+
+// TimeoutContextKey context key of setting timeout
+var TimeoutContextKey string = contextPrefix + timeoutContextKey
+
+// SetTimeout set timeout from context
+func SetTimeout(ctx context.Context, timeout uint) context.Context {
+	return context.WithValue(ctx, TimeoutContextKey, timeout)
+}
+
+func getTimeout(ctx context.Context) (uint, bool) {
+	val, ok := ctx.Value(TimeoutContextKey).(uint)
+	return val, ok
+}
+
+/*
+ * catalog
+ */
+
+const catalogContextKey string = "catalog_key"
+
+// CatalogContextKey context key of setting catalog
+var CatalogContextKey string = contextPrefix + catalogContextKey
+
+// SetCatalog set catalog from context
+func SetTimout(ctx context.Context, catalog string) context.Context {
+	return context.WithValue(ctx, CatalogContextKey, catalog)
+}
+
+func getCatalog(ctx context.Context) (string, bool) {
+	val, ok := ctx.Value(CatalogContextKey).(string)
+	return val, ok
+}