amazon-web-servicesgoaws-sdkamazon-athenaaws-sdk-go

Wait for AWS Athena query execution in Go SDK


I have a working code that runs an Athena Query and waits for the query to finish by polling the error return from GetQueryResults using the following code:

func GetQueryResults(client *athena.Client, QueryID *string) []types.Row {

    params := &athena.GetQueryResultsInput{
        QueryExecutionId: QueryID,
    }

    data, err := client.GetQueryResults(context.TODO(), params)

    for err != nil {
        println(err.Error())
        time.Sleep(time.Second)
        data, err = client.GetQueryResults(context.TODO(), params)
    }

    return data.ResultSet.Rows
}

The problem is that in case the query fails, I have absolutely no way to break the loop.

For example, in Python I can do something like:

    while athena.get_query_execution(QueryExecutionId=execution_id)["QueryExecution"][
        "Status"
    ]["State"] in ["RUNNING", "QUEUED"]:
        sleep(2)

I can do a check like strings.Contains(err.Error(),"FAILED") inside the for loop, but I am looking for a cleaner way.

I tried looking for an equivalent for Go, but I wasn't successful. Is there a function for Go SDK that can return the execution status? Is there a better way to further examine an error in Go instead of err != nil?


Solution

  • The SDK has already provided retries.

    Here is an example for you, using aws-sdk-go-v2.

    package main
    
    import (
        "context"
        "fmt"
        "time"
    
        "github.com/aws/aws-sdk-go-v2/aws"
        "github.com/aws/aws-sdk-go-v2/service/athena"
        "github.com/aws/aws-sdk-go-v2/service/athena/types"
    )
    
    func main() {
        cfg := aws.NewConfig()
        ath := athena.NewFromConfig(*cfg)
    
        ctx, cancelFunc := context.WithTimeout(context.Background(), time.Minute*5)
        defer cancelFunc()
    
        rows, err := GetQueryResults(ctx, ath, aws.String("query-id"), 10)
        if err != nil {
            panic(err) // TODO: handle error
        }
    
        fmt.Println(rows)
    }
    
    func GetQueryResults(ctx context.Context, client *athena.Client, QueryID *string, attempts int) ([]types.Row, error) {
        t := time.NewTicker(time.Second * 5)
        defer t.Stop()
    
        attemptsFunc := func(o *athena.Options) { o.RetryMaxAttempts = attempts }
    
    WAIT:
        for {
            select {
            case <-t.C:
                out, err := client.GetQueryExecution(ctx, &athena.GetQueryExecutionInput{
                    QueryExecutionId: QueryID,
                }, attemptsFunc)
                if err != nil {
                    return nil, err
                }
    
                switch out.QueryExecution.Status.State {
                case types.QueryExecutionStateCancelled,
                    types.QueryExecutionStateFailed,
                    types.QueryExecutionStateSucceeded:
                    break WAIT
                }
    
            case <-ctx.Done():
                break WAIT
            }
        }
    
        data, err := client.GetQueryResults(ctx, &athena.GetQueryResultsInput{
            QueryExecutionId: QueryID,
        })
        if err != nil {
            return nil, err
        }
    
        return data.ResultSet.Rows, nil
    }