UBERF-9608 Fix bucket selection in datalake (#8210)

Signed-off-by: Alexander Onnikov <[email protected]>

Author: aonnikov

services/datalake/pod-datalake/src/datalake/datalake.ts

@@ -25,7 +25,7 @@ import { type S3Bucket } from '../s3'
 export class DatalakeImpl implements Datalake {
   constructor (
     private readonly db: BlobDB,
-    private readonly buckets: Partial<Record<Location, S3Bucket>>,
+    private readonly buckets: Array<{ location: Location, bucket: S3Bucket }>,
     private readonly options: {
       cacheControl: string
     }
@@ -44,13 +44,12 @@ export class DatalakeImpl implements Datalake {
   }
 
   async head (ctx: MeasureContext, workspace: string, name: string): Promise<BlobHead | null> {
-    const { bucket } = await this.selectStorage(ctx, workspace)
-
     const blob = await this.db.getBlob(ctx, { workspace, name })
     if (blob === null) {
       return null
     }
 
+    const { bucket } = await this.selectStorage(ctx, workspace, blob.location)
     const head = await bucket.head(ctx, blob.filename)
     if (head == null) {
       return null
@@ -72,13 +71,13 @@ export class DatalakeImpl implements Datalake {
     name: string,
     options: { range?: string }
   ): Promise<BlobBody | null> {
-    const { bucket } = await this.selectStorage(ctx, workspace)
-
     const blob = await this.db.getBlob(ctx, { workspace, name })
     if (blob === null) {
       return null
     }
 
+    const { bucket } = await this.selectStorage(ctx, workspace, blob.location)
+
     const range = options.range
     const object = await bucket.get(ctx, blob.filename, { range })
     if (object == null) {
@@ -176,17 +175,17 @@ export class DatalakeImpl implements Datalake {
     await this.db.setParent(ctx, { workspace, name }, parent !== null ? { workspace, name: parent } : null)
   }
 
-  async selectStorage (ctx: MeasureContext, workspace: string): Promise<BlobStorage> {
-    const location = this.selectLocation(workspace)
-    const bucket = this.buckets[location]
+  async selectStorage (ctx: MeasureContext, workspace: string, location?: Location): Promise<BlobStorage> {
+    location ??= this.selectLocation(workspace)
+
+    const bucket = this.buckets.find((b) => b.location === location)?.bucket
     if (bucket == null) {
       throw new Error(`Unsupported location: ${location}`)
     }
     return { location, bucket }
   }
 
   selectLocation (workspace: string): Location {
-    // TODO select location based on workspace
-    return 'weur'
+    return this.buckets[0].location
   }
 }

services/datalake/pod-datalake/src/datalake/types.ts

@@ -17,7 +17,7 @@ import { MeasureContext } from '@hcengineering/core'
 import { type Readable } from 'stream'
 import { S3Bucket } from '../s3'
 
-export type Location = 'weur' | 'eeur' | 'wnam' | 'enam' | 'apac'
+export type Location = 'eu' | 'weur' | 'eeur' | 'wnam' | 'enam' | 'apac'
 
 export type UUID = string & { __uuid: true }
 

services/datalake/pod-datalake/src/server.ts

@@ -74,17 +74,20 @@ const wrapRequest =
     }
 
 export function createServer (ctx: MeasureContext, config: Config): { app: Express, close: () => void } {
-  const buckets: Partial<Record<Location, S3Bucket>> = {}
+  const buckets: Array<{ location: Location, bucket: S3Bucket }> = []
   for (const bucket of config.Buckets) {
     const location = bucket.location as Location
     if (
+      location === 'eu' ||
       location === 'weur' ||
       location === 'eeur' ||
       location === 'wnam' ||
       location === 'enam' ||
       location === 'apac'
     ) {
-      buckets[location] = createBucket(createClient(bucket), bucket.bucket)
+      buckets.push({ location, bucket: createBucket(createClient(bucket), bucket.bucket) })
+    } else {
+      ctx.warn('invalid bucket location', { location, bucket })
     }
   }