From e1c74a3c8e9f9a066259f1310409d91fa4930940 Mon Sep 17 00:00:00 2001
From: Aris Tzoumas <atzoumas@rudderstack.com>
Date: Fri, 11 Oct 2024 16:32:50 +0300
Subject: [PATCH] chore(redshift): retry schema does not exist errors in list
 schemas

---
 go.mod                                        |  2 +-
 sqlconnect/internal/redshift/schemaadmin.go   | 32 +++++++++++++++++++
 .../internal/redshift/schemaadmin_test.go     | 11 +++++++
 3 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 sqlconnect/internal/redshift/schemaadmin.go
 create mode 100644 sqlconnect/internal/redshift/schemaadmin_test.go

diff --git a/go.mod b/go.mod
index 97db235..1807de5 100644
--- a/go.mod
+++ b/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.17.39
 	github.com/aws/aws-sdk-go-v2/service/redshiftdata v1.30.1
 	github.com/aws/aws-sdk-go-v2/service/sts v1.32.0
+	github.com/cenkalti/backoff/v4 v4.3.0
 	github.com/databricks/databricks-sql-go v1.6.1
 	github.com/dlclark/regexp2 v1.11.4
 	github.com/gliderlabs/ssh v0.3.7
@@ -70,7 +71,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sso v1.24.0 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.0 // indirect
 	github.com/aws/smithy-go v1.22.0 // indirect
-	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/containerd/continuity v0.4.3 // indirect
 	github.com/coreos/go-oidc/v3 v3.5.0 // indirect
 	github.com/danieljoos/wincred v1.1.2 // indirect
diff --git a/sqlconnect/internal/redshift/schemaadmin.go b/sqlconnect/internal/redshift/schemaadmin.go
new file mode 100644
index 0000000..e6613a7
--- /dev/null
+++ b/sqlconnect/internal/redshift/schemaadmin.go
@@ -0,0 +1,32 @@
+package redshift
+
+import (
+	"context"
+	"regexp"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+
+	"github.com/rudderlabs/sqlconnect-go/sqlconnect"
+)
+
+var schemaDoesNotExistRegex = regexp.MustCompile(`schema "(.*)" does not exist`)
+
+func (db *DB) ListSchemas(ctx context.Context) ([]sqlconnect.SchemaRef, error) {
+	// If the list schemas query is executed while a schema is being deleted, the query will fail with a schema does not exist error.
+	retryableError := func(err error) bool {
+		return schemaDoesNotExistRegex.MatchString(err.Error())
+	}
+	return retryOperationWithData(ctx, func() ([]sqlconnect.SchemaRef, error) {
+		schemas, err := db.DB.ListSchemas(ctx)
+		if err != nil && !retryableError(err) {
+			return nil, backoff.Permanent(err)
+		}
+		return schemas, err
+	})
+}
+
+// retryOperationWithData retries the given operation with a constant backoff policy of 100ms for 10 times.
+func retryOperationWithData[T any](ctx context.Context, o backoff.OperationWithData[T]) (T, error) {
+	return backoff.RetryWithData(o, backoff.WithMaxRetries(backoff.WithContext(backoff.NewConstantBackOff(100*time.Millisecond), ctx), 10))
+}
diff --git a/sqlconnect/internal/redshift/schemaadmin_test.go b/sqlconnect/internal/redshift/schemaadmin_test.go
new file mode 100644
index 0000000..59d982c
--- /dev/null
+++ b/sqlconnect/internal/redshift/schemaadmin_test.go
@@ -0,0 +1,11 @@
+package redshift
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSchemaDoesNotExistRegex(t *testing.T) {
+	require.True(t, schemaDoesNotExistRegex.MatchString(`ERROR: schema "tsqlcon_zvtquuiqulpz_1728651180" does not exist`))
+}