merge master into pigfoot-feature-workgroup

cmd/query.go

@@ -18,7 +18,6 @@ var QueryCmd = &cobra.Command{
 	Use:   "query",
 	Short: "query Athena",
 	Run: func(cmd *cobra.Command, args []string) {
-
 		_, err := lib.ParseFormat(q.Format)
 		if err != nil {
 			log.Fatal(err)
@@ -55,6 +54,7 @@ func init() {
 	QueryCmd.PersistentFlags().StringVarP(&q.SQL, "sql", "s", "", "SQL query to execute. Can be a file or raw query")
 	QueryCmd.PersistentFlags().StringVarP(&q.Format, "format", "f", "csv", "format the output as either json, csv, or table")
 	QueryCmd.PersistentFlags().StringVarP(&q.OutputFile, "output", "o", "", "(optional) file name to write this content to (defaults to standard output)")
+	QueryCmd.PersistentFlags().StringVarP(&q.WorkGroup, "workgroup", "w", "", "(optional) WorkGroup (defaults to primary)")
 	QueryCmd.PersistentFlags().BoolVar(&q.Statistics, "statistics", false, "print query statistics to stderr")
 	// QueryCmd.PersistentFlags().StringVar(&q.JMESPath, "jmespath", "", "optional JMESPath to further filter or format results. See jmespath.org for more.")
 }

go.mod

@@ -1,19 +1,17 @@
 module github.com/justmiles/athena-cli
 
-go 1.12
+go 1.16
 
 require (
-	github.com/360EntSecGroup-Skylar/excelize/v2 v2.3.2
 	github.com/aws/aws-sdk-go v1.40.50
 	github.com/dustin/go-humanize v1.0.0
-	github.com/elliotchance/orderedmap v1.3.0
-	github.com/gocarina/gocsv v0.0.0-20200827134620-49f5c3fa2b3e
-	github.com/olekukonko/tablewriter v0.0.4
+	github.com/elliotchance/orderedmap v1.4.0
+	github.com/gocarina/gocsv v0.0.0-20210516172204-ca9e8a8ddea8
+	github.com/olekukonko/tablewriter v0.0.5
 	github.com/recursionpharma/go-csv-map v0.0.0-20160524001940-792523c65ae9
-	github.com/sirupsen/logrus v1.6.0
-	github.com/smartystreets/goconvey v1.6.4 // indirect
-	github.com/spf13/cobra v1.0.0
-	github.com/spf13/pflag v1.0.5 // indirect
-	github.com/xitongsys/parquet-go v1.5.3
-	github.com/xitongsys/parquet-go-source v0.0.0-20200817004010-026bad9b25d0
+	github.com/sirupsen/logrus v1.8.1
+	github.com/spf13/cobra v1.2.1
+	github.com/xitongsys/parquet-go v1.6.0
+	github.com/xitongsys/parquet-go-source v0.0.0-20201108113611-f372b7d813be
+	github.com/xuri/excelize/v2 v2.4.1
 )

lib/output.go

@@ -5,9 +5,9 @@ import (
 	"fmt"
 	"os"
 
-	"github.com/360EntSecGroup-Skylar/excelize/v2"
 	"github.com/gocarina/gocsv"
 	"github.com/olekukonko/tablewriter"
+	"github.com/xuri/excelize/v2"
 )
 
 func values(c []string, m map[string]string) []string {
@@ -22,12 +22,10 @@ func values(c []string, m map[string]string) []string {
 
 // RenderAsTable will render an interfact to table.
 func RenderAsTable(i interface{}) error {
-
 	data, _ := json.Marshal(i)
 
 	var d []map[string]interface{}
 	err := json.Unmarshal(data, &d)
-
 	if err != nil {
 		return err
 	}
@@ -56,33 +54,27 @@ func RenderAsTable(i interface{}) error {
 
 	table.Render()
 	return nil
-
 }
 
 // RenderAsCSV will render an interfact to table
 func RenderAsCSV(i interface{}) error {
-
 	data, err := gocsv.MarshalBytes(i)
 	if err != nil {
 		return err
 	}
 	fmt.Print(string(data))
 	return nil
-
 }
 
 // RenderAsJSON will render an interface as json
 func RenderAsJSON(i interface{}) error {
-
 	data, err := json.MarshalIndent(i, "", "  ")
-
 	if err != nil {
 		return err
 	}
 
 	fmt.Print(string(data))
 	return nil
-
 }
 
 // RenderAsXLSX will render an interface as XLSX

lib/query.go

@@ -8,7 +8,6 @@ import (
 	"path"
 	"time"
 
-	"github.com/360EntSecGroup-Skylar/excelize/v2"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/athena"
@@ -17,6 +16,7 @@ import (
 	"github.com/olekukonko/tablewriter"
 	csvmap "github.com/recursionpharma/go-csv-map"
 	"github.com/sirupsen/logrus"
+	"github.com/xuri/excelize/v2"
 )
 
 // Query ...
@@ -29,6 +29,7 @@ type Query struct {
 	Format             string
 	JMESPath           string
 	Statistics         bool
+	WorkGroup          string
 }
 
 // Format is an enumeration of available query output formats
@@ -39,7 +40,6 @@ type Format int
 
 // Execute a SQL query against Athena
 func (q *Query) Execute() (*os.File, error) {
-
 	// Check to see if `--sql` points to a file
 	if _, err := os.Stat(q.SQL); err == nil {
 		queryFromFile, err := ioutil.ReadFile(q.SQL)
@@ -49,16 +49,21 @@ func (q *Query) Execute() (*os.File, error) {
 		q.SQL = string(queryFromFile)
 	}
 
-	result, err := svc.StartQueryExecution(&athena.StartQueryExecutionInput{
+	startQueryExecutionInput := athena.StartQueryExecutionInput{
 		QueryString: &q.SQL,
 		QueryExecutionContext: &athena.QueryExecutionContext{
 			Database: &q.Database,
 		},
 		ResultConfiguration: &athena.ResultConfiguration{
 			OutputLocation: aws.String("s3://" + path.Join(q.QueryResultsBucket, q.QueryResultsPrefix)),
 		},
-	})
+	}
 
+	if q.WorkGroup != "" {
+		startQueryExecutionInput.WorkGroup = aws.String(q.WorkGroup)
+	}
+
+	result, err := svc.StartQueryExecution(&startQueryExecutionInput)
 	if err != nil {
 		return nil, err
 	}
@@ -105,7 +110,6 @@ func (q *Query) Execute() (*os.File, error) {
 			Bucket: aws.String(q.QueryResultsBucket),
 			Key:    aws.String(*result.QueryExecutionId + ".csv"),
 		})
-
 		if err != nil {
 			if aerr, ok := err.(awserr.Error); ok {
 				switch aerr.Code() {