gcs: add error classification

Author: jgao54

flow/alerting/classifier.go

@@ -88,6 +88,7 @@ const (
 	ErrorSourceMySQL           ErrorSource = "mysql"
 	ErrorSourceMongoDB         ErrorSource = "mongodb"
 	ErrorSourceBigQuery        ErrorSource = "bigquery"
+	ErrorSourceGCS             ErrorSource = "gcs"
 	ErrorSourcePostgresCatalog ErrorSource = "postgres_catalog"
 	ErrorSourceSSH             ErrorSource = "ssh_tunnel"
 	ErrorSourceNet             ErrorSource = "net"
@@ -943,6 +944,23 @@ func GetErrorClass(ctx context.Context, err error) (ErrorClass, ErrorInfo) {
 		return ErrorOther, bqErrorInfo
 	}
 
+	if _, ok := errors.AsType[*exceptions.GCSError](err); ok {
+		gcsErrorInfo := ErrorInfo{
+			Source: ErrorSourceGCS,
+			Code:   "UNKNOWN",
+		}
+		if apiErr, ok := errors.AsType[*googleapi.Error](err); ok {
+			gcsErrorInfo.Code = strconv.Itoa(apiErr.Code)
+			switch apiErr.Code {
+			case 503: // Service Unavailable
+				return ErrorRetryRecoverable, gcsErrorInfo
+			default:
+				return ErrorOther, gcsErrorInfo
+			}
+		}
+		return ErrorOther, gcsErrorInfo
+	}
+
 	if chException, ok := errors.AsType[*clickhouse.Exception](err); ok {
 		chErrorInfo := ErrorInfo{
 			Source: ErrorSourceClickHouse,

flow/alerting/classifier_test.go

@@ -1188,6 +1188,22 @@ func TestBigQueryGCSBucketNotExistShouldBeConnectivity(t *testing.T) {
 	}, errInfo)
 }
 
+func TestGCSTransientErrorsShouldBeRecoverable(t *testing.T) {
+	apiErr := &googleapi.Error{
+		Code:    503,
+		Message: "We encountered an internal error. Please try again.",
+		Errors:  []googleapi.ErrorItem{{Reason: "backendError"}},
+	}
+	err := exceptions.NewGCSError(fmt.Errorf("failed to finalize GCS upload for a.avro: %w", apiErr))
+	errorClass, errInfo := GetErrorClass(t.Context(), fmt.Errorf(
+		"failed to push records: failed to upload to staging: %w", err))
+	assert.Equal(t, ErrorRetryRecoverable, errorClass)
+	assert.Equal(t, ErrorInfo{
+		Source: ErrorSourceGCS,
+		Code:   strconv.Itoa(503),
+	}, errInfo)
+}
+
 func TestBigQueryUnclassifiedCodeShouldBeOther(t *testing.T) {
 	t.Parallel()
 

flow/connectors/clickhouse/staging_gcs.go

@@ -17,6 +17,7 @@ import (
 	peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse"
 	"github.com/PeerDB-io/peerdb/flow/pkg/objectstore"
 	"github.com/PeerDB-io/peerdb/flow/shared"
+	"github.com/PeerDB-io/peerdb/flow/shared/exceptions"
 )
 
 const gcsSignedURLExpiry = 1 * time.Hour
@@ -50,7 +51,7 @@ func newGCSStagingStore(ctx context.Context, bucketName string) (StagingStore, e
 	// Uses Application Default Credentials (Workload Identity on GKE)
 	client, err := storage.NewClient(ctx)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create GCS client: %w", err)
+		return nil, exceptions.NewGCSError(fmt.Errorf("failed to create GCS client: %w", err))
 	}
 
 	return &gcsStagingStore{
@@ -69,10 +70,10 @@ func (g *gcsStagingStore) Upload(ctx context.Context, env map[string]string, key
 
 	if _, err := io.Copy(w, body); err != nil {
 		w.Close()
-		return fmt.Errorf("failed to write to GCS object %s: %w", key, err)
+		return exceptions.NewGCSError(fmt.Errorf("failed to write to GCS object %s: %w", key, err))
 	}
 	if err := w.Close(); err != nil {
-		return fmt.Errorf("failed to finalize GCS upload for %s: %w", key, err)
+		return exceptions.NewGCSError(fmt.Errorf("failed to finalize GCS upload for %s: %w", key, err))
 	}
 
 	logger.Info("finished GCS upload", slog.String("key", key))
@@ -85,7 +86,7 @@ func (g *gcsStagingStore) TableFunctionExpr(ctx context.Context, key string, for
 		Expires: time.Now().Add(gcsSignedURLExpiry),
 	})
 	if err != nil {
-		return "", fmt.Errorf("failed to generate signed URL for gs://%s/%s: %w", g.bucket, key, err)
+		return "", exceptions.NewGCSError(fmt.Errorf("failed to generate signed URL for gs://%s/%s: %w", g.bucket, key, err))
 	}
 
 	var expr strings.Builder
@@ -109,10 +110,10 @@ func (g *gcsStagingStore) DeletePrefix(ctx context.Context, prefix string) error
 			break
 		}
 		if err != nil {
-			return fmt.Errorf("failed to list GCS objects: %w", err)
+			return exceptions.NewGCSError(fmt.Errorf("failed to list GCS objects: %w", err))
 		}
 		if err := g.client.Bucket(g.bucket).Object(attrs.Name).Delete(ctx); err != nil {
-			return fmt.Errorf("failed to delete GCS object %s: %w", attrs.Name, err)
+			return exceptions.NewGCSError(fmt.Errorf("failed to delete GCS object %s: %w", attrs.Name, err))
 		}
 	}
 

flow/shared/exceptions/gcs.go

@@ -0,0 +1,17 @@
+package exceptions
+
+type GCSError struct {
+	error
+}
+
+func NewGCSError(err error) *GCSError {
+	return &GCSError{err}
+}
+
+func (e *GCSError) Error() string {
+	return "GCS Error: " + e.error.Error()
+}
+
+func (e *GCSError) Unwrap() error {
+	return e.error
+}