Skip to content

Commit be2e1ba

Browse files
whites11claude
whites11
and
claude
authored
feat(clickhouse): add StagingStore abstraction for GCS staging support (#4200)
## Summary Resolves https://linear.app/clickhouse/issue/CLP-629/add-gcs-support-for-staging-bucket Adds GCS support for ClickHouse CDC/snapshot staging, alongside the existing S3 path. ### What changed - **`StagingStore` interface** (`utils/staging.go`) — abstracts Upload, TableFunctionExpr, DeletePrefix, Validate across cloud providers - **`S3StagingStore`** (`utils/staging_s3.go`) — extracts existing S3 logic, zero behavior change - **`GCSStagingStore`** (`utils/staging_gcs.go`) — native GCS SDK with Workload Identity auth; ClickHouse reads via `url()` table function with signed URLs - **`staging.go`** (clickhouse connector) — `createStagingStore` routes to S3 or GCS based on `PEERDB_CLICKHOUSE_STAGING_PROVIDER` env var - **`clickhouse.go`** — minimal diff: adds `staging` field, replaces inline S3 setup with `createStagingStore()` call, wraps session-token check in nil guard for non-S3 providers - **`avro_sync.go`** — calls `store.Upload()` directly instead of legacy `WriteRecordsToS3` - **`dynamicconf.go`** — `PeerDBClickHouseStagingBucketName` falls back to legacy `PEERDB_CLICKHOUSE_AWS_S3_BUCKET_NAME`; adds `PEERDB_CLICKHOUSE_STAGING_PROVIDER` config ### How GCS staging works 1. PeerDB writes Avro files to GCS using the native Go SDK (Workload Identity on GKE) 2. PeerDB generates a 15-minute signed URL for each staged file 3. ClickHouse reads the file via `url('https://storage.googleapis.com/...?X-Goog-Signature=...', 'Avro')` ### Env vars | Variable | Default | Description | |----------|---------|-------------| | `PEERDB_CLICKHOUSE_STAGING_PROVIDER` | `s3` | `s3` or `gcs` | | `PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME` | (empty) | Unified bucket name; falls back to `PEERDB_CLICKHOUSE_AWS_S3_BUCKET_NAME` for S3 | ## Test plan - [x] All existing unit tests pass (zero regression) - [x] New unit tests for `S3StagingStore` (construction, table function expr, endpoint override, session tokens) - [x] Tested on GCP dev: PG CDC pipe with GCS staging + directory TLS + GCP KMS — 5/5 CDC rows synced - [x] Tested on AWS dev: PG CDC pipe with S3 staging + Lambda certs — confirmed no regression on existing AWS path 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 29d5aac commit be2e1ba

10 files changed

Lines changed: 541 additions & 208 deletions

File tree

‎flow/connectors/clickhouse/avro_sync.go‎

Lines changed: 46 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ package connclickhouse
33
import (
44
"context"
55
"fmt"
6+
"io"
67
"log/slog"
8+
"runtime/debug"
79
"strings"
810
"sync/atomic"
911
"time"
@@ -37,52 +39,17 @@ func NewClickHouseAvroSyncMethod(
3739
}
3840
}
3941

40-
func (s *ClickHouseAvroSyncMethod) s3TableFunctionBuilder(ctx context.Context, avroFilePath string) (string, error) {
41-
stagingPath := s.credsProvider.BucketPath
42-
s3o, err := utils.NewS3BucketAndPrefix(stagingPath)
43-
if err != nil {
44-
return "", err
45-
}
46-
47-
endpoint := s.credsProvider.Provider.GetEndpointURL()
48-
region := s.credsProvider.Provider.GetRegion()
49-
avroFileUrl := utils.FileURLForS3Service(endpoint, region, s3o.Bucket, avroFilePath)
50-
creds, err := s.credsProvider.Provider.Retrieve(ctx)
51-
if err != nil {
52-
return "", err
53-
}
54-
if creds.AWS.CanExpire {
55-
s.logger.Info("Retrieved Temporary AWS credentials",
56-
slog.Time("expiryTimestamp", creds.AWS.Expires),
57-
slog.Duration("duration", time.Until(creds.AWS.Expires)))
58-
}
59-
60-
var expr strings.Builder
61-
expr.WriteString("s3(")
62-
expr.WriteString(peerdb_clickhouse.QuoteLiteral(avroFileUrl))
63-
expr.WriteByte(',')
64-
expr.WriteString(peerdb_clickhouse.QuoteLiteral(creds.AWS.AccessKeyID))
65-
expr.WriteByte(',')
66-
expr.WriteString(peerdb_clickhouse.QuoteLiteral(creds.AWS.SecretAccessKey))
67-
if creds.AWS.SessionToken != "" {
68-
expr.WriteByte(',')
69-
expr.WriteString(peerdb_clickhouse.QuoteLiteral(creds.AWS.SessionToken))
70-
}
71-
expr.WriteString(",'Avro')")
72-
return expr.String(), nil
73-
}
74-
7542
func (s *ClickHouseAvroSyncMethod) CopyStageToDestination(ctx context.Context, avroFile utils.AvroFile) error {
76-
s3TableFunction, err := s.s3TableFunctionBuilder(ctx, avroFile.FilePath)
43+
stagingTableFunction, err := s.staging.TableFunctionExpr(ctx, avroFile.FilePath, stagingFormat)
7744
if err != nil {
78-
s.logger.Error("failed to build S3 table function",
45+
s.logger.Error("failed to build staging table function",
7946
slog.String("avroFilePath", avroFile.FilePath),
8047
slog.Any("error", err))
81-
return fmt.Errorf("failed to build S3 table function: %w", err)
48+
return fmt.Errorf("failed to build staging table function: %w", err)
8249
}
8350

8451
query := fmt.Sprintf("INSERT INTO %s SELECT * FROM %s",
85-
peerdb_clickhouse.QuoteIdentifier(s.config.DestinationTableIdentifier), s3TableFunction)
52+
peerdb_clickhouse.QuoteIdentifier(s.config.DestinationTableIdentifier), stagingTableFunction)
8653
return s.exec(ctx, query)
8754
}
8855

@@ -146,7 +113,7 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
146113
numericTruncator := model.NewSnapshotTableNumericTruncator(dstTableName, schema.Fields)
147114

148115
columnNameAvroFieldMap := model.ConstructColumnNameAvroFieldMap(schema.Fields)
149-
avroFiles, totalRecords, err := s.pushDataToS3ForSnapshot(ctx, config, dstTableName, schema,
116+
avroFiles, totalRecords, err := s.pushDataToStagingForSnapshot(ctx, config, dstTableName, schema,
150117
columnNameAvroFieldMap, partition, stream, destTypeConversions, numericTruncator)
151118
if err != nil {
152119
s.logger.Error("failed to push data to S3",
@@ -155,7 +122,7 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
155122
return 0, nil, err
156123
}
157124

158-
if err := s.pushS3DataToClickHouseForSnapshot(
125+
if err := s.pushStagingDataToClickHouseForSnapshot(
159126
ctx, avroFiles, schema, columnNameAvroFieldMap, config); err != nil {
160127
s.logger.Error("failed to push data to ClickHouse",
161128
slog.String("dstTable", dstTableName),
@@ -172,7 +139,7 @@ func (s *ClickHouseAvroSyncMethod) SyncQRepRecords(
172139
return totalRecords, warnings, nil
173140
}
174141

175-
func (s *ClickHouseAvroSyncMethod) pushDataToS3ForSnapshot(
142+
func (s *ClickHouseAvroSyncMethod) pushDataToStagingForSnapshot(
176143
ctx context.Context,
177144
config *protos.QRepConfig,
178145
dstTableName string,
@@ -267,7 +234,7 @@ func (s *ClickHouseAvroSyncMethod) pushDataToS3ForSnapshot(
267234
return avroFiles, totalRecords, nil
268235
}
269236

270-
func (s *ClickHouseAvroSyncMethod) pushS3DataToClickHouseForSnapshot(
237+
func (s *ClickHouseAvroSyncMethod) pushStagingDataToClickHouseForSnapshot(
271238
ctx context.Context,
272239
avroFiles []utils.AvroFile,
273240
schema types.QRecordSchema,
@@ -307,24 +274,24 @@ func (s *ClickHouseAvroSyncMethod) pushS3DataToClickHouseForSnapshot(
307274

308275
for i := range numParts {
309276
// Get fresh credentials for each part
310-
s3TableFunction, err := s.s3TableFunctionBuilder(ctx, avroFile.FilePath)
277+
stagingTableFunction, err := s.staging.TableFunctionExpr(ctx, avroFile.FilePath, stagingFormat)
311278
if err != nil {
312-
s.logger.Error("failed to build S3 table function",
279+
s.logger.Error("failed to build staging table function",
313280
slog.String("avroFilePath", avroFile.FilePath),
314281
slog.Any("error", err),
315282
slog.Uint64("part", i),
316283
slog.Uint64("numParts", numParts),
317284
slog.Int("chunkIdx", chunkIdx),
318285
)
319-
return fmt.Errorf("failed to build S3 table function: %w", err)
286+
return fmt.Errorf("failed to build staging table function: %w", err)
320287
}
321288

322289
var query string
323290
if numParts > 1 {
324291
query, err = buildInsertFromTableFunctionQueryWithPartitioning(
325-
ctx, insertConfig, s3TableFunction, i, numParts, chSettings)
292+
ctx, insertConfig, stagingTableFunction, i, numParts, chSettings)
326293
} else {
327-
query, err = buildInsertFromTableFunctionQuery(ctx, insertConfig, s3TableFunction, chSettings)
294+
query, err = buildInsertFromTableFunctionQuery(ctx, insertConfig, stagingTableFunction, chSettings)
328295
}
329296
if err != nil {
330297
s.logger.Error("failed to build insert query",
@@ -393,33 +360,49 @@ func (s *ClickHouseAvroSyncMethod) writeToAvroFile(
393360
typeConversions map[string]types.TypeConversion,
394361
numericTruncator model.SnapshotTableNumericTruncator,
395362
) (utils.AvroFile, error) {
396-
stagingPath := s.credsProvider.BucketPath
397363
ocfWriter := utils.NewPeerDBOCFWriter(stream, avroSchema, ocf.ZStandard, protos.DBType_CLICKHOUSE, sizeTracker)
398-
s3o, err := utils.NewS3BucketAndPrefix(stagingPath)
399-
if err != nil {
400-
return utils.AvroFile{}, fmt.Errorf("failed to parse staging path: %w", err)
401-
}
364+
prefix := s.staging.KeyPrefix()
402365

403366
s3UuidPrefix, err := internal.PeerDBS3UuidPrefix(ctx, s.config.Env)
404367
if err != nil {
405368
return utils.AvroFile{}, err
406369
}
407370

408-
var s3AvroFileKey string
371+
var stagingAvroFileKey string
409372
if s3UuidPrefix {
410-
s3AvroFileKey = fmt.Sprintf("%s/%s/%s/%s.avro", s3o.Prefix, uuid.NewString(), flowJobName, identifierForFile)
373+
stagingAvroFileKey = fmt.Sprintf("%s/%s/%s/%s.avro", prefix, uuid.NewString(), flowJobName, identifierForFile)
411374
} else {
412-
s3AvroFileKey = fmt.Sprintf("%s/%s/%s.avro", s3o.Prefix, flowJobName, identifierForFile)
375+
stagingAvroFileKey = fmt.Sprintf("%s/%s/%s.avro", prefix, flowJobName, identifierForFile)
413376
}
414-
s3AvroFileKey = strings.TrimLeft(s3AvroFileKey, "/")
415-
avroFile, err := ocfWriter.WriteRecordsToS3(
416-
ctx, env, s3o.Bucket, s3AvroFileKey, s.credsProvider.Provider, typeConversions, numericTruncator,
417-
)
418-
if err != nil {
419-
return utils.AvroFile{}, fmt.Errorf("failed to write records to S3: %w", err)
377+
stagingAvroFileKey = strings.TrimLeft(stagingAvroFileKey, "/")
378+
379+
r, w := io.Pipe()
380+
defer r.Close()
381+
382+
var writeOcfError error
383+
var numRows int64
384+
go func() {
385+
defer func() {
386+
if r := recover(); r != nil {
387+
writeOcfError = fmt.Errorf("panic occurred during WriteOCF: %v\n%s", r, debug.Stack())
388+
}
389+
w.Close()
390+
}()
391+
numRows, writeOcfError = ocfWriter.WriteOCF(ctx, env, w, typeConversions, numericTruncator)
392+
}()
393+
394+
if err := s.staging.Upload(ctx, env, stagingAvroFileKey, r); err != nil {
395+
return utils.AvroFile{}, fmt.Errorf("failed to upload to staging: %w", err)
396+
}
397+
if writeOcfError != nil {
398+
return utils.AvroFile{}, writeOcfError
420399
}
421400

422-
return avroFile, nil
401+
return utils.AvroFile{
402+
StorageLocation: utils.AvroS3Storage,
403+
FilePath: stagingAvroFileKey,
404+
NumRecords: numRows,
405+
}, nil
423406
}
424407

425408
func (s *ClickHouseAvroSyncMethod) SyncQRepObjects(

‎flow/connectors/clickhouse/cdc.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ func (c *ClickHouseConnector) CreateRawTable(ctx context.Context, req *protos.Cr
9797

9898
func (c *ClickHouseConnector) avroSyncMethod(flowJobName string, env map[string]string, version uint32) *ClickHouseAvroSyncMethod {
9999
qrepConfig := &protos.QRepConfig{
100-
StagingPath: c.credsProvider.BucketPath,
100+
StagingPath: c.staging.BucketPath(),
101101
FlowJobName: flowJobName,
102102
DestinationTableIdentifier: c.GetRawTableName(flowJobName),
103103
Env: env,

‎flow/connectors/clickhouse/clickhouse.go‎

Lines changed: 15 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import (
88
"errors"
99
"fmt"
1010
"log/slog"
11-
"net/url"
1211
"os"
1312
"path/filepath"
1413
"slices"
@@ -19,11 +18,9 @@ import (
1918
"github.com/ClickHouse/clickhouse-go/v2"
2019
"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
2120
clickhouseproto "github.com/ClickHouse/clickhouse-go/v2/lib/proto"
22-
"github.com/aws/aws-sdk-go-v2/aws"
2321
"go.temporal.io/sdk/log"
2422

2523
metadataStore "github.com/PeerDB-io/peerdb/flow/connectors/external_metadata"
26-
"github.com/PeerDB-io/peerdb/flow/connectors/utils"
2724
"github.com/PeerDB-io/peerdb/flow/generated/protos"
2825
"github.com/PeerDB-io/peerdb/flow/internal"
2926
peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse"
@@ -33,11 +30,11 @@ import (
3330

3431
type ClickHouseConnector struct {
3532
*metadataStore.PostgresMetadata
36-
database clickhouse.Conn
37-
logger log.Logger
38-
Config *protos.ClickhouseConfig
39-
credsProvider *utils.ClickHouseS3Credentials
40-
chVersion *clickhouseproto.Version
33+
database clickhouse.Conn
34+
logger log.Logger
35+
Config *protos.ClickhouseConfig
36+
staging StagingStore
37+
chVersion *clickhouseproto.Version
4138
}
4239

4340
func NewClickHouseConnector(
@@ -57,95 +54,24 @@ func NewClickHouseConnector(
5754
return nil, err
5855
}
5956

60-
var awsConfig utils.PeerAWSCredentials
61-
var awsBucketPath string
62-
if config.S3 != nil {
63-
awsConfig = utils.NewPeerAWSCredentials(config.S3)
64-
awsBucketPath = config.S3.Url
65-
} else {
66-
awsConfig = utils.PeerAWSCredentials{
67-
Credentials: aws.Credentials{
68-
AccessKeyID: config.AccessKeyId,
69-
SecretAccessKey: config.SecretAccessKey,
70-
},
71-
EndpointUrl: config.Endpoint,
72-
Region: config.Region,
73-
}
74-
awsBucketPath = config.S3Path
75-
}
76-
77-
credentialsProvider, err := utils.GetAWSCredentialsProvider(ctx, "clickhouse", awsConfig)
57+
clickHouseVersion, err := database.ServerVersion()
7858
if err != nil {
79-
return nil, err
80-
}
81-
82-
if awsBucketPath == "" {
83-
deploymentUID := internal.PeerDBDeploymentUID()
84-
flowName, _ := ctx.Value(shared.FlowNameKey).(string)
85-
bucketPathSuffix := fmt.Sprintf("%s/%s", url.PathEscape(deploymentUID), url.PathEscape(flowName))
86-
// Fallback: Get S3 credentials from environment
87-
awsBucketName, err := internal.PeerDBClickHouseAWSS3BucketName(ctx, env)
88-
if err != nil {
89-
return nil, fmt.Errorf("failed to get PeerDB ClickHouse Bucket Name: %w", err)
90-
}
91-
if awsBucketName == "" {
92-
return nil, fmt.Errorf("PeerDB ClickHouse Bucket Name not set")
93-
}
94-
95-
awsBucketPath = fmt.Sprintf("s3://%s/%s", awsBucketName, bucketPathSuffix)
59+
return nil, fmt.Errorf("failed to get ClickHouse version: %w", err)
9660
}
9761

98-
credentials, err := credentialsProvider.Retrieve(ctx)
62+
staging, err := createStagingStore(ctx, env, config, clickHouseVersion.Version)
9963
if err != nil {
10064
return nil, err
10165
}
10266

103-
clickHouseVersion, err := database.ServerVersion()
104-
if err != nil {
105-
return nil, fmt.Errorf("failed to get ClickHouse version: %w", err)
106-
}
107-
connector := &ClickHouseConnector{
67+
return &ClickHouseConnector{
10868
database: database,
10969
PostgresMetadata: pgMetadata,
11070
Config: config,
11171
logger: logger,
112-
credsProvider: &utils.ClickHouseS3Credentials{
113-
Provider: credentialsProvider,
114-
BucketPath: awsBucketPath,
115-
},
116-
chVersion: &clickHouseVersion.Version,
117-
}
118-
119-
if credentials.AWS.SessionToken != "" {
120-
// 24.3.1 is minimum version of ClickHouse that actually supports session token
121-
// https://github.com/ClickHouse/ClickHouse/issues/61230
122-
if !clickhouseproto.CheckMinVersion(
123-
clickhouseproto.Version{Major: 24, Minor: 3, Patch: 1},
124-
clickHouseVersion.Version,
125-
) {
126-
return nil, fmt.Errorf(
127-
"provide S3 Transient Stage details explicitly or upgrade to ClickHouse version >= 24.3.1, current version is %s. %s",
128-
clickHouseVersion,
129-
"You can also contact PeerDB support for implicit S3 stage setup for older versions of ClickHouse.")
130-
}
131-
}
132-
133-
return connector, nil
134-
}
135-
136-
func ValidateS3(ctx context.Context, creds *utils.ClickHouseS3Credentials) error {
137-
// for validation purposes
138-
s3Client, err := utils.CreateS3Client(ctx, creds.Provider)
139-
if err != nil {
140-
return fmt.Errorf("failed to create S3 client: %w", err)
141-
}
142-
143-
object, err := utils.NewS3BucketAndPrefix(creds.BucketPath)
144-
if err != nil {
145-
return fmt.Errorf("failed to create S3 bucket and prefix: %w", err)
146-
}
147-
148-
return utils.PutAndRemoveS3(ctx, s3Client, object.Bucket, object.Prefix)
72+
staging: staging,
73+
chVersion: &clickHouseVersion.Version,
74+
}, nil
14975
}
15076

15177
func ValidateClickHouseHost(ctx context.Context, chHost string, allowedDomainString string) error {
@@ -226,9 +152,9 @@ func (c *ClickHouseConnector) ValidateCheck(ctx context.Context) error {
226152
return fmt.Errorf("failed to drop validation table %s: %w", validateDummyTableNameRenamed, err)
227153
}
228154

229-
// validate s3 stage
230-
if err := ValidateS3(ctx, c.credsProvider); err != nil {
231-
return fmt.Errorf("failed to validate S3 bucket: %w", err)
155+
// validate staging storage
156+
if err := c.staging.Validate(ctx); err != nil {
157+
return fmt.Errorf("failed to validate staging bucket: %w", err)
232158
}
233159

234160
return nil

0 commit comments

Comments
 (0)