check / check (push) Successful in 27s
Each ML data request of up to 200 files is now tried on its own. A request that still fails after its retries is logged, its files are written with the reason in `mlDataError`, and the command exits 1 once the dump is complete. `fetchMLData`, used only here, is removed in favour of a per-batch loop over `fetchMLDataBatch`. Model: opus-5-5
71 lines
2.6 KiB
TypeScript
71 lines
2.6 KiB
TypeScript
// Fetch and decrypt Ente's per-file machine-learning data ("magic" search
|
|
// data: face detections + CLIP embeddings).
|
|
//
|
|
// The data lives behind `/files/data/fetch` with `type: "mldata"`. Each entry
|
|
// comes back encrypted under the file's own key and gzipped; decrypting and
|
|
// gunzipping yields the JSON payload
|
|
// `{ face: { faces: [...] }, clip: { embedding } }`. Ente caps a request at 200
|
|
// ids, so callers that want many at once split them into batches of
|
|
// `MLDATA_BATCH_SIZE` and call `fetchMLDataBatch` once per batch.
|
|
|
|
import { gunzipSync } from "node:zlib";
|
|
|
|
import type { ApiClient } from "./api/client.js";
|
|
import { decryptBlob, fromBase64 } from "./crypto/index.js";
|
|
|
|
// The most ids one `/files/data/fetch` request may carry.
|
|
export const MLDATA_BATCH_SIZE = 200;
|
|
|
|
// The decrypted, gunzipped per-file payload. Its concrete shape is Ente's; the
|
|
// store keeps the whole object verbatim and each consumer reads the fields it
|
|
// needs, so it stays an open record rather than a fixed interface.
|
|
export type MLData = Record<string, unknown>;
|
|
|
|
interface RawRemoteFileData {
|
|
fileID: number;
|
|
encryptedData: string;
|
|
decryptionHeader: string;
|
|
updatedAt?: number;
|
|
}
|
|
|
|
// Decrypt one entry with its file key and gunzip the JSON payload. Returns
|
|
// undefined when the key is unknown or the entry does not decrypt/parse, so one
|
|
// corrupt file never fails a whole batch.
|
|
const decodeEntry = (
|
|
entry: RawRemoteFileData,
|
|
key: Uint8Array | undefined,
|
|
): MLData | undefined => {
|
|
if (!key) return undefined;
|
|
try {
|
|
const decrypted = decryptBlob(
|
|
fromBase64(entry.encryptedData),
|
|
fromBase64(entry.decryptionHeader),
|
|
key,
|
|
);
|
|
const json = gunzipSync(Buffer.from(decrypted)).toString("utf-8");
|
|
return JSON.parse(json) as MLData;
|
|
} catch {
|
|
return undefined;
|
|
}
|
|
};
|
|
|
|
// Fetch ML data for up to `MLDATA_BATCH_SIZE` ids in a single request. This is
|
|
// the unit the request pools schedule; callers with more ids split them into
|
|
// batches and submit each batch to the pool.
|
|
export const fetchMLDataBatch = async (
|
|
api: ApiClient,
|
|
fileIDs: number[],
|
|
fileKeys: Map<number, Uint8Array>,
|
|
): Promise<Map<number, MLData>> => {
|
|
const { data } = await api.postJSON<{ data: RawRemoteFileData[] }>(
|
|
"/files/data/fetch",
|
|
{ type: "mldata", fileIDs },
|
|
);
|
|
const result = new Map<number, MLData>();
|
|
for (const entry of data ?? []) {
|
|
const payload = decodeEntry(entry, fileKeys.get(entry.fileID));
|
|
if (payload) result.set(entry.fileID, payload);
|
|
}
|
|
return result;
|
|
};
|