src/r2.ts

Functions

Imports

It imports src/types.ts. Imported by src/dump.ts, src/runs.ts.

Source

1import type { R2Like, R2UploadedPartLike } from "./types.js";
2
3// R2 requires every part except the last to be the same size, and at least 5 MiB.
4const MULTIPART_PART_BYTES = 8 * 1024 * 1024;
5// R2's bulk delete refuses more than 1000 keys in one call rather than truncating.
6const R2_DELETE_MAX = 1000;
7
8export const JSON_TYPE = { httpMetadata: { contentType: "application/json" } };
9
10// Writes `head`, then the rows a page at a time, then "]}", as one multipart object.
11// The bytes equal JSON.stringify of the whole object, holding at most one part and one
12// page in memory.
13export async function putJsonStreamed(bucket: R2Like, key: string, head: string, pages: AsyncIterable<unknown[]>): Promise<number> {
14  const upload = await bucket.createMultipartUpload(key, JSON_TYPE);
15  const parts: R2UploadedPartLike[] = [];
16  const encoder = new TextEncoder();
17  let part = new Uint8Array(MULTIPART_PART_BYTES);
18  let filled = 0;
19  let rows = 0;
20  const flush = async () => {
21    parts.push(await upload.uploadPart(parts.length + 1, part.slice(0, filled)));
22    part = new Uint8Array(MULTIPART_PART_BYTES);
23    filled = 0;
24  };
25  const write = async (text: string) => {
26    let bytes = encoder.encode(text);
27    while (bytes.length > 0) {
28      const n = Math.min(bytes.length, MULTIPART_PART_BYTES - filled);
29      part.set(bytes.subarray(0, n), filled);
30      filled += n;
31      bytes = bytes.subarray(n);
32      if (filled === MULTIPART_PART_BYTES) await flush();
33    }
34  };
35  try {
36    await write(head);
37    for await (const page of pages) {
38      for (const row of page) {
39        await write((rows === 0 ? "" : ",") + JSON.stringify(row));
40        rows++;
41      }
42    }
43    await write("]}");
44    if (filled > 0 || parts.length === 0) await flush();
45    await upload.complete(parts);
46    return rows;
47  } catch (err) {
48    // The abort's own failure is reported with the cause, never in place of it.
49    await upload.abort().catch((abortErr: unknown) => {
50      throw new AggregateError([err, abortErr], `d1-dump: writing ${key} failed, and so did aborting its upload`);
51    });
52    throw err;
53  }
54}
55
56export async function listAllKeys(bucket: R2Like, prefix: string): Promise<string[]> {
57  const keys: string[] = [];
58  let cursor: string | undefined;
59  do {
60    const page = await bucket.list({ prefix, cursor });
61    for (const obj of page.objects) keys.push(obj.key);
62    cursor = page.truncated ? page.cursor : undefined;
63  } while (cursor);
64  return keys;
65}
66
67export async function deleteInChunks(bucket: R2Like, keys: string[]): Promise<number> {
68  for (let i = 0; i < keys.length; i += R2_DELETE_MAX) await bucket.delete(keys.slice(i, i + R2_DELETE_MAX));
69  return keys.length;
70}