support showing query history

cmd/history.go

@@ -0,0 +1,43 @@
+package cmd
+
+import (
+	"log"
+
+	lib "github.com/justmiles/athena-cli/lib"
+
+	"github.com/spf13/cobra"
+)
+
+var (
+	historyCmdMaxResults int
+	historyCmdWorkgroup  string
+)
+
+// HistoryCmd ...
+var HistoryCmd = &cobra.Command{
+	Use:   "history",
+	Short: "export Athena execution history",
+	Run: func(cmd *cobra.Command, args []string) {
+
+		h, err := lib.ListHistory(historyCmdMaxResults)
+
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		err = lib.RenderHistoryResults(h, q.Format)
+		// err = lib.RenderAsTable(h)
+		// err = lib.RenderAsCSV(h)
+
+		if err != nil {
+			log.Fatal(err)
+		}
+	},
+}
+
+func init() {
+	HistoryCmd.PersistentFlags().IntVarP(&historyCmdMaxResults, "max", "m", 100, "maximum results")
+	HistoryCmd.PersistentFlags().StringVarP(&q.Format, "format", "f", "csv", "format the output as either json, csv, or table")
+	// TODO: Add support for selecting a workgroup
+	// HistoryCmd.PersistentFlags().StringVarP(&historyCmdWorkgroup, "workgroup", "w", "default", "workgroup")
+}

cmd/import.go

@@ -4,5 +4,5 @@ import "github.com/spf13/cobra"
 
 // Import is used to import all of these package's commands
 func Import(rootCmd *cobra.Command) {
-	rootCmd.AddCommand(QueryCmd)
+	rootCmd.AddCommand(QueryCmd, HistoryCmd)
 }

cmd/query.go

@@ -36,7 +36,7 @@ var QueryCmd = &cobra.Command{
 		// Clean up
 		defer lib.CleanCache(file.Name())
 
-		err = lib.OutputData(q.Format, q.JMESPath, file)
+		err = lib.RenderQueryResults(q.Format, q.JMESPath, file)
 
 		if err != nil {
 			log.Fatal(err)

go.mod

@@ -5,6 +5,7 @@ go 1.12
 require (
 	github.com/aws/aws-sdk-go v1.34.13
 	github.com/dustin/go-humanize v1.0.0
+	github.com/gocarina/gocsv v0.0.0-20200827134620-49f5c3fa2b3e
 	github.com/olekukonko/tablewriter v0.0.4
 	github.com/recursionpharma/go-csv-map v0.0.0-20160524001940-792523c65ae9
 	github.com/sirupsen/logrus v1.6.0

go.sum

@@ -30,6 +30,8 @@ github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9
 github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk=
 github.com/go-sql-driver/mysql v1.5.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
+github.com/gocarina/gocsv v0.0.0-20200827134620-49f5c3fa2b3e h1:f9zU2ojLUYe8f/uWnetnr0p6TnAGQBCV/WdPKxodBqA=
+github.com/gocarina/gocsv v0.0.0-20200827134620-49f5c3fa2b3e/go.mod h1:5YoVOkjYAQumqlV356Hj3xeYh4BdZuLE0/nRkf2NKkI=
 github.com/gogo/protobuf v1.1.1/go.mod h1:r8qH/GZQm5c6nD/R0oafs1akxWv10x8SbQlK7atdtwQ=
 github.com/gogo/protobuf v1.2.1/go.mod h1:hp+jE20tsWTFYpLwKvXlhS1hjn+gTNwPg2I6zVXpSg4=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=

lib/history.go

@@ -0,0 +1,103 @@
+package lib
+
+import (
+	"log"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/athena"
+)
+
+// HistoricalExecution ..
+type HistoricalExecution struct {
+	Query                      string
+	Catalog                    string
+	Database                   string
+	QueryExecutionID           string
+	OutputLocation             string
+	State                      string
+	WorkGroup                  string
+	TotalExecutionTimeInMillis int64
+	DataScannedInBytes         int64
+	Cost                       float64
+	SubmissionDateTime         time.Time
+}
+
+// ListHistory returns query results per workgroup
+func ListHistory(maxResults int) (h []HistoricalExecution, err error) {
+	params := &athena.ListQueryExecutionsInput{
+		// WorkGroup
+	}
+
+	if maxResults < 50 {
+		params.MaxResults = aws.Int64(int64(maxResults))
+	}
+
+	err = svc.ListQueryExecutionsPages(params, func(page *athena.ListQueryExecutionsOutput, lastPage bool) bool {
+		h = append(h, getQueryExecution(page.QueryExecutionIds)...)
+		return len(h) <= maxResults
+	})
+	if err != nil {
+		return h, err
+	}
+
+	if len(h) > maxResults {
+		return h[0:maxResults], nil
+	}
+	return h, nil
+}
+
+func getQueryExecution(ids []*string) (h []HistoricalExecution) {
+	o, err := svc.BatchGetQueryExecution(&athena.BatchGetQueryExecutionInput{
+		QueryExecutionIds: ids,
+	})
+
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// fmt.Println(o.GoString(), err)
+	for _, e := range o.QueryExecutions {
+
+		he := HistoricalExecution{
+			Query:                      *e.Query,
+			QueryExecutionID:           *e.QueryExecutionId,
+			OutputLocation:             *e.ResultConfiguration.OutputLocation,
+			State:                      *e.Status.State,
+			WorkGroup:                  *e.WorkGroup,
+			TotalExecutionTimeInMillis: *e.Statistics.TotalExecutionTimeInMillis,
+			SubmissionDateTime:         *e.Status.SubmissionDateTime,
+		}
+
+		if e.Statistics.DataScannedInBytes != nil {
+			he.DataScannedInBytes = *e.Statistics.DataScannedInBytes
+			he.Cost = float64((float64(*e.Statistics.DataScannedInBytes) / 1024 / 1024 / 1024 / 1024) * 5)
+		}
+
+		if e.QueryExecutionContext.Database != nil {
+			he.Database = *e.QueryExecutionContext.Database
+		}
+
+		if e.QueryExecutionContext.Catalog != nil {
+			he.Catalog = *e.QueryExecutionContext.Catalog
+		} else {
+			he.Catalog = "AwsDataCatalog"
+		}
+
+		h = append(h, he)
+	}
+
+	return h
+}
+
+// RenderHistoryResults ..
+func RenderHistoryResults(h []HistoricalExecution, outputFormat string) error {
+	switch outputFormat {
+	case "csv":
+		return RenderAsCSV(h)
+	case "table":
+		return RenderAsTable(h)
+	default:
+		return RenderAsJSON(h)
+	}
+}

lib/output.go

@@ -6,18 +6,18 @@ import (
 	"io/ioutil"
 	"os"
 
+	"github.com/gocarina/gocsv"
 	"github.com/olekukonko/tablewriter"
 	csvmap "github.com/recursionpharma/go-csv-map"
 )
 
-// OutputData formats json data as the desired output format
-// and sends the results to stdout
-func OutputData(outputFormat string, query string, file *os.File) error {
+// RenderQueryResults formats query results a in the
+// desired format and sends to stdout
+func RenderQueryResults(outputFormat string, query string, file *os.File) error {
 
 	if outputFormat == "json" {
 
 		reader := csvmap.NewReader(file)
-		reader.Columns, _ = reader.ReadHeader()
 		records, err := reader.ReadAll()
 		if err != nil {
 			return fmt.Errorf("Unable to read query results from %q, %v", file.Name(), err)
@@ -73,3 +73,68 @@ func values(c []string, m map[string]string) []string {
 
 	return values
 }
+
+// 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
+	}
+
+	if len(d) == 0 {
+		return nil
+	}
+
+	table := tablewriter.NewWriter(os.Stdout)
+
+	var headers []string
+	for key := range d[0] {
+		headers = append(headers, key)
+	}
+	table.SetHeader(headers)
+	table.SetBorders(tablewriter.Border{Left: true, Top: false, Right: true, Bottom: false})
+	table.SetCenterSeparator("|")
+
+	for _, record := range d {
+		var row []string
+		for _, key := range headers {
+			row = append(row, fmt.Sprintln(record[key]))
+		}
+		table.Append(row)
+	}
+
+	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.Println(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.Println(string(data))
+	return nil
+
+}

lib/query.go

@@ -135,7 +135,7 @@ func (q *Query) ExecuteToStdout() error {
 	// Clean up
 	defer CleanCache(file.Name())
 
-	OutputData(q.Format, q.JMESPath, file)
+	RenderQueryResults(q.Format, q.JMESPath, file)
 
 	return nil
 }