Refactor agent runtime and task handling

Author: marcusschiesser

.changeset/smooth-sync-projection.md

@@ -0,0 +1,5 @@
+---
+"@agentpond/duckdb": patch
+---
+
+Speed up DuckDB sync by batching raw event writes and timestamp-ordered projection, and rename the cache class to AgentPondCache.

apps/cli/src/index.ts

@@ -9,7 +9,7 @@ import {
 	loadEnvFile,
 	S3ObjectStore,
 } from "@agentpond/core";
-import { AgentPondDuckDb } from "@agentpond/duckdb";
+import { AgentPondCache } from "@agentpond/duckdb";
 import { manualTraceResourceSpans } from "./otel-trace.js";
 import {
 	writeEventsAndSyncCache,
@@ -39,7 +39,7 @@ export async function main(argv = process.argv): Promise<void> {
 	try {
 		if (!resource || parsed.flags.help || parsed.flags.h) return printHelp();
 		if (resource === "sync") {
-			const db = new AgentPondDuckDb(config.dbPath);
+			const db = new AgentPondCache(config.dbPath);
 			const result = await db.syncFromStore({
 				store: new S3ObjectStore(config.s3),
 				projectId: config.projectId,
@@ -51,7 +51,7 @@ export async function main(argv = process.argv): Promise<void> {
 		if (resource === "sql") {
 			const query = rest.length > 0 ? [action, ...rest].join(" ") : action;
 			if (!query) throw new CliError("Missing SQL query");
-			const db = new AgentPondDuckDb(config.dbPath);
+			const db = new AgentPondCache(config.dbPath);
 			const rows = await db.query(query);
 			await db.close();
 			return print(rows, json);
@@ -63,7 +63,7 @@ export async function main(argv = process.argv): Promise<void> {
 			return createTrace(parsed, config, json);
 		}
 
-		const db = new AgentPondDuckDb(config.dbPath);
+		const db = new AgentPondCache(config.dbPath);
 		const rows = await runReadCommand(db, resource, action, rest, parsed);
 		await db.close();
 		return print(rows, json);
@@ -141,7 +141,7 @@ async function createTrace(
 }
 
 async function runReadCommand(
-	db: AgentPondDuckDb,
+	db: AgentPondCache,
 	resource: string,
 	action: string | undefined,
 	rest: string[],

apps/cli/src/sync-write.ts

@@ -5,7 +5,7 @@ import {
 	type IngestionEvent,
 	type ObjectStore,
 } from "@agentpond/core";
-import { AgentPondDuckDb } from "@agentpond/duckdb";
+import { AgentPondCache } from "@agentpond/duckdb";
 
 export async function writeOtelAndSyncCache(
 	config: Pick<AgentPondConfig, "dbPath" | "projectId" | "s3">,
@@ -18,7 +18,7 @@ export async function writeOtelAndSyncCache(
 		prefix: config.s3.prefix,
 	});
 	const object = await writer.writeOtelResourceSpans(resourceSpans);
-	const db = new AgentPondDuckDb(config.dbPath);
+	const db = new AgentPondCache(config.dbPath);
 	try {
 		await db.syncFromStore({
 			store,
@@ -42,7 +42,7 @@ export async function writeEventsAndSyncCache(
 		prefix: config.s3.prefix,
 	});
 	const manifest = await writer.writeAcceptedEvents(events);
-	const db = new AgentPondDuckDb(config.dbPath);
+	const db = new AgentPondCache(config.dbPath);
 	try {
 		await db.syncFromStore({
 			store,

packages/duckdb/src/cache/db-operations.ts

@@ -0,0 +1,79 @@
+import { mkdirSync } from "node:fs";
+import { dirname } from "node:path";
+import { type DuckDBConnection, DuckDBInstance } from "@duckdb/node-api";
+
+type ProcessedTable = "processed_manifests" | "processed_objects";
+
+export class DuckDbOperations {
+	private instance?: DuckDBInstance;
+	private connection?: DuckDBConnection;
+
+	constructor(readonly dbPath: string) {
+		mkdirSync(dirname(dbPath), { recursive: true });
+	}
+
+	async exec(sqlText: string): Promise<void> {
+		await (await this.getConnection()).run(sqlText);
+	}
+
+	async all<T = Record<string, unknown>>(sqlText: string): Promise<T[]> {
+		const reader = await (await this.getConnection()).runAndReadAll(sqlText);
+		return reader.getRowObjectsJS() as T[];
+	}
+
+	sql(value: unknown): string {
+		return sql(value);
+	}
+
+	async processedKeyExists(
+		table: ProcessedTable,
+		key: string,
+	): Promise<boolean> {
+		const rows = await this.all(
+			`SELECT key FROM ${table} WHERE key = ${this.sql(key)} LIMIT 1`,
+		);
+		return rows.length > 0;
+	}
+
+	async insertProcessedKey(
+		table: ProcessedTable,
+		key: string,
+		manifestKey?: string,
+	): Promise<void> {
+		if (table === "processed_objects") {
+			await this.exec(
+				`INSERT INTO processed_objects (key, manifest_key) VALUES (${this.sql(key)}, ${this.sql(manifestKey)})`,
+			);
+			return;
+		}
+		await this.exec(
+			`INSERT INTO processed_manifests (key) VALUES (${this.sql(key)})`,
+		);
+	}
+
+	async close(): Promise<void> {
+		if (this.connection) {
+			const connection = this.connection;
+			this.connection = undefined;
+			connection.closeSync();
+		}
+		if (this.instance) {
+			const instance = this.instance;
+			this.instance = undefined;
+			instance.closeSync();
+		}
+	}
+
+	private async getConnection(): Promise<DuckDBConnection> {
+		if (!this.connection) {
+			this.instance = await DuckDBInstance.create(this.dbPath);
+			this.connection = await this.instance.connect();
+		}
+		return this.connection;
+	}
+}
+
+export function sql(value: unknown): string {
+	if (value === null || value === undefined) return "NULL";
+	return `'${String(value).replaceAll("'", "''")}'`;
+}