@squiz/job-runner-lib 2.3.0 → 2.4.0

package/src/domains/job/repository/MultiTenantJobExecutionRepository/MultiTenantJobExecutionRepository.ts

@@ -0,0 +1,42 @@
+import { AttributeValue, DynamoDBClient, QueryCommand, QueryCommandOutput } from '@aws-sdk/client-dynamodb';
+import { JobExecutionStatus } from '../../constants/JobExecutionStatus.constants';
+import { JobExecutionDto, JobExecutionsByTenant } from '../../model/JobExecution';
+import { parsePartitionKey } from '../../ParsePartitionKey';
+import { convertDynamoDbItemToJobExecution } from '../JobExecutionRepository/utils';
+import { ENTITY_TYPES } from '../../constants/Repository.constants';
+
+export class MultiTenantJobExecutionRepository {
+  public constructor(
+    private readonly dynamoDbClient: DynamoDBClient,
+    private readonly tableName: string,
+  ) {}
+
+  public async findByStatus(status: JobExecutionStatus): Promise<JobExecutionsByTenant> {
+    const result: JobExecutionsByTenant = {};
+    let startKey: Record<string, AttributeValue> | undefined = undefined;
+
+    do {
+      const res: QueryCommandOutput = await this.dynamoDbClient.send(
+        new QueryCommand({
+          KeyConditionExpression: 'gsi1pk = :gsi1pk',
+          ExpressionAttributeValues: {
+            ':gsi1pk': { S: `${ENTITY_TYPES.jobExecution}~${status}` },
+          },
+          TableName: this.tableName,
+          IndexName: 'gsi1',
+          ExclusiveStartKey: startKey,
+        }),
+      );
+
+      res.Items?.forEach((item) => {
+        const { tenant } = parsePartitionKey(item.pk.S || '');
+        result[tenant] ||= [];
+        result[tenant].push(convertDynamoDbItemToJobExecution(item as JobExecutionDto));
+      });
+
+      startKey = res.LastEvaluatedKey;
+    } while (startKey);
+
+    return result;
+  }
+}

package/src/domains/job/repository/index.ts

@@ -2,4 +2,5 @@ export * from './JobContextRepository/JobContextRepository';
 export * from './JobExecutionRepository/JobExecutionRepository';
 export * from './JobManifestRepository/JobManifestRepository';
 export * from './JobRepository/JobRepository';
+export * from './MultiTenantJobExecutionRepository/MultiTenantJobExecutionRepository';
 export * from './AbstractRepository';

package/src/domains/job/services/JobContextService/JobContextService.spec.ts

@@ -6,7 +6,6 @@ import {
   AttributeValue,
   DeleteItemCommandOutput,
   DynamoDBClient,
-  DynamoDBClientConfig,
   GetItemCommandOutput,
   PutItemCommandOutput,
   ScanCommandOutput,
@@ -15,7 +14,7 @@ import {
 import * as PaginationUtils from '../../core/pagination/utils/PaginationUtils';
 import { faker } from '@faker-js/faker/locale/en';
 import { JobContextService } from './JobContextService';
-import { PAGINATION_DIRECTION_TYPE_AFTER, PaginationService } from '../../core/pagination';
+import { PAGINATION_DIRECTION_TYPE_AFTER } from '../../core/pagination';
 import { JobContextRepository } from '../../repository';
 import {
   createMockCreateJobContextRequest,
@@ -31,41 +30,45 @@ import {
   createMockClientPutResponse,
   createMockClientScanResponse,
 } from '../../core/dynamoClient.mock';
-import { PARTITION_KEYS } from '../../constants/Repository.constants';
+import { ENTITY_TYPES } from '../../constants/Repository.constants';
+import { initIocContainer, iocContainer } from '../../../../mocks/mockIoc';
+import { testLogger } from '../../../../jest-global-setup';
 
 // Mock DynamoDB Client
-const mockOptions = {} as DynamoDBClientConfig;
-const mockClient = new DynamoDBClient(mockOptions);
+const mockSend = jest.spyOn(DynamoDBClient.prototype, 'send');
 
-const mockSend = jest.fn();
-mockClient.send = mockSend;
+initIocContainer(testLogger);
 
-let jobContextService: JobContextService;
-let paginationService: PaginationService<any>;
 describe('JobContextService', (): void => {
+  let jobContextService: JobContextService;
+
   beforeAll((): void => {
-    const jobContextRepository = new JobContextRepository(mockClient);
-    paginationService = new PaginationService();
-    jobContextService = new JobContextService(jobContextRepository, paginationService);
+    jobContextService = iocContainer.get(JobContextService);
   });
 
   describe(`createJobContext`, (): void => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should return a response for createJobContext`, async (): Promise<void> => {
-      const newJobContextToCreate = createMockCreateJobContextRequest();
+      const newJobContextToCreate = createMockCreateJobContextRequest({
+        environment: {
+          secret1: faker.string.uuid(),
+          secret2: faker.string.uuid(),
+        },
+      });
       const environmentMap: Record<string, AttributeValue> = {};
       for (const key of Object.keys(newJobContextToCreate.environment)) {
         environmentMap[key] = { S: newJobContextToCreate.environment[key] };
       }
       const mockResponse = createMockClientPutResponse({
         Attributes: {
-          pk: { S: PARTITION_KEYS.context },
+          pk: { S: ENTITY_TYPES.context },
           sk: { S: newJobContextToCreate.contextName },
           name: { S: newJobContextToCreate.contextName },
-          version: { M: environmentMap },
+          version: {
+            M: {
+              secret1: { S: newJobContextToCreate.environment.secret1 },
+              secret2: { S: newJobContextToCreate.environment.secret2 },
+            },
+          },
         },
       });
 
@@ -75,16 +78,27 @@ describe('JobContextService', (): void => {
       expect(mockSend).toBeCalledTimes(1);
       expect(response).toStrictEqual({
         ...newJobContextToCreate,
-        type: PARTITION_KEYS.context,
+        type: ENTITY_TYPES.context,
+      });
+      expect(testLogger.audit).toHaveBeenCalledWith(`Created job context ${newJobContextToCreate.contextName}`, {
+        context: expect.any(String),
+        processable: true,
+        remote: true,
+        details: {
+          context: {
+            contextName: newJobContextToCreate.contextName,
+            description: newJobContextToCreate.description,
+            environment: {
+              secret1: '*redacted*',
+              secret2: '*redacted*',
+            },
+          },
+        },
       });
     });
   });
 
   describe(`deleteJobContext`, (): void => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should return an empty response for job context deletion`, async (): Promise<void> => {
       mockSend.mockImplementationOnce((): DeleteItemCommandOutput => {
         return {
@@ -92,17 +106,32 @@ describe('JobContextService', (): void => {
         };
       });
 
-      const response = await jobContextService.deleteJobContext(`name`);
+      const context = createMockJobContext({
+        environment: {
+          secret1: faker.string.uuid(),
+        },
+      });
+      const response = await jobContextService.deleteJobContext(context);
       expect(response).toBe(undefined);
       expect(mockSend).toBeCalledTimes(1);
+      expect(testLogger.audit).toHaveBeenCalledWith(`Deleted job context ${context.contextName}`, {
+        context: expect.any(String),
+        processable: true,
+        remote: true,
+        details: {
+          context: {
+            contextName: context.contextName,
+            description: context.description,
+            environment: {
+              secret1: '*redacted*',
+            },
+          },
+        },
+      });
     });
   });
 
   describe(`findJobContext`, (): void => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should return job context by name`, async (): Promise<void> => {
       const mockJobContextRecord = createMockJobContextDto({
         environment: { key: 'value' },
@@ -140,10 +169,6 @@ describe('JobContextService', (): void => {
   });
 
   describe(`listAllJobContexts`, (): void => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should return an empty list`, async (): Promise<void> => {
       mockSend.mockImplementationOnce((): ScanCommandOutput => createMockClientScanResponse());
 
@@ -202,8 +227,6 @@ describe('JobContextService', (): void => {
     let getContextsSpy: any;
 
     beforeEach(() => {
-      jest.resetAllMocks();
-
       mockJobContexts = createMockJobContexts();
 
       mockPaginationInfo = {
@@ -235,10 +258,6 @@ describe('JobContextService', (): void => {
   });
 
   describe(`updateJobContext`, (): void => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should update a job context record`, async (): Promise<void> => {
       const jobContextToUpdate = createMockJobContext();
       const updatedJobContext: JobContext = {
@@ -267,6 +286,21 @@ describe('JobContextService', (): void => {
       );
       expect(response).toStrictEqual(updatedJobContext);
       expect(mockSend).toBeCalledTimes(1);
+      expect(testLogger.audit).toHaveBeenCalledWith(`Updated job context ${updatedJobContext.contextName}`, {
+        context: expect.any(String),
+        processable: true,
+        remote: true,
+        details: {
+          context: {
+            contextName: updatedJobContext.contextName,
+            description: updatedJobContext.description,
+            environment: {
+              name: '*redacted*',
+              age: '*redacted*',
+            },
+          },
+        },
+      });
     });
   });
 });

package/src/domains/job/services/JobContextService/JobContextService.ts

@@ -4,6 +4,7 @@
  */
 
 import { inject } from 'inversify';
+import { produce } from 'immer';
 import { JobContextRepository } from '../../repository';
 import { provideSingleton } from '../../singletonProvider';
 import {
@@ -16,20 +17,27 @@ import {
   PaginationService,
 } from '../../core/pagination';
 import { CreateJobContextRequest, JobContext, JobContextDto, UpdateJobContextRequest } from '../../model/JobContext';
+import { CreateLog, Logger } from '@squiz/optimization-logger';
+import { InjectTokens } from '../../constants/InjectTokens';
 
 @provideSingleton(JobContextService)
 export class JobContextService {
   constructor(
     @inject(JobContextRepository) protected jobContextRepository: JobContextRepository,
     @inject(PaginationService) protected paginationService: PaginationService<JobContext>,
+    @inject(InjectTokens.Logger) protected logger: Logger,
   ) {}
 
   public async createJobContext(input: CreateJobContextRequest): Promise<JobContext> {
-    return this.jobContextRepository.create(input);
+    const context = await this.jobContextRepository.create(input);
+    this.logContextAction(`Created job context ${context.contextName}`, context);
+    return context;
   }
 
-  public async updateJobContext(name: string, input: UpdateJobContextRequest): Promise<JobContext | null> {
-    return this.jobContextRepository.update(name, input);
+  public async updateJobContext(contextName: string, input: UpdateJobContextRequest): Promise<JobContext | null> {
+    const context = await this.jobContextRepository.update({ contextName }, input);
+    this.logContextAction(`Updated job context ${context.contextName}`, context);
+    return context;
   }
 
   public async listAllJobContexts(): Promise<Array<JobContext>> {
@@ -68,10 +76,33 @@ export class JobContextService {
   }
 
   public async findJobContext(contextName: string): Promise<JobContext | null> {
-    return this.jobContextRepository.findOne(contextName);
+    return this.jobContextRepository.findOne({ contextName });
   }
 
-  public async deleteJobContext(contextName: string): Promise<void> {
-    return await this.jobContextRepository.delete(contextName);
+  public async deleteJobContext(context: JobContext): Promise<void> {
+    await this.jobContextRepository.delete({ contextName: context.contextName });
+    this.logContextAction(`Deleted job context ${context.contextName}`, context);
+  }
+
+  private logContextAction(message: string, context: JobContext) {
+    const { type: _, ...maskedContext } = produce(context, (draft) => {
+      for (const key in draft.environment) {
+        draft.environment[key] = '*redacted*';
+      }
+
+      return draft;
+    });
+
+    const log = CreateLog.create()
+      .attachMetadata({
+        details: {
+          context: maskedContext,
+        },
+      })
+      .processable()
+      .remote()
+      .create();
+
+    this.logger.audit(...log(message));
   }
 }

package/src/domains/job/services/JobExecutionTaskService/JobExecutionTaskService.defs.ts

@@ -3,9 +3,7 @@ import { JobExecution } from '../../model/JobExecution';
 /**
  * The required parameters for retrieving the number of running ECS tasks
  */
-export interface ListRunningTasksParams {
-  clusterArn: string;
-  runtimeTaskFamily: string;
+export interface HasReachedMaxTenantConcurrencyParams {
   maxTenantConcurrentTasks: number;
 }
 

package/src/domains/job/services/JobExecutionTaskService/JobExecutionTaskService.spec.ts

@@ -1,22 +1,25 @@
 import 'aws-sdk-client-mock-jest';
 import { AwsClientStub, mockClient } from 'aws-sdk-client-mock';
-import { ECSClient, ListTasksCommand, RunTaskCommand } from '@aws-sdk/client-ecs';
+import { ConflictException, ResourceNotFoundException, ECSClient, RunTaskCommand } from '@aws-sdk/client-ecs';
 import { faker } from '@faker-js/faker';
 import { JobExecutionTaskService } from './JobExecutionTaskService';
 import { JobManifestService } from '../JobManifestService/JobManifestService';
 import { JobContextService } from '../JobContextService/JobContextService';
-import { JobDetailsRepository, S3Repository } from '../../S3Repository';
+import { JobDetailsRepository, JobUploadsRepository } from '../../S3Repository';
 import { createMockJobExecution } from '../../model/JobExecution';
 import { createMockJobContext } from '../../model/JobContext';
 import { createMockJobManifest } from '../../manifest';
 import { RunJobTaskParams } from './JobExecutionTaskService.defs';
+import { JobExecutionRepository } from '../../repository';
+import { JOB_EXECUTION_STATUS } from '../../constants/JobExecutionStatus.constants';
 
 describe('JobExecutionTaskService', () => {
   let mockEcsClient: AwsClientStub<ECSClient>;
   let mockJobManifestService: jest.Mocked<JobManifestService>;
   let mockJobContextService: jest.Mocked<JobContextService>;
+  let mockJobExecutionRepository: jest.Mocked<JobExecutionRepository>;
   let mockJobDetailsRepository: jest.Mocked<JobDetailsRepository>;
-  let mockJobUploadRepository: jest.Mocked<S3Repository>;
+  let mockJobUploadRepository: jest.Mocked<JobUploadsRepository>;
   let jobExecutionTaskService: JobExecutionTaskService;
 
   const jobDetailsSignedUrl = faker.internet.url();
@@ -33,8 +36,6 @@ describe('JobExecutionTaskService', () => {
   };
 
   beforeEach(() => {
-    jest.resetAllMocks();
-
     mockEcsClient = mockClient(ECSClient);
     mockJobManifestService = {
       findJobManifest: jest.fn(),
@@ -42,59 +43,60 @@ describe('JobExecutionTaskService', () => {
     mockJobContextService = {
       findJobContext: jest.fn(),
     } as unknown as jest.Mocked<JobContextService>;
+    mockJobExecutionRepository = {
+      list: jest.fn(),
+    } as unknown as jest.Mocked<JobExecutionRepository>;
     mockJobDetailsRepository = {
       createJobDetailsFile: jest.fn(),
       getS3SignedUrl: jest.fn().mockResolvedValue(jobDetailsSignedUrl),
     } as unknown as jest.Mocked<JobDetailsRepository>;
     mockJobUploadRepository = {
       getS3SignedUrl: jest.fn().mockResolvedValue(jobCodeArtifactSignedUrl),
-    } as unknown as jest.Mocked<S3Repository>;
+    } as unknown as jest.Mocked<JobUploadsRepository>;
 
     jobExecutionTaskService = new JobExecutionTaskService(
       new ECSClient({}),
       mockJobManifestService,
       mockJobContextService,
+      mockJobExecutionRepository,
       mockJobDetailsRepository,
       mockJobUploadRepository,
     );
   });
 
-  describe(`getTenantRunningTasks`, () => {
-    it('should 0 if no tasks were found', async (): Promise<void> => {
-      mockEcsClient.on(ListTasksCommand).resolves({});
-      const response = await (jobExecutionTaskService as any).getTenantRunningTasks({
-        clusterArn: `clusterArn`,
-        runtimeTaskFamily: 'some_family',
-        maxTenantConcurrentTasks: 3,
-      });
-
-      expect(response).toBe(0);
-    });
-  });
-
   describe(`hasReachedMaxTenantConcurrency`, () => {
     it('should return false if number of running jobs is less than max', async (): Promise<void> => {
-      mockEcsClient.on(ListTasksCommand).resolves({ taskArns: ['12345', '67890'] });
+      const maxTenantConcurrentTasks = 3;
+
+      mockJobExecutionRepository.list.mockResolvedValue([createMockJobExecution(), createMockJobExecution()]);
 
       const response = await jobExecutionTaskService.hasReachedMaxTenantConcurrency({
-        clusterArn: `clusterArn`,
-        runtimeTaskFamily: 'some_family',
-        maxTenantConcurrentTasks: 3,
+        maxTenantConcurrentTasks,
       });
 
       expect(response).toBe(false);
+      expect(mockJobExecutionRepository.list).toHaveBeenCalledWith({
+        limit: maxTenantConcurrentTasks,
+        consistentRead: true,
+        filters: { status: JOB_EXECUTION_STATUS.running },
+      });
     });
 
     it('should return true if number of running jobs is more than max', async (): Promise<void> => {
-      mockEcsClient.on(ListTasksCommand).resolves({ taskArns: ['12345', '67890'] });
+      const maxTenantConcurrentTasks = 1;
+
+      mockJobExecutionRepository.list.mockResolvedValue([createMockJobExecution(), createMockJobExecution()]);
 
       const response = await jobExecutionTaskService.hasReachedMaxTenantConcurrency({
-        clusterArn: `clusterArn`,
-        runtimeTaskFamily: 'some_family',
-        maxTenantConcurrentTasks: 1,
+        maxTenantConcurrentTasks,
       });
 
       expect(response).toBe(true);
+      expect(mockJobExecutionRepository.list).toHaveBeenCalledWith({
+        limit: maxTenantConcurrentTasks,
+        consistentRead: true,
+        filters: { status: JOB_EXECUTION_STATUS.running },
+      });
     });
   });
 
@@ -105,7 +107,7 @@ describe('JobExecutionTaskService', () => {
       const jobContext = createMockJobContext();
       const jobManifest = createMockJobManifest();
       const jobExecution = createMockJobExecution({
-        id: `${new Date().toISOString()}-jobExecutionId`,
+        id: `20250917T005949395Z-9602417782`,
         jobName: 'jobName',
         version: 'jobVersion',
       });
@@ -175,15 +177,17 @@ describe('JobExecutionTaskService', () => {
           { key: 'cmp:tenant', value: 'some_tenant' },
         ],
         taskDefinition: defaultRunJobTaskParams.taskDefinitionArn,
+        // sha1 of execution job name/id + tenant + retryCount.
+        clientToken: '4725755e08138edde031a177c36a269ecdbed5a3',
       });
       expect(mockJobDetailsRepository.createJobDetailsFile).toHaveBeenCalledWith(
-        `${defaultRunJobTaskParams.tenantId}/${jobExecution.id}`,
+        defaultRunJobTaskParams.tenantId,
         jobContext,
         jobManifest,
         jobExecution,
       );
       expect(mockJobUploadRepository.getS3SignedUrl).toHaveBeenCalledWith(
-        `${jobExecution.jobName}/${jobExecution.version}/export.zip`,
+        `${defaultRunJobTaskParams.tenantId}/${jobExecution.jobName}/${jobExecution.version}/export.zip`,
       );
     });
 
@@ -201,11 +205,89 @@ describe('JobExecutionTaskService', () => {
       await jobExecutionTaskService.runJobExecutionTask(defaultRunJobTaskParams);
 
       expect(mockJobDetailsRepository.createJobDetailsFile).toHaveBeenCalledWith(
-        `${defaultRunJobTaskParams.tenantId}/${defaultRunJobTaskParams.jobExecution.id}`,
+        defaultRunJobTaskParams.tenantId,
         null,
         null,
         defaultRunJobTaskParams.jobExecution,
       );
     });
+
+    it('should handle ConflictException with a previously started task arn', async () => {
+      const mockArn = `mock-arn`;
+      const jobContext = createMockJobContext();
+      const jobManifest = createMockJobManifest();
+      const jobExecution = createMockJobExecution({
+        id: `20250917T005949395Z-9602417782`,
+        jobName: 'jobName',
+        version: 'jobVersion',
+      });
+
+      mockJobContextService.findJobContext.mockResolvedValue(jobContext);
+      mockJobManifestService.findJobManifest.mockResolvedValue(jobManifest);
+      mockEcsClient.on(RunTaskCommand).rejects(
+        new ConflictException({
+          message: 'The RunTask request could not be processed due to conflicts.',
+          resourceIds: [mockArn],
+          $metadata: {},
+        }),
+      );
+
+      const response = await jobExecutionTaskService.runJobExecutionTask({
+        ...defaultRunJobTaskParams,
+        jobExecution: jobExecution,
+      });
+
+      expect(response).toBe(mockArn);
+    });
+
+    it('should handle ConflictException with no previously started task arn', async () => {
+      const jobContext = createMockJobContext();
+      const jobManifest = createMockJobManifest();
+      const jobExecution = createMockJobExecution({
+        id: `20250917T005949395Z-9602417782`,
+        jobName: 'jobName',
+        version: 'jobVersion',
+      });
+
+      mockJobContextService.findJobContext.mockResolvedValue(jobContext);
+      mockJobManifestService.findJobManifest.mockResolvedValue(jobManifest);
+      mockEcsClient.on(RunTaskCommand).rejects(
+        new ConflictException({
+          message: 'The RunTask request could not be processed due to conflicts.',
+          resourceIds: [],
+          $metadata: {},
+        }),
+      );
+
+      const response = await jobExecutionTaskService.runJobExecutionTask({
+        ...defaultRunJobTaskParams,
+        jobExecution: jobExecution,
+      });
+
+      // A ConflictException may be thrown without returning the ARN of the task.
+      expect(response).toBe(undefined);
+    });
+
+    it('should re-throw unhandled errors when starting a task', async () => {
+      const expectedError = new ResourceNotFoundException({ message: faker.lorem.words(), $metadata: {} });
+      const jobContext = createMockJobContext();
+      const jobManifest = createMockJobManifest();
+      const jobExecution = createMockJobExecution({
+        id: `20250917T005949395Z-9602417782`,
+        jobName: 'jobName',
+        version: 'jobVersion',
+      });
+
+      mockJobContextService.findJobContext.mockResolvedValue(jobContext);
+      mockJobManifestService.findJobManifest.mockResolvedValue(jobManifest);
+      mockEcsClient.on(RunTaskCommand).rejects(expectedError);
+
+      await expect(
+        jobExecutionTaskService.runJobExecutionTask({
+          ...defaultRunJobTaskParams,
+          jobExecution: jobExecution,
+        }),
+      ).rejects.toThrow(expectedError);
+    });
   });
 });