Add internal route to count the number of objects in a bucket (#48)

Author: AdiRishi

src/routes/internal.test.ts

@@ -113,4 +113,53 @@ describe('/internal Routes', () => {
       expect(list.objects.length).toBe(0);
     });
   });
+
+  describe('/internal/count-objects route', () => {
+    beforeEach(() => {
+      workerEnv = getMiniflareBindings();
+      ctx = new ExecutionContext();
+    });
+
+    afterEach(() => {
+      vi.restoreAllMocks();
+    });
+
+    test('should return the number of objects in R2 when bucket is empty', async () => {
+      const request = new Request('http://localhost/internal/count-objects', {
+        method: 'GET',
+        headers: {
+          'Content-Type': 'application/json',
+          Authorization: `Bearer ${workerEnv.TURBO_TOKEN}`,
+        },
+      });
+      const response = await app.fetch(request, workerEnv, ctx);
+      expect(response.status).toBe(200);
+      expect(await response.json()).toEqual({ count: 0 });
+    });
+
+    test('should return the number of objects in R2 when bucket is not empty', async () => {
+      await workerEnv.R2_STORE.put('key', 'value');
+      const request = new Request('http://localhost/internal/count-objects', {
+        method: 'GET',
+        headers: {
+          'Content-Type': 'application/json',
+          Authorization: `Bearer ${workerEnv.TURBO_TOKEN}`,
+        },
+      });
+      const response = await app.fetch(request, workerEnv, ctx);
+      expect(response.status).toBe(200);
+      expect(await response.json()).toEqual({ count: 1 });
+    });
+
+    test('should return 401 if no auth token is provided', async () => {
+      const request = new Request('http://localhost/internal/count-objects', {
+        method: 'GET',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+      });
+      const response = await app.fetch(request, workerEnv, ctx);
+      expect(response.status).toBe(401);
+    });
+  });
 });

src/routes/internal.ts

@@ -48,3 +48,18 @@ internalRouter.post(
     return c.json({ success: true });
   }
 );
+
+internalRouter.get('/count-objects', async (c) => {
+  let truncated = false;
+  let cursor: string | undefined;
+  let list: R2Objects;
+  let count = 0;
+  do {
+    list = await c.env.R2_STORE.list({ limit: 999, cursor });
+    truncated = list.truncated;
+    cursor = list.truncated ? list.cursor : undefined;
+    count += list.objects.length;
+  } while (truncated);
+
+  return c.json({ count });
+});