feat(clickhouse): add StagingStore abstraction for GCS staging support

flow/connectors/clickhouse/avro_sync.go

@@ -3,7 +3,9 @@ package connclickhouse
 import (
 	"context"
 	"fmt"
+	"io"
 	"log/slog"
+	"runtime/debug"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -37,52 +39,17 @@ func NewClickHouseAvroSyncMethod(
 	}
 }
 
-func (s *ClickHouseAvroSyncMethod) s3TableFunctionBuilder(ctx context.Context, avroFilePath string) (string, error) {
-	stagingPath := s.credsProvider.BucketPath
-	s3o, err := utils.NewS3BucketAndPrefix(stagingPath)
-	if err != nil {
-		return "", err
-	}
-
-	endpoint := s.credsProvider.Provider.GetEndpointURL()
-	region := s.credsProvider.Provider.GetRegion()
-	avroFileUrl := utils.FileURLForS3Service(endpoint, region, s3o.Bucket, avroFilePath)
-	creds, err := s.credsProvider.Provider.Retrieve(ctx)
-	if err != nil {
-		return "", err
-	}
-	if creds.AWS.CanExpire {
-		s.logger.Info("Retrieved Temporary AWS credentials",
-			slog.Time("expiryTimestamp", creds.AWS.Expires),
-			slog.Duration("duration", time.Until(creds.AWS.Expires)))
-	}
-
-	var expr strings.Builder
-	expr.WriteString("s3(")
-	expr.WriteString(peerdb_clickhouse.QuoteLiteral(avroFileUrl))
-	expr.WriteByte(',')
-	expr.WriteString(peerdb_clickhouse.QuoteLiteral(creds.AWS.AccessKeyID))
-	expr.WriteByte(',')
-	expr.WriteString(peerdb_clickhouse.QuoteLiteral(creds.AWS.SecretAccessKey))
-	if creds.AWS.SessionToken != "" {
-		expr.WriteByte(',')
-		expr.WriteString(peerdb_clickhouse.QuoteLiteral(creds.AWS.SessionToken))
-	}
-	expr.WriteString(",'Avro')")
-	return expr.String(), nil
-}
-
 func (s *ClickHouseAvroSyncMethod) CopyStageToDestination(ctx context.Context, avroFile utils.AvroFile) error {
-	s3TableFunction, err := s.s3TableFunctionBuilder(ctx, avroFile.FilePath)
+	stagingTableFunction, err := s.staging.TableFunctionExpr(ctx, avroFile.FilePath, stagingFormat)
 	if err != nil {
-		s.logger.Error("failed to build S3 table function",
+		s.logger.Error("failed to build staging table function",
 			slog.String("avroFilePath", avroFile.FilePath),
 			slog.Any("error", err))
-		return fmt.Errorf("failed to build S3 table function: %w", err)
+		return fmt.Errorf("failed to build staging table function: %w", err)
 	}
 
 	query := fmt.Sprintf("INSERT INTO %s SELECT * FROM %s",
-		peerdb_clickhouse.QuoteIdentifier(s.config.DestinationTableIdentifier), s3TableFunction)
+		peerdb_clickhouse.QuoteIdentifier(s.config.DestinationTableIdentifier), stagingTableFunction)
 	return s.exec(ctx, query)
 }
 
@@ -146,7 +113,7 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
 	numericTruncator := model.NewSnapshotTableNumericTruncator(dstTableName, schema.Fields)
 
 	columnNameAvroFieldMap := model.ConstructColumnNameAvroFieldMap(schema.Fields)
-	avroFiles, totalRecords, err := s.pushDataToS3ForSnapshot(ctx, config, dstTableName, schema,
+	avroFiles, totalRecords, err := s.pushDataToStagingForSnapshot(ctx, config, dstTableName, schema,
 		columnNameAvroFieldMap, partition, stream, destTypeConversions, numericTruncator)
 	if err != nil {
 		s.logger.Error("failed to push data to S3",
@@ -155,7 +122,7 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
 		return 0, nil, err
 	}
 
-	if err := s.pushS3DataToClickHouseForSnapshot(
+	if err := s.pushStagingDataToClickHouseForSnapshot(
 		ctx, avroFiles, schema, columnNameAvroFieldMap, config); err != nil {
 		s.logger.Error("failed to push data to ClickHouse",
 			slog.String("dstTable", dstTableName),
@@ -172,7 +139,7 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
 	return totalRecords, warnings, nil
 }
 
-func (s *ClickHouseAvroSyncMethod) pushDataToS3ForSnapshot(
+func (s *ClickHouseAvroSyncMethod) pushDataToStagingForSnapshot(
 	ctx context.Context,
 	config *protos.QRepConfig,
 	dstTableName string,
@@ -267,7 +234,7 @@ func (s *ClickHouseAvroSyncMethod) pushDataToS3ForSnapshot(
 	return avroFiles, totalRecords, nil
 }
 
-func (s *ClickHouseAvroSyncMethod) pushS3DataToClickHouseForSnapshot(
+func (s *ClickHouseAvroSyncMethod) pushStagingDataToClickHouseForSnapshot(
 	ctx context.Context,
 	avroFiles []utils.AvroFile,
 	schema types.QRecordSchema,
@@ -307,24 +274,24 @@ func (s *ClickHouseAvroSyncMethod) pushS3DataToClickHouseForSnapshot(
 
 		for i := range numParts {
 			// Get fresh credentials for each part
-			s3TableFunction, err := s.s3TableFunctionBuilder(ctx, avroFile.FilePath)
+			stagingTableFunction, err := s.staging.TableFunctionExpr(ctx, avroFile.FilePath, stagingFormat)
 			if err != nil {
-				s.logger.Error("failed to build S3 table function",
+				s.logger.Error("failed to build staging table function",
 					slog.String("avroFilePath", avroFile.FilePath),
 					slog.Any("error", err),
 					slog.Uint64("part", i),
 					slog.Uint64("numParts", numParts),
 					slog.Int("chunkIdx", chunkIdx),
 				)
-				return fmt.Errorf("failed to build S3 table function: %w", err)
+				return fmt.Errorf("failed to build staging table function: %w", err)
 			}
 
 			var query string
 			if numParts > 1 {
 				query, err = buildInsertFromTableFunctionQueryWithPartitioning(
-					ctx, insertConfig, s3TableFunction, i, numParts, chSettings)
+					ctx, insertConfig, stagingTableFunction, i, numParts, chSettings)
 			} else {
-				query, err = buildInsertFromTableFunctionQuery(ctx, insertConfig, s3TableFunction, chSettings)
+				query, err = buildInsertFromTableFunctionQuery(ctx, insertConfig, stagingTableFunction, chSettings)
 			}
 			if err != nil {
 				s.logger.Error("failed to build insert query",
@@ -393,33 +360,49 @@ func (s *ClickHouseAvroSyncMethod) writeToAvroFile(
 	typeConversions map[string]types.TypeConversion,
 	numericTruncator model.SnapshotTableNumericTruncator,
 ) (utils.AvroFile, error) {
-	stagingPath := s.credsProvider.BucketPath
 	ocfWriter := utils.NewPeerDBOCFWriter(stream, avroSchema, ocf.ZStandard, protos.DBType_CLICKHOUSE, sizeTracker)
-	s3o, err := utils.NewS3BucketAndPrefix(stagingPath)
-	if err != nil {
-		return utils.AvroFile{}, fmt.Errorf("failed to parse staging path: %w", err)
-	}
+	prefix := s.staging.KeyPrefix()
 
 	s3UuidPrefix, err := internal.PeerDBS3UuidPrefix(ctx, s.config.Env)
 	if err != nil {
 		return utils.AvroFile{}, err
 	}
 
-	var s3AvroFileKey string
+	var stagingAvroFileKey string
 	if s3UuidPrefix {
-		s3AvroFileKey = fmt.Sprintf("%s/%s/%s/%s.avro", s3o.Prefix, uuid.NewString(), flowJobName, identifierForFile)
+		stagingAvroFileKey = fmt.Sprintf("%s/%s/%s/%s.avro", prefix, uuid.NewString(), flowJobName, identifierForFile)
 	} else {
-		s3AvroFileKey = fmt.Sprintf("%s/%s/%s.avro", s3o.Prefix, flowJobName, identifierForFile)
+		stagingAvroFileKey = fmt.Sprintf("%s/%s/%s.avro", prefix, flowJobName, identifierForFile)
 	}
-	s3AvroFileKey = strings.TrimLeft(s3AvroFileKey, "/")
-	avroFile, err := ocfWriter.WriteRecordsToS3(
-		ctx, env, s3o.Bucket, s3AvroFileKey, s.credsProvider.Provider, typeConversions, numericTruncator,
-	)
-	if err != nil {
-		return utils.AvroFile{}, fmt.Errorf("failed to write records to S3: %w", err)
+	stagingAvroFileKey = strings.TrimLeft(stagingAvroFileKey, "/")
+
+	r, w := io.Pipe()
+	defer r.Close()
+
+	var writeOcfError error
+	var numRows int64
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				writeOcfError = fmt.Errorf("panic occurred during WriteOCF: %v\n%s", r, debug.Stack())
+			}
+			w.Close()
+		}()
+		numRows, writeOcfError = ocfWriter.WriteOCF(ctx, env, w, typeConversions, numericTruncator)
+	}()
+
+	if err := s.staging.Upload(ctx, env, stagingAvroFileKey, r); err != nil {
+		return utils.AvroFile{}, fmt.Errorf("failed to upload to staging: %w", err)
+	}
+	if writeOcfError != nil {
+		return utils.AvroFile{}, writeOcfError
 	}
 
-	return avroFile, nil
+	return utils.AvroFile{
+		StorageLocation: utils.AvroS3Storage,
+		FilePath:        stagingAvroFileKey,
+		NumRecords:      numRows,
+	}, nil
 }
 
 func (s *ClickHouseAvroSyncMethod) SyncQRepObjects(

flow/connectors/clickhouse/cdc.go

@@ -97,7 +97,7 @@ func (c *ClickHouseConnector) CreateRawTable(ctx context.Context, req *protos.Cr
 
 func (c *ClickHouseConnector) avroSyncMethod(flowJobName string, env map[string]string, version uint32) *ClickHouseAvroSyncMethod {
 	qrepConfig := &protos.QRepConfig{
-		StagingPath:                c.credsProvider.BucketPath,
+		StagingPath:                c.staging.BucketPath(),
 		FlowJobName:                flowJobName,
 		DestinationTableIdentifier: c.GetRawTableName(flowJobName),
 		Env:                        env,

flow/connectors/clickhouse/clickhouse.go

@@ -8,7 +8,6 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
-	"net/url"
 	"os"
 	"path/filepath"
 	"slices"
@@ -19,11 +18,9 @@ import (
 	"github.com/ClickHouse/clickhouse-go/v2"
 	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
 	clickhouseproto "github.com/ClickHouse/clickhouse-go/v2/lib/proto"
-	"github.com/aws/aws-sdk-go-v2/aws"
 	"go.temporal.io/sdk/log"
 
 	metadataStore "github.com/PeerDB-io/peerdb/flow/connectors/external_metadata"
-	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
 	"github.com/PeerDB-io/peerdb/flow/internal"
 	peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse"
@@ -33,11 +30,11 @@ import (
 
 type ClickHouseConnector struct {
 	*metadataStore.PostgresMetadata
-	database      clickhouse.Conn
-	logger        log.Logger
-	Config        *protos.ClickhouseConfig
-	credsProvider *utils.ClickHouseS3Credentials
-	chVersion     *clickhouseproto.Version
+	database  clickhouse.Conn
+	logger    log.Logger
+	Config    *protos.ClickhouseConfig
+	staging   StagingStore
+	chVersion *clickhouseproto.Version
 }
 
 func NewClickHouseConnector(
@@ -57,95 +54,24 @@ func NewClickHouseConnector(
 		return nil, err
 	}
 
-	var awsConfig utils.PeerAWSCredentials
-	var awsBucketPath string
-	if config.S3 != nil {
-		awsConfig = utils.NewPeerAWSCredentials(config.S3)
-		awsBucketPath = config.S3.Url
-	} else {
-		awsConfig = utils.PeerAWSCredentials{
-			Credentials: aws.Credentials{
-				AccessKeyID:     config.AccessKeyId,
-				SecretAccessKey: config.SecretAccessKey,
-			},
-			EndpointUrl: config.Endpoint,
-			Region:      config.Region,
-		}
-		awsBucketPath = config.S3Path
-	}
-
-	credentialsProvider, err := utils.GetAWSCredentialsProvider(ctx, "clickhouse", awsConfig)
+	clickHouseVersion, err := database.ServerVersion()
 	if err != nil {
-		return nil, err
-	}
-
-	if awsBucketPath == "" {
-		deploymentUID := internal.PeerDBDeploymentUID()
-		flowName, _ := ctx.Value(shared.FlowNameKey).(string)
-		bucketPathSuffix := fmt.Sprintf("%s/%s", url.PathEscape(deploymentUID), url.PathEscape(flowName))
-		// Fallback: Get S3 credentials from environment
-		awsBucketName, err := internal.PeerDBClickHouseAWSS3BucketName(ctx, env)
-		if err != nil {
-			return nil, fmt.Errorf("failed to get PeerDB ClickHouse Bucket Name: %w", err)
-		}
-		if awsBucketName == "" {
-			return nil, fmt.Errorf("PeerDB ClickHouse Bucket Name not set")
-		}
-
-		awsBucketPath = fmt.Sprintf("s3://%s/%s", awsBucketName, bucketPathSuffix)
+		return nil, fmt.Errorf("failed to get ClickHouse version: %w", err)
 	}
 
-	credentials, err := credentialsProvider.Retrieve(ctx)
+	staging, err := createStagingStore(ctx, env, config, clickHouseVersion.Version)
 	if err != nil {
 		return nil, err
 	}
 
-	clickHouseVersion, err := database.ServerVersion()
-	if err != nil {
-		return nil, fmt.Errorf("failed to get ClickHouse version: %w", err)
-	}
-	connector := &ClickHouseConnector{
+	return &ClickHouseConnector{
 		database:         database,
 		PostgresMetadata: pgMetadata,
 		Config:           config,
 		logger:           logger,
-		credsProvider: &utils.ClickHouseS3Credentials{
-			Provider:   credentialsProvider,
-			BucketPath: awsBucketPath,
-		},
-		chVersion: &clickHouseVersion.Version,
-	}
-
-	if credentials.AWS.SessionToken != "" {
-		// 24.3.1 is minimum version of ClickHouse that actually supports session token
-		// https://github.com/ClickHouse/ClickHouse/issues/61230
-		if !clickhouseproto.CheckMinVersion(
-			clickhouseproto.Version{Major: 24, Minor: 3, Patch: 1},
-			clickHouseVersion.Version,
-		) {
-			return nil, fmt.Errorf(
-				"provide S3 Transient Stage details explicitly or upgrade to ClickHouse version >= 24.3.1, current version is %s. %s",
-				clickHouseVersion,
-				"You can also contact PeerDB support for implicit S3 stage setup for older versions of ClickHouse.")
-		}
-	}
-
-	return connector, nil
-}
-
-func ValidateS3(ctx context.Context, creds *utils.ClickHouseS3Credentials) error {
-	// for validation purposes
-	s3Client, err := utils.CreateS3Client(ctx, creds.Provider)
-	if err != nil {
-		return fmt.Errorf("failed to create S3 client: %w", err)
-	}
-
-	object, err := utils.NewS3BucketAndPrefix(creds.BucketPath)
-	if err != nil {
-		return fmt.Errorf("failed to create S3 bucket and prefix: %w", err)
-	}
-
-	return utils.PutAndRemoveS3(ctx, s3Client, object.Bucket, object.Prefix)
+		staging:          staging,
+		chVersion:        &clickHouseVersion.Version,
+	}, nil
 }
 
 func ValidateClickHouseHost(ctx context.Context, chHost string, allowedDomainString string) error {
@@ -226,9 +152,9 @@ func (c *ClickHouseConnector) ValidateCheck(ctx context.Context) error {
 		return fmt.Errorf("failed to drop validation table %s: %w", validateDummyTableNameRenamed, err)
 	}
 
-	// validate s3 stage
-	if err := ValidateS3(ctx, c.credsProvider); err != nil {
-		return fmt.Errorf("failed to validate S3 bucket: %w", err)
+	// validate staging storage
+	if err := c.staging.Validate(ctx); err != nil {
+		return fmt.Errorf("failed to validate staging bucket: %w", err)
 	}
 
 	return nil