Skip to content

Commit a8c3800

Browse files
whites11whites11
committed
refactor(staging): unify factory functions, struct at top of file
Merge newS3StagingStoreFromConfig + newS3StagingStore into one newS3StagingStore. Same for GCS. Each staging_*.go file now has: struct → constructor → methods. staging.go is just the router.
1 parent 501bec8 commit a8c3800

3 files changed

Lines changed: 90 additions & 94 deletions

File tree

‎flow/connectors/clickhouse/staging.go‎

Lines changed: 3 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,13 @@ package connclickhouse
22

33
import (
44
"context"
5-
"errors"
65
"fmt"
7-
"net/url"
86
"strings"
97

10-
clickhouseproto "github.com/ClickHouse/clickhouse-go/v2/lib/proto"
11-
"github.com/aws/aws-sdk-go-v2/aws"
12-
13-
"github.com/PeerDB-io/peerdb/flow/connectors/utils"
148
"github.com/PeerDB-io/peerdb/flow/generated/protos"
159
"github.com/PeerDB-io/peerdb/flow/internal"
16-
"github.com/PeerDB-io/peerdb/flow/shared"
10+
11+
clickhouseproto "github.com/ClickHouse/clickhouse-go/v2/lib/proto"
1712
)
1813

1914
func createStagingStore(
@@ -41,84 +36,10 @@ func createStagingStore(
4136

4237
switch strings.ToLower(provider) {
4338
case "s3", "":
44-
return newS3StagingStoreFromConfig(ctx, config, bucketName, chVersion)
39+
return newS3StagingStore(ctx, config, bucketName, chVersion)
4540
case "gcs":
4641
return newGCSStagingStore(ctx, bucketName)
4742
default:
4843
return nil, fmt.Errorf("unsupported staging provider %q (expected s3 or gcs)", provider)
4944
}
5045
}
51-
52-
//nolint:iface // factory function intentionally returns interface
53-
func newS3StagingStoreFromConfig(
54-
ctx context.Context,
55-
config *protos.ClickhouseConfig,
56-
unifiedBucketName string,
57-
chVersion clickhouseproto.Version,
58-
) (StagingStore, error) {
59-
var awsConfig utils.PeerAWSCredentials
60-
var awsBucketPath string
61-
if config.S3 != nil {
62-
awsConfig = utils.NewPeerAWSCredentials(config.S3)
63-
awsBucketPath = config.S3.Url
64-
} else {
65-
awsConfig = utils.PeerAWSCredentials{
66-
Credentials: aws.Credentials{
67-
AccessKeyID: config.AccessKeyId,
68-
SecretAccessKey: config.SecretAccessKey,
69-
},
70-
EndpointUrl: config.Endpoint,
71-
Region: config.Region,
72-
}
73-
awsBucketPath = config.S3Path
74-
}
75-
76-
credentialsProvider, err := utils.GetAWSCredentialsProvider(ctx, "clickhouse", awsConfig)
77-
if err != nil {
78-
return nil, err
79-
}
80-
81-
if awsBucketPath == "" {
82-
if unifiedBucketName == "" {
83-
return nil, errors.New("PeerDB ClickHouse Bucket Name not set")
84-
}
85-
deploymentUID := internal.PeerDBDeploymentUID()
86-
flowName, _ := ctx.Value(shared.FlowNameKey).(string)
87-
bucketPathSuffix := fmt.Sprintf("%s/%s", url.PathEscape(deploymentUID), url.PathEscape(flowName))
88-
awsBucketPath = fmt.Sprintf("s3://%s/%s", unifiedBucketName, bucketPathSuffix)
89-
}
90-
91-
// S3 with session tokens requires ClickHouse >= 24.3.1
92-
// https://github.com/ClickHouse/ClickHouse/issues/61230
93-
credentials, err := credentialsProvider.Retrieve(ctx)
94-
if err != nil {
95-
return nil, err
96-
}
97-
if credentials.AWS.SessionToken != "" {
98-
if !clickhouseproto.CheckMinVersion(
99-
clickhouseproto.Version{Major: 24, Minor: 3, Patch: 1},
100-
chVersion,
101-
) {
102-
return nil, fmt.Errorf(
103-
"provide S3 Transient Stage details explicitly or upgrade to ClickHouse version >= 24.3.1, current version is %s. %s",
104-
chVersion,
105-
"You can also contact PeerDB support for implicit S3 stage setup for older versions of ClickHouse.")
106-
}
107-
}
108-
109-
return newS3StagingStore(awsBucketPath, credentialsProvider)
110-
}
111-
112-
//nolint:iface // factory function intentionally returns interface
113-
func newGCSStagingStore(ctx context.Context, bucketName string) (StagingStore, error) {
114-
if bucketName == "" {
115-
return nil, errors.New("PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME must be set when staging provider is gcs")
116-
}
117-
118-
deploymentUID := internal.PeerDBDeploymentUID()
119-
flowName, _ := ctx.Value(shared.FlowNameKey).(string)
120-
bucketPath := fmt.Sprintf("gs://%s/%s/%s",
121-
bucketName, url.PathEscape(deploymentUID), url.PathEscape(flowName))
122-
123-
return newGCSStagingStoreFromPath(ctx, bucketPath)
124-
}

‎flow/connectors/clickhouse/staging_gcs.go‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,21 @@ package connclickhouse
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"io"
78
"log/slog"
9+
"net/url"
810
"strings"
911
"time"
1012

13+
"github.com/PeerDB-io/peerdb/flow/internal"
14+
peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse"
15+
"github.com/PeerDB-io/peerdb/flow/shared"
16+
1117
"cloud.google.com/go/storage"
1218
"github.com/google/uuid"
1319
"google.golang.org/api/iterator"
14-
15-
"github.com/PeerDB-io/peerdb/flow/internal"
16-
peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse"
1720
)
1821

1922
const gcsSignedURLExpiry = 1 * time.Hour
@@ -27,7 +30,17 @@ type gcsStagingStore struct {
2730
fullPath string
2831
}
2932

30-
func newGCSStagingStoreFromPath(ctx context.Context, bucketPath string) (*gcsStagingStore, error) {
33+
//nolint:iface // factory function intentionally returns interface
34+
func newGCSStagingStore(ctx context.Context, bucketName string) (StagingStore, error) {
35+
if bucketName == "" {
36+
return nil, errors.New("PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME must be set when staging provider is gcs")
37+
}
38+
39+
deploymentUID := internal.PeerDBDeploymentUID()
40+
flowName, _ := ctx.Value(shared.FlowNameKey).(string)
41+
bucketPath := fmt.Sprintf("gs://%s/%s/%s",
42+
bucketName, url.PathEscape(deploymentUID), url.PathEscape(flowName))
43+
3144
// bucketPath may use either "gs://" (standard GCS URI scheme) or "gcs://" prefix.
3245
path := strings.TrimPrefix(bucketPath, "gcs://")
3346
path = strings.TrimPrefix(path, "gs://")

‎flow/connectors/clickhouse/staging_s3.go‎

Lines changed: 70 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,24 @@ package connclickhouse
22

33
import (
44
"context"
5+
"errors"
56
"fmt"
67
"io"
78
"log/slog"
9+
"net/url"
810
"strings"
911
"time"
1012

11-
"github.com/aws/aws-sdk-go-v2/aws"
12-
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
13-
"github.com/aws/aws-sdk-go-v2/service/s3"
14-
1513
"github.com/PeerDB-io/peerdb/flow/connectors/utils"
14+
"github.com/PeerDB-io/peerdb/flow/generated/protos"
1615
"github.com/PeerDB-io/peerdb/flow/internal"
1716
peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse"
17+
"github.com/PeerDB-io/peerdb/flow/shared"
18+
19+
clickhouseproto "github.com/ClickHouse/clickhouse-go/v2/lib/proto"
20+
"github.com/aws/aws-sdk-go-v2/aws"
21+
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
22+
"github.com/aws/aws-sdk-go-v2/service/s3"
1823
)
1924

2025
// s3StagingStore implements StagingStore for AWS S3 (and S3-compatible services).
@@ -25,16 +30,73 @@ type s3StagingStore struct {
2530
fullPath string // original "s3://bucket/prefix" for logging
2631
}
2732

28-
func newS3StagingStore(bucketPath string, creds utils.AWSCredentialsProvider) (*s3StagingStore, error) {
29-
s3o, err := utils.NewS3BucketAndPrefix(bucketPath)
33+
//nolint:iface // factory function intentionally returns interface
34+
func newS3StagingStore(
35+
ctx context.Context,
36+
config *protos.ClickhouseConfig,
37+
unifiedBucketName string,
38+
chVersion clickhouseproto.Version,
39+
) (StagingStore, error) {
40+
var awsConfig utils.PeerAWSCredentials
41+
var awsBucketPath string
42+
if config.S3 != nil {
43+
awsConfig = utils.NewPeerAWSCredentials(config.S3)
44+
awsBucketPath = config.S3.Url
45+
} else {
46+
awsConfig = utils.PeerAWSCredentials{
47+
Credentials: aws.Credentials{
48+
AccessKeyID: config.AccessKeyId,
49+
SecretAccessKey: config.SecretAccessKey,
50+
},
51+
EndpointUrl: config.Endpoint,
52+
Region: config.Region,
53+
}
54+
awsBucketPath = config.S3Path
55+
}
56+
57+
credentialsProvider, err := utils.GetAWSCredentialsProvider(ctx, "clickhouse", awsConfig)
58+
if err != nil {
59+
return nil, err
60+
}
61+
62+
if awsBucketPath == "" {
63+
if unifiedBucketName == "" {
64+
return nil, errors.New("PeerDB ClickHouse Bucket Name not set")
65+
}
66+
deploymentUID := internal.PeerDBDeploymentUID()
67+
flowName, _ := ctx.Value(shared.FlowNameKey).(string)
68+
bucketPathSuffix := fmt.Sprintf("%s/%s", url.PathEscape(deploymentUID), url.PathEscape(flowName))
69+
awsBucketPath = fmt.Sprintf("s3://%s/%s", unifiedBucketName, bucketPathSuffix)
70+
}
71+
72+
// S3 with session tokens requires ClickHouse >= 24.3.1
73+
// https://github.com/ClickHouse/ClickHouse/issues/61230
74+
credentials, err := credentialsProvider.Retrieve(ctx)
75+
if err != nil {
76+
return nil, err
77+
}
78+
if credentials.AWS.SessionToken != "" {
79+
if !clickhouseproto.CheckMinVersion(
80+
clickhouseproto.Version{Major: 24, Minor: 3, Patch: 1},
81+
chVersion,
82+
) {
83+
return nil, fmt.Errorf(
84+
"provide S3 Transient Stage details explicitly or upgrade to ClickHouse version >= 24.3.1, current version is %s. %s",
85+
chVersion,
86+
"You can also contact PeerDB support for implicit S3 stage setup for older versions of ClickHouse.")
87+
}
88+
}
89+
90+
s3o, err := utils.NewS3BucketAndPrefix(awsBucketPath)
3091
if err != nil {
3192
return nil, fmt.Errorf("failed to parse S3 bucket path: %w", err)
3293
}
94+
3395
return &s3StagingStore{
3496
bucket: s3o.Bucket,
3597
prefix: s3o.Prefix,
36-
fullPath: bucketPath,
37-
creds: creds,
98+
fullPath: awsBucketPath,
99+
creds: credentialsProvider,
38100
}, nil
39101
}
40102

0 commit comments

Comments
 (0)