Skip to content

Commit 4286c43

Browse files
whites11whites11
committed
refactor(staging): move staging store types from utils/ to clickhouse/
Move StagingStore interface, S3StagingStore, and GCSStagingStore into the clickhouse connector package — their only consumer. This eliminates the iface lint workaround and keeps staging logic colocated with the connector that uses it. Types are now unexported (s3StagingStore, gcsStagingStore) since they don't need to be visible outside the package.
1 parent 5c65d6a commit 4286c43

6 files changed

Lines changed: 43 additions & 180 deletions

File tree

‎flow/connectors/clickhouse/clickhouse.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ type ClickHouseConnector struct {
3535
logger log.Logger
3636
Config *protos.ClickhouseConfig
3737
credsProvider *utils.ClickHouseS3Credentials
38-
staging utils.StagingStore
38+
staging StagingStore
3939
chVersion *clickhouseproto.Version
4040
}
4141

‎flow/connectors/clickhouse/staging.go‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ func createStagingStore(
2424
env map[string]string,
2525
config *protos.ClickhouseConfig,
2626
chVersion clickhouseproto.Version,
27-
) (utils.StagingStore, *utils.ClickHouseS3Credentials, error) {
27+
) (StagingStore, *utils.ClickHouseS3Credentials, error) {
2828
provider, err := internal.PeerDBClickHouseStagingProvider(ctx, env)
2929
if err != nil {
3030
return nil, nil, fmt.Errorf("failed to get staging provider: %w", err)
@@ -37,19 +37,19 @@ func createStagingStore(
3737

3838
switch strings.ToLower(provider) {
3939
case "gcs":
40-
store, err := createGCSStagingStore(ctx, bucketName)
40+
store, err := newGCSStagingStore(ctx, bucketName)
4141
return store, nil, err
4242
default:
43-
return createS3StagingStore(ctx, config, bucketName, chVersion)
43+
return newS3StagingStoreFromConfig(ctx, config, bucketName, chVersion)
4444
}
4545
}
4646

47-
func createS3StagingStore(
47+
func newS3StagingStoreFromConfig(
4848
ctx context.Context,
4949
config *protos.ClickhouseConfig,
5050
unifiedBucketName string,
5151
chVersion clickhouseproto.Version,
52-
) (utils.StagingStore, *utils.ClickHouseS3Credentials, error) {
52+
) (StagingStore, *utils.ClickHouseS3Credentials, error) {
5353
var awsConfig utils.PeerAWSCredentials
5454
var awsBucketPath string
5555
if config.S3 != nil {
@@ -105,18 +105,15 @@ func createS3StagingStore(
105105
BucketPath: awsBucketPath,
106106
}
107107

108-
store, err := utils.NewS3StagingStore(awsBucketPath, credentialsProvider)
108+
store, err := newS3StagingStore(awsBucketPath, credentialsProvider)
109109
if err != nil {
110110
return nil, nil, err
111111
}
112112

113113
return store, creds, nil
114114
}
115115

116-
func createGCSStagingStore(
117-
ctx context.Context,
118-
bucketName string,
119-
) (utils.StagingStore, error) {
116+
func newGCSStagingStore(ctx context.Context, bucketName string) (StagingStore, error) {
120117
if bucketName == "" {
121118
return nil, errors.New("PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME must be set when staging provider is gcs")
122119
}
@@ -126,5 +123,5 @@ func createGCSStagingStore(
126123
bucketPath := fmt.Sprintf("gs://%s/%s/%s",
127124
bucketName, url.PathEscape(deploymentUID), url.PathEscape(flowName))
128125

129-
return utils.NewGCSStagingStore(ctx, bucketPath)
126+
return newGCSStagingStoreFromPath(ctx, bucketPath)
130127
}

flow/connectors/utils/staging_gcs.go renamed to flow/connectors/clickhouse/staging_gcs.go

Lines changed: 11 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package utils
1+
package connclickhouse
22

33
import (
44
"context"
@@ -17,16 +17,16 @@ import (
1717

1818
const gcsSignedURLExpiry = 15 * time.Minute
1919

20-
// GCSStagingStore implements StagingStore for Google Cloud Storage using
20+
// gcsStagingStore implements StagingStore for Google Cloud Storage using
2121
// native GCS SDK and signed URLs for ClickHouse reads.
22-
type GCSStagingStore struct {
22+
type gcsStagingStore struct {
2323
client *storage.Client
2424
bucket string
2525
prefix string
2626
fullPath string
2727
}
2828

29-
func NewGCSStagingStore(ctx context.Context, bucketPath string) (*GCSStagingStore, error) {
29+
func newGCSStagingStoreFromPath(ctx context.Context, bucketPath string) (*gcsStagingStore, error) {
3030
// bucketPath may use either "gs://" (standard GCS URI scheme) or "gcs://" prefix.
3131
path := strings.TrimPrefix(bucketPath, "gcs://")
3232
path = strings.TrimPrefix(path, "gs://")
@@ -39,21 +39,19 @@ func NewGCSStagingStore(ctx context.Context, bucketPath string) (*GCSStagingStor
3939
return nil, fmt.Errorf("failed to create GCS client: %w", err)
4040
}
4141

42-
return &GCSStagingStore{
42+
return &gcsStagingStore{
4343
client: client,
4444
bucket: bucket,
4545
prefix: prefix,
4646
fullPath: bucketPath,
4747
}, nil
4848
}
4949

50-
func (g *GCSStagingStore) Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error {
50+
func (g *gcsStagingStore) Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error {
5151
logger := internal.LoggerFromCtx(ctx)
5252

5353
obj := g.client.Bucket(g.bucket).Object(key)
5454
w := obj.NewWriter(ctx)
55-
// GCS supports up to 5 TiB objects; the client handles chunked uploads internally.
56-
// Default chunk size is 16 MiB which is fine for most staging files.
5755

5856
if _, err := io.Copy(w, body); err != nil {
5957
w.Close()
@@ -67,9 +65,7 @@ func (g *GCSStagingStore) Upload(ctx context.Context, env map[string]string, key
6765
return nil
6866
}
6967

70-
func (g *GCSStagingStore) TableFunctionExpr(ctx context.Context, key string, format string) (string, error) {
71-
// Generate a signed URL that ClickHouse can use to read the staged file.
72-
// Uses Application Default Credentials for signing.
68+
func (g *gcsStagingStore) TableFunctionExpr(ctx context.Context, key string, format string) (string, error) {
7369
signedURL, err := g.client.Bucket(g.bucket).SignedURL(key, &storage.SignedURLOptions{
7470
Method: "GET",
7571
Expires: time.Now().Add(gcsSignedURLExpiry),
@@ -87,7 +83,7 @@ func (g *GCSStagingStore) TableFunctionExpr(ctx context.Context, key string, for
8783
return expr.String(), nil
8884
}
8985

90-
func (g *GCSStagingStore) DeletePrefix(ctx context.Context, prefix string) error {
86+
func (g *gcsStagingStore) DeletePrefix(ctx context.Context, prefix string) error {
9187
logger := internal.LoggerFromCtx(ctx)
9288
logger.Info("Deleting objects from GCS",
9389
slog.String("bucket", g.bucket), slog.String("prefix", prefix))
@@ -111,7 +107,7 @@ func (g *GCSStagingStore) DeletePrefix(ctx context.Context, prefix string) error
111107
return nil
112108
}
113109

114-
func (g *GCSStagingStore) Validate(ctx context.Context) error {
110+
func (g *gcsStagingStore) Validate(ctx context.Context) error {
115111
testKey := g.prefix + "/_peerdb_check_" + time.Now().Format("20060102150405")
116112
obj := g.client.Bucket(g.bucket).Object(testKey)
117113

@@ -131,20 +127,10 @@ func (g *GCSStagingStore) Validate(ctx context.Context) error {
131127
return nil
132128
}
133129

134-
func (g *GCSStagingStore) BucketPath() string {
130+
func (g *gcsStagingStore) BucketPath() string {
135131
return g.fullPath
136132
}
137133

138-
func (g *GCSStagingStore) KeyPrefix() string {
139-
return g.prefix
140-
}
141-
142-
// Bucket returns the bucket name.
143-
func (g *GCSStagingStore) Bucket() string {
144-
return g.bucket
145-
}
146-
147-
// Prefix returns the key prefix.
148-
func (g *GCSStagingStore) Prefix() string {
134+
func (g *gcsStagingStore) KeyPrefix() string {
149135
return g.prefix
150136
}

flow/connectors/utils/staging_s3.go renamed to flow/connectors/clickhouse/staging_s3.go

Lines changed: 19 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package utils
1+
package connclickhouse
22

33
import (
44
"context"
@@ -12,35 +12,36 @@ import (
1212
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
1313
"github.com/aws/aws-sdk-go-v2/service/s3"
1414

15+
"github.com/PeerDB-io/peerdb/flow/connectors/utils"
1516
"github.com/PeerDB-io/peerdb/flow/internal"
1617
peerdb_clickhouse "github.com/PeerDB-io/peerdb/flow/pkg/clickhouse"
1718
)
1819

19-
// S3StagingStore implements StagingStore for AWS S3 (and S3-compatible services).
20-
type S3StagingStore struct {
21-
creds AWSCredentialsProvider
20+
// s3StagingStore implements StagingStore for AWS S3 (and S3-compatible services).
21+
type s3StagingStore struct {
22+
creds utils.AWSCredentialsProvider
2223
bucket string
2324
prefix string
2425
fullPath string // original "s3://bucket/prefix" for logging
2526
}
2627

27-
func NewS3StagingStore(bucketPath string, creds AWSCredentialsProvider) (*S3StagingStore, error) {
28-
s3o, err := NewS3BucketAndPrefix(bucketPath)
28+
func newS3StagingStore(bucketPath string, creds utils.AWSCredentialsProvider) (*s3StagingStore, error) {
29+
s3o, err := utils.NewS3BucketAndPrefix(bucketPath)
2930
if err != nil {
3031
return nil, fmt.Errorf("failed to parse S3 bucket path: %w", err)
3132
}
32-
return &S3StagingStore{
33+
return &s3StagingStore{
3334
bucket: s3o.Bucket,
3435
prefix: s3o.Prefix,
3536
fullPath: bucketPath,
3637
creds: creds,
3738
}, nil
3839
}
3940

40-
func (s *S3StagingStore) Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error {
41+
func (s *s3StagingStore) Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error {
4142
logger := internal.LoggerFromCtx(ctx)
4243

43-
s3svc, err := CreateS3Client(ctx, s.creds)
44+
s3svc, err := utils.CreateS3Client(ctx, s.creds)
4445
if err != nil {
4546
return fmt.Errorf("failed to create S3 client: %w", err)
4647
}
@@ -73,10 +74,10 @@ func (s *S3StagingStore) Upload(ctx context.Context, env map[string]string, key
7374
return nil
7475
}
7576

76-
func (s *S3StagingStore) TableFunctionExpr(ctx context.Context, key string, format string) (string, error) {
77+
func (s *s3StagingStore) TableFunctionExpr(ctx context.Context, key string, format string) (string, error) {
7778
endpoint := s.creds.GetEndpointURL()
7879
region := s.creds.GetRegion()
79-
fileURL := FileURLForS3Service(endpoint, region, s.bucket, key)
80+
fileURL := utils.FileURLForS3Service(endpoint, region, s.bucket, key)
8081

8182
creds, err := s.creds.Retrieve(ctx)
8283
if err != nil {
@@ -106,10 +107,10 @@ func (s *S3StagingStore) TableFunctionExpr(ctx context.Context, key string, form
106107
return expr.String(), nil
107108
}
108109

109-
func (s *S3StagingStore) DeletePrefix(ctx context.Context, prefix string) error {
110+
func (s *s3StagingStore) DeletePrefix(ctx context.Context, prefix string) error {
110111
logger := internal.LoggerFromCtx(ctx)
111112

112-
s3svc, err := CreateS3Client(ctx, s.creds)
113+
s3svc, err := utils.CreateS3Client(ctx, s.creds)
113114
if err != nil {
114115
return fmt.Errorf("failed to create S3 client: %w", err)
115116
}
@@ -142,35 +143,18 @@ func (s *S3StagingStore) DeletePrefix(ctx context.Context, prefix string) error
142143
return nil
143144
}
144145

145-
func (s *S3StagingStore) Validate(ctx context.Context) error {
146-
s3Client, err := CreateS3Client(ctx, s.creds)
146+
func (s *s3StagingStore) Validate(ctx context.Context) error {
147+
s3Client, err := utils.CreateS3Client(ctx, s.creds)
147148
if err != nil {
148149
return fmt.Errorf("failed to create S3 client: %w", err)
149150
}
150-
return PutAndRemoveS3(ctx, s3Client, s.bucket, s.prefix)
151+
return utils.PutAndRemoveS3(ctx, s3Client, s.bucket, s.prefix)
151152
}
152153

153-
func (s *S3StagingStore) BucketPath() string {
154+
func (s *s3StagingStore) BucketPath() string {
154155
return s.fullPath
155156
}
156157

157-
func (s *S3StagingStore) KeyPrefix() string {
158-
return s.prefix
159-
}
160-
161-
// CredentialsProvider returns the underlying AWS credentials provider.
162-
// This is needed for backward compatibility (e.g., ClickHouse version checks
163-
// that inspect whether credentials use session tokens).
164-
func (s *S3StagingStore) CredentialsProvider() AWSCredentialsProvider {
165-
return s.creds
166-
}
167-
168-
// Bucket returns the bucket name.
169-
func (s *S3StagingStore) Bucket() string {
170-
return s.bucket
171-
}
172-
173-
// Prefix returns the key prefix within the bucket.
174-
func (s *S3StagingStore) Prefix() string {
158+
func (s *s3StagingStore) KeyPrefix() string {
175159
return s.prefix
176160
}

flow/connectors/utils/staging.go renamed to flow/connectors/clickhouse/staging_store.go

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,13 @@
1-
package utils
1+
package connclickhouse
22

33
import (
44
"context"
55
"io"
66
)
77

8-
var (
9-
_ StagingStore = (*S3StagingStore)(nil)
10-
_ StagingStore = (*GCSStagingStore)(nil)
11-
)
12-
13-
// StagingStore abstracts cloud storage used by PeerDB's ClickHouse connector
14-
// for staging Avro files. Files are written by PeerDB and read by ClickHouse
15-
// via table functions (s3(), url(), etc.).
8+
// StagingStore abstracts cloud storage used for staging Avro files.
9+
// Files are written by PeerDB and read by ClickHouse via table functions
10+
// (s3(), url(), etc.).
1611
type StagingStore interface {
1712
// Upload streams data from body to the given key in the staging bucket.
1813
Upload(ctx context.Context, env map[string]string, key string, body io.Reader) error

0 commit comments

Comments
 (0)