Line data Source code
1 : import * as arrow from 'apache-arrow';
2 : import {
3 : HyperDB as HyperDBEngine,
4 : type HyperDBConnection as HyperDBEngineConnection,
5 : type MountedResource,
6 : type QueryInfo,
7 : } from 'hyperdb-wasm';
8 :
9 : import type {
10 : EmbeddedComputeDatabase,
11 : EmbeddedConnectionOptions,
12 : EmbeddedDatabaseAlias,
13 : EmbeddedExternalDirectoryDatabase,
14 : EmbeddedExternalFileDatabase,
15 : EmbeddedResultStream,
16 : ExternalDirectoryMode,
17 : EmbeddedTableImportConnection,
18 : EmbeddedTableInsertOptions,
19 : } from '../database/embedded_database.js';
20 : import { prepareHyperArrowInsert } from './hyperdb_arrow_insert.js';
21 :
22 : export { HyperDBEngine };
23 : export type { HyperDBSettings } from 'hyperdb-wasm';
24 : /// The client of the engine worker. Tests share one between test frames, every frame wraps it in its own HyperDBEngine.
25 : export type HyperDBClient = ConstructorParameters<typeof HyperDBEngine>[0];
26 : /// A response of the engine client
27 : export type HyperDBResult = Awaited<ReturnType<HyperDBClient['connect']>>;
28 :
29 : /// The result of statements without result, Hyper sends no Arrow IPC for them.
30 1 : const EMPTY_RESULT_IPC = arrow.tableToIPC(arrow.tableFromArrays({}), 'stream');
31 :
32 : export interface HyperDBOptions {
33 : /// The database of the compute tables, created on startup and dropped on termination
34 : computeDatabasePath?: string;
35 : /// The database attached as "hyper", created on startup and dropped on termination
36 : defaultDatabasePath?: string;
37 : /// The database of the snippet search, attached as "__dashql_snippets", created on startup and dropped on termination
38 : snippetsDatabasePath?: string;
39 : /// Terminate the engine with the database. Tests share one engine client between databases.
40 : terminateEngine?: boolean;
41 : }
42 :
43 : /// The rows a statement affected.
44 : /// Hyper reports a row count for every statement, but it only means affected rows for these kinds.
45 1 : function readAffectedRowCount(info: QueryInfo | undefined): bigint | null {
46 1 : switch (info?.statementKind) {
47 1 : case 'COPY':
48 1 : case 'DELETE':
49 1 : case 'FETCH':
50 1 : case 'INSERT':
51 1 : case 'MERGE':
52 1 : case 'UPDATE':
53 1 : return info.rowCount;
54 1 : default:
55 1 : return null;
56 1 : }
57 1 : }
58 :
59 1 : async function* prependChunk(first: Uint8Array, rest: AsyncGenerator<Uint8Array, unknown>): AsyncGenerator<Uint8Array> {
60 1 : yield first;
61 1 : yield* rest;
62 1 : }
63 :
64 1 : async function* singleChunk(chunk: Uint8Array): AsyncGenerator<Uint8Array> {
65 1 : yield chunk;
66 1 : }
67 :
68 0 : function yieldToEventLoop(): Promise<void> {
69 0 : return new Promise(resolve => setTimeout(resolve, 0));
70 0 : }
71 :
72 : // Pooled connections stay attached to their default database when closed and are reused.
73 : // ATTACH is expensive since Hyper loads the statistics of every table in the database, and the compute code
74 : // opens a connection per statement. Connections are only pooled when the caller opts in.
75 : // A session keeps the client context it was connected with, the sessions are pooled per client context.
76 1 : const MAX_IDLE_SESSIONS = 4;
77 :
78 : /// The key of the idle sessions with a default database and a client context
79 1 : function poolKey(alias: string, clientContext: string | null): string {
80 1 : return clientContext == null ? alias : `${alias}:${clientContext}`;
81 1 : }
82 :
83 1 : export class HyperDB implements EmbeddedComputeDatabase, EmbeddedExternalFileDatabase, EmbeddedExternalDirectoryDatabase {
84 1 : private readonly sessions = new Set<HyperDBSession>();
85 : /// The idle pooled sessions by default database alias and client context, see poolKey()
86 1 : private readonly idleSessions = new Map<string, HyperDBSession[]>();
87 : /// The sessions of connect calls that are still attaching their database
88 1 : private readonly openingSessions = new Set<Promise<HyperDBSession>>();
89 1 : private initialization: Promise<void> | null = null;
90 1 : private termination: Promise<void> | null = null;
91 1 : private terminated = false;
92 1 : private databaseOperation: Promise<void> = Promise.resolve();
93 1 : private readonly terminateEngine: boolean;
94 1 : private readonly computeDatabasePath: string | null;
95 1 : private readonly defaultDatabasePath: string | null;
96 1 : private readonly snippetsDatabasePath: string | null;
97 1 : private readonly externalFiles = new Map<string, MountedResource>();
98 1 : private readonly externalDirectories = new Map<string, MountedResource>();
99 :
100 1 : constructor(
101 1 : private readonly engine: HyperDBEngine,
102 1 : options: HyperDBOptions = {},
103 1 : ) {
104 1 : this.computeDatabasePath = options.computeDatabasePath ?? null;
105 1 : this.defaultDatabasePath = options.defaultDatabasePath ?? null;
106 1 : this.snippetsDatabasePath = options.snippetsDatabasePath ?? null;
107 1 : this.terminateEngine = options.terminateEngine ?? true;
108 1 : }
109 :
110 1 : static async create(engine: HyperDBEngine, options?: HyperDBOptions): Promise<HyperDB> {
111 1 : const database = new HyperDB(engine, options);
112 1 : await database.initialize();
113 1 : return database;
114 1 : }
115 :
116 1 : async connect(options: EmbeddedConnectionOptions = {}): Promise<HyperDBConnection> {
117 1 : await this.initialize();
118 1 : if (this.termination || this.terminated) {
119 1 : throw new Error('database is terminated');
120 : }
121 1 : const alias = options.defaultDatabase ?? '__dashql_compute';
122 1 : const clientContext = options.clientContext == null ? null : JSON.stringify(options.clientContext);
123 1 : let pool: HyperDBSession[] | null = null;
124 1 : if (options.pooled) {
125 1 : const key = poolKey(alias, clientContext);
126 1 : pool = this.idleSessions.get(key) ?? [];
127 1 : this.idleSessions.set(key, pool);
128 1 : }
129 1 : const release = (session: HyperDBSession) => this.releaseSession(session, pool);
130 :
131 1 : const idle = pool?.pop();
132 1 : if (idle != null) {
133 1 : return new HyperDBConnection(idle, release);
134 : }
135 :
136 1 : const opening = this.openSession(this.databasePath(alias), alias, clientContext);
137 1 : this.openingSessions.add(opening);
138 1 : try {
139 1 : return new HyperDBConnection(await opening, release);
140 1 : } finally {
141 1 : this.openingSessions.delete(opening);
142 1 : }
143 1 : }
144 :
145 : /// The path of the database that is attached under an alias
146 1 : private databasePath(alias: EmbeddedDatabaseAlias): string | null {
147 1 : switch (alias) {
148 1 : case '__dashql_compute':
149 1 : return this.computeDatabasePath;
150 1 : case 'hyper':
151 1 : return this.defaultDatabasePath;
152 1 : case '__dashql_snippets':
153 1 : return this.snippetsDatabasePath;
154 1 : }
155 1 : }
156 :
157 : /// The paths of the databases that are created on startup and dropped on termination
158 1 : private databasePaths(): string[] {
159 1 : return [this.computeDatabasePath, this.defaultDatabasePath, this.snippetsDatabasePath].filter(path => path != null);
160 1 : }
161 :
162 : /// Open an engine connection with the database attached under the alias.
163 : /// A termination waits for the opening sessions, they close themselves if it started while they were attaching.
164 1 : private async openSession(databasePath: string | null, alias: string, clientContext: string | null): Promise<HyperDBSession> {
165 : // Hyper logs the client context as ctx.client of the query events of the session
166 1 : const connection = await this.engine.connect(clientContext == null
167 1 : ? undefined
168 1 : : { settings: { external_client_context: clientContext } });
169 1 : if (this.termination || this.terminated) {
170 1 : await connection.close();
171 1 : throw new Error('database is terminated');
172 1 : }
173 :
174 1 : const session = new HyperDBSession(connection, this.mountFile, () => {
175 1 : this.sessions.delete(session);
176 1 : }, databasePath == null ? undefined : alias);
177 1 : try {
178 1 : if (databasePath != null) {
179 1 : await session.execute(`ATTACH DATABASE ${quoteIdentifier(databasePath)} AS ${quoteIdentifier(alias)}`);
180 : }
181 1 : if (this.termination || this.terminated) {
182 1 : throw new Error('database is terminated');
183 : }
184 1 : } catch (error) {
185 1 : await session.close();
186 1 : throw error;
187 1 : }
188 1 : this.sessions.add(session);
189 1 : return session;
190 1 : }
191 :
192 : /// Mount a file for a statement of a session
193 1 : private readonly mountFile = async (file: File, path: string): Promise<MountedResource> =>
194 1 : await this.engine.mount(file, { path });
195 :
196 1 : private async releaseSession(session: HyperDBSession, pool: HyperDBSession[] | null): Promise<void> {
197 1 : if (pool != null && session.reusable && !this.termination && !this.terminated
198 1 : && pool.length < MAX_IDLE_SESSIONS) {
199 1 : pool.push(session);
200 1 : return;
201 1 : }
202 1 : await session.close();
203 1 : }
204 :
205 1 : async getVersion(): Promise<string> {
206 1 : const connection = await this.connect();
207 1 : try {
208 1 : const result = await connection.query('SELECT version() AS version');
209 1 : const version = result.getChild('version')?.get(0);
210 1 : if (typeof version !== 'string') {
211 0 : throw new Error('version query returned an invalid result');
212 : }
213 1 : return version;
214 1 : } finally {
215 1 : await connection.close();
216 1 : }
217 1 : }
218 :
219 1 : async registerExternalFile(name: string, blob: Blob, path?: string, signal?: AbortSignal): Promise<string> {
220 1 : return await this.runDatabaseOperation(async () => {
221 1 : throwIfAborted(signal);
222 1 : path ??= externalMountPath(name);
223 1 : validateExternalMountPath(path);
224 : // The engine wants a file name without separators, the last path component is one.
225 1 : const file = new File([blob], path.slice(path.lastIndexOf('/') + 1), { type: blob.type });
226 1 : const resource = await this.engine.mount(file, { path });
227 1 : if (signal?.aborted) {
228 1 : try {
229 1 : await resource.close();
230 1 : } catch {
231 : // Termination retries the removal.
232 0 : this.externalFiles.set(path, resource);
233 0 : }
234 1 : throw signal.reason ?? new DOMException('The operation was aborted', 'AbortError');
235 1 : }
236 1 : this.externalFiles.set(path, resource);
237 1 : return path;
238 1 : });
239 1 : }
240 :
241 1 : async removeExternalFile(path: string): Promise<void> {
242 1 : await this.runDatabaseOperation(async () => {
243 1 : const resource = this.externalFiles.get(path);
244 1 : if (resource == null) throw new Error(`external file is not registered: ${path}`);
245 1 : await resource.close();
246 1 : this.externalFiles.delete(path);
247 1 : });
248 1 : }
249 :
250 1 : async registerExternalDirectory(
251 1 : handle: FileSystemDirectoryHandle,
252 1 : path: string | undefined,
253 1 : mode: ExternalDirectoryMode,
254 1 : ): Promise<string> {
255 1 : return await this.runDatabaseOperation(async () => {
256 1 : path ??= externalMountPath(handle.name);
257 1 : validateExternalMountPath(path);
258 1 : const resource = await this.engine.mount(handle, { path, mode });
259 1 : this.externalDirectories.set(path, resource);
260 1 : return path;
261 1 : });
262 1 : }
263 :
264 1 : async removeExternalDirectory(path: string): Promise<void> {
265 1 : await this.runDatabaseOperation(async () => {
266 1 : const resource = this.externalDirectories.get(path);
267 1 : if (resource == null) throw new Error(`external directory is not registered: ${path}`);
268 1 : await resource.close();
269 1 : this.externalDirectories.delete(path);
270 1 : });
271 1 : }
272 :
273 1 : async refreshExternalDirectory(path: string): Promise<void> {
274 1 : await this.runDatabaseOperation(async () => {
275 1 : const resource = this.externalDirectories.get(path);
276 1 : if (resource == null) throw new Error(`external directory is not registered: ${path}`);
277 1 : await resource.refresh();
278 1 : });
279 1 : }
280 :
281 1 : async terminate(): Promise<void> {
282 1 : if (this.terminated) {
283 0 : return;
284 : }
285 1 : if (!this.termination) {
286 1 : this.termination = (async () => {
287 1 : try {
288 1 : await this.databaseOperation;
289 1 : await Promise.allSettled([...this.openingSessions]);
290 1 : this.idleSessions.clear();
291 : // A failed close must not leave the other closes running into a retried termination
292 1 : const closes = await Promise.allSettled([...this.sessions].map(session => session.closeForTermination()));
293 1 : const failed = closes.find(close => close.status === 'rejected');
294 1 : if (failed) throw failed.reason;
295 1 : for (const mounts of [this.externalFiles, this.externalDirectories]) {
296 1 : for (const [path, resource] of mounts) {
297 1 : await resource.close();
298 1 : mounts.delete(path);
299 1 : }
300 : }
301 1 : for (const path of this.databasePaths()) {
302 1 : await this.executeDatabaseStatement(`DROP DATABASE ${quoteIdentifier(path)}`);
303 : }
304 1 : if (this.terminateEngine) {
305 1 : await this.engine.terminate();
306 : }
307 1 : this.terminated = true;
308 1 : } catch (error) {
309 0 : this.termination = null;
310 0 : throw error;
311 0 : }
312 1 : })();
313 : }
314 1 : await this.termination;
315 1 : }
316 :
317 1 : private async initialize(): Promise<void> {
318 1 : if (this.terminated) {
319 1 : throw new Error('database is terminated');
320 : }
321 1 : if (!this.initialization) {
322 1 : this.initialization = (async () => {
323 1 : await this.engine.ready();
324 1 : for (const path of this.databasePaths()) {
325 1 : await this.executeDatabaseStatement(`CREATE DATABASE ${quoteIdentifier(path)}`);
326 : }
327 1 : })().catch(error => {
328 0 : this.initialization = null;
329 0 : throw error;
330 1 : });
331 : }
332 1 : await this.initialization;
333 1 : }
334 :
335 1 : private async runDatabaseOperation<T>(operation: () => Promise<T>): Promise<T> {
336 1 : const previous = this.databaseOperation;
337 1 : let release!: () => void;
338 1 : this.databaseOperation = new Promise(resolve => { release = resolve; });
339 1 : await previous;
340 1 : try {
341 1 : if (this.termination || this.terminated) throw new Error('database is terminated');
342 1 : return await operation();
343 1 : } finally {
344 1 : release();
345 1 : }
346 1 : }
347 :
348 1 : private async executeDatabaseStatement(statement: string): Promise<void> {
349 1 : const session = new HyperDBSession(await this.engine.connect(), this.mountFile, () => {});
350 1 : try {
351 1 : await session.execute(statement);
352 1 : } finally {
353 1 : await session.close();
354 1 : }
355 1 : }
356 1 : }
357 :
358 1 : function externalMountPath(name: string): string {
359 : // Mount paths allow only ASCII components and at most 200 bytes below /mnt/external/.
360 1 : const component = name.replace(/[^a-zA-Z0-9._-]/g, '_').slice(0, 200);
361 1 : return `/mnt/external/${!component || component === '.' || component === '..' ? 'file' : component}`;
362 1 : }
363 :
364 1 : function validateExternalMountPath(path: string): void {
365 : const prefix = '/mnt/external/';
366 1 : const relativePath = path.startsWith(prefix) ? path.slice(prefix.length) : '';
367 1 : if (
368 1 : relativePath.length === 0 ||
369 1 : new TextEncoder().encode(relativePath).byteLength > 200 ||
370 1 : relativePath.split('/').some(component => !/^[A-Za-z0-9._-]+$/.test(component) || component === '.' || component === '..')
371 : ) {
372 1 : throw new Error('external mount path must be below /mnt/external/ with ASCII path components');
373 : }
374 1 : }
375 :
376 1 : function throwIfAborted(signal?: AbortSignal): void {
377 1 : if (signal?.aborted) throw signal.reason ?? new DOMException('The operation was aborted', 'AbortError');
378 1 : }
379 :
380 : /// A connection lease handed out by HyperDB.connect().
381 : /// Closing it releases the physical session, which HyperDB may hand to the next lease.
382 : /// The lease stays closed, so a stale reference cannot run statements on a reused session.
383 1 : export class HyperDBConnection implements EmbeddedTableImportConnection {
384 1 : private closed = false;
385 :
386 1 : constructor(
387 1 : private readonly session: HyperDBSession,
388 1 : private readonly release: (session: HyperDBSession) => Promise<void>,
389 1 : ) {}
390 :
391 1 : async close(): Promise<void> {
392 1 : if (this.closed) {
393 1 : return;
394 : }
395 1 : if (this.session.active) {
396 1 : throw new Error('connection has an active operation');
397 : }
398 1 : this.closed = true;
399 1 : await this.release(this.session);
400 1 : }
401 :
402 1 : async query(query: string, abort?: AbortSignal): Promise<arrow.Table> {
403 1 : this.checkOpen();
404 1 : return await this.session.query(query, abort);
405 1 : }
406 :
407 1 : async queryStream(query: string, abort?: AbortSignal): Promise<EmbeddedResultStream> {
408 1 : this.checkOpen();
409 1 : return await this.session.queryStream(query, abort);
410 1 : }
411 :
412 1 : async execute(statement: string, abort?: AbortSignal): Promise<bigint | null> {
413 1 : this.checkOpen();
414 1 : return await this.session.execute(statement, abort);
415 1 : }
416 :
417 1 : async insertArrowTable(table: arrow.Table, options: EmbeddedTableInsertOptions): Promise<void> {
418 1 : this.checkOpen();
419 1 : await this.session.insertArrowTable(table, options);
420 1 : }
421 :
422 1 : async createTableAs(name: string, query: string, abort?: AbortSignal): Promise<void> {
423 1 : this.checkOpen();
424 1 : await this.session.createTableAs(name, query, abort);
425 1 : }
426 :
427 1 : private checkOpen(): void {
428 1 : if (this.closed) {
429 1 : throw new Error('connection is closed');
430 : }
431 1 : }
432 1 : }
433 :
434 : /// The result batches of a session query.
435 : /// Arrow readers never close their source, so the stream stops the query of the session itself.
436 1 : class HyperDBResultStream implements EmbeddedResultStream {
437 1 : constructor(
438 : /// The reader of the Arrow IPC chunks
439 1 : private readonly reader: arrow.AsyncRecordBatchStreamReader,
440 : /// Stops the query and waits until the session is free
441 1 : private readonly stop: () => Promise<void>,
442 1 : ) {}
443 :
444 1 : get schema(): arrow.Schema {
445 1 : return this.reader.schema;
446 1 : }
447 :
448 1 : async next(): Promise<IteratorResult<arrow.RecordBatch>> {
449 1 : try {
450 1 : const next = await this.reader.next();
451 1 : if (next.done) await this.stop();
452 1 : return next;
453 1 : } catch (error) {
454 1 : await this.stop();
455 1 : throw error;
456 1 : }
457 1 : }
458 :
459 1 : async cancel(): Promise<void> {
460 1 : await this.stop();
461 1 : }
462 :
463 1 : async *[Symbol.asyncIterator](): AsyncGenerator<arrow.RecordBatch> {
464 1 : try {
465 1 : for (let next = await this.next(); !next.done; next = await this.next()) {
466 1 : yield next.value;
467 : }
468 1 : } finally {
469 1 : await this.stop();
470 1 : }
471 1 : }
472 1 : }
473 :
474 : /// The running query of a session
475 : interface HyperDBSessionQuery {
476 : /// The Arrow IPC chunks of the query
477 : chunks: AsyncGenerator<Uint8Array, QueryInfo | undefined>;
478 : /// Cancels the query when the session terminates
479 : cancellation: AbortController;
480 : }
481 :
482 : /// A physical engine connection.
483 1 : class HyperDBSession {
484 1 : private closed = false;
485 : /// Set while an operation runs, the engine connection serves one at a time.
486 1 : active = false;
487 : /// Cleared when a query could not be cancelled or finished, the session must not be reused then.
488 1 : reusable = true;
489 1 : private closing: Promise<void> | null = null;
490 : /// The running query, if any
491 1 : private current: HyperDBSessionQuery | null = null;
492 :
493 1 : constructor(
494 1 : private readonly connection: HyperDBEngineConnection,
495 1 : private readonly mountFile: (file: File, path: string) => Promise<MountedResource>,
496 1 : private readonly onClose: () => void,
497 1 : private readonly databaseAlias?: string,
498 1 : ) {}
499 :
500 1 : async close(): Promise<void> {
501 1 : if (this.closed) {
502 1 : return;
503 : }
504 1 : if (this.active) {
505 1 : throw new Error('connection has an active operation');
506 : }
507 1 : await this.closeImpl();
508 1 : }
509 :
510 1 : async query(query: string, abort?: AbortSignal): Promise<arrow.Table> {
511 1 : const stream = await this.queryStream(query, abort);
512 1 : const batches: arrow.RecordBatch[] = [];
513 1 : for await (const batch of stream) {
514 1 : batches.push(batch);
515 : }
516 1 : return new arrow.Table(stream.schema, batches);
517 1 : }
518 :
519 1 : async queryStream(query: string, abort?: AbortSignal): Promise<EmbeddedResultStream> {
520 1 : this.begin();
521 1 : const chunks = this.startQuery(query, abort);
522 1 : const cancellation = this.current!.cancellation;
523 1 : const stop = async () => {
524 1 : cancellation.abort(new DOMException('The result stream was cancelled', 'AbortError'));
525 1 : await chunks.return(undefined).catch(() => {});
526 1 : };
527 1 : const first = await chunks.next();
528 1 : const source = first.done ? singleChunk(EMPTY_RESULT_IPC) : prependChunk(first.value, chunks);
529 1 : try {
530 1 : const reader = await arrow.RecordBatchReader.from(source);
531 1 : if (!reader.isAsync() || !reader.isStream()) {
532 1 : throw new Error('query did not produce an Arrow stream');
533 : }
534 1 : await reader.open();
535 1 : return new HyperDBResultStream(reader, stop);
536 1 : } catch (error) {
537 1 : await stop();
538 1 : throw error;
539 1 : }
540 1 : }
541 :
542 1 : async execute(statement: string, abort?: AbortSignal): Promise<bigint | null> {
543 1 : this.begin();
544 1 : return await this.drain(this.startQuery(statement, abort));
545 1 : }
546 :
547 : /// Insert an Arrow table through an external scan of its Arrow IPC stream, mounted as a file.
548 1 : async insertArrowTable(table: arrow.Table, options: EmbeddedTableInsertOptions): Promise<void> {
549 1 : if (!options.name) {
550 0 : throw new Error('Arrow insert table name must not be empty');
551 : }
552 1 : const insert = prepareHyperArrowInsert(table);
553 1 : this.begin();
554 1 : let resource: MountedResource | null = null;
555 1 : try {
556 1 : const name = `__dashql_insert_${crypto.randomUUID().replace(/-/g, '_')}.arrows`;
557 1 : const path = `/mnt/external/${name}`;
558 : // tableToIPC writes into a fresh ArrayBuffer.
559 1 : const bytes = arrow.tableToIPC(insert.table, 'stream') as Uint8Array<ArrayBuffer>;
560 1 : resource = await this.mountFile(new File([bytes], name), path);
561 1 : const target = options.schema == null
562 1 : ? quoteIdentifier(options.name)
563 1 : : `${quoteIdentifier(options.schema)}.${quoteIdentifier(options.name)}`;
564 1 : const source = `SELECT * FROM external(${quoteLiteral(path)}, format => 'arrowstream', columns => descriptor(${insert.columns}))`;
565 1 : await this.drain(this.startQuery(
566 1 : options.create ?? true ? `CREATE TABLE ${target} AS ${source}` : `INSERT INTO ${target} ${source}`,
567 1 : ));
568 1 : } finally {
569 1 : this.active = false;
570 1 : await resource?.close();
571 1 : }
572 1 : }
573 :
574 : /// Read the chunks of a statement to the end and return the rows it affected.
575 1 : private async drain(chunks: AsyncGenerator<Uint8Array, QueryInfo | undefined>): Promise<bigint | null> {
576 1 : let next = await chunks.next();
577 1 : while (!next.done) {
578 1 : next = await chunks.next();
579 : }
580 1 : return readAffectedRowCount(next.value);
581 1 : }
582 :
583 1 : async createTableAs(name: string, query: string, abort?: AbortSignal): Promise<void> {
584 1 : await this.execute(`CREATE TABLE ${quoteIdentifier(name)} AS ${query}`, abort);
585 1 : }
586 :
587 1 : async closeForTermination(): Promise<void> {
588 : // The running query of a close is its detach, aborting it would keep the database attached
589 1 : if (this.closing) {
590 1 : await this.closing;
591 1 : return;
592 1 : }
593 1 : const current = this.current;
594 1 : if (current != null) {
595 1 : current.cancellation.abort(new Error('database is terminating'));
596 1 : await current.chunks.return(undefined).catch(() => {});
597 : // A query that never started has nothing running in the engine.
598 1 : if (this.current === current) this.finishQuery();
599 1 : }
600 1 : while (this.active) {
601 0 : await yieldToEventLoop();
602 : }
603 1 : await this.closeImpl();
604 1 : }
605 :
606 : /// Start a query of the running operation, the returned chunks run it as they are read.
607 1 : private startQuery(query: string, abort?: AbortSignal): AsyncGenerator<Uint8Array, QueryInfo | undefined> {
608 1 : const cancellation = new AbortController();
609 1 : const signal = abort == null ? cancellation.signal : AbortSignal.any([abort, cancellation.signal]);
610 1 : const chunks = this.readChunks(query, signal);
611 1 : this.current = { chunks, cancellation };
612 1 : return chunks;
613 1 : }
614 :
615 1 : private async *readChunks(query: string, signal: AbortSignal): AsyncGenerator<Uint8Array, QueryInfo | undefined> {
616 1 : const stream = this.connection.queryStreamIPC(query, { signal });
617 1 : const chunks = stream[Symbol.asyncIterator]();
618 1 : try {
619 1 : for (;;) {
620 1 : const next = await chunks.next();
621 1 : if (next.done) break;
622 1 : yield next.value;
623 1 : }
624 1 : return await stream.info;
625 1 : } finally {
626 : // An aborted scan rejects at once, return() waits until the engine connection finished the query.
627 1 : await chunks.return?.(undefined).catch(() => {});
628 : // The query information is only missing when the query could not be cancelled or finished.
629 1 : if (!signal.aborted && !await stream.info.then(() => true, () => false)) {
630 1 : this.reusable = false;
631 : }
632 1 : this.finishQuery();
633 1 : }
634 1 : }
635 :
636 1 : private begin(): void {
637 1 : this.checkOpen();
638 1 : if (this.active) {
639 1 : throw new Error('connection has an active operation');
640 : }
641 1 : this.active = true;
642 1 : }
643 :
644 1 : private finishQuery(): void {
645 1 : this.current = null;
646 1 : this.active = false;
647 1 : }
648 :
649 1 : private async closeImpl(): Promise<void> {
650 1 : if (!this.closing) {
651 1 : this.closing = (async () => {
652 1 : if (this.databaseAlias != null) {
653 1 : await this.execute(`DETACH DATABASE ${quoteIdentifier(this.databaseAlias)}`);
654 : }
655 1 : await this.connection.close();
656 1 : this.closed = true;
657 1 : this.onClose();
658 1 : })().catch(error => {
659 0 : this.closing = null;
660 0 : throw error;
661 1 : });
662 : }
663 1 : await this.closing;
664 1 : }
665 :
666 1 : private checkOpen(): void {
667 1 : if (this.closed) {
668 1 : throw new Error('connection is closed');
669 : }
670 1 : }
671 1 : }
672 :
673 1 : function quoteLiteral(value: string): string {
674 1 : return `'${value.replace(/'/g, "''")}'`;
675 1 : }
676 :
677 1 : export function quoteIdentifier(value: string): string {
678 1 : return `"${value.replace(/"/g, '""')}"`;
679 1 : }
|
