Add support for pod annotation service account. (#10415)

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);

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);
 

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,14 +11,17 @@ 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
 ): Promise<string> {
   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,15 +35,19 @@ export async function writeToRelocationStorage(
 export async function readFromRelocationStorage<T = unknown>(
   dataPath: string
 ): Promise<T> {
-  const relocationBucket = getBucketInstance(config.getGcsRelocationBucket());
+  const relocationBucket = getBucketInstance(config.getGcsRelocationBucket(), {
+    useServiceAccount: isDevelopment(),
+  });
 
   const content = await relocationBucket.fetchFileContent(dataPath);
 
   return JSON.parse(content) as T;
 }
 
 export async function deleteFromRelocationStorage(dataPath: string) {
-  const relocationBucket = getBucketInstance(config.getGcsRelocationBucket());
+  const relocationBucket = getBucketInstance(config.getGcsRelocationBucket(), {
+    useServiceAccount: isDevelopment(),
+  });
 
   await relocationBucket.delete(dataPath, { ignoreNotFound: true });
 }

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,

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";