Spaces:
Running
Running
| 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; | |
| })); | |
| } |