From 4df96964b2014956d8b0f1dcbc12a50be4ef925f Mon Sep 17 00:00:00 2001
From: bnechyporenko <bnechyporenko@bol.com>
Date: Tue, 4 Jun 2024 10:54:22 +0200
Subject: [PATCH] Storing the serialized workspace into google cloud bucket

Signed-off-by: bnechyporenko <bnechyporenko@bol.com>
---
 .../src/scaffolder/tasks/StorageTaskBroker.ts | 28 ++----
 .../src/scaffolder/tasks/WorkspaceService.ts  | 90 +++++++++++++++++++
 2 files changed, 96 insertions(+), 22 deletions(-)
 create mode 100644 plugins/scaffolder-backend/src/scaffolder/tasks/WorkspaceService.ts

diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
index 03a2833938..4177780e06 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
@@ -34,6 +34,7 @@ import {
   AuthService,
   BackstageCredentials,
 } from '@backstage/backend-plugin-api';
+import { DefaultWorkspaceService, WorkspaceService } from './WorkspaceService';
 
 type TaskState = {
   checkpoints: {
@@ -71,8 +72,8 @@ export class TaskManager implements TaskContext {
       storage,
       abortSignal,
       logger,
+      DefaultWorkspaceService.create(task, storage, config),
       auth,
-      config,
     );
     agent.startTimeout();
     return agent;
@@ -84,8 +85,8 @@ export class TaskManager implements TaskContext {
     private readonly storage: TaskStore,
     private readonly signal: AbortSignal,
     private readonly logger: Logger,
+    private readonly workspaceService: WorkspaceService,
     private readonly auth?: AuthService,
-    private readonly config?: Config,
   ) {}
 
   get spec() {
@@ -112,9 +113,7 @@ export class TaskManager implements TaskContext {
     taskId: string;
     targetPath: string;
   }): Promise<void> {
-    if (this.isWorkspaceSerializationEnabled()) {
-      this.storage.rehydrateWorkspace?.(options);
-    }
+    await this.workspaceService.rehydrateWorkspace(options);
   }
 
   get done() {
@@ -163,18 +162,11 @@ export class TaskManager implements TaskContext {
   }
 
   async serializeWorkspace?(options: { path: string }): Promise<void> {
-    if (this.isWorkspaceSerializationEnabled()) {
-      await this.storage.serializeWorkspace?.({
-        path: options.path,
-        taskId: this.task.taskId,
-      });
-    }
+    await this.workspaceService.serializeWorkspace(options);
   }
 
   async cleanWorkspace?(): Promise<void> {
-    if (this.isWorkspaceSerializationEnabled()) {
-      await this.storage.cleanWorkspace?.({ taskId: this.task.taskId });
-    }
+    await this.workspaceService.cleanWorkspace();
   }
 
   async complete(
@@ -211,14 +203,6 @@ export class TaskManager implements TaskContext {
     }, 1000);
   }
 
-  private isWorkspaceSerializationEnabled(): boolean {
-    return (
-      this.config?.getOptionalBoolean(
-        'scaffolder.EXPERIMENTAL_workspaceSerialization',
-      ) ?? false
-    );
-  }
-
   async getInitiatorCredentials(): Promise<BackstageCredentials> {
     const secrets = this.task.secrets as InternalTaskSecrets;
 
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/WorkspaceService.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/WorkspaceService.ts
new file mode 100644
index 0000000000..15de15859d
--- /dev/null
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/WorkspaceService.ts
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { TaskStore } from './types';
+
+import { CurrentClaimedTask } from './StorageTaskBroker';
+import { Config } from '@backstage/config';
+
+type WorkspaceSerializationProvider = 'database' | 'googleCloudBucket';
+
+export interface WorkspaceService {
+  serializeWorkspace(options: { path: string }): Promise<void>;
+
+  cleanWorkspace(): Promise<void>;
+
+  rehydrateWorkspace(options: {
+    taskId: string;
+    targetPath: string;
+  }): Promise<void>;
+}
+
+export class DefaultWorkspaceService implements WorkspaceService {
+  static create(task: CurrentClaimedTask, storage: TaskStore, config?: Config) {
+    return new DefaultWorkspaceService(task, storage, config);
+  }
+
+  private constructor(
+    private readonly task: CurrentClaimedTask,
+    private readonly storage: TaskStore,
+    private readonly config?: Config,
+  ) {}
+
+  public async serializeWorkspace(options: { path: string }): Promise<void> {
+    if (this.isWorkspaceSerializationEnabled()) {
+      const provider = this.getWorkspaceSerializationProvider();
+      switch (provider) {
+        case 'database':
+          this.storage.serializeWorkspace?.({
+            path: options.path,
+            taskId: this.task.taskId,
+          });
+          return;
+        case 'googleCloudBucket':
+          return;
+        default:
+          throw new Error(
+            `Workspace serialization provider ${provider} is not supported`,
+          );
+      }
+    }
+  }
+
+  public async rehydrateWorkspace(options: {
+    taskId: string;
+    targetPath: string;
+  }): Promise<void> {
+    if (this.isWorkspaceSerializationEnabled()) {
+      this.storage.rehydrateWorkspace?.(options);
+    }
+  }
+
+  public async cleanWorkspace(): Promise<void> {}
+
+  private isWorkspaceSerializationEnabled(): boolean {
+    return (
+      this.config?.getOptionalBoolean(
+        'scaffolder.EXPERIMENTAL_workspaceSerialization',
+      ) ?? false
+    );
+  }
+
+  private getWorkspaceSerializationProvider(): WorkspaceSerializationProvider {
+    return (this.config?.getOptionalString(
+      'scaffolder.EXPERIMENTAL_workspaceSerializationProvider',
+    ) ?? 'database') as WorkspaceSerializationProvider;
+  }
+}