@@ -117,19 +117,26 @@ func createStagingStore(
117117 return nil , nil , fmt .Errorf ("failed to get staging provider: %w" , err )
118118 }
119119
120+ // Prefer unified PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME, fall back to provider-specific env vars.
121+ bucketName , err := internal .PeerDBClickHouseStagingBucketName (ctx , env )
122+ if err != nil {
123+ return nil , nil , fmt .Errorf ("failed to get staging bucket name: %w" , err )
124+ }
125+
120126 switch strings .ToLower (provider ) {
121127 case "gcs" :
122- store , err := createGCSStagingStore (ctx , env )
128+ store , err := createGCSStagingStore (ctx , bucketName )
123129 return store , nil , err
124130 default :
125- return createS3StagingStore (ctx , env , config )
131+ return createS3StagingStore (ctx , env , config , bucketName )
126132 }
127133}
128134
129135func createS3StagingStore (
130136 ctx context.Context ,
131137 env map [string ]string ,
132138 config * protos.ClickhouseConfig ,
139+ unifiedBucketName string ,
133140) (utils.StagingStore , * utils.ClickHouseS3Credentials , error ) {
134141 var awsConfig utils.PeerAWSCredentials
135142 var awsBucketPath string
@@ -157,9 +164,13 @@ func createS3StagingStore(
157164 deploymentUID := internal .PeerDBDeploymentUID ()
158165 flowName , _ := ctx .Value (shared .FlowNameKey ).(string )
159166 bucketPathSuffix := fmt .Sprintf ("%s/%s" , url .PathEscape (deploymentUID ), url .PathEscape (flowName ))
160- awsBucketName , err := internal .PeerDBClickHouseAWSS3BucketName (ctx , env )
161- if err != nil {
162- return nil , nil , fmt .Errorf ("failed to get PeerDB ClickHouse Bucket Name: %w" , err )
167+ // Prefer unified bucket name, fall back to legacy S3-specific env var.
168+ awsBucketName := unifiedBucketName
169+ if awsBucketName == "" {
170+ awsBucketName , err = internal .PeerDBClickHouseAWSS3BucketName (ctx , env )
171+ if err != nil {
172+ return nil , nil , fmt .Errorf ("failed to get PeerDB ClickHouse Bucket Name: %w" , err )
173+ }
163174 }
164175 if awsBucketName == "" {
165176 return nil , nil , errors .New ("PeerDB ClickHouse Bucket Name not set" )
@@ -182,20 +193,16 @@ func createS3StagingStore(
182193
183194func createGCSStagingStore (
184195 ctx context.Context ,
185- env map [ string ] string ,
196+ bucketName string ,
186197) (utils.StagingStore , error ) {
187- gcsBucketName , err := internal .PeerDBClickHouseGCSBucketName (ctx , env )
188- if err != nil {
189- return nil , fmt .Errorf ("failed to get GCS bucket name: %w" , err )
190- }
191- if gcsBucketName == "" {
192- return nil , errors .New ("PEERDB_CLICKHOUSE_GCS_BUCKET_NAME not set but staging provider is gcs" )
198+ if bucketName == "" {
199+ return nil , errors .New ("PEERDB_CLICKHOUSE_STAGING_BUCKET_NAME must be set when staging provider is gcs" )
193200 }
194201
195202 deploymentUID := internal .PeerDBDeploymentUID ()
196203 flowName , _ := ctx .Value (shared .FlowNameKey ).(string )
197204 bucketPath := fmt .Sprintf ("gs://%s/%s/%s" ,
198- gcsBucketName , url .PathEscape (deploymentUID ), url .PathEscape (flowName ))
205+ bucketName , url .PathEscape (deploymentUID ), url .PathEscape (flowName ))
199206
200207 return utils .NewGCSStagingStore (ctx , bucketPath )
201208}
0 commit comments