Added a query status check to account for race conditions creating the table.

Author: beqqrry-aws

gov2/redshift/actions/redshiftdata_actions.go

@@ -87,7 +87,7 @@ func (actor RedshiftDataActions) ListDatabases(ctx context.Context, clusterId st
 // snippet-start:[gov2.redshift.CreateTable]
 
 // CreateTable creates a table named <tableName> in the <databaseName> database with the given arguments.
-func (actor RedshiftDataActions) CreateTable(ctx context.Context, clusterId string, databaseName string, tableName string, userName string, args []string) error {
+func (actor RedshiftDataActions) CreateTable(ctx context.Context, clusterId string, databaseName string, tableName string, userName string, pauser demotools.IPausable, args []string) (*redshiftdata.ExecuteStatementOutput, error) {
 	sql := "CREATE TABLE " + tableName + " (" +
 		"id bigint identity(1, 1), " +
 		"PRIMARY KEY (id)"
@@ -108,11 +108,25 @@ func (actor RedshiftDataActions) CreateTable(ctx context.Context, clusterId stri
 		log.Printf("Could not connect to the database.")
 	} else if err != nil {
 		log.Printf("Failed to create table: %v\n", err)
-		return err
+		return nil, err
+	}
+
+	describeStatementInput := &redshiftdata.DescribeStatementInput{Id: output.Id}
+	query := RedshiftQuery{
+		Context: ctx,
+		Input:   *describeStatementInput,
+		Result:  output,
 	}
 
+	err = actor.WaitForQueryStatus(query, pauser, true)
+	if err != nil {
+		log.Printf("Failed to execute query: %v\n", err)
+		panic(err)
+	}
+	log.Printf("Successfully executed query\n")
+
 	log.Println("Table created:", *output.Id)
-	return nil
+	return output, nil
 }
 
 // snippet-end:[gov2.redshift.CreateTable]

gov2/redshift/scenarios/redshift_basics.go

@@ -195,12 +195,28 @@ func (runner *RedshiftBasicsScenario) Run() {
 	// Create the "Movies" table
 	log.Println("Now you will create a table named " + tableName + ".")
 	runner.questioner.Ask("Press Enter to continue...")
-	err = runner.redshiftDataActor.CreateTable(ctx, clusterId, databaseName, tableName, user.Username, []string{"title VARCHAR(256)", "year INT"})
+	err = nil
+	result, err := runner.redshiftDataActor.CreateTable(ctx, clusterId, databaseName, tableName, user.Username, runner.pauser, []string{"title VARCHAR(256)", "year INT"})
 	if err != nil {
 		log.Printf("Failed to create table: %v\n", err)
 		panic(err)
 	}
 
+	describeInput := redshiftdata.DescribeStatementInput{
+		Id: result.Id,
+	}
+	query := actions.RedshiftQuery{
+		Context: ctx,
+		Input:   describeInput,
+		Result:  result,
+	}
+	err = runner.redshiftDataActor.WaitForQueryStatus(query, runner.pauser, true)
+	if err != nil {
+		log.Printf("Failed to execute query: %v\n", err)
+		panic(err)
+	}
+	log.Printf("Successfully executed query\n")
+
 	// Populate the "Movies" table
 	runner.PopulateMoviesTable(ctx, clusterId, databaseName, tableName, user.Username, fileName)
 

gov2/redshift/scenarios/redshift_basics_test.go

@@ -102,6 +102,8 @@ func (scenarioTest *BasicsScenarioTest) SetupDataAndStubs() []testtools.Stub {
 	stubList = append(stubList, stubs.StubDescribeClusters(clusterId, nil))
 	stubList = append(stubList, stubs.StubListDatabases(clusterId, databaseName, user.Username, nil))
 	stubList = append(stubList, stubs.StubExecuteStatement(clusterId, databaseName, user.Username, sql, testId, nil))
+	stubList = append(stubList, stubs.StubDescribeStatement(testId, nil))
+	stubList = append(stubList, stubs.StubDescribeStatement(testId, nil))
 	stubList = append(stubList, stubs.StubBatchExecuteStatement(clusterId, databaseName, user.Username, sqls, testId, nil))
 	stubList = append(stubList, stubs.StubDescribeStatement(testId, nil))
 	stubList = append(stubList, stubs.StubExecuteStatement(clusterId, databaseName, user.Username, sql, testId, nil))