From ccd8623be563b72c1bbd23e2feecf278c3300db7 Mon Sep 17 00:00:00 2001 From: Flavien David Date: Fri, 31 Jan 2025 10:37:33 +0100 Subject: [PATCH] Add support for pod annotation service account. (#10415) --- front/lib/file_storage/index.ts | 28 ++++++++++++------- front/start_worker.ts | 2 +- .../relocation/lib/file_storage/relocation.ts | 15 ++++++++-- .../relocation/lib/file_storage/transfer.ts | 8 ++++-- front/temporal/relocation/worker.ts | 6 ++-- 5 files changed, 40 insertions(+), 19 deletions(-) diff --git a/front/lib/file_storage/index.ts b/front/lib/file_storage/index.ts index 626da6f6a7ce..c7e61c9c5af6 100644 --- a/front/lib/file_storage/index.ts +++ b/front/lib/file_storage/index.ts @@ -9,13 +9,20 @@ import { isGCSNotFoundError } from "@app/lib/file_storage/types"; const DEFAULT_SIGNED_URL_EXPIRATION_DELAY_MS = 5 * 60 * 1000; // 5 minutes. +interface FileStorageOptions { + useServiceAccount?: boolean; +} + export class FileStorage { private readonly bucket: Bucket; private readonly storage: Storage; - constructor(bucketKey: string) { + constructor( + bucketKey: string, + { useServiceAccount }: FileStorageOptions = { useServiceAccount: true } + ) { this.storage = new Storage({ - keyFilename: config.getServiceAccount(), + keyFilename: useServiceAccount ? config.getServiceAccount() : undefined, }); this.bucket = this.storage.bucket(bucketKey); @@ -126,17 +133,18 @@ export class FileStorage { const bucketInstances = new Map(); -export const getBucketInstance: (bucketConfig: string) => FileStorage = ( - bucketConfig: string -) => { +export const getBucketInstance: ( + bucketConfig: string, + options?: FileStorageOptions +) => FileStorage = (bucketConfig, options) => { if (!bucketInstances.has(bucketConfig)) { - bucketInstances.set(bucketConfig, new FileStorage(bucketConfig)); + bucketInstances.set(bucketConfig, new FileStorage(bucketConfig, options)); } return bucketInstances.get(bucketConfig); }; -export const getPrivateUploadBucket = () => - getBucketInstance(config.getGcsPrivateUploadsBucket()); +export const getPrivateUploadBucket = (options?: FileStorageOptions) => + getBucketInstance(config.getGcsPrivateUploadsBucket(), options); -export const getPublicUploadBucket = () => - getBucketInstance(config.getGcsPublicUploadBucket()); +export const getPublicUploadBucket = (options?: FileStorageOptions) => + getBucketInstance(config.getGcsPublicUploadBucket(), options); diff --git a/front/start_worker.ts b/front/start_worker.ts index 85a0fd52d365..4d043085632b 100644 --- a/front/start_worker.ts +++ b/front/start_worker.ts @@ -10,6 +10,7 @@ import { runLabsWorker } from "@app/temporal/labs/worker"; import { runMentionsCountWorker } from "@app/temporal/mentions_count_queue/worker"; import { runPermissionsWorker } from "@app/temporal/permissions_queue/worker"; import { runProductionChecksWorker } from "@app/temporal/production_checks/worker"; +import { runRelocationWorker } from "@app/temporal/relocation/worker"; import { runScrubWorkspaceQueueWorker } from "@app/temporal/scrub_workspace/worker"; import { runTrackerNotificationWorker, @@ -18,7 +19,6 @@ import { import { runUpsertQueueWorker } from "@app/temporal/upsert_queue/worker"; import { runUpsertTableQueueWorker } from "@app/temporal/upsert_tables/worker"; import { runUpdateWorkspaceUsageWorker } from "@app/temporal/usage_queue/worker"; -import { runRelocationWorker } from "@app/temporal/relocation/worker"; setupGlobalErrorHandler(logger); diff --git a/front/temporal/relocation/lib/file_storage/relocation.ts b/front/temporal/relocation/lib/file_storage/relocation.ts index d5523d1c05bb..bd654d3ea3f6 100644 --- a/front/temporal/relocation/lib/file_storage/relocation.ts +++ b/front/temporal/relocation/lib/file_storage/relocation.ts @@ -1,3 +1,5 @@ +import { isDevelopment } from "@dust-tt/types"; + import { getBucketInstance } from "@app/lib/file_storage"; import config from "@app/temporal/relocation/activities/config"; @@ -9,6 +11,7 @@ interface RelocationStorageOptions { operation: string; } +// In prod, we use pod annotations to set the service account. export async function writeToRelocationStorage( data: unknown, { workspaceId, type, operation }: RelocationStorageOptions @@ -16,7 +19,9 @@ export async function writeToRelocationStorage( const timestamp = Date.now(); const path = `${RELOCATION_PATH_PREFIX}/${workspaceId}/${type}/${operation}/${timestamp}.json`; - const relocationBucket = getBucketInstance(config.getGcsRelocationBucket()); + const relocationBucket = getBucketInstance(config.getGcsRelocationBucket(), { + useServiceAccount: isDevelopment(), + }); await relocationBucket.uploadRawContentToBucket({ content: JSON.stringify(data), @@ -30,7 +35,9 @@ export async function writeToRelocationStorage( export async function readFromRelocationStorage( dataPath: string ): Promise { - const relocationBucket = getBucketInstance(config.getGcsRelocationBucket()); + const relocationBucket = getBucketInstance(config.getGcsRelocationBucket(), { + useServiceAccount: isDevelopment(), + }); const content = await relocationBucket.fetchFileContent(dataPath); @@ -38,7 +45,9 @@ export async function readFromRelocationStorage( } export async function deleteFromRelocationStorage(dataPath: string) { - const relocationBucket = getBucketInstance(config.getGcsRelocationBucket()); + const relocationBucket = getBucketInstance(config.getGcsRelocationBucket(), { + useServiceAccount: isDevelopment(), + }); await relocationBucket.delete(dataPath, { ignoreNotFound: true }); } diff --git a/front/temporal/relocation/lib/file_storage/transfer.ts b/front/temporal/relocation/lib/file_storage/transfer.ts index c195b631e1ad..72e79aa887b6 100644 --- a/front/temporal/relocation/lib/file_storage/transfer.ts +++ b/front/temporal/relocation/lib/file_storage/transfer.ts @@ -1,5 +1,5 @@ import type { Result } from "@dust-tt/types"; -import { Err, Ok } from "@dust-tt/types"; +import { Err, isDevelopment, Ok } from "@dust-tt/types"; import { protos } from "@google-cloud/storage-transfer"; import { StorageTransferServiceClient } from "@google-cloud/storage-transfer"; import type { google } from "@google-cloud/storage-transfer/build/protos/protos"; @@ -23,7 +23,11 @@ export class StorageTransferService { private transferClient: StorageTransferServiceClient; constructor() { - const serviceAccountPath = config.getServiceAccount(); + // Only use service account in dev. In prod, we use pod annotations to set the + // service account. + const serviceAccountPath = isDevelopment() + ? config.getServiceAccount() + : undefined; this.transferClient = new StorageTransferServiceClient({ keyFilename: serviceAccountPath, diff --git a/front/temporal/relocation/worker.ts b/front/temporal/relocation/worker.ts index d487c4b5f98d..419204367796 100644 --- a/front/temporal/relocation/worker.ts +++ b/front/temporal/relocation/worker.ts @@ -5,12 +5,12 @@ import TsconfigPathsPlugin from "tsconfig-paths-webpack-plugin"; import { config } from "@app/lib/api/regions/config"; import { ActivityInboundLogInterceptor } from "@app/lib/temporal_monitoring"; import logger from "@app/logger/logger"; -import * as frontDestinationActivities from "@app/temporal/relocation/activities/destination_region/front"; -import * as frontSourceActivities from "@app/temporal/relocation/activities/source_region/front"; import * as connectorsDestinationActivities from "@app/temporal/relocation/activities/destination_region/connectors/sql"; +import * as coreDestinationActivities from "@app/temporal/relocation/activities/destination_region/core"; +import * as frontDestinationActivities from "@app/temporal/relocation/activities/destination_region/front"; import * as connectorsSourceActivities from "@app/temporal/relocation/activities/source_region/connectors/sql"; import * as coreSourceActivities from "@app/temporal/relocation/activities/source_region/core"; -import * as coreDestinationActivities from "@app/temporal/relocation/activities/destination_region/core"; +import * as frontSourceActivities from "@app/temporal/relocation/activities/source_region/front"; import { RELOCATION_QUEUES_PER_REGION } from "@app/temporal/relocation/config"; import { getTemporalWorkerConnection } from "@app/temporal/relocation/temporal";