@@ -8,15 +8,20 @@ import (
88 "errors"
99 "github.com/aws/aws-sdk-go-v2/aws"
1010 "github.com/aws/aws-sdk-go-v2/service/redshiftdata"
11+ "github.com/aws/aws-sdk-go-v2/service/redshiftdata/types"
1112 "github.com/awsdocs/aws-doc-sdk-examples/gov2/demotools"
1213 "log"
1314)
1415
16+ // snippet-start:[gov2.redshift.DataActionsStruct]
17+
1518// RedshiftDataActions wraps RedshiftData actions.
1619type RedshiftDataActions struct {
1720 RedshiftDataClient * redshiftdata.Client
1821}
1922
23+ // snippet-end:[gov2.redshift.DataActionsStruct]
24+
2025// snippet-start:[gov2.redshift.RedshiftQuery.struct]
2126
2227// RedshiftQuery makes it easier to deal with RedshiftQuery objects.
@@ -32,9 +37,7 @@ type RedshiftQuery struct {
3237
3338// ExecuteStatement calls the ExecuteStatement operation from the RedshiftDataClient
3439func (actor RedshiftDataActions ) ExecuteStatement (ctx context.Context , input redshiftdata.ExecuteStatementInput ) (* redshiftdata.ExecuteStatementOutput , error ) {
35-
3640 return actor .RedshiftDataClient .ExecuteStatement (ctx , & input )
37-
3841}
3942
4043// snippet-end:[gov2.redshift.ExecuteStatement
@@ -52,19 +55,28 @@ func (actor RedshiftDataActions) ExecuteBatchStatement(ctx context.Context, inpu
5255
5356// ListDatabases lists all databases in the given cluster.
5457func (actor RedshiftDataActions ) ListDatabases (ctx context.Context , clusterId string , databaseName string , userName string ) error {
55- input := redshiftdata.ListDatabasesInput {
56- ClusterIdentifier : aws .String (clusterId ),
58+
59+ var opErr * types.DatabaseConnectionException
60+ var databaseNames []string
61+ paginator := redshiftdata .NewListDatabasesPaginator (actor .RedshiftDataClient , & redshiftdata.ListDatabasesInput {
5762 Database : aws .String (databaseName ),
63+ ClusterIdentifier : aws .String (clusterId ),
5864 DbUser : aws .String (userName ),
65+ })
66+ for paginator .HasMorePages () {
67+ output , err := paginator .NextPage (ctx )
68+ if err != nil && errors .As (err , & opErr ) {
69+ log .Printf ("Could not connect to the database." )
70+ panic (err )
71+ } else if err != nil {
72+ log .Printf ("Couldn't finish listing the tables. Here's why: %v\n " , err )
73+ return err
74+ } else {
75+ databaseNames = append (databaseNames , output .Databases ... )
76+ }
5977 }
6078
61- output , err := actor .RedshiftDataClient .ListDatabases (ctx , & input )
62- if err != nil {
63- log .Printf ("Failed to list databases: %v\n " , err )
64- return err
65- }
66-
67- for _ , database := range output .Databases {
79+ for _ , database := range databaseNames {
6880 log .Printf ("The database name is : %s\n " , database )
6981 }
7082 return nil
@@ -90,8 +102,11 @@ func (actor RedshiftDataActions) CreateTable(ctx context.Context, clusterId stri
90102 Sql : aws .String (sql ),
91103 }
92104
105+ var opErr * types.DatabaseConnectionException
93106 output , err := actor .RedshiftDataClient .ExecuteStatement (ctx , createTableInput )
94- if err != nil {
107+ if err != nil && errors .As (err , & opErr ) {
108+ log .Printf ("Could not connect to the database." )
109+ } else if err != nil {
95110 log .Printf ("Failed to create table: %v\n " , err )
96111 return err
97112 }
@@ -114,8 +129,11 @@ func (actor RedshiftDataActions) DeleteTable(ctx context.Context, clusterId stri
114129 Sql : aws .String (sql ),
115130 }
116131
132+ var opErr * types.DatabaseConnectionException
117133 output , err := actor .RedshiftDataClient .ExecuteStatement (ctx , deleteTableInput )
118- if err != nil {
134+ if err != nil && errors .As (err , & opErr ) {
135+ log .Printf ("Could not connect to the database." )
136+ } else if err != nil {
119137 log .Printf ("Failed to delete table " + tableName + " from " + databaseName + " database: %v\n " , err )
120138 return false , err
121139 }
@@ -162,6 +180,24 @@ func (actor RedshiftDataActions) DeleteDataRows(ctx context.Context, clusterId s
162180
163181// snippet-end:[gov2.redshift.DeleteRows]
164182
183+ // snippet-start:[gov2.redshift.DescribeStatement]
184+
185+ // DescribeStatement gets information about the given statement.
186+ func (actor RedshiftDataActions ) DescribeStatement (query RedshiftQuery ) (* redshiftdata.DescribeStatementOutput , error ) {
187+ describeOutput , err := actor .RedshiftDataClient .DescribeStatement (query .Context , & query .Input )
188+ var opErr * types.QueryTimeoutException
189+ if errors .As (err , & opErr ) {
190+ println ("The connection to the redshift data request timed out." )
191+ panic (err )
192+ } else if err != nil {
193+ println ("Failed to execute describe statement" )
194+ return nil , err
195+ }
196+ return describeOutput , nil
197+ }
198+
199+ // snippet-end:[gov2.redshift.DescribeStatement]
200+
165201// snippet-start:[gov2.redshift.WaitForQueryStatus]
166202
167203// WaitForQueryStatus waits until the given RedshiftQuery object has succeeded or failed.
@@ -170,8 +206,7 @@ func (actor RedshiftDataActions) WaitForQueryStatus(query RedshiftQuery, pauser
170206 attempts := 0
171207 maxWaitCycles := 30
172208 for done == false {
173- // snippet-start:[gov2.redshift.DescribeStatement]
174- describeOutput , err := actor .RedshiftDataClient .DescribeStatement (query .Context , & query .Input )
209+ describeOutput , err := actor .DescribeStatement (query )
175210 if err != nil {
176211 return err
177212 }
@@ -187,7 +222,6 @@ func (actor RedshiftDataActions) WaitForQueryStatus(query RedshiftQuery, pauser
187222 if describeOutput .Status == "FINISHED" {
188223 done = true
189224 }
190- // snippet-end:[gov2.redshift.DescribeStatement]
191225 attempts ++
192226 pauser .Pause (attempts )
193227 }
@@ -200,10 +234,15 @@ func (actor RedshiftDataActions) WaitForQueryStatus(query RedshiftQuery, pauser
200234
201235// GetStatementResult returns the result of the statement with the given id.
202236func (actor RedshiftDataActions ) GetStatementResult (ctx context.Context , statementId string ) (* redshiftdata.GetStatementResultOutput , error ) {
237+
238+ var opErr * types.QueryTimeoutException
203239 getStatementResultOutput , err := actor .RedshiftDataClient .GetStatementResult (ctx , & redshiftdata.GetStatementResultInput {
204240 Id : aws .String (statementId ),
205241 })
206- if err != nil {
242+ if err != nil && errors .As (err , & opErr ) {
243+ log .Printf ("Query timed out: %v\n " , err )
244+ return nil , err
245+ } else if err != nil {
207246 return nil , err
208247 }
209248 return getStatementResultOutput , nil
0 commit comments