Line data Source code
1 : import * as arrow from 'apache-arrow';
2 : import * as pb from "../../../proto.js";
3 : import * as connection from '@ankoh/dashql-jsonschema/connection.js';
4 :
5 : import {
6 : HyperDatabaseChannel,
7 : HyperDatabaseClient,
8 : HyperDatabaseConnectionContext,
9 : HyperQueryResultStream,
10 : } from '../hyperdb_grpc_client.js';
11 : import {
12 : HyperDatabaseHttpClient,
13 : HyperHttpAuthProvider,
14 : HyperHttpError,
15 : mergeQueryStatus,
16 : QueryStatus,
17 : } from '../hyperdb_http_client.js';
18 : import { fromHttpErrorResponse, HyperQueryError } from '../hyper_error_info.js';
19 : import {
20 : createQueryResponseStreamMetrics,
21 : QueryExecutionMetrics,
22 : QueryExecutionProgress,
23 : QueryExecutionResponseStream,
24 : QueryExecutionStatus,
25 : } from '../../../query/query_execution_state.js';
26 : import { AsyncConsumer } from '../../../utils/async_consumer.js';
27 : import { AsyncValue } from '../../../utils/async_value.js';
28 : import { HttpClient } from '../../../platform/http/http_client.js';
29 : import { Logger } from '../../../platform/logger/logger.js';
30 : import { yieldToBrowser } from '../../../utils/yield_to_browser.js';
31 :
32 1 : const LOG_CTX = "web_hyperdb_http_client";
33 :
34 : /// Convert a HyperHttpError (v3 JSON error body) into a structured HyperQueryError.
35 : /// Other errors are rethrown unchanged.
36 1 : function rethrowAsHyperQueryError(e: unknown): never {
37 1 : if (e instanceof HyperHttpError && e.errorResponse) {
38 1 : throw new HyperQueryError(fromHttpErrorResponse(e.errorResponse), LOG_CTX);
39 : }
40 1 : throw e;
41 1 : }
42 :
43 1 : const CHUNK_POLL_DELAY_MS = 250;
44 : export const DEFAULT_PARALLEL_CHUNKS = 4;
45 1 : const STREAM_YIELD_INTERVAL = 8;
46 :
47 : /// Bytes-reader that feeds Arrow IPC stream chunks from v3 HTTP into an AsyncRecordBatchStreamReader.
48 : /// Chunks are prefetched in parallel up to `parallelChunks`, but yielded in-order so that the
49 : /// Arrow stream framing stays valid.
50 1 : class WebHyperResultReader implements AsyncIterator<Uint8Array>, AsyncIterable<Uint8Array> {
51 1 : httpClient: HyperDatabaseHttpClient;
52 1 : logger: Logger;
53 1 : queryId: string;
54 : /// The initial response body from POST /api/v3/query (already the first/only chunk when RESULTS_PRODUCED).
55 1 : initialBytes: Uint8Array | null;
56 : /// Current status (from POST or last poll / chunk response).
57 1 : status: QueryStatus;
58 : /// Maximum number of in-flight chunk fetches.
59 1 : parallelChunks: number;
60 : /// Index of the next chunk to enqueue for prefetch.
61 1 : nextPrefetchIndex: number;
62 : /// FIFO of in-flight chunk byte Promises; the head is the next chunk to return.
63 1 : inflight: Array<Promise<Uint8Array>>;
64 : /// Has the Arrow schema prefix already been emitted to the reader (via POST body or a prior chunk)?
65 1 : schemaEmitted: boolean;
66 : /// Public reader state.
67 1 : currentStatus: QueryExecutionStatus;
68 1 : metrics: QueryExecutionMetrics;
69 1 : metadata: Map<string, string>;
70 : /// Callback invoked whenever a new QueryStatus is observed.
71 1 : progressCallback: ((progress: QueryExecutionProgress) => void) | null;
72 1 : abort: AbortSignal | undefined;
73 :
74 1 : constructor(httpClient: HyperDatabaseHttpClient, logger: Logger, queryId: string, initialStatus: QueryStatus, initialBytes: Uint8Array | null, parallelChunks: number) {
75 1 : this.httpClient = httpClient;
76 1 : this.logger = logger;
77 1 : this.queryId = queryId;
78 1 : this.initialBytes = initialBytes;
79 1 : this.status = initialStatus;
80 1 : this.parallelChunks = Math.max(1, parallelChunks);
81 1 : this.nextPrefetchIndex = initialBytes && initialBytes.byteLength > 0 ? 1 : 0;
82 1 : this.inflight = [];
83 1 : this.schemaEmitted = !!(initialBytes && initialBytes.byteLength > 0);
84 1 : this.currentStatus = QueryExecutionStatus.RUNNING;
85 1 : this.metrics = createQueryResponseStreamMetrics();
86 1 : this.metadata = new Map();
87 1 : this.progressCallback = null;
88 1 : this.abort = undefined;
89 1 : if (initialBytes) {
90 1 : this.metrics.totalDataBytesReceived += initialBytes.byteLength;
91 : }
92 1 : }
93 :
94 1 : emitProgress(): void {
95 1 : if (!this.progressCallback) return;
96 1 : const stats = this.status.executionStats;
97 1 : this.progressCallback({
98 1 : isQueued: null,
99 1 : progressPercent: this.status.progress,
100 1 : resultRows: this.status.rowCount,
101 1 : rowsProcessed: stats?.rowsProcessed ?? null,
102 1 : executionTimeMs: stats?.wallClockTime ?? null,
103 1 : metrics: { ...this.metrics },
104 1 : engineProgress: { kind: 'hyper_query_status', status: this.status, receivedAt: Date.now() },
105 1 : });
106 1 : }
107 :
108 1 : private updateStatus(status: QueryStatus): void {
109 1 : this.status = mergeQueryStatus(this.status, status);
110 1 : this.emitProgress();
111 1 : }
112 :
113 1 : private isCompleted(): boolean {
114 1 : return this.status.completionStatus === "RESULTS_PRODUCED" || this.status.completionStatus === "FINISHED";
115 1 : }
116 :
117 1 : private fillPrefetchWindow(): void {
118 1 : const chunkCount = this.status.chunkCount ?? 0;
119 1 : while (this.inflight.length < this.parallelChunks && this.nextPrefetchIndex < chunkCount) {
120 1 : const chunkId = this.nextPrefetchIndex++;
121 1 : const omitSchema = this.schemaEmitted;
122 1 : this.schemaEmitted = true;
123 1 : this.inflight.push(this.fetchChunk(chunkId, omitSchema));
124 1 : }
125 1 : }
126 :
127 1 : private async fetchChunk(chunkId: number, omitSchema: boolean): Promise<Uint8Array> {
128 1 : const { status, response } = await this.httpClient.getQueryChunk({
129 1 : queryId: this.queryId,
130 1 : chunkId,
131 1 : omitSchema,
132 1 : }, undefined, this.abort).catch(rethrowAsHyperQueryError);
133 1 : if (status) {
134 1 : this.updateStatus(status);
135 : }
136 1 : const body = await response.arrayBuffer();
137 1 : return new Uint8Array(body);
138 1 : }
139 :
140 1 : async next(): Promise<IteratorResult<Uint8Array>> {
141 1 : if (this.initialBytes && this.initialBytes.byteLength > 0) {
142 1 : const value = this.initialBytes;
143 1 : this.initialBytes = null;
144 1 : this.fillPrefetchWindow();
145 1 : return { done: false, value };
146 1 : }
147 1 : this.initialBytes = null;
148 :
149 1 : while (true) {
150 1 : this.fillPrefetchWindow();
151 :
152 1 : if (this.inflight.length > 0) {
153 1 : const bytes = await this.inflight.shift()!;
154 : // Top up the window as soon as one slot frees.
155 1 : this.fillPrefetchWindow();
156 1 : if (bytes.byteLength > 0) {
157 1 : this.metrics.totalDataBytesReceived += bytes.byteLength;
158 1 : return { done: false, value: bytes };
159 1 : }
160 0 : continue;
161 0 : }
162 :
163 1 : if (this.isCompleted()) {
164 1 : return { done: true, value: null };
165 : }
166 :
167 : // No chunks available yet and query is still running — poll for progress.
168 0 : this.updateStatus(await this.httpClient.getQueryStatus({
169 0 : queryId: this.queryId,
170 0 : waitTimeMs: CHUNK_POLL_DELAY_MS,
171 0 : }, this.abort).catch(rethrowAsHyperQueryError));
172 0 : }
173 1 : }
174 :
175 1 : [Symbol.asyncIterator]() {
176 1 : return this;
177 1 : }
178 1 : }
179 :
180 : /// A web HyperDatabaseChannel backed by the v3 HTTP API.
181 1 : export class WebHyperQueryResultStream implements QueryExecutionResponseStream {
182 1 : reader: WebHyperResultReader;
183 1 : resultSchema: AsyncValue<arrow.Schema<any>, Error>;
184 :
185 1 : constructor(reader: WebHyperResultReader) {
186 1 : this.reader = reader;
187 1 : this.resultSchema = new AsyncValue();
188 1 : }
189 :
190 1 : async cancel(): Promise<void> {
191 1 : await this.reader.httpClient.cancelQuery(this.reader.queryId);
192 1 : }
193 :
194 0 : getMetadata(): Map<string, string> {
195 0 : return this.reader.metadata;
196 0 : }
197 1 : getMetrics(): QueryExecutionMetrics {
198 1 : return this.reader.metrics;
199 1 : }
200 0 : getStatus(): QueryExecutionStatus {
201 0 : return this.reader.currentStatus;
202 0 : }
203 0 : getSchema(): Promise<arrow.Schema> {
204 0 : return this.resultSchema.getValue();
205 0 : }
206 :
207 1 : async produce(batches: AsyncConsumer<QueryExecutionResponseStream, arrow.RecordBatch>, progress: AsyncConsumer<QueryExecutionResponseStream, QueryExecutionProgress>, abort?: AbortSignal): Promise<void> {
208 1 : this.reader.abort = abort;
209 1 : this.reader.progressCallback = (update: QueryExecutionProgress) => {
210 1 : progress.resolve(this, update);
211 1 : };
212 : // Report the status of the query response
213 1 : this.reader.emitProgress();
214 1 : const arrowReader = await arrow.AsyncRecordBatchStreamReader.from(this.reader);
215 1 : abort?.throwIfAborted();
216 :
217 1 : await arrowReader.open();
218 1 : abort?.throwIfAborted();
219 1 : this.resultSchema.resolve(arrowReader.schema);
220 :
221 1 : let batchesRead = 0;
222 1 : while (true) {
223 1 : const iter = await arrowReader.next();
224 1 : abort?.throwIfAborted();
225 1 : if (iter.done) {
226 1 : if (iter.value !== undefined) {
227 0 : batches.resolve(this, iter.value);
228 : }
229 1 : break;
230 1 : }
231 1 : batches.resolve(this, iter.value);
232 1 : if (++batchesRead % STREAM_YIELD_INTERVAL === 0) await yieldToBrowser();
233 1 : }
234 1 : }
235 1 : }
236 :
237 1 : class WebHyperDatabaseChannel implements HyperDatabaseChannel {
238 1 : httpClient: HyperDatabaseHttpClient;
239 1 : connection: HyperDatabaseConnectionContext;
240 1 : logger: Logger;
241 1 : parallelChunks: number;
242 :
243 1 : constructor(httpClient: HyperDatabaseHttpClient, connection: HyperDatabaseConnectionContext, logger: Logger, parallelChunks: number) {
244 1 : this.httpClient = httpClient;
245 1 : this.connection = connection;
246 1 : this.logger = logger;
247 1 : this.parallelChunks = parallelChunks;
248 1 : }
249 :
250 1 : async executeQuery(params: pb.salesforce_hyperdb_grpc_v1.pb.QueryParam, abort?: AbortSignal): Promise<HyperQueryResultStream> {
251 1 : const sql = params.query;
252 1 : const queryParameters = this.connection.getQueryParameters();
253 1 : const settings = Object.keys(queryParameters).length > 0 ? queryParameters : undefined;
254 1 : const attachedDatabases = this.connection.getAttachedDatabases();
255 1 : const databases = attachedDatabases.length > 0 ? attachedDatabases.map(db => ({ path: db.path, alias: db.alias || undefined })) : undefined;
256 1 : const clientContext = this.connection.getClientContext?.();
257 1 : const headers = clientContext == null ? undefined : { externalClientContext: JSON.stringify(clientContext) };
258 1 : const { status, response } = await this.httpClient
259 1 : .executeQuery({ sql, settings, databases }, headers, abort)
260 1 : .catch(rethrowAsHyperQueryError);
261 1 : if (!status || !status.queryId) {
262 1 : throw new Error("v3 query response missing status header");
263 : }
264 1 : const bodyBuffer = await response.arrayBuffer();
265 1 : const bytes = new Uint8Array(bodyBuffer);
266 1 : const reader = new WebHyperResultReader(this.httpClient, this.logger, status.queryId, status, bytes.byteLength > 0 ? bytes : null, this.parallelChunks);
267 1 : return new WebHyperQueryResultStream(reader);
268 1 : }
269 :
270 0 : async close(): Promise<void> { }
271 1 : }
272 :
273 : /// HyperHttpAuthProvider that delegates to the connection context's getRequestMetadata.
274 1 : class ContextAuthProvider implements HyperHttpAuthProvider {
275 1 : context: HyperDatabaseConnectionContext;
276 1 : constructor(context: HyperDatabaseConnectionContext) {
277 1 : this.context = context;
278 1 : }
279 1 : async getAuthHeaders(): Promise<Record<string, string>> {
280 1 : return await this.context.getRequestMetadata();
281 1 : }
282 1 : }
283 :
284 : export interface WebHyperDatabaseClientConfig {
285 : /// Maximum number of chunk fetches in flight per query result stream.
286 : parallelChunks?: number;
287 : }
288 :
289 1 : export class WebHyperDatabaseClient implements HyperDatabaseClient {
290 1 : httpClient: HttpClient;
291 1 : logger: Logger;
292 1 : parallelChunks: number;
293 :
294 1 : constructor(httpClient: HttpClient, logger: Logger, config: WebHyperDatabaseClientConfig = {}) {
295 1 : this.httpClient = httpClient;
296 1 : this.logger = logger;
297 1 : this.parallelChunks = Math.max(1, config.parallelChunks ?? DEFAULT_PARALLEL_CHUNKS);
298 1 : }
299 :
300 1 : async connect(hyperArgs: connection.HyperConnectionParams, context: HyperDatabaseConnectionContext): Promise<HyperDatabaseChannel> {
301 1 : if (!hyperArgs.endpoint) {
302 1 : throw new Error("missing hyper endpoint");
303 : }
304 1 : const endpointUrl = new URL(hyperArgs.endpoint);
305 1 : const auth = new ContextAuthProvider(context);
306 1 : const client = new HyperDatabaseHttpClient(this.httpClient, endpointUrl, auth, this.logger);
307 1 : return new WebHyperDatabaseChannel(client, context, this.logger, this.parallelChunks);
308 1 : }
309 1 : }
|
