Skip to content

Commit 66abd8c

Browse files
marcusschiessermarcusschiesser
committed
Extract DuckDB store sync logic
1 parent 4381fa5 commit 66abd8c

3 files changed

Lines changed: 281 additions & 258 deletions

File tree

‎packages/duckdb/src/cache/index.ts‎

Lines changed: 9 additions & 258 deletions
Original file line numberDiff line numberDiff line change
@@ -1,54 +1,12 @@
1-
import {
2-
type BatchManifest,
3-
type IngestionEvent,
4-
type ObjectStore,
5-
otelResourceSpansToEvents,
6-
} from "@agentpond/core";
71
import { DuckDbDirectIngestion } from "../ingestion/direct-ingestion.js";
82
import { type DuckDbAccessMode, DuckDbOperations } from "./db-operations.js";
9-
import { BatchProjection, rawEventRow, stringValue } from "./projection.js";
10-
import type { SyncStateStore } from "./sync-state.js";
11-
import { createSyncStateStore } from "./sync-state.js";
12-
13-
export type SyncResult = {
14-
manifestsProcessed: number;
15-
objectsProcessed: number;
16-
eventsProcessed: number;
17-
};
18-
19-
export type SyncProgress = SyncResult & {
20-
manifestsTotal: number;
21-
manifestsSeen: number;
22-
manifestsSkipped: number;
23-
objectsSkipped: number;
24-
phase:
25-
| "listed"
26-
| "manifest-skipped"
27-
| "manifest-processed"
28-
| "object-skipped"
29-
| "object-processed"
30-
| "events-processed"
31-
| "complete";
32-
currentManifestKey?: string;
33-
currentObjectKey?: string;
34-
};
35-
36-
type PendingObject = {
37-
manifestKey: string | null;
38-
objectKey: string;
39-
progress?: Pick<SyncProgress, "currentManifestKey">;
40-
loadEvents: () => Promise<IngestionEvent[]>;
41-
entityIdForEvent: (event: IngestionEvent) => string;
42-
};
43-
44-
type ProcessedObject = {
45-
objectKey: string;
46-
manifestKey: string | null;
47-
};
48-
49-
// Keeps the 100k-trace perf fixture (~650k events) below the default V8 heap
50-
// limit by flushing raw JSON strings and projection maps about 33 times.
51-
const MAX_PENDING_EVENTS_PER_COMMIT = 20_000;
3+
import { DuckDbStoreSync } from "./store-sync.js";
4+
import type { SyncFromStoreParams, SyncResult } from "./sync-types.js";
5+
export type {
6+
SyncFromStoreParams,
7+
SyncProgress,
8+
SyncResult,
9+
} from "./sync-types.js";
5210

5311
export class AgentPondCache {
5412
private readonly db: DuckDbOperations;
@@ -65,180 +23,9 @@ export class AgentPondCache {
6523
await this.db.createSchema();
6624
}
6725

68-
async syncFromStore(params: {
69-
store: ObjectStore;
70-
projectId: string;
71-
prefix: string;
72-
onProgress?: (progress: SyncProgress) => void;
73-
}): Promise<SyncResult> {
26+
async syncFromStore(params: SyncFromStoreParams): Promise<SyncResult> {
7427
await this.init();
75-
const result: SyncResult = {
76-
manifestsProcessed: 0,
77-
objectsProcessed: 0,
78-
eventsProcessed: 0,
79-
};
80-
let manifestsSeen = 0;
81-
let manifestsSkipped = 0;
82-
let objectsSkipped = 0;
83-
let projection = new BatchProjection(this.db);
84-
let processedObjects: ProcessedObject[] = [];
85-
let processedManifestKeys: string[] = [];
86-
const syncState = createSyncStateStore({
87-
store: params.store,
88-
prefix: params.prefix,
89-
projectId: params.projectId,
90-
db: this.db,
91-
});
92-
const otelKeys = await syncState.listKeysForScanWindow("otel");
93-
const manifestKeys = await syncState.listKeysForScanWindow("manifests");
94-
const processedObjectKeys = await this.db.processedKeys(
95-
"processed_objects",
96-
otelKeys,
97-
);
98-
const processedManifestKeysSeen = await this.db.processedKeys(
99-
"processed_manifests",
100-
manifestKeys,
101-
);
102-
const emitProgress = (
103-
phase: SyncProgress["phase"],
104-
current?: Pick<SyncProgress, "currentManifestKey" | "currentObjectKey">,
105-
) => {
106-
params.onProgress?.({
107-
...result,
108-
manifestsTotal: manifestKeys.length,
109-
manifestsSeen,
110-
manifestsSkipped,
111-
objectsSkipped,
112-
phase,
113-
...current,
114-
});
115-
};
116-
const processObject = async (
117-
object: PendingObject,
118-
): Promise<"processed" | "skipped"> => {
119-
const progress = {
120-
...object.progress,
121-
currentObjectKey: object.objectKey,
122-
};
123-
if (processedObjectKeys.has(object.objectKey)) {
124-
emitProgress("object-skipped", progress);
125-
return "skipped";
126-
}
127-
128-
const events = await object.loadEvents();
129-
for (const event of events) {
130-
const row = rawEventRow({
131-
projectId: params.projectId,
132-
manifestKey: object.manifestKey,
133-
objectKey: object.objectKey,
134-
entityId: object.entityIdForEvent(event),
135-
event,
136-
});
137-
projection.addRawEvent(row);
138-
result.eventsProcessed += 1;
139-
if (result.eventsProcessed % 1000 === 0) {
140-
emitProgress("events-processed", progress);
141-
}
142-
}
143-
144-
processedObjects.push({
145-
objectKey: object.objectKey,
146-
manifestKey: object.manifestKey,
147-
});
148-
processedObjectKeys.add(object.objectKey);
149-
result.objectsProcessed += 1;
150-
emitProgress("object-processed", progress);
151-
return "processed";
152-
};
153-
const flushPending = async (): Promise<void> => {
154-
await this.commitSyncBatch({
155-
projectId: params.projectId,
156-
projection,
157-
processedObjects,
158-
processedManifestKeys,
159-
syncState,
160-
advanceSyncState: false,
161-
});
162-
projection = new BatchProjection(this.db);
163-
processedObjects = [];
164-
processedManifestKeys = [];
165-
};
166-
const flushFinal = async (): Promise<void> => {
167-
await this.commitSyncBatch({
168-
projectId: params.projectId,
169-
projection,
170-
processedObjects,
171-
processedManifestKeys,
172-
syncState,
173-
advanceSyncState: true,
174-
});
175-
};
176-
const flushIfNeeded = async (): Promise<void> => {
177-
if (projection.pendingEventCount < MAX_PENDING_EVENTS_PER_COMMIT) return;
178-
await flushPending();
179-
};
180-
emitProgress("listed");
181-
182-
const otelObjects: PendingObject[] = otelKeys.map((objectKey) => ({
183-
manifestKey: null,
184-
objectKey,
185-
loadEvents: async () =>
186-
otelResourceSpansToEvents(
187-
await params.store.getJson<unknown[]>(objectKey),
188-
),
189-
entityIdForEvent: (event) =>
190-
stringValue((event.body as Record<string, unknown>).id) ?? objectKey,
191-
}));
192-
for (const object of otelObjects) {
193-
const outcome = await processObject(object);
194-
if (outcome === "skipped") {
195-
objectsSkipped += 1;
196-
}
197-
await flushIfNeeded();
198-
}
199-
200-
for (const manifestKey of manifestKeys) {
201-
manifestsSeen += 1;
202-
if (processedManifestKeysSeen.has(manifestKey)) {
203-
manifestsSkipped += 1;
204-
emitProgress("manifest-skipped", {
205-
currentManifestKey: manifestKey,
206-
});
207-
continue;
208-
}
209-
const manifest = await params.store.getJson<BatchManifest>(manifestKey);
210-
const processedManifestObjectKeys = await this.db.processedKeys(
211-
"processed_objects",
212-
manifest.objects.map((object) => object.key),
213-
);
214-
for (const objectKey of processedManifestObjectKeys) {
215-
processedObjectKeys.add(objectKey);
216-
}
217-
for (const object of manifest.objects) {
218-
const outcome = await processObject({
219-
manifestKey,
220-
objectKey: object.key,
221-
progress: { currentManifestKey: manifestKey },
222-
loadEvents: async () =>
223-
params.store.getJson<IngestionEvent[]>(object.key),
224-
entityIdForEvent: () => object.entityId,
225-
});
226-
if (outcome === "skipped") {
227-
objectsSkipped += 1;
228-
}
229-
await flushIfNeeded();
230-
}
231-
processedManifestKeys.push(manifestKey);
232-
result.manifestsProcessed += 1;
233-
emitProgress("manifest-processed", {
234-
currentManifestKey: manifestKey,
235-
});
236-
await flushIfNeeded();
237-
}
238-
await flushFinal();
239-
emitProgress("complete");
240-
241-
return result;
28+
return new DuckDbStoreSync(this.db).syncFromStore(params);
24229
}
24330

24431
async query<T = Record<string, unknown>>(sql: string): Promise<T[]> {
@@ -253,40 +40,4 @@ export class AgentPondCache {
25340
async close(): Promise<void> {
25441
await this.db.close();
25542
}
256-
257-
private async commitSyncBatch(params: {
258-
projectId: string;
259-
projection: BatchProjection;
260-
processedObjects: ProcessedObject[];
261-
processedManifestKeys: string[];
262-
syncState: SyncStateStore;
263-
advanceSyncState: boolean;
264-
}): Promise<void> {
265-
if (
266-
params.processedObjects.length === 0 &&
267-
params.processedManifestKeys.length === 0 &&
268-
params.projection.pendingEventCount === 0
269-
) {
270-
if (params.advanceSyncState) {
271-
await params.syncState.advanceLastFinalized("otel");
272-
await params.syncState.advanceLastFinalized("manifests");
273-
}
274-
return;
275-
}
276-
277-
await this.db.exec("BEGIN TRANSACTION");
278-
try {
279-
await params.projection.commit(params.projectId);
280-
await this.db.insertProcessedObjects(params.processedObjects);
281-
await this.db.insertProcessedManifests(params.processedManifestKeys);
282-
if (params.advanceSyncState) {
283-
await params.syncState.advanceLastFinalized("otel");
284-
await params.syncState.advanceLastFinalized("manifests");
285-
}
286-
await this.db.exec("COMMIT");
287-
} catch (error) {
288-
await this.db.exec("ROLLBACK");
289-
throw error;
290-
}
291-
}
29243
}

0 commit comments

Comments
 (0)