Spaces:
Running
Running
File size: 2,452 Bytes
e5d9748 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import {
AsyncDuckDB,
ConsoleLogger,
DuckDBDataProtocol,
getJsDelivrBundles,
selectBundle,
} from '@duckdb/duckdb-wasm';
export type QueryMonthFile = {
month: string;
path: string;
};
type QueryEngine = {
database: AsyncDuckDB;
connection: Awaited<ReturnType<AsyncDuckDB['connect']>>;
worker: Worker;
registeredFiles: Map<string, string>;
};
let queryEnginePromise: Promise<QueryEngine> | null = null;
async function createQueryEngine(): Promise<QueryEngine> {
const bundle = await selectBundle(getJsDelivrBundles());
const workerUrl = URL.createObjectURL(
new Blob([`importScripts("${bundle.mainWorker}");`], {
type: 'text/javascript',
}),
);
const worker = new Worker(workerUrl);
const database = new AsyncDuckDB(new ConsoleLogger(), worker);
try {
await database.instantiate(bundle.mainModule, bundle.pthreadWorker);
URL.revokeObjectURL(workerUrl);
const connection = await database.connect();
const engine = {
database,
connection,
worker,
registeredFiles: new Map<string, string>(),
};
window.addEventListener('pagehide', () => {
void connection.close().catch(() => undefined);
void database.terminate().catch(() => undefined);
worker.terminate();
}, { once: true });
return engine;
} catch (error) {
URL.revokeObjectURL(workerUrl);
worker.terminate();
await database.terminate().catch(() => undefined);
throw error;
}
}
export function getQueryEngine(): Promise<QueryEngine> {
if (!queryEnginePromise) {
queryEnginePromise = createQueryEngine().catch((error: unknown) => {
queryEnginePromise = null;
throw error;
});
}
return queryEnginePromise;
}
export async function registerMonthFiles(
engine: QueryEngine,
months: QueryMonthFile[],
): Promise<string[]> {
return Promise.all(months.map(async (month) => {
const filename = `${month.month}.slim.parquet`;
const url = new URL(month.path, window.location.href).toString();
const registeredUrl = engine.registeredFiles.get(filename);
if (registeredUrl === url) return filename;
if (registeredUrl) {
throw new Error(`A different snapshot file is already registered as ${filename}.`);
}
await engine.database.registerFileURL(
filename,
url,
DuckDBDataProtocol.HTTP,
false,
);
engine.registeredFiles.set(filename, url);
return filename;
}));
} |