Line data Source code
1 : import type * as arrow from 'apache-arrow';
2 : import type * as pb from '../../proto.js';
3 :
4 : import type { EmbeddedConnection, EmbeddedResultStream } from '../../platform/database/embedded_database.js';
5 : import type { AsyncConsumer } from '../../utils/async_consumer.js';
6 : import {
7 : createQueryResponseStreamMetrics,
8 : QueryExecutionProgress,
9 : QueryExecutionResponseStream,
10 : QueryExecutionStatus,
11 : type QueryExecutionMetrics,
12 : } from '../../query/query_execution_state.js';
13 : import type { HyperDatabaseChannel } from '../hyper/hyperdb_grpc_client.js';
14 : import { yieldToBrowser } from '../../utils/yield_to_browser.js';
15 :
16 1 : const STREAM_YIELD_INTERVAL = 8;
17 :
18 : /// A query result that streams the record batches as the embedded database produces them.
19 1 : class EmbeddedQueryResultStream implements QueryExecutionResponseStream {
20 1 : private readonly metrics: QueryExecutionMetrics;
21 1 : private status = QueryExecutionStatus.RUNNING;
22 : /// Resolves once the result is read to the end or cancelled, the connection is free again then
23 1 : readonly finished: Promise<void>;
24 1 : private resolveFinished!: () => void;
25 :
26 1 : constructor(
27 : /// The result batches
28 1 : private readonly reader: EmbeddedResultStream,
29 : /// The time at which the query was started
30 1 : private readonly startedAt: number,
31 1 : ) {
32 1 : this.metrics = createQueryResponseStreamMetrics();
33 1 : this.metrics.totalQueryRequestsStarted = 1;
34 1 : this.finished = new Promise(resolve => { this.resolveFinished = resolve; });
35 1 : }
36 :
37 1 : async cancel(): Promise<void> {
38 1 : try {
39 1 : await this.reader.cancel();
40 1 : } finally {
41 1 : if (this.status === QueryExecutionStatus.RUNNING) this.status = QueryExecutionStatus.CANCELLED;
42 1 : this.resolveFinished();
43 1 : }
44 1 : }
45 :
46 1 : getMetadata(): Map<string, string> {
47 1 : return new Map();
48 1 : }
49 :
50 1 : getMetrics(): QueryExecutionMetrics {
51 1 : return this.metrics;
52 1 : }
53 :
54 1 : getStatus(): QueryExecutionStatus {
55 1 : return this.status;
56 1 : }
57 :
58 1 : async getSchema(): Promise<arrow.Schema> {
59 1 : return this.reader.schema;
60 1 : }
61 :
62 : /// Produce the batches, the abort signal of the query cancels it in the database
63 1 : async produce(
64 1 : batches: AsyncConsumer<QueryExecutionResponseStream, arrow.RecordBatch>,
65 1 : _progress: AsyncConsumer<QueryExecutionResponseStream, QueryExecutionProgress>,
66 1 : abort?: AbortSignal,
67 1 : ): Promise<void> {
68 1 : try {
69 1 : for (;;) {
70 1 : abort?.throwIfAborted();
71 1 : const next = await this.reader.next();
72 1 : if (next.done) break;
73 1 : const batch = next.value;
74 1 : this.metrics.durationUntilFirstBatchMs ??= performance.now() - this.startedAt;
75 1 : this.metrics.totalDataBytesReceived += batch.data.byteLength;
76 1 : this.metrics.totalBatchesReceived += 1;
77 1 : this.metrics.totalRowsReceived += batch.numRows;
78 1 : batches.resolve(this, batch);
79 1 : if (this.metrics.totalBatchesReceived % STREAM_YIELD_INTERVAL === 0) await yieldToBrowser();
80 1 : }
81 1 : this.metrics.totalQueryRequestsSucceeded = 1;
82 1 : this.status = QueryExecutionStatus.SUCCEEDED;
83 1 : } catch (error) {
84 1 : this.metrics.totalQueryRequestsFailed = 1;
85 1 : this.status = abort?.aborted ? QueryExecutionStatus.CANCELLED : QueryExecutionStatus.FAILED;
86 : // Stop the query when the batches are not read to the end
87 1 : await this.reader.cancel().catch(() => {});
88 1 : throw error;
89 1 : } finally {
90 1 : this.metrics.totalQueryRequestDurationMs = performance.now() - this.startedAt;
91 1 : this.resolveFinished();
92 1 : }
93 1 : }
94 1 : }
95 :
96 : /// Wait for a promise, an abort rejects right away
97 1 : function waitUnlessAborted(promise: Promise<void>, abort?: AbortSignal): Promise<void> {
98 1 : if (abort == null) return promise;
99 1 : return new Promise((resolve, reject) => {
100 1 : const onAbort = () => reject(abort.reason);
101 1 : abort.addEventListener('abort', onAbort, { once: true });
102 1 : void promise.then(() => {
103 1 : abort.removeEventListener('abort', onAbort);
104 1 : resolve();
105 1 : });
106 1 : });
107 1 : }
108 :
109 1 : export class EmbeddedDatabaseChannel {
110 : /// The connection runs one statement at a time, so queries wait until the result of the previous one is read.
111 : /// One connection keeps the session state of the scripts, such as settings and attached databases.
112 1 : private previousQuery: Promise<void> = Promise.resolve();
113 :
114 1 : constructor(private readonly connection: EmbeddedConnection) {}
115 :
116 1 : async executeQuery(query: string, abort?: AbortSignal): Promise<QueryExecutionResponseStream> {
117 1 : abort?.throwIfAborted();
118 1 : const previous = this.previousQuery;
119 1 : let release!: () => void;
120 1 : const finished = new Promise<void>(resolve => { release = resolve; });
121 : // The next query waits for the previous one as well, this query may be cancelled while it waits
122 1 : this.previousQuery = previous.then(() => finished);
123 1 : try {
124 1 : await waitUnlessAborted(previous, abort);
125 1 : abort?.throwIfAborted();
126 1 : const startedAt = performance.now();
127 1 : const stream = new EmbeddedQueryResultStream(await this.connection.queryStream(query, abort), startedAt);
128 1 : void stream.finished.then(release);
129 1 : return stream;
130 1 : } catch (error) {
131 1 : release();
132 1 : throw error;
133 1 : }
134 1 : }
135 :
136 1 : async close(): Promise<void> {
137 1 : await this.connection.close();
138 1 : }
139 1 : }
140 :
141 1 : export class EmbeddedHyperDatabaseChannel implements HyperDatabaseChannel {
142 1 : constructor(
143 1 : private readonly channel: EmbeddedDatabaseChannel,
144 1 : ) {}
145 :
146 1 : async executeQuery(
147 1 : params: pb.salesforce_hyperdb_grpc_v1.pb.QueryParam,
148 1 : abort?: AbortSignal,
149 1 : ): Promise<QueryExecutionResponseStream> {
150 1 : return await this.channel.executeQuery(params.query, abort);
151 1 : }
152 :
153 0 : async close(): Promise<void> {
154 0 : await this.channel.close();
155 0 : }
156 1 : }
|
