Merge pull request #22773 from acierto/task-action-idempotency

Introducing checkpoints for scaffolder task action idempotency

.changeset/sixty-queens-mix.md

@@ -0,0 +1,6 @@
+---
+'@backstage/plugin-scaffolder-backend': minor
+'@backstage/plugin-scaffolder-node': patch
+---
+
+Introducing checkpoints for scaffolder task action idempotency

packages/backend-next/package.json

@@ -57,6 +57,7 @@
     "@backstage/plugin-playlist-backend": "workspace:^",
     "@backstage/plugin-proxy-backend": "workspace:^",
     "@backstage/plugin-scaffolder-backend": "workspace:^",
+    "@backstage/plugin-scaffolder-backend-module-github": "workspace:^",
     "@backstage/plugin-search-backend": "workspace:^",
     "@backstage/plugin-search-backend-module-catalog": "workspace:^",
     "@backstage/plugin-search-backend-module-explore": "workspace:^",

packages/backend-next/src/index.ts

@@ -45,6 +45,7 @@ backend.add(
 backend.add(import('@backstage/plugin-permission-backend/alpha'));
 backend.add(import('@backstage/plugin-proxy-backend/alpha'));
 backend.add(import('@backstage/plugin-scaffolder-backend/alpha'));
+backend.add(import('@backstage/plugin-scaffolder-backend-module-github'));
 backend.add(import('@backstage/plugin-search-backend-module-catalog/alpha'));
 backend.add(import('@backstage/plugin-search-backend-module-explore/alpha'));
 backend.add(import('@backstage/plugin-search-backend-module-techdocs/alpha'));

plugins/scaffolder-backend/api-report.md

@@ -21,6 +21,7 @@ import * as gitlab from '@backstage/plugin-scaffolder-backend-module-gitlab';
 import { HumanDuration } from '@backstage/types';
 import { IdentityApi } from '@backstage/plugin-auth-node';
 import { JsonObject } from '@backstage/types';
+import { JsonValue } from '@backstage/types';
 import { Knex } from 'knex';
 import { LifecycleService } from '@backstage/backend-plugin-api';
 import { Logger } from 'winston';
@@ -359,6 +360,7 @@ export interface CurrentClaimedTask {
   createdBy?: string;
   secrets?: TaskSecrets_2;
   spec: TaskSpec;
+  state?: JsonObject;
   taskId: string;
 }
 
@@ -397,6 +399,13 @@ export class DatabaseTaskStore implements TaskStore {
   // (undocumented)
   getTask(taskId: string): Promise<SerializedTask_2>;
   // (undocumented)
+  getTaskState({ taskId }: { taskId: string }): Promise<
+    | {
+        state: JsonObject;
+      }
+    | undefined
+  >;
+  // (undocumented)
   heartbeatTask(taskId: string): Promise<void>;
   // (undocumented)
   list(options: { createdBy?: string }): Promise<{
@@ -418,6 +427,8 @@ export class DatabaseTaskStore implements TaskStore {
     ids: string[];
   }>;
   // (undocumented)
+  saveTaskState(options: { taskId: string; state?: JsonObject }): Promise<void>;
+  // (undocumented)
   shutdownTask(options: TaskStoreShutDownTaskOptions): Promise<void>;
 }
 
@@ -519,11 +530,32 @@ export class TaskManager implements TaskContext_2 {
   // (undocumented)
   emitLog(message: string, logMetadata?: JsonObject): Promise<void>;
   // (undocumented)
+  getTaskState?(): Promise<
+    | {
+        state?: JsonObject;
+      }
+    | undefined
+  >;
+  // (undocumented)
   getWorkspaceName(): Promise<string>;
   // (undocumented)
   get secrets(): TaskSecrets_2 | undefined;
   // (undocumented)
   get spec(): TaskSpecV1beta3;
+  // (undocumented)
+  updateCheckpoint?(
+    options:
+      | {
+          key: string;
+          status: 'success';
+          value: JsonValue;
+        }
+      | {
+          key: string;
+          status: 'failed';
+          reason: string;
+        },
+  ): Promise<void>;
 }
 
 // @public @deprecated (undocumented)
@@ -553,6 +585,13 @@ export interface TaskStore {
   // (undocumented)
   getTask(taskId: string): Promise<SerializedTask>;
   // (undocumented)
+  getTaskState?({ taskId }: { taskId: string }): Promise<
+    | {
+        state: JsonObject;
+      }
+    | undefined
+  >;
+  // (undocumented)
   heartbeatTask(taskId: string): Promise<void>;
   // (undocumented)
   list?(options: { createdBy?: string }): Promise<{
@@ -573,6 +612,11 @@ export interface TaskStore {
     ids: string[];
   }>;
   // (undocumented)
+  saveTaskState?(options: {
+    taskId: string;
+    state?: JsonObject;
+  }): Promise<void>;
+  // (undocumented)
   shutdownTask?(options: TaskStoreShutDownTaskOptions): Promise<void>;
 }
 

plugins/scaffolder-backend/migrations/20240203232000_state.js

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2020 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.
+ */
+
+// @ts-check
+
+/**
+ * @param {import('knex').Knex} knex
+ */
+exports.up = async function up(knex) {
+  await knex.schema.alterTable('tasks', table => {
+    table.text('state').nullable().comment('A state of the checkpoints');
+  });
+};
+
+/**
+ * @param {import('knex').Knex} knex
+ */
+exports.down = async function down(knex) {
+  await knex.schema.alterTable('tasks', table => {
+    table.dropColumn('state');
+  });
+};

plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.test.ts

@@ -188,4 +188,37 @@ describe('DatabaseTaskStore', () => {
       await store.shutdownTask({ taskId });
     }).rejects.toThrow(ConflictError);
   });
+
+  it('should store checkpoints and retrieve task state', async () => {
+    const { store } = await createStore();
+    const { taskId } = await store.createTask({
+      spec: {} as TaskSpec,
+      createdBy: 'me',
+    });
+
+    await store.saveTaskState({
+      taskId,
+      state: {
+        checkpoints: {
+          'repo.create': {
+            status: 'success',
+            value: { repoUrl: 'https://github.com/backstage/backstage.git' },
+          },
+        },
+      },
+    });
+
+    const state = await store.getTaskState({ taskId });
+
+    expect(state).toStrictEqual({
+      state: {
+        checkpoints: {
+          'repo.create': {
+            status: 'success',
+            value: { repoUrl: 'https://github.com/backstage/backstage.git' },
+          },
+        },
+      },
+    });
+  });
 });

plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts

@@ -52,6 +52,7 @@ export type RawDbTaskRow = {
   id: string;
   spec: string;
   status: TaskStatus;
+  state?: string;
   last_heartbeat_at?: string;
   created_at: string;
   created_by: string | null;
@@ -394,6 +395,32 @@ export class DatabaseTaskStore implements TaskStore {
     });
   }
 
+  async getTaskState({ taskId }: { taskId: string }): Promise<
+    | {
+        state: JsonObject;
+      }
+    | undefined
+  > {
+    const [result] = await this.db<RawDbTaskRow>('tasks')
+      .where({ id: taskId })
+      .select('state');
+    return result.state ? JSON.parse(result.state) : undefined;
+  }
+
+  async saveTaskState(options: {
+    taskId: string;
+    state?: JsonObject;
+  }): Promise<void> {
+    if (options.state) {
+      const serializedState = JSON.stringify({ state: options.state });
+      await this.db<RawDbTaskRow>('tasks')
+        .where({ id: options.taskId })
+        .update({
+          state: serializedState,
+        });
+    }
+  }
+
   async listEvents(
     options: TaskStoreListEventsOptions,
   ): Promise<{ events: SerializedTaskEvent[] }> {

plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts

@@ -18,6 +18,7 @@ import { getVoidLogger } from '@backstage/backend-common';
 import { NunjucksWorkflowRunner } from './NunjucksWorkflowRunner';
 import { TemplateActionRegistry } from '../actions';
 import { ScmIntegrations } from '@backstage/integration';
+import { JsonObject } from '@backstage/types';
 import { ConfigReader } from '@backstage/config';
 import { TaskSpec } from '@backstage/plugin-scaffolder-common';
 import {
@@ -133,6 +134,26 @@ describe('NunjucksWorkflowRunner', () => {
       },
     });
 
+    actionRegistry.register({
+      id: 'checkpoints-action',
+      description: 'Mock action with checkpoints',
+      handler: async ctx => {
+        const key1 = await ctx.checkpoint?.('key1', async () => {
+          return 'updated';
+        });
+        const key2 = await ctx.checkpoint?.('key2', async () => {
+          return 'updated';
+        });
+        const key3 = await ctx.checkpoint?.('key3', async () => {
+          return 'updated';
+        });
+
+        ctx.output('key1', key1);
+        ctx.output('key2', key2);
+        ctx.output('key3', key3);
+      },
+    });
+
     mockedPermissionApi.authorizeConditional.mockResolvedValue([
       { result: AuthorizeResult.ALLOW },
     ]);
@@ -538,6 +559,54 @@ describe('NunjucksWorkflowRunner', () => {
       );
     });
 
+    it('should deal with checkpoints', async () => {
+      const task = {
+        ...createMockTaskWithSpec({
+          apiVersion: 'scaffolder.backstage.io/v1beta3',
+          parameters: {},
+          steps: [
+            {
+              id: 'test',
+              name: 'name',
+              action: 'checkpoints-action',
+              input: { foo: 1 },
+            },
+          ],
+          output: {
+            key1: '${{steps.test.output.key1}}',
+            key2: '${{steps.test.output.key2}}',
+            key3: '${{steps.test.output.key3}}',
+          },
+        }),
+        getTaskState: (): Promise<
+          | {
+              state: JsonObject;
+            }
+          | undefined
+        > => {
+          return Promise.resolve({
+            state: {
+              checkpoints: {
+                ['v1.task.checkpoint.key1']: {
+                  status: 'success',
+                  value: 'initial',
+                },
+                ['v1.task.checkpoint.key2']: {
+                  status: 'failed',
+                  reason: 'fatal error',
+                },
+              },
+            },
+          });
+        },
+      };
+      const result = await runner.execute(task);
+
+      expect(result.output.key1).toEqual('initial');
+      expect(result.output.key2).toEqual('updated');
+      expect(result.output.key3).toEqual('updated');
+    });
+
     it('should template the output from simple actions', async () => {
       const task = createMockTaskWithSpec({
         apiVersion: 'scaffolder.backstage.io/v1beta3',

plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts

@@ -21,7 +21,7 @@ import fs from 'fs-extra';
 import path from 'path';
 import nunjucks from 'nunjucks';
 import { JsonArray, JsonObject, JsonValue } from '@backstage/types';
-import { InputError, NotAllowedError } from '@backstage/errors';
+import { InputError, NotAllowedError, stringifyError } from '@backstage/errors';
 import { PassThrough } from 'stream';
 import { generateExampleOutput, isTruthy } from './helper';
 import { validate as validateJsonSchema } from 'jsonschema';
@@ -79,6 +79,16 @@ type TemplateContext = {
   each?: JsonValue;
 };
 
+type CheckpointState =
+  | {
+      status: 'failed';
+      reason: string;
+    }
+  | {
+      status: 'success';
+      value: JsonValue;
+    };
+
 const isValidTaskSpec = (taskSpec: TaskSpec): taskSpec is TaskSpecV1beta3 => {
   return taskSpec.apiVersion === 'scaffolder.backstage.io/v1beta3';
 };
@@ -330,6 +340,7 @@ export class NunjucksWorkflowRunner implements WorkflowRunner {
       }
       const tmpDirs = new Array<string>();
       const stepOutput: { [outputName: string]: JsonValue } = {};
+      const prevTaskState = await task.getTaskState?.();
 
       for (const iteration of iterations) {
         if (iteration.each) {
@@ -347,6 +358,43 @@ export class NunjucksWorkflowRunner implements WorkflowRunner {
           logger: taskLogger,
           logStream: streamLogger,
           workspacePath,
+          async checkpoint<U extends JsonValue>(
+            keySuffix: string,
+            fn: () => Promise<U>,
+          ) {
+            const key = `v1.task.checkpoint.${keySuffix}`;
+            try {
+              let prevValue: U | undefined;
+              if (prevTaskState) {
+                const prevState = (
+                  prevTaskState.state?.checkpoints as {
+                    [key: string]: CheckpointState;
+                  }
+                )?.[key];
+                if (prevState && prevState.status === 'success') {
+                  prevValue = prevState.value as U;
+                }
+              }
+
+              const value = prevValue ? prevValue : await fn();
+
+              if (!prevValue) {
+                task.updateCheckpoint?.({
+                  key,
+                  status: 'success',
+                  value,
+                });
+              }
+              return value;
+            } catch (err) {
+              task.updateCheckpoint?.({
+                key,
+                status: 'failed',
+                reason: stringifyError(err),
+              });
+              throw err;
+            }
+          },
           createTemporaryDirectory: async () => {
             const tmpDir = await fs.mkdtemp(
               `${workspacePath}_step-${step.id}-`,

plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.test.ts

@@ -242,4 +242,31 @@ describe('StorageTaskBroker', () => {
     const promise = broker.list({ createdBy: 'user:default/foo' });
     await expect(promise).resolves.toEqual({ tasks: [task] });
   });
+
+  it('should handle checkpoints in task state', async () => {
+    const broker = new StorageTaskBroker(storage, logger);
+
+    await broker.dispatch({
+      spec: { steps: [] } as unknown as TaskSpec,
+      createdBy: 'user:default/foo',
+    });
+
+    const taskA = await broker.claim();
+    await taskA.updateCheckpoint?.({
+      key: 'repo.create',
+      status: 'success',
+      value: 'https://github.com/backstage/backstage.git',
+    });
+
+    expect(await taskA.getTaskState?.()).toEqual({
+      state: {
+        checkpoints: {
+          'repo.create': {
+            status: 'success',
+            value: 'https://github.com/backstage/backstage.git',
+          },
+        },
+      },
+    });
+  });
 });

plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts

@@ -16,6 +16,9 @@
 
 import { Config } from '@backstage/config';
 import { TaskSpec } from '@backstage/plugin-scaffolder-common';
+import { JsonObject, JsonValue, Observable } from '@backstage/types';
+import { Logger } from 'winston';
+import ObservableImpl from 'zen-observable';
 import {
   TaskSecrets,
   SerializedTask,
@@ -25,12 +28,22 @@ import {
   TaskCompletionState,
   TaskContext,
 } from '@backstage/plugin-scaffolder-node';
-import { JsonObject, Observable } from '@backstage/types';
-import { Logger } from 'winston';
-import ObservableImpl from 'zen-observable';
 import { TaskStore } from './types';
 import { readDuration } from './helper';
 
+type TaskState = {
+  checkpoints: {
+    [key: string]:
+      | {
+          status: 'failed';
+          reason: string;
+        }
+      | {
+          status: 'success';
+          value: JsonValue;
+        };
+  };
+};
 /**
  * TaskManager
  *
@@ -91,6 +104,40 @@ export class TaskManager implements TaskContext {
     });
   }
 
+  async getTaskState?(): Promise<
+    | {
+        state?: JsonObject;
+      }
+    | undefined
+  > {
+    return this.storage.getTaskState?.({ taskId: this.task.taskId });
+  }
+
+  async updateCheckpoint?(
+    options:
+      | {
+          key: string;
+          status: 'success';
+          value: JsonValue;
+        }
+      | {
+          key: string;
+          status: 'failed';
+          reason: string;
+        },
+  ): Promise<void> {
+    const { key, ...value } = options;
+    if (this.task.state) {
+      (this.task.state as TaskState).checkpoints[key] = value;
+    } else {
+      this.task.state = { checkpoints: { [key]: value } };
+    }
+    await this.storage.saveTaskState?.({
+      taskId: this.task.taskId,
+      state: this.task.state,
+    });
+  }
+
   async complete(
     result: TaskCompletionState,
     metadata?: JsonObject,
@@ -144,6 +191,10 @@ export interface CurrentClaimedTask {
    * The secrets that are stored with the task.
    */
   secrets?: TaskSecrets;
+  /**
+   * The state of checkpoints of the task.
+   */
+  state?: JsonObject;
   /**
    * The creator of the task.
    */

plugins/scaffolder-backend/src/scaffolder/tasks/types.ts

@@ -194,6 +194,18 @@ export interface TaskStore {
 
   emitLogEvent(options: TaskStoreEmitOptions): Promise<void>;
 
+  getTaskState?({ taskId }: { taskId: string }): Promise<
+    | {
+        state: JsonObject;
+      }
+    | undefined
+  >;
+
+  saveTaskState?(options: {
+    taskId: string;
+    state?: JsonObject;
+  }): Promise<void>;
+
   listEvents(
     options: TaskStoreListEventsOptions,
   ): Promise<{ events: SerializedTaskEvent[] }>;

plugins/scaffolder-node/api-report.md

@@ -30,6 +30,10 @@ export type ActionContext<
   secrets?: TaskSecrets;
   workspacePath: string;
   input: TActionInput;
+  checkpoint?<U extends JsonValue>(
+    key: string,
+    fn: () => Promise<U>,
+  ): Promise<U>;
   output(
     name: keyof TActionOutput,
     value: TActionOutput[keyof TActionOutput],
@@ -345,6 +349,13 @@ export interface TaskContext {
   // (undocumented)
   emitLog(message: string, logMetadata?: JsonObject): Promise<void>;
   // (undocumented)
+  getTaskState?(): Promise<
+    | {
+        state?: JsonObject;
+      }
+    | undefined
+  >;
+  // (undocumented)
   getWorkspaceName(): Promise<string>;
   // (undocumented)
   isDryRun?: boolean;
@@ -352,6 +363,20 @@ export interface TaskContext {
   secrets?: TaskSecrets;
   // (undocumented)
   spec: TaskSpec;
+  // (undocumented)
+  updateCheckpoint?(
+    options:
+      | {
+          key: string;
+          status: 'success';
+          value: JsonValue;
+        }
+      | {
+          key: string;
+          status: 'failed';
+          reason: string;
+        },
+  ): Promise<void>;
 }
 
 // @public

plugins/scaffolder-node/src/actions/types.ts

@@ -16,7 +16,7 @@
 
 import { Logger } from 'winston';
 import { Writable } from 'stream';
-import { JsonObject } from '@backstage/types';
+import { JsonObject, JsonValue } from '@backstage/types';
 import { TaskSecrets } from '../tasks';
 import { TemplateInfo } from '@backstage/plugin-scaffolder-common';
 import { UserEntity } from '@backstage/catalog-model';
@@ -35,6 +35,10 @@ export type ActionContext<
   secrets?: TaskSecrets;
   workspacePath: string;
   input: TActionInput;
+  checkpoint?<U extends JsonValue>(
+    key: string,
+    fn: () => Promise<U>,
+  ): Promise<U>;
   output(
     name: keyof TActionOutput,
     value: TActionOutput[keyof TActionOutput],

plugins/scaffolder-node/src/tasks/types.ts

@@ -15,7 +15,7 @@
  */
 
 import { TaskSpec } from '@backstage/plugin-scaffolder-common';
-import { JsonObject, Observable } from '@backstage/types';
+import { JsonObject, JsonValue, Observable } from '@backstage/types';
 
 /**
  * TaskSecrets
@@ -118,6 +118,27 @@ export interface TaskContext {
 
   emitLog(message: string, logMetadata?: JsonObject): Promise<void>;
 
+  getTaskState?(): Promise<
+    | {
+        state?: JsonObject;
+      }
+    | undefined
+  >;
+
+  updateCheckpoint?(
+    options:
+      | {
+          key: string;
+          status: 'success';
+          value: JsonValue;
+        }
+      | {
+          key: string;
+          status: 'failed';
+          reason: string;
+        },
+  ): Promise<void>;
+
   getWorkspaceName(): Promise<string>;
 }
 

yarn.lock

@@ -27474,6 +27474,7 @@ __metadata:
     "@backstage/plugin-playlist-backend": "workspace:^"
     "@backstage/plugin-proxy-backend": "workspace:^"
     "@backstage/plugin-scaffolder-backend": "workspace:^"
+    "@backstage/plugin-scaffolder-backend-module-github": "workspace:^"
     "@backstage/plugin-search-backend": "workspace:^"
     "@backstage/plugin-search-backend-module-catalog": "workspace:^"
     "@backstage/plugin-search-backend-module-explore": "workspace:^"