diff --git a/flow/connectors/clickhouse/avro_sync.go b/flow/connectors/clickhouse/avro_sync.go index fe2151820..39c8f94e1 100644 --- a/flow/connectors/clickhouse/avro_sync.go +++ b/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( diff --git a/flow/connectors/clickhouse/cdc.go b/flow/connectors/clickhouse/cdc.go index f58239a98..c50442e95 100644 --- a/flow/connectors/clickhouse/cdc.go +++ b/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, diff --git a/flow/connectors/clickhouse/clickhouse.go b/flow/connectors/clickhouse/clickhouse.go index 3e732a3a4..6011e972a 100644 --- a/flow/connectors/clickhouse/clickhouse.go +++ b/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 diff --git a/flow/connectors/clickhouse/qrep.go b/flow/connectors/clickhouse/qrep.go index 6295a8ba4..f14c181e7 100644 --- a/flow/connectors/clickhouse/qrep.go +++ b/flow/connectors/clickhouse/qrep.go @@ -4,12 +4,7 @@ import ( "context" "fmt" "log/slog" - "strings" - "github.com/aws/aws-sdk-go-v2/aws" - "github.com/aws/aws-sdk-go-v2/service/s3" - - "github.com/PeerDB-io/peerdb/flow/connectors/utils" "github.com/PeerDB-io/peerdb/flow/generated/protos" "github.com/PeerDB-io/peerdb/flow/model" "github.com/PeerDB-io/peerdb/flow/shared" @@ -49,53 +44,13 @@ func (c *ClickHouseConnector) ConsolidateQRepPartitions(_ context.Context, confi // CleanupQRepFlow function for clickhouse connector func (c *ClickHouseConnector) CleanupQRepFlow(ctx context.Context, config *protos.QRepConfig) error { flowName := config.FlowJobName - stagingPath := config.StagingPath c.logger.Info("Cleaning up stage after QRepFlow", - slog.String("stagingPath", stagingPath), slog.String("flowName", flowName)) - - // if s3 we need to delete the contents of the bucket - if strings.HasPrefix(stagingPath, "s3://") { - s3o, err := utils.NewS3BucketAndPrefix(stagingPath) - if err != nil { - c.logger.Error("failed to create S3 bucket and prefix", slog.Any("error", err)) - return fmt.Errorf("failed to create S3 bucket and prefix: %w", err) - } - - prefix := fmt.Sprintf("%s/%s", s3o.Prefix, flowName) - c.logger.Info("Deleting contents of bucket", slog.String("bucket", s3o.Bucket), slog.String("prefix", prefix)) - - // deleting the contents of the bucket with prefix - s3svc, err := utils.CreateS3Client(ctx, c.credsProvider.Provider) - if err != nil { - c.logger.Error("failed to create S3 client", slog.Any("error", err)) - return fmt.Errorf("failed to create S3 client: %w", err) - } - - // Create a list of all objects with the defined prefix in the bucket - pages := s3.NewListObjectsV2Paginator(s3svc, &s3.ListObjectsV2Input{ - Bucket: aws.String(s3o.Bucket), - Prefix: aws.String(prefix), - }) - for pages.HasMorePages() { - page, err := pages.NextPage(ctx) - if err != nil { - c.logger.Error("failed to list objects from bucket", slog.Any("error", err)) - return fmt.Errorf("failed to list objects from bucket: %w", err) - } - - for _, object := range page.Contents { - _, err = s3svc.DeleteObject(ctx, &s3.DeleteObjectInput{ - Bucket: aws.String(s3o.Bucket), - Key: object.Key, - }) - if err != nil { - c.logger.Error("failed to delete objects from bucket", slog.Any("error", err)) - return fmt.Errorf("failed to delete objects from bucket: %w", err) - } - } - } + slog.String("stagingPath", c.staging.BucketPath()), slog.String("flowName", flowName)) - c.logger.Info("Deleted contents of bucket", slog.String("bucket", s3o.Bucket), slog.String("prefix", prefix)) + prefix := fmt.Sprintf("%s/%s", c.staging.KeyPrefix(), flowName) + if err := c.staging.DeletePrefix(ctx, prefix); err != nil { + c.logger.Error("failed to clean up staging", slog.Any("error", err)) + return fmt.Errorf("failed to clean up staging: %w", err) } return nil diff --git a/flow/connectors/clickhouse/staging.go b/flow/connectors/clickhouse/staging.go new file mode 100644 index 000000000..46dbed94c --- /dev/null +++ b/flow/connectors/clickhouse/staging.go @@ -0,0 +1,45 @@ +package connclickhouse + +import ( + "context" + "fmt" + "strings" + + clickhouseproto "github.com/ClickHouse/clickhouse-go/v2/lib/proto" + + "github.com/PeerDB-io/peerdb/flow/generated/protos" + "github.com/PeerDB-io/peerdb/flow/internal" +) + +func createStagingStore( + ctx context.Context, + env map[string]string, + config *protos.ClickhouseConfig, + chVersion clickhouseproto.Version, +) (StagingStore, error) { + provider, err := internal.PeerDBClickHouseStagingProvider(ctx, env) + if err != nil { + return nil, fmt.Errorf("failed to get staging provider: %w", err) + } + + // Prefer unified PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME, fall back to legacy S3-specific env var. + bucketName, err := internal.PeerDBClickHouseStagingBucketName(ctx, env) + if err != nil { + return nil, fmt.Errorf("failed to get staging bucket name: %w", err) + } + if bucketName == "" { + bucketName, err = internal.PeerDBClickHouseAWSS3BucketName(ctx, env) + if err != nil { + return nil, fmt.Errorf("failed to get S3 bucket name: %w", err) + } + } + + switch strings.ToLower(provider) { + case "s3", "": + return newS3StagingStore(ctx, config, bucketName, chVersion) + case "gcs": + return newGCSStagingStore(ctx, bucketName) + default: + return nil, fmt.Errorf("unsupported staging provider %q (expected s3 or gcs)", provider) + } +} diff --git a/flow/connectors/clickhouse/staging_gcs.go b/flow/connectors/clickhouse/staging_gcs.go new file mode 100644 index 000000000..b71bcd7f1 --- /dev/null +++ b/flow/connectors/clickhouse/staging_gcs.go @@ -0,0 +1,150 @@ +package connclickhouse + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "net/url" + "strings" + "time" + + "cloud.google.com/go/storage" + "github.com/google/uuid" + "google.golang.org/api/iterator" + + "github.com/PeerDB-io/peerdb/flow/internal" + peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse" + "github.com/PeerDB-io/peerdb/flow/shared" +) + +const gcsSignedURLExpiry = 1 * time.Hour + +// gcsStagingStore implements StagingStore for Google Cloud Storage using +// native GCS SDK and signed URLs for ClickHouse reads. +type gcsStagingStore struct { + client *storage.Client + bucket string + prefix string + fullPath string +} + +//nolint:iface // factory function intentionally returns interface +func newGCSStagingStore(ctx context.Context, bucketName string) (StagingStore, error) { + if bucketName == "" { + return nil, errors.New("PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME must be set when staging provider is gcs") + } + + deploymentUID := internal.PeerDBDeploymentUID() + flowName, _ := ctx.Value(shared.FlowNameKey).(string) + bucketPath := fmt.Sprintf("gs://%s/%s/%s", + bucketName, url.PathEscape(deploymentUID), url.PathEscape(flowName)) + + // bucketPath may use either "gs://" (standard GCS URI scheme) or "gcs://" prefix. + path := strings.TrimPrefix(bucketPath, "gcs://") + path = strings.TrimPrefix(path, "gs://") + bucket, prefix, _ := strings.Cut(path, "/") + prefix = strings.Trim(prefix, "/") + + // 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 &gcsStagingStore{ + client: client, + bucket: bucket, + prefix: prefix, + fullPath: bucketPath, + }, nil +} + +func (g *gcsStagingStore) Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error { + logger := internal.LoggerFromCtx(ctx) + + obj := g.client.Bucket(g.bucket).Object(key) + w := obj.NewWriter(ctx) + + if _, err := io.Copy(w, body); err != nil { + w.Close() + return 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) + } + + logger.Info("finished GCS upload", slog.String("key", key)) + return nil +} + +func (g *gcsStagingStore) TableFunctionExpr(ctx context.Context, key string, format string) (string, error) { + signedURL, err := g.client.Bucket(g.bucket).SignedURL(key, &storage.SignedURLOptions{ + Method: "GET", + 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) + } + + var expr strings.Builder + expr.WriteString("url(") + expr.WriteString(peerdb_clickhouse.QuoteLiteral(signedURL)) + expr.WriteString(",") + expr.WriteString(peerdb_clickhouse.QuoteLiteral(format)) + expr.WriteByte(')') + return expr.String(), nil +} + +func (g *gcsStagingStore) DeletePrefix(ctx context.Context, prefix string) error { + logger := internal.LoggerFromCtx(ctx) + logger.Info("Deleting objects from GCS", + slog.String("bucket", g.bucket), slog.String("prefix", prefix)) + + it := g.client.Bucket(g.bucket).Objects(ctx, &storage.Query{Prefix: prefix}) + for { + attrs, err := it.Next() + if err == iterator.Done { + break + } + if err != nil { + return 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) + } + } + + logger.Info("Deleted objects from GCS", + slog.String("bucket", g.bucket), slog.String("prefix", prefix)) + return nil +} + +func (g *gcsStagingStore) Validate(ctx context.Context) error { + testKey := strings.TrimPrefix(g.prefix+"/_peerdb_check_"+uuid.NewString(), "/") + obj := g.client.Bucket(g.bucket).Object(testKey) + + w := obj.NewWriter(ctx) + if _, err := w.Write([]byte(time.Now().Format(time.RFC3339))); err != nil { + w.Close() + return fmt.Errorf("failed to write test object to GCS: %w", err) + } + if err := w.Close(); err != nil { + return fmt.Errorf("failed to finalize test object in GCS: %w", err) + } + + if err := obj.Delete(ctx); err != nil { + return fmt.Errorf("failed to delete test object from GCS: %w", err) + } + + return nil +} + +func (g *gcsStagingStore) BucketPath() string { + return g.fullPath +} + +func (g *gcsStagingStore) KeyPrefix() string { + return g.prefix +} diff --git a/flow/connectors/clickhouse/staging_s3.go b/flow/connectors/clickhouse/staging_s3.go new file mode 100644 index 000000000..f2ee614d5 --- /dev/null +++ b/flow/connectors/clickhouse/staging_s3.go @@ -0,0 +1,222 @@ +package connclickhouse + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "net/url" + "strings" + "time" + + clickhouseproto "github.com/ClickHouse/clickhouse-go/v2/lib/proto" + "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/feature/s3/manager" + "github.com/aws/aws-sdk-go-v2/service/s3" + + "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" + "github.com/PeerDB-io/peerdb/flow/shared" +) + +// s3StagingStore implements StagingStore for AWS S3 (and S3-compatible services). +type s3StagingStore struct { + creds utils.AWSCredentialsProvider + bucket string + prefix string + fullPath string // original "s3://bucket/prefix" for logging +} + +//nolint:iface // factory function intentionally returns interface +func newS3StagingStore( + ctx context.Context, + config *protos.ClickhouseConfig, + unifiedBucketName string, + chVersion clickhouseproto.Version, +) (StagingStore, error) { + 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) + if err != nil { + return nil, err + } + + if awsBucketPath == "" { + if unifiedBucketName == "" { + return nil, errors.New("PeerDB ClickHouse Bucket Name not set") + } + deploymentUID := internal.PeerDBDeploymentUID() + flowName, _ := ctx.Value(shared.FlowNameKey).(string) + bucketPathSuffix := fmt.Sprintf("%s/%s", url.PathEscape(deploymentUID), url.PathEscape(flowName)) + awsBucketPath = fmt.Sprintf("s3://%s/%s", unifiedBucketName, bucketPathSuffix) + } + + // S3 with session tokens requires ClickHouse >= 24.3.1 + // https://github.com/ClickHouse/ClickHouse/issues/61230 + credentials, err := credentialsProvider.Retrieve(ctx) + if err != nil { + return nil, err + } + if credentials.AWS.SessionToken != "" { + if !clickhouseproto.CheckMinVersion( + clickhouseproto.Version{Major: 24, Minor: 3, Patch: 1}, + chVersion, + ) { + return nil, fmt.Errorf( + "provide S3 Transient Stage details explicitly or upgrade to ClickHouse version >= 24.3.1, current version is %s. %s", + chVersion, + "You can also contact PeerDB support for implicit S3 stage setup for older versions of ClickHouse.") + } + } + + s3o, err := utils.NewS3BucketAndPrefix(awsBucketPath) + if err != nil { + return nil, fmt.Errorf("failed to parse S3 bucket path: %w", err) + } + + return &s3StagingStore{ + bucket: s3o.Bucket, + prefix: s3o.Prefix, + fullPath: awsBucketPath, + creds: credentialsProvider, + }, nil +} + +func (s *s3StagingStore) Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error { + logger := internal.LoggerFromCtx(ctx) + + s3svc, err := utils.CreateS3Client(ctx, s.creds) + if err != nil { + return fmt.Errorf("failed to create S3 client: %w", err) + } + + partSize, err := internal.PeerDBS3PartSize(ctx, env) + if err != nil { + return fmt.Errorf("could not get s3 part size config: %w", err) + } + + uploader := manager.NewUploader(s3svc, func(u *manager.Uploader) { + if partSize > 0 { + u.PartSize = partSize + if partSize > 256*1024*1024 { + u.Concurrency = 1 + } + } + }) + + if _, err := uploader.Upload(ctx, &s3.PutObjectInput{ + Bucket: aws.String(s.bucket), + Key: aws.String(key), + Body: body, + }); err != nil { + s3Path := "s3://" + s.bucket + "/" + key + logger.Error("failed to upload file", slog.Any("error", err), slog.String("s3_path", s3Path)) + return fmt.Errorf("failed to upload file to S3: %w", err) + } + + logger.Info("finished S3 upload", slog.String("key", key)) + return nil +} + +func (s *s3StagingStore) TableFunctionExpr(ctx context.Context, key string, format string) (string, error) { + endpoint := s.creds.GetEndpointURL() + region := s.creds.GetRegion() + fileURL := utils.FileURLForS3Service(endpoint, region, s.bucket, key) + + creds, err := s.creds.Retrieve(ctx) + if err != nil { + return "", fmt.Errorf("failed to retrieve S3 credentials: %w", err) + } + if creds.AWS.CanExpire { + logger := internal.LoggerFromCtx(ctx) + logger.Info("Retrieved temporary AWS credentials for table function", + 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(fileURL)) + 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(",") + expr.WriteString(peerdb_clickhouse.QuoteLiteral(format)) + expr.WriteByte(')') + return expr.String(), nil +} + +func (s *s3StagingStore) DeletePrefix(ctx context.Context, prefix string) error { + logger := internal.LoggerFromCtx(ctx) + + s3svc, err := utils.CreateS3Client(ctx, s.creds) + if err != nil { + return fmt.Errorf("failed to create S3 client: %w", err) + } + + logger.Info("Deleting objects from S3", + slog.String("bucket", s.bucket), slog.String("prefix", prefix)) + + pages := s3.NewListObjectsV2Paginator(s3svc, &s3.ListObjectsV2Input{ + Bucket: aws.String(s.bucket), + Prefix: aws.String(prefix), + }) + for pages.HasMorePages() { + page, err := pages.NextPage(ctx) + if err != nil { + return fmt.Errorf("failed to list objects from bucket: %w", err) + } + + for _, object := range page.Contents { + if _, err = s3svc.DeleteObject(ctx, &s3.DeleteObjectInput{ + Bucket: aws.String(s.bucket), + Key: object.Key, + }); err != nil { + return fmt.Errorf("failed to delete object from bucket: %w", err) + } + } + } + + logger.Info("Deleted objects from S3", + slog.String("bucket", s.bucket), slog.String("prefix", prefix)) + return nil +} + +func (s *s3StagingStore) Validate(ctx context.Context) error { + s3Client, err := utils.CreateS3Client(ctx, s.creds) + if err != nil { + return fmt.Errorf("failed to create S3 client: %w", err) + } + return utils.PutAndRemoveS3(ctx, s3Client, s.bucket, s.prefix) +} + +func (s *s3StagingStore) BucketPath() string { + return s.fullPath +} + +func (s *s3StagingStore) KeyPrefix() string { + return s.prefix +} diff --git a/flow/connectors/clickhouse/staging_store.go b/flow/connectors/clickhouse/staging_store.go new file mode 100644 index 000000000..ea06ca5f5 --- /dev/null +++ b/flow/connectors/clickhouse/staging_store.go @@ -0,0 +1,33 @@ +package connclickhouse + +import ( + "context" + "io" +) + +const stagingFormat = "Avro" + +// StagingStore abstracts cloud storage used for staging Avro files. +// Files are written by PeerDB and read by ClickHouse via table functions +// (s3(), url(), etc.). +type StagingStore interface { + // Upload streams data from body to the given key in the staging bucket. + Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error + + // TableFunctionExpr returns a ClickHouse SQL expression that reads the staged file. + // For S3: s3('url', 'access_key', 'secret_key', 'Avro') + // For GCS: url('signed_url', 'Avro') + TableFunctionExpr(ctx context.Context, key string, format string) (string, error) + + // DeletePrefix removes all objects whose key starts with prefix. + DeletePrefix(ctx context.Context, prefix string) error + + // Validate checks that the store is writable by putting and removing a test object. + Validate(ctx context.Context) error + + // BucketPath returns the full staging path (e.g. "s3://bucket/prefix") for logging. + BucketPath() string + + // KeyPrefix returns the object key prefix within the bucket. + KeyPrefix() string +} diff --git a/flow/connectors/utils/aws.go b/flow/connectors/utils/aws.go index 564d07207..6d65ff1fa 100644 --- a/flow/connectors/utils/aws.go +++ b/flow/connectors/utils/aws.go @@ -60,11 +60,6 @@ func NewPeerAWSCredentials(s3 *protos.S3Config) PeerAWSCredentials { } } -type ClickHouseS3Credentials struct { - Provider AWSCredentialsProvider - BucketPath string -} - type AWSCredentials struct { EndpointUrl *string AWS aws.Credentials diff --git a/flow/internal/dynamicconf.go b/flow/internal/dynamicconf.go index b3f0e2386..685180620 100644 --- a/flow/internal/dynamicconf.go +++ b/flow/internal/dynamicconf.go @@ -174,6 +174,22 @@ var DynamicSettings = [...]*protos.DynamicSetting{ ApplyMode: protos.DynconfApplyMode_APPLY_MODE_IMMEDIATE, TargetForSetting: protos.DynconfTarget_CLICKHOUSE, }, + { + Name: "PEERDB_CLICKHOUSE_STAGING_PROVIDER", + Description: "Cloud storage provider for ClickHouse staging: s3 (default) or gcs", + DefaultValue: "s3", + ValueType: protos.DynconfValueType_STRING, + ApplyMode: protos.DynconfApplyMode_APPLY_MODE_AFTER_RESUME, + TargetForSetting: protos.DynconfTarget_CLICKHOUSE, + }, + { + Name: "PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME", + Description: "Staging bucket name for ClickHouse mirrors (provider-agnostic, preferred over legacy env vars)", + DefaultValue: "", + ValueType: protos.DynconfValueType_STRING, + ApplyMode: protos.DynconfApplyMode_APPLY_MODE_IMMEDIATE, + TargetForSetting: protos.DynconfTarget_CLICKHOUSE, + }, { Name: "PEERDB_S3_UUID_PREFIX", Description: "Use random UUID as prefix instead of flow name, can help partitioning on non-AWS based s3 providers", @@ -706,6 +722,14 @@ func PeerDBSnowflakeAutoCompress(ctx context.Context, env map[string]string) (bo return dynamicConfBool(ctx, env, "PEERDB_SNOWFLAKE_AUTO_COMPRESS") } +func PeerDBClickHouseStagingProvider(ctx context.Context, env map[string]string) (string, error) { + return dynLookup(ctx, env, "PEERDB_CLICKHOUSE_STAGING_PROVIDER") +} + +func PeerDBClickHouseStagingBucketName(ctx context.Context, env map[string]string) (string, error) { + return dynLookup(ctx, env, "PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME") +} + func PeerDBClickHouseAWSS3BucketName(ctx context.Context, env map[string]string) (string, error) { return dynLookup(ctx, env, "PEERDB_CLICKHOUSE_AWS_S3_BUCKET_NAME") }