Line data Source code
1 : import * as arrow from 'apache-arrow';
2 : import * as buf from "@bufbuild/protobuf";
3 : import * as pb from "../../../proto.js";
4 : import * as connection from '@ankoh/dashql-jsonschema/connection.js';
5 :
6 : import {
7 : HyperDatabaseChannel,
8 : HyperDatabaseClient,
9 : HyperDatabaseConnectionContext,
10 : HyperQueryExecutionOptions,
11 : HyperQueryResultStream,
12 : } from '../hyperdb_grpc_client.js';
13 : import {
14 : NativeGrpcChannel,
15 : NativeGrpcClient,
16 : NativeGrpcProxyConfig,
17 : NativeGrpcServerStream,
18 : NativeGrpcServerStreamMessageIterator,
19 : } from '../../../platform/grpc/native_grpc_client.js';
20 : import {
21 : createQueryResponseStreamMetrics,
22 : QueryExecutionProgress,
23 : QueryExecutionResponseStream, QueryExecutionMetrics,
24 : QueryExecutionStatus,
25 : } from '../../../query/query_execution_state.js';
26 : import { ChannelArgs, ChannelError } from '../../../platform/channel_common.js';
27 : import { decodeGrpcErrorHeaders, HyperQueryError } from '../hyper_error_info.js';
28 : import { Logger } from '../../../platform/logger/logger.js';
29 : import { AsyncConsumer } from '../../../utils/async_consumer.js';
30 : import { AsyncValue } from '../../../utils/async_value.js';
31 : import { yieldToBrowser } from '../../../utils/yield_to_browser.js';
32 :
33 1 : const LOG_CTX = "native_hyperdb_grpc_client";
34 1 : const STREAM_YIELD_INTERVAL = 8;
35 :
36 1 : export class QueryResultReader implements AsyncIterator<Uint8Array>, AsyncIterable<Uint8Array> {
37 : /// The gRPC stream
38 1 : grpcStream: NativeGrpcServerStream;
39 : /// The logger
40 1 : logger: Logger;
41 : /// The message iterator
42 1 : messageIterator: NativeGrpcServerStreamMessageIterator;
43 : /// The current status
44 1 : currentStatus: QueryExecutionStatus;
45 : /// The metrics
46 1 : metrics: QueryExecutionMetrics;
47 :
48 1 : constructor(stream: NativeGrpcServerStream, logger: Logger) {
49 1 : this.grpcStream = stream;
50 1 : this.logger = logger;
51 1 : this.messageIterator = new NativeGrpcServerStreamMessageIterator(this.grpcStream, logger);
52 1 : this.currentStatus = QueryExecutionStatus.RUNNING;
53 1 : this.metrics = createQueryResponseStreamMetrics();
54 1 : }
55 :
56 : /// Get the result metadata (if any)
57 0 : get metadata() { return this.messageIterator.metadata; }
58 : /// Get the next binary result chunk
59 1 : async next(): Promise<IteratorResult<Uint8Array>> {
60 1 : while (true) {
61 1 : const next = await this.messageIterator.next();
62 1 : if (next.value == null) {
63 1 : return { done: true, value: null };
64 : }
65 1 : const resultMessage = buf.fromBinary(pb.salesforce_hyperdb_grpc_v1.pb.QueryResultSchema$, next.value);
66 1 : switch (resultMessage.result.case) {
67 : // We skip any dedicated header prefix
68 1 : case "header":
69 1 : continue;
70 : // Skip qsv1 chunks
71 1 : case "qsv1Chunk":
72 1 : throw new Error("invalid result data message. expected arrowChunk, received qsv1Chunk");
73 : // Unpack an arrow chunk
74 1 : case "arrowChunk": {
75 1 : const buffer = resultMessage.result.value.data;
76 1 : this.metrics.totalDataBytesReceived += buffer.byteLength;
77 1 : return { done: false, value: buffer };
78 1 : }
79 1 : }
80 1 : }
81 1 : }
82 :
83 : /// Get the async iterator
84 1 : [Symbol.asyncIterator]() {
85 1 : return this;
86 1 : }
87 1 : }
88 :
89 : /// A native Hyper query result stream
90 1 : export class NativeHyperQueryResultStream implements QueryExecutionResponseStream {
91 : /// The query result iterator
92 1 : resultReader: QueryResultReader;
93 : /// The schema Promise
94 1 : resultSchema: AsyncValue<arrow.Schema<any>, Error>;
95 :
96 1 : constructor(stream: NativeGrpcServerStream, _connection: HyperDatabaseConnectionContext, logger: Logger) {
97 1 : this.resultReader = new QueryResultReader(stream, logger);
98 1 : this.resultSchema = new AsyncValue();
99 1 : }
100 0 : async cancel(): Promise<void> {
101 0 : await this.resultReader.grpcStream.destroy();
102 0 : }
103 :
104 : /// Get the metadata
105 0 : getMetadata(): Map<string, string> {
106 0 : return this.resultReader.metadata;
107 0 : }
108 : /// Get the metrics
109 1 : getMetrics(): QueryExecutionMetrics {
110 1 : return this.resultReader.metrics;
111 1 : }
112 : /// Get the current query status
113 0 : getStatus(): QueryExecutionStatus {
114 0 : return this.resultReader.currentStatus;
115 0 : }
116 : /// Await the Arrow schema
117 1 : getSchema(): Promise<arrow.Schema> {
118 1 : return this.resultSchema.getValue();
119 1 : }
120 : /// Produce the result batches
121 1 : async produce(batches: AsyncConsumer<QueryExecutionResponseStream, arrow.RecordBatch>, _progress: AsyncConsumer<QueryExecutionResponseStream, QueryExecutionProgress>, abort?: AbortSignal): Promise<void> {
122 1 : const cancel = () => { void this.cancel(); };
123 1 : abort?.addEventListener('abort', cancel, { once: true });
124 1 : try {
125 1 : const arrowReader = await arrow.AsyncRecordBatchStreamReader.from(this.resultReader);
126 1 : abort?.throwIfAborted();
127 :
128 1 : await arrowReader.open();
129 1 : abort?.throwIfAborted();
130 1 : this.resultSchema.resolve(arrowReader.schema);
131 :
132 1 : let batchesRead = 0;
133 1 : while (true) {
134 1 : const iter = await arrowReader!.next();
135 1 : abort?.throwIfAborted();
136 1 : if (iter.done) {
137 1 : if (iter.value !== undefined) {
138 0 : batches.resolve(this, iter.value);
139 : }
140 1 : return;
141 1 : }
142 1 : batches.resolve(this, iter.value);
143 1 : if (++batchesRead % STREAM_YIELD_INTERVAL === 0) await yieldToBrowser();
144 1 : }
145 1 : } catch (e: unknown) {
146 : // A failed gRPC stream surfaces as a ChannelError whose headers carry
147 : // the forwarded gRPC richer error model. Decode it into a structured
148 : // HyperQueryError so the rest of the app sees Hyper's ErrorInfo.
149 1 : if (e instanceof ChannelError) {
150 0 : const info = decodeGrpcErrorHeaders(e.headers);
151 0 : if (info) {
152 0 : throw new HyperQueryError(info, LOG_CTX);
153 : }
154 0 : }
155 1 : throw e;
156 1 : } finally {
157 1 : abort?.removeEventListener('abort', cancel);
158 1 : }
159 1 : }
160 1 : }
161 :
162 : /// A native Hyper database connection
163 1 : class NativeHyperDatabaseChannel implements HyperDatabaseChannel {
164 : /// A gRPC channel
165 1 : grpcChannel: NativeGrpcChannel;
166 : /// The connection context
167 1 : connection: HyperDatabaseConnectionContext;
168 : /// A logger
169 1 : logger: Logger;
170 :
171 1 : constructor(channel: NativeGrpcChannel, connection: HyperDatabaseConnectionContext, logger: Logger) {
172 1 : this.grpcChannel = channel;
173 1 : this.connection = connection;
174 1 : this.logger = logger;
175 1 : }
176 :
177 : /// Execute a query against Hyper
178 1 : public async executeQuery(params: pb.salesforce_hyperdb_grpc_v1.pb.QueryParam, abort?: AbortSignal, options?: HyperQueryExecutionOptions): Promise<HyperQueryResultStream> {
179 1 : abort?.throwIfAborted();
180 1 : params.outputFormat = pb.salesforce_hyperdb_grpc_v1.pb.QueryParam_OutputFormat.ARROW_STREAM;
181 1 : for (const db of this.connection.getAttachedDatabases()) {
182 1 : params.database.push(buf.create(pb.salesforce_hyperdb_grpc_v1.pb.AttachedDatabaseSchema, db));
183 : }
184 1 : for (const [key, value] of Object.entries(this.connection.getQueryParameters())) {
185 1 : params.params[key] = value;
186 : }
187 1 : const clientContext = this.connection.getClientContext?.();
188 1 : if (clientContext != null) params.params.external_client_context = JSON.stringify(clientContext);
189 1 : const stream = await this.grpcChannel.startServerStream({
190 1 : path: "/salesforce.hyperdb.grpc.v1.HyperService/ExecuteQuery",
191 1 : body: buf.toBinary(pb.salesforce_hyperdb_grpc_v1.pb.QueryParamSchema, params),
192 1 : readTimeoutMs: options?.readTimeoutMs,
193 1 : });
194 1 : if (abort?.aborted) {
195 0 : await stream.destroy();
196 0 : abort.throwIfAborted();
197 0 : }
198 1 : return new NativeHyperQueryResultStream(stream, this.connection, this.logger);
199 1 : }
200 :
201 : /// Close the connection
202 1 : public async close(): Promise<void> {
203 1 : await this.grpcChannel.close();
204 1 : }
205 1 : }
206 :
207 : /// A native Hyper database client
208 1 : export class NativeHyperDatabaseClient implements HyperDatabaseClient {
209 : /// A logger
210 1 : logger: Logger;
211 : /// A native Hyper client
212 1 : client: NativeGrpcClient;
213 :
214 1 : constructor(config: NativeGrpcProxyConfig, logger: Logger) {
215 1 : this.logger = logger;
216 1 : this.client = new NativeGrpcClient(config, logger);
217 1 : }
218 :
219 : /// Create a database connection
220 1 : public async connect(hyperArgs: connection.HyperConnectionParams, connection: HyperDatabaseConnectionContext): Promise<NativeHyperDatabaseChannel> {
221 1 : const tls = hyperArgs.tls;
222 1 : const isHttps = new URL(hyperArgs.endpoint).protocol === "https:";
223 1 : const useTls = isHttps || Boolean(
224 1 : tls?.clientKeyPath || tls?.clientCertPath || tls?.caCertsPath,
225 1 : );
226 1 : if (useTls && !isHttps) {
227 1 : throw new Error("TLS certificate paths require an https:// endpoint");
228 : }
229 1 : const args: ChannelArgs = {
230 1 : endpoint: hyperArgs.endpoint,
231 1 : tls: useTls ? {
232 1 : keyPath: tls?.clientKeyPath,
233 1 : pubPath: tls?.clientCertPath,
234 1 : caPath: tls?.caCertsPath,
235 1 : } : undefined
236 1 : };
237 1 : const channel = await this.client.connect(args, connection);
238 1 : return new NativeHyperDatabaseChannel(channel, connection, this.logger);
239 1 : }
240 1 : }
|
