Address co-pilot review comments

Author: stevebux

internal/datastore/src/__test__/letter-queue-repository.test.ts

@@ -174,21 +174,27 @@ describe("LetterQueueRepository", () => {
     });
 
     it("returns letters in timestamp order", async () => {
+      jest.useFakeTimers().setSystemTime(new Date());
       await letterQueueRepository.putLetter(
         createLetter({ letterId: "first-letter" }),
       );
+      jest.advanceTimersByTime(1);
       await letterQueueRepository.putLetter(
         createLetter({ letterId: "second-letter" }),
       );
+      jest.advanceTimersByTime(1);
       await letterQueueRepository.putLetter(
         createLetter({ letterId: "third-letter" }),
       );
+      jest.advanceTimersByTime(1);
       await letterQueueRepository.putLetter(
         createLetter({ letterId: "fourth-letter" }),
       );
+      jest.advanceTimersByTime(1);
       await letterQueueRepository.putLetter(
         createLetter({ letterId: "fifth-letter" }),
       );
+      jest.advanceTimersByTime(1);
 
       const letters = await letterQueueRepository.getLetters("supplier1", 5);
 
@@ -218,6 +224,80 @@ describe("LetterQueueRepository", () => {
       expect(letters).toHaveLength(3);
       expect(letters[2].letterId).toBe("third-letter");
     });
+
+    it("applies the limit after filtering on supplier", async () => {
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "first-letter" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "second-letter", supplierId: "supplier2" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "third-letter" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "fourth-letter" }),
+      );
+
+      const letters = await letterQueueRepository.getLetters("supplier1", 3);
+
+      expect(letters).toHaveLength(3);
+      expect(letters[2].letterId).toBe("fourth-letter");
+    });
+
+    it("applies the limit after filtering on visibilityTimestamp", async () => {
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "first-letter" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "second-letter" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "third-letter" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "fourth-letter" }),
+      );
+      await letterQueueRepository.updateVisibilityTimestamp(
+        createLetter({ letterId: "second-letter" }),
+        new Date(Date.now() + 600_000),
+      );
+
+      const letters = await letterQueueRepository.getLetters("supplier1", 3);
+
+      expect(letters).toHaveLength(3);
+      expect(letters[2].letterId).toBe("fourth-letter");
+    });
+
+    it("paginates through multiple DynamoDB pages to reach the limit", async () => {
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "first-letter" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "second-letter" }),
+      );
+      await letterQueueRepository.putLetter(
+        createLetter({ letterId: "third-letter" }),
+      );
+
+      const pagedRepository = new LetterQueueRepository(db.docClient, logger, {
+        ...db.config,
+        queryPageSize: 1,
+      });
+
+      const letters = await pagedRepository.getLetters("supplier1", 3);
+
+      expect(letters).toHaveLength(3);
+      expect(letters[0].letterId).toBe("first-letter");
+      expect(letters[1].letterId).toBe("second-letter");
+      expect(letters[2].letterId).toBe("third-letter");
+    });
+
+    it("returns an empty array if no items found", async () => {
+      const letters = await letterQueueRepository.getLetters("supplier1", 3);
+
+      expect(letters).toHaveLength(0);
+    });
   });
 
   describe("updateVisibilityTimestamp", () => {

internal/datastore/src/letter-queue-repository.ts

@@ -19,6 +19,8 @@ import { LetterDoesNotExistError } from "./letter-does-not-exist-error";
 type LetterQueueRepositoryConfig = {
   letterQueueTableName: string;
   letterQueueTtlHours: number;
+  /** Maximum number of items to fetch per DynamoDB page. Defaults to 100. */
+  queryPageSize?: number;
 };
 
 export default class LetterQueueRepository {
@@ -95,21 +97,34 @@ export default class LetterQueueRepository {
     supplierId: string,
     limit: number,
   ): Promise<PendingLetter[]> {
-    const result = await this.ddbClient.send(
-      new QueryCommand({
-        TableName: this.config.letterQueueTableName,
-        IndexName: "queueSortOrder-index",
-        KeyConditionExpression: "supplierId = :supplierId",
-        FilterExpression: "visibilityTimestamp < :now",
-        ExpressionAttributeValues: {
-          ":supplierId": supplierId,
-          ":now": new Date().toISOString(),
-        },
-        Limit: limit,
-      }),
-    );
+    const letters: PendingLetter[] = [];
+    let lastEvaluatedKey: Record<string, unknown> | undefined;
+
+    do {
+      const result = await this.ddbClient.send(
+        new QueryCommand({
+          TableName: this.config.letterQueueTableName,
+          IndexName: "queueSortOrder-index",
+          KeyConditionExpression: "supplierId = :supplierId",
+          FilterExpression: "visibilityTimestamp < :now",
+          ExpressionAttributeValues: {
+            ":supplierId": supplierId,
+            ":now": new Date().toISOString(),
+          },
+          // 1000 is a compromise - a smaller number might result in a lot of round trips, a larger one might
+          // entail fetching and then throwing away a lot of data
+          Limit: this.config.queryPageSize ?? 1000,
+          ExclusiveStartKey: lastEvaluatedKey,
+        }),
+      );
+
+      const page = z.array(PendingLetterSchema).parse(result.Items);
+      letters.push(...page);
+
+      lastEvaluatedKey = result.LastEvaluatedKey;
+    } while (lastEvaluatedKey !== undefined && letters.length < limit);
 
-    return z.array(PendingLetterSchema).parse(result.Items);
+    return letters.slice(0, limit);
   }
 
   async updateVisibilityTimestamp(

internal/datastore/src/types.ts

@@ -78,28 +78,28 @@ export type UpdateLetter = {
   reasonText?: string;
 };
 
-export const PendingLetterSchema = z.object({
+export const PendingLetterSchemaBase = z.object({
   supplierId: idRef(SupplierSchema, "id"),
   letterId: idRef(LetterSchema, "id"),
+  specificationId: z.string(),
+  groupId: z.string(),
+});
+
+export const PendingLetterSchema = PendingLetterSchemaBase.extend({
   queueTimestamp: z.string(),
   visibilityTimestamp: z.string(),
   queueSortOrderSk: z.string().describe("Secondary index SK"),
-  specificationId: z.string(),
-  groupId: z.string(),
   priority: z.int().min(0).max(99).optional(),
   ttl: z.int(),
 });
 
-export const PendingLetterSchemaBase = z.object({
-  supplierId: idRef(SupplierSchema, "id"),
-  letterId: idRef(LetterSchema, "id"),
-  specificationId: z.string(),
-  groupId: z.string(),
+export const InsertPendingLetter = PendingLetterSchemaBase.extend({
+  priority: z.int().min(0).max(99).optional(),
 });
 
 export type PendingLetter = z.infer<typeof PendingLetterSchema>;
 export type PendingLetterBase = z.infer<typeof PendingLetterSchemaBase>;
-export type InsertPendingLetter = PendingLetterBase & { priority?: number };
+export type InsertPendingLetter = z.infer<typeof InsertPendingLetter>;
 
 export const MISchemaBase = z.object({
   id: z.string(),

lambdas/api-handler/src/services/__tests__/letter-operations.test.ts

@@ -49,10 +49,6 @@ afterEach(async () => {
   jest.useRealTimers();
 });
 
-afterEach(async () => {
-  jest.useRealTimers();
-});
-
 describe("getPendingLetters", () => {
   beforeEach(() => {
     jest.clearAllMocks();