@query-farm/vgi-rpc 0.2.4 → 0.3.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +3 -3
- package/dist/client/connect.d.ts.map +1 -1
- package/dist/client/index.d.ts +3 -3
- package/dist/client/index.d.ts.map +1 -1
- package/dist/client/introspect.d.ts +1 -1
- package/dist/client/introspect.d.ts.map +1 -1
- package/dist/client/ipc.d.ts +1 -1
- package/dist/client/ipc.d.ts.map +1 -1
- package/dist/client/pipe.d.ts +2 -2
- package/dist/client/pipe.d.ts.map +1 -1
- package/dist/client/stream.d.ts +1 -1
- package/dist/client/stream.d.ts.map +1 -1
- package/dist/dispatch/describe.d.ts +1 -1
- package/dist/dispatch/describe.d.ts.map +1 -1
- package/dist/dispatch/stream.d.ts +1 -1
- package/dist/dispatch/stream.d.ts.map +1 -1
- package/dist/dispatch/unary.d.ts.map +1 -1
- package/dist/http/common.d.ts +1 -1
- package/dist/http/common.d.ts.map +1 -1
- package/dist/http/dispatch.d.ts.map +1 -1
- package/dist/http/handler.d.ts.map +1 -1
- package/dist/http/index.d.ts +3 -1
- package/dist/http/index.d.ts.map +1 -1
- package/dist/http/token.d.ts.map +1 -1
- package/dist/http/types.d.ts.map +1 -1
- package/dist/index.d.ts +7 -7
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +2162 -2119
- package/dist/index.js.map +26 -26
- package/dist/protocol.d.ts +1 -1
- package/dist/protocol.d.ts.map +1 -1
- package/dist/schema.d.ts +1 -1
- package/dist/schema.d.ts.map +1 -1
- package/dist/server.d.ts +1 -1
- package/dist/server.d.ts.map +1 -1
- package/dist/types.d.ts +1 -1
- package/dist/types.d.ts.map +1 -1
- package/dist/util/schema.d.ts +1 -1
- package/dist/util/schema.d.ts.map +1 -1
- package/dist/util/zstd.d.ts.map +1 -1
- package/dist/wire/reader.d.ts +1 -1
- package/dist/wire/reader.d.ts.map +1 -1
- package/dist/wire/request.d.ts +1 -1
- package/dist/wire/request.d.ts.map +1 -1
- package/dist/wire/response.d.ts +1 -1
- package/dist/wire/response.d.ts.map +1 -1
- package/dist/wire/writer.d.ts +1 -1
- package/dist/wire/writer.d.ts.map +1 -1
- package/package.json +9 -5
- package/src/client/connect.ts +12 -20
- package/src/client/index.ts +8 -8
- package/src/client/introspect.ts +11 -15
- package/src/client/ipc.ts +18 -32
- package/src/client/pipe.ts +17 -37
- package/src/client/stream.ts +20 -46
- package/src/dispatch/describe.ts +15 -27
- package/src/dispatch/stream.ts +7 -21
- package/src/dispatch/unary.ts +1 -2
- package/src/http/common.ts +9 -19
- package/src/http/dispatch.ts +115 -110
- package/src/http/handler.ts +20 -55
- package/src/http/index.ts +3 -1
- package/src/http/token.ts +2 -7
- package/src/http/types.ts +2 -6
- package/src/index.ts +44 -41
- package/src/protocol.ts +8 -8
- package/src/schema.ts +12 -16
- package/src/server.ts +16 -36
- package/src/types.ts +9 -36
- package/src/util/schema.ts +1 -1
- package/src/util/zstd.ts +2 -8
- package/src/wire/reader.ts +2 -4
- package/src/wire/request.ts +4 -15
- package/src/wire/response.ts +9 -25
- package/src/wire/writer.ts +1 -5
package/dist/index.js
CHANGED
|
@@ -50,11 +50,122 @@ var init_zstd = __esm(() => {
|
|
|
50
50
|
isBun = typeof globalThis.Bun !== "undefined";
|
|
51
51
|
});
|
|
52
52
|
|
|
53
|
-
// src/
|
|
54
|
-
|
|
53
|
+
// src/constants.ts
|
|
54
|
+
var RPC_METHOD_KEY = "vgi_rpc.method";
|
|
55
|
+
var LOG_LEVEL_KEY = "vgi_rpc.log_level";
|
|
56
|
+
var LOG_MESSAGE_KEY = "vgi_rpc.log_message";
|
|
57
|
+
var LOG_EXTRA_KEY = "vgi_rpc.log_extra";
|
|
58
|
+
var REQUEST_VERSION_KEY = "vgi_rpc.request_version";
|
|
59
|
+
var REQUEST_VERSION = "1";
|
|
60
|
+
var SERVER_ID_KEY = "vgi_rpc.server_id";
|
|
61
|
+
var REQUEST_ID_KEY = "vgi_rpc.request_id";
|
|
62
|
+
var PROTOCOL_NAME_KEY = "vgi_rpc.protocol_name";
|
|
63
|
+
var DESCRIBE_VERSION_KEY = "vgi_rpc.describe_version";
|
|
64
|
+
var DESCRIBE_VERSION = "2";
|
|
65
|
+
var DESCRIBE_METHOD_NAME = "__describe__";
|
|
66
|
+
var STATE_KEY = "vgi_rpc.stream_state#b64";
|
|
67
|
+
|
|
68
|
+
// src/http/common.ts
|
|
69
|
+
import {
|
|
70
|
+
makeData,
|
|
71
|
+
RecordBatch,
|
|
72
|
+
RecordBatchReader,
|
|
73
|
+
RecordBatchStreamWriter,
|
|
74
|
+
Struct
|
|
75
|
+
} from "@query-farm/apache-arrow";
|
|
76
|
+
var ARROW_CONTENT_TYPE = "application/vnd.apache.arrow.stream";
|
|
77
|
+
|
|
78
|
+
class HttpRpcError extends Error {
|
|
79
|
+
statusCode;
|
|
80
|
+
constructor(message, statusCode) {
|
|
81
|
+
super(message);
|
|
82
|
+
this.statusCode = statusCode;
|
|
83
|
+
this.name = "HttpRpcError";
|
|
84
|
+
}
|
|
85
|
+
}
|
|
86
|
+
function conformBatchToSchema(batch, schema) {
|
|
87
|
+
if (batch.numRows === 0)
|
|
88
|
+
return batch;
|
|
89
|
+
const children = schema.fields.map((f, i) => batch.data.children[i].clone(f.type));
|
|
90
|
+
const structType = new Struct(schema.fields);
|
|
91
|
+
const data = makeData({
|
|
92
|
+
type: structType,
|
|
93
|
+
length: batch.numRows,
|
|
94
|
+
children,
|
|
95
|
+
nullCount: batch.data.nullCount,
|
|
96
|
+
nullBitmap: batch.data.nullBitmap
|
|
97
|
+
});
|
|
98
|
+
return new RecordBatch(schema, data, batch.metadata);
|
|
99
|
+
}
|
|
100
|
+
function serializeIpcStream(schema, batches) {
|
|
101
|
+
const writer = new RecordBatchStreamWriter;
|
|
102
|
+
writer.reset(undefined, schema);
|
|
103
|
+
for (const batch of batches) {
|
|
104
|
+
writer.write(conformBatchToSchema(batch, schema));
|
|
105
|
+
}
|
|
106
|
+
writer.close();
|
|
107
|
+
return writer.toUint8Array(true);
|
|
108
|
+
}
|
|
109
|
+
function arrowResponse(body, status = 200, extraHeaders) {
|
|
110
|
+
const headers = extraHeaders ?? new Headers;
|
|
111
|
+
headers.set("Content-Type", ARROW_CONTENT_TYPE);
|
|
112
|
+
return new Response(body, { status, headers });
|
|
113
|
+
}
|
|
114
|
+
async function readRequestFromBody(body) {
|
|
115
|
+
const reader = await RecordBatchReader.from(body);
|
|
116
|
+
await reader.open();
|
|
117
|
+
const schema = reader.schema;
|
|
118
|
+
if (!schema) {
|
|
119
|
+
throw new HttpRpcError("Empty IPC stream: no schema", 400);
|
|
120
|
+
}
|
|
121
|
+
const batches = reader.readAll();
|
|
122
|
+
if (batches.length === 0) {
|
|
123
|
+
throw new HttpRpcError("IPC stream contains no batches", 400);
|
|
124
|
+
}
|
|
125
|
+
return { schema, batch: batches[0] };
|
|
126
|
+
}
|
|
127
|
+
|
|
128
|
+
// src/client/introspect.ts
|
|
129
|
+
import { Schema as ArrowSchema, RecordBatchReader as RecordBatchReader4 } from "@query-farm/apache-arrow";
|
|
130
|
+
|
|
131
|
+
// src/client/ipc.ts
|
|
132
|
+
import {
|
|
133
|
+
Binary,
|
|
134
|
+
Bool,
|
|
135
|
+
DataType,
|
|
136
|
+
Float64,
|
|
137
|
+
Int64,
|
|
138
|
+
makeData as makeData2,
|
|
139
|
+
RecordBatch as RecordBatch2,
|
|
140
|
+
RecordBatchReader as RecordBatchReader3,
|
|
141
|
+
Struct as Struct2,
|
|
142
|
+
Utf8,
|
|
143
|
+
vectorFromArray
|
|
144
|
+
} from "@query-farm/apache-arrow";
|
|
145
|
+
|
|
146
|
+
// src/errors.ts
|
|
147
|
+
class RpcError extends Error {
|
|
148
|
+
errorType;
|
|
149
|
+
errorMessage;
|
|
150
|
+
remoteTraceback;
|
|
151
|
+
constructor(errorType, errorMessage, remoteTraceback) {
|
|
152
|
+
super(`${errorType}: ${errorMessage}`);
|
|
153
|
+
this.errorType = errorType;
|
|
154
|
+
this.errorMessage = errorMessage;
|
|
155
|
+
this.remoteTraceback = remoteTraceback;
|
|
156
|
+
this.name = "RpcError";
|
|
157
|
+
}
|
|
158
|
+
}
|
|
159
|
+
|
|
160
|
+
class VersionError extends Error {
|
|
161
|
+
constructor(message) {
|
|
162
|
+
super(message);
|
|
163
|
+
this.name = "VersionError";
|
|
164
|
+
}
|
|
165
|
+
}
|
|
55
166
|
|
|
56
167
|
// src/wire/reader.ts
|
|
57
|
-
import { RecordBatchReader } from "apache-arrow";
|
|
168
|
+
import { RecordBatchReader as RecordBatchReader2 } from "@query-farm/apache-arrow";
|
|
58
169
|
|
|
59
170
|
class IpcStreamReader {
|
|
60
171
|
reader;
|
|
@@ -64,7 +175,7 @@ class IpcStreamReader {
|
|
|
64
175
|
this.reader = reader;
|
|
65
176
|
}
|
|
66
177
|
static async create(input) {
|
|
67
|
-
const reader = await
|
|
178
|
+
const reader = await RecordBatchReader2.from(input);
|
|
68
179
|
await reader.open({ autoDestroy: false });
|
|
69
180
|
if (reader.closed) {
|
|
70
181
|
throw new Error("Input stream closed before first IPC message");
|
|
@@ -124,2342 +235,2274 @@ class IpcStreamReader {
|
|
|
124
235
|
}
|
|
125
236
|
}
|
|
126
237
|
|
|
127
|
-
// src/
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
offset += written;
|
|
141
|
-
} catch (e) {
|
|
142
|
-
if (e.code === "EAGAIN") {
|
|
143
|
-
Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 1);
|
|
144
|
-
continue;
|
|
145
|
-
}
|
|
146
|
-
throw e;
|
|
147
|
-
}
|
|
148
|
-
}
|
|
238
|
+
// src/client/ipc.ts
|
|
239
|
+
function inferArrowType(value) {
|
|
240
|
+
if (typeof value === "string")
|
|
241
|
+
return new Utf8;
|
|
242
|
+
if (typeof value === "boolean")
|
|
243
|
+
return new Bool;
|
|
244
|
+
if (typeof value === "bigint")
|
|
245
|
+
return new Int64;
|
|
246
|
+
if (typeof value === "number")
|
|
247
|
+
return new Float64;
|
|
248
|
+
if (value instanceof Uint8Array)
|
|
249
|
+
return new Binary;
|
|
250
|
+
return new Utf8;
|
|
149
251
|
}
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
252
|
+
function coerceForArrow(type, value) {
|
|
253
|
+
if (value == null)
|
|
254
|
+
return value;
|
|
255
|
+
if (DataType.isInt(type) && type.bitWidth === 64) {
|
|
256
|
+
if (typeof value === "number")
|
|
257
|
+
return BigInt(value);
|
|
258
|
+
return value;
|
|
155
259
|
}
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
260
|
+
if (DataType.isMap(type)) {
|
|
261
|
+
if (value instanceof Map) {
|
|
262
|
+
const entriesField = type.children[0];
|
|
263
|
+
const valueType = entriesField.type.children[1].type;
|
|
264
|
+
const coerced = new Map;
|
|
265
|
+
for (const [k, v] of value) {
|
|
266
|
+
coerced.set(k, coerceForArrow(valueType, v));
|
|
267
|
+
}
|
|
268
|
+
return coerced;
|
|
161
269
|
}
|
|
162
|
-
|
|
163
|
-
const bytes = writer.toUint8Array(true);
|
|
164
|
-
writeAll(this.fd, bytes);
|
|
270
|
+
return value;
|
|
165
271
|
}
|
|
166
|
-
|
|
167
|
-
|
|
272
|
+
if (DataType.isList(type)) {
|
|
273
|
+
if (Array.isArray(value)) {
|
|
274
|
+
const elemType = type.children[0].type;
|
|
275
|
+
return value.map((v) => coerceForArrow(elemType, v));
|
|
276
|
+
}
|
|
277
|
+
return value;
|
|
168
278
|
}
|
|
279
|
+
return value;
|
|
169
280
|
}
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
281
|
+
function buildRequestIpc(schema, params, method) {
|
|
282
|
+
const metadata = new Map;
|
|
283
|
+
metadata.set(RPC_METHOD_KEY, method);
|
|
284
|
+
metadata.set(REQUEST_VERSION_KEY, REQUEST_VERSION);
|
|
285
|
+
if (schema.fields.length === 0) {
|
|
286
|
+
const structType2 = new Struct2(schema.fields);
|
|
287
|
+
const data2 = makeData2({
|
|
288
|
+
type: structType2,
|
|
289
|
+
length: 1,
|
|
290
|
+
children: [],
|
|
291
|
+
nullCount: 0
|
|
292
|
+
});
|
|
293
|
+
const batch2 = new RecordBatch2(schema, data2, metadata);
|
|
294
|
+
return serializeIpcStream(schema, [batch2]);
|
|
180
295
|
}
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
296
|
+
const children = schema.fields.map((f) => {
|
|
297
|
+
const val = coerceForArrow(f.type, params[f.name]);
|
|
298
|
+
return vectorFromArray([val], f.type).data[0];
|
|
299
|
+
});
|
|
300
|
+
const structType = new Struct2(schema.fields);
|
|
301
|
+
const data = makeData2({
|
|
302
|
+
type: structType,
|
|
303
|
+
length: 1,
|
|
304
|
+
children,
|
|
305
|
+
nullCount: 0
|
|
306
|
+
});
|
|
307
|
+
const batch = new RecordBatch2(schema, data, metadata);
|
|
308
|
+
return serializeIpcStream(schema, [batch]);
|
|
309
|
+
}
|
|
310
|
+
async function readResponseBatches(body) {
|
|
311
|
+
const reader = await RecordBatchReader3.from(body);
|
|
312
|
+
await reader.open();
|
|
313
|
+
const schema = reader.schema;
|
|
314
|
+
if (!schema) {
|
|
315
|
+
throw new RpcError("ProtocolError", "Empty IPC stream: no schema", "");
|
|
186
316
|
}
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
317
|
+
const batches = reader.readAll();
|
|
318
|
+
return { schema, batches };
|
|
319
|
+
}
|
|
320
|
+
function dispatchLogOrError(batch, onLog) {
|
|
321
|
+
const meta = batch.metadata;
|
|
322
|
+
if (!meta)
|
|
323
|
+
return false;
|
|
324
|
+
const level = meta.get(LOG_LEVEL_KEY);
|
|
325
|
+
if (!level)
|
|
326
|
+
return false;
|
|
327
|
+
const message = meta.get(LOG_MESSAGE_KEY) ?? "";
|
|
328
|
+
if (level === "EXCEPTION") {
|
|
329
|
+
const extraStr = meta.get(LOG_EXTRA_KEY);
|
|
330
|
+
let errorType = "RpcError";
|
|
331
|
+
let errorMessage = message;
|
|
332
|
+
let traceback = "";
|
|
333
|
+
if (extraStr) {
|
|
334
|
+
try {
|
|
335
|
+
const extra = JSON.parse(extraStr);
|
|
336
|
+
errorType = extra.exception_type ?? "RpcError";
|
|
337
|
+
errorMessage = extra.exception_message ?? message;
|
|
338
|
+
traceback = extra.traceback ?? "";
|
|
339
|
+
} catch {}
|
|
340
|
+
}
|
|
341
|
+
throw new RpcError(errorType, errorMessage, traceback);
|
|
193
342
|
}
|
|
194
|
-
|
|
195
|
-
const
|
|
196
|
-
|
|
197
|
-
|
|
343
|
+
if (onLog) {
|
|
344
|
+
const extraStr = meta.get(LOG_EXTRA_KEY);
|
|
345
|
+
let extra;
|
|
346
|
+
if (extraStr) {
|
|
347
|
+
try {
|
|
348
|
+
extra = JSON.parse(extraStr);
|
|
349
|
+
} catch {}
|
|
198
350
|
}
|
|
199
|
-
|
|
351
|
+
onLog({ level, message, extra });
|
|
200
352
|
}
|
|
353
|
+
return true;
|
|
201
354
|
}
|
|
202
|
-
|
|
203
|
-
|
|
204
|
-
|
|
205
|
-
|
|
206
|
-
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
|
|
210
|
-
|
|
211
|
-
|
|
212
|
-
|
|
213
|
-
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
|
|
217
|
-
var DESCRIBE_VERSION = "2";
|
|
218
|
-
var DESCRIBE_METHOD_NAME = "__describe__";
|
|
219
|
-
var STATE_KEY = "vgi_rpc.stream_state#b64";
|
|
220
|
-
|
|
221
|
-
// src/errors.ts
|
|
222
|
-
class RpcError extends Error {
|
|
223
|
-
errorType;
|
|
224
|
-
errorMessage;
|
|
225
|
-
remoteTraceback;
|
|
226
|
-
constructor(errorType, errorMessage, remoteTraceback) {
|
|
227
|
-
super(`${errorType}: ${errorMessage}`);
|
|
228
|
-
this.errorType = errorType;
|
|
229
|
-
this.errorMessage = errorMessage;
|
|
230
|
-
this.remoteTraceback = remoteTraceback;
|
|
231
|
-
this.name = "RpcError";
|
|
355
|
+
function extractBatchRows(batch) {
|
|
356
|
+
const rows = [];
|
|
357
|
+
for (let r = 0;r < batch.numRows; r++) {
|
|
358
|
+
const row = {};
|
|
359
|
+
for (let i = 0;i < batch.schema.fields.length; i++) {
|
|
360
|
+
const field = batch.schema.fields[i];
|
|
361
|
+
let value = batch.getChildAt(i)?.get(r);
|
|
362
|
+
if (typeof value === "bigint") {
|
|
363
|
+
if (value >= BigInt(Number.MIN_SAFE_INTEGER) && value <= BigInt(Number.MAX_SAFE_INTEGER)) {
|
|
364
|
+
value = Number(value);
|
|
365
|
+
}
|
|
366
|
+
}
|
|
367
|
+
row[field.name] = value;
|
|
368
|
+
}
|
|
369
|
+
rows.push(row);
|
|
232
370
|
}
|
|
371
|
+
return rows;
|
|
233
372
|
}
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
|
|
239
|
-
|
|
373
|
+
async function readSequentialStreams(body) {
|
|
374
|
+
const stream = new ReadableStream({
|
|
375
|
+
start(controller) {
|
|
376
|
+
controller.enqueue(body);
|
|
377
|
+
controller.close();
|
|
378
|
+
}
|
|
379
|
+
});
|
|
380
|
+
return IpcStreamReader.create(stream);
|
|
240
381
|
}
|
|
241
382
|
|
|
242
|
-
// src/
|
|
243
|
-
function
|
|
244
|
-
const
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
|
|
252
|
-
|
|
253
|
-
|
|
254
|
-
|
|
383
|
+
// src/client/introspect.ts
|
|
384
|
+
async function deserializeSchema(bytes) {
|
|
385
|
+
const reader = await RecordBatchReader4.from(bytes);
|
|
386
|
+
await reader.open();
|
|
387
|
+
return reader.schema;
|
|
388
|
+
}
|
|
389
|
+
async function parseDescribeResponse(batches, onLog) {
|
|
390
|
+
let dataBatch = null;
|
|
391
|
+
for (const batch of batches) {
|
|
392
|
+
if (batch.numRows === 0) {
|
|
393
|
+
dispatchLogOrError(batch, onLog);
|
|
394
|
+
continue;
|
|
395
|
+
}
|
|
396
|
+
dataBatch = batch;
|
|
255
397
|
}
|
|
256
|
-
|
|
257
|
-
|
|
258
|
-
if (schema.fields.length > 0 && batch.numRows !== 1) {
|
|
259
|
-
throw new RpcError("ProtocolError", `Expected 1 row in request batch, got ${batch.numRows}. ` + "Each parameter is a column (not a row). The batch should have exactly 1 row.", "");
|
|
398
|
+
if (!dataBatch) {
|
|
399
|
+
throw new Error("Empty __describe__ response");
|
|
260
400
|
}
|
|
261
|
-
|
|
262
|
-
|
|
263
|
-
|
|
264
|
-
|
|
265
|
-
|
|
401
|
+
const meta = dataBatch.metadata;
|
|
402
|
+
const protocolName = meta?.get(PROTOCOL_NAME_KEY) ?? "";
|
|
403
|
+
const methods = [];
|
|
404
|
+
for (let i = 0;i < dataBatch.numRows; i++) {
|
|
405
|
+
const name = dataBatch.getChildAt(0).get(i);
|
|
406
|
+
const methodType = dataBatch.getChildAt(1).get(i);
|
|
407
|
+
const doc = dataBatch.getChildAt(2)?.get(i);
|
|
408
|
+
const _hasReturn = dataBatch.getChildAt(3).get(i);
|
|
409
|
+
const paramsIpc = dataBatch.getChildAt(4).get(i);
|
|
410
|
+
const resultIpc = dataBatch.getChildAt(5).get(i);
|
|
411
|
+
const paramTypesJson = dataBatch.getChildAt(6)?.get(i);
|
|
412
|
+
const paramDefaultsJson = dataBatch.getChildAt(7)?.get(i);
|
|
413
|
+
const hasHeader = dataBatch.getChildAt(8).get(i);
|
|
414
|
+
const headerIpc = dataBatch.getChildAt(9)?.get(i);
|
|
415
|
+
const paramsSchema = await deserializeSchema(paramsIpc);
|
|
416
|
+
const resultSchema = await deserializeSchema(resultIpc);
|
|
417
|
+
let paramTypes;
|
|
418
|
+
if (paramTypesJson) {
|
|
419
|
+
try {
|
|
420
|
+
paramTypes = JSON.parse(paramTypesJson);
|
|
421
|
+
} catch {}
|
|
266
422
|
}
|
|
267
|
-
let
|
|
268
|
-
if (
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
}
|
|
423
|
+
let defaults;
|
|
424
|
+
if (paramDefaultsJson) {
|
|
425
|
+
try {
|
|
426
|
+
defaults = JSON.parse(paramDefaultsJson);
|
|
427
|
+
} catch {}
|
|
272
428
|
}
|
|
273
|
-
|
|
429
|
+
const info = {
|
|
430
|
+
name,
|
|
431
|
+
type: methodType,
|
|
432
|
+
paramsSchema,
|
|
433
|
+
resultSchema,
|
|
434
|
+
doc: doc ?? undefined,
|
|
435
|
+
paramTypes,
|
|
436
|
+
defaults
|
|
437
|
+
};
|
|
438
|
+
if (methodType === "stream") {
|
|
439
|
+
info.outputSchema = resultSchema;
|
|
440
|
+
}
|
|
441
|
+
if (hasHeader && headerIpc) {
|
|
442
|
+
info.headerSchema = await deserializeSchema(headerIpc);
|
|
443
|
+
}
|
|
444
|
+
methods.push(info);
|
|
274
445
|
}
|
|
275
|
-
return {
|
|
276
|
-
|
|
277
|
-
|
|
278
|
-
|
|
279
|
-
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
|
|
446
|
+
return { protocolName, methods };
|
|
447
|
+
}
|
|
448
|
+
async function httpIntrospect(baseUrl, options) {
|
|
449
|
+
const prefix = options?.prefix ?? "/vgi";
|
|
450
|
+
const emptySchema = new ArrowSchema([]);
|
|
451
|
+
const body = buildRequestIpc(emptySchema, {}, DESCRIBE_METHOD_NAME);
|
|
452
|
+
const response = await fetch(`${baseUrl}${prefix}/${DESCRIBE_METHOD_NAME}`, {
|
|
453
|
+
method: "POST",
|
|
454
|
+
headers: { "Content-Type": ARROW_CONTENT_TYPE },
|
|
455
|
+
body
|
|
456
|
+
});
|
|
457
|
+
const responseBody = new Uint8Array(await response.arrayBuffer());
|
|
458
|
+
const { batches } = await readResponseBatches(responseBody);
|
|
459
|
+
return parseDescribeResponse(batches);
|
|
283
460
|
}
|
|
284
461
|
|
|
285
|
-
// src/
|
|
286
|
-
import {
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
291
|
-
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
|
|
301
|
-
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
462
|
+
// src/client/stream.ts
|
|
463
|
+
import { Field, makeData as makeData3, RecordBatch as RecordBatch3, Schema, Struct as Struct3, vectorFromArray as vectorFromArray2 } from "@query-farm/apache-arrow";
|
|
464
|
+
class HttpStreamSession {
|
|
465
|
+
_baseUrl;
|
|
466
|
+
_prefix;
|
|
467
|
+
_method;
|
|
468
|
+
_stateToken;
|
|
469
|
+
_outputSchema;
|
|
470
|
+
_inputSchema;
|
|
471
|
+
_onLog;
|
|
472
|
+
_pendingBatches;
|
|
473
|
+
_finished;
|
|
474
|
+
_header;
|
|
475
|
+
_compressionLevel;
|
|
476
|
+
_compressFn;
|
|
477
|
+
_decompressFn;
|
|
478
|
+
constructor(opts) {
|
|
479
|
+
this._baseUrl = opts.baseUrl;
|
|
480
|
+
this._prefix = opts.prefix;
|
|
481
|
+
this._method = opts.method;
|
|
482
|
+
this._stateToken = opts.stateToken;
|
|
483
|
+
this._outputSchema = opts.outputSchema;
|
|
484
|
+
this._inputSchema = opts.inputSchema;
|
|
485
|
+
this._onLog = opts.onLog;
|
|
486
|
+
this._pendingBatches = opts.pendingBatches;
|
|
487
|
+
this._finished = opts.finished;
|
|
488
|
+
this._header = opts.header;
|
|
489
|
+
this._compressionLevel = opts.compressionLevel;
|
|
490
|
+
this._compressFn = opts.compressFn;
|
|
491
|
+
this._decompressFn = opts.decompressFn;
|
|
307
492
|
}
|
|
308
|
-
|
|
309
|
-
|
|
310
|
-
function buildResultBatch(schema, values, serverId, requestId) {
|
|
311
|
-
const metadata = new Map;
|
|
312
|
-
metadata.set(SERVER_ID_KEY, serverId);
|
|
313
|
-
if (requestId !== null) {
|
|
314
|
-
metadata.set(REQUEST_ID_KEY, requestId);
|
|
493
|
+
get header() {
|
|
494
|
+
return this._header;
|
|
315
495
|
}
|
|
316
|
-
|
|
317
|
-
|
|
496
|
+
_buildHeaders() {
|
|
497
|
+
const headers = {
|
|
498
|
+
"Content-Type": ARROW_CONTENT_TYPE
|
|
499
|
+
};
|
|
500
|
+
if (this._compressionLevel != null) {
|
|
501
|
+
headers["Content-Encoding"] = "zstd";
|
|
502
|
+
headers["Accept-Encoding"] = "zstd";
|
|
503
|
+
}
|
|
504
|
+
return headers;
|
|
318
505
|
}
|
|
319
|
-
|
|
320
|
-
if (
|
|
321
|
-
|
|
322
|
-
throw new TypeError(`Handler result missing required field '${field.name}'. Got keys: [${got.join(", ")}]`);
|
|
506
|
+
_prepareBody(content) {
|
|
507
|
+
if (this._compressionLevel != null && this._compressFn) {
|
|
508
|
+
return this._compressFn(content, this._compressionLevel);
|
|
323
509
|
}
|
|
510
|
+
return content;
|
|
324
511
|
}
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
return val;
|
|
512
|
+
async _readResponse(resp) {
|
|
513
|
+
let body = new Uint8Array(await resp.arrayBuffer());
|
|
514
|
+
if (resp.headers.get("Content-Encoding") === "zstd" && this._decompressFn) {
|
|
515
|
+
body = new Uint8Array(this._decompressFn(body));
|
|
330
516
|
}
|
|
331
|
-
|
|
332
|
-
return arr.data[0];
|
|
333
|
-
});
|
|
334
|
-
const structType = new Struct(schema.fields);
|
|
335
|
-
const data = makeData({
|
|
336
|
-
type: structType,
|
|
337
|
-
length: 1,
|
|
338
|
-
children,
|
|
339
|
-
nullCount: 0
|
|
340
|
-
});
|
|
341
|
-
return new RecordBatch(schema, data, metadata);
|
|
342
|
-
}
|
|
343
|
-
function buildErrorBatch(schema, error, serverId, requestId) {
|
|
344
|
-
const metadata = new Map;
|
|
345
|
-
metadata.set(LOG_LEVEL_KEY, "EXCEPTION");
|
|
346
|
-
metadata.set(LOG_MESSAGE_KEY, `${error.constructor.name}: ${error.message}`);
|
|
347
|
-
const extra = {
|
|
348
|
-
exception_type: error.constructor.name,
|
|
349
|
-
exception_message: error.message,
|
|
350
|
-
traceback: error.stack ?? ""
|
|
351
|
-
};
|
|
352
|
-
metadata.set(LOG_EXTRA_KEY, JSON.stringify(extra));
|
|
353
|
-
metadata.set(SERVER_ID_KEY, serverId);
|
|
354
|
-
if (requestId !== null) {
|
|
355
|
-
metadata.set(REQUEST_ID_KEY, requestId);
|
|
356
|
-
}
|
|
357
|
-
return buildEmptyBatch(schema, metadata);
|
|
358
|
-
}
|
|
359
|
-
function buildLogBatch(schema, level, message, extra, serverId, requestId) {
|
|
360
|
-
const metadata = new Map;
|
|
361
|
-
metadata.set(LOG_LEVEL_KEY, level);
|
|
362
|
-
metadata.set(LOG_MESSAGE_KEY, message);
|
|
363
|
-
if (extra) {
|
|
364
|
-
metadata.set(LOG_EXTRA_KEY, JSON.stringify(extra));
|
|
517
|
+
return body;
|
|
365
518
|
}
|
|
366
|
-
|
|
367
|
-
|
|
519
|
+
async exchange(input) {
|
|
520
|
+
if (this._stateToken === null) {
|
|
521
|
+
throw new RpcError("ProtocolError", "Stream has finished — no state token available", "");
|
|
522
|
+
}
|
|
523
|
+
if (input.length === 0) {
|
|
524
|
+
const zeroSchema = this._inputSchema ?? this._outputSchema;
|
|
525
|
+
const emptyBatch = this._buildEmptyBatch(zeroSchema);
|
|
526
|
+
const metadata2 = new Map;
|
|
527
|
+
metadata2.set(STATE_KEY, this._stateToken);
|
|
528
|
+
const batchWithMeta = new RecordBatch3(zeroSchema, emptyBatch.data, metadata2);
|
|
529
|
+
return this._doExchange(zeroSchema, [batchWithMeta]);
|
|
530
|
+
}
|
|
531
|
+
const keys = Object.keys(input[0]);
|
|
532
|
+
const fields = keys.map((key) => {
|
|
533
|
+
let sample;
|
|
534
|
+
for (const row of input) {
|
|
535
|
+
if (row[key] != null) {
|
|
536
|
+
sample = row[key];
|
|
537
|
+
break;
|
|
538
|
+
}
|
|
539
|
+
}
|
|
540
|
+
const arrowType = inferArrowType(sample);
|
|
541
|
+
const nullable = input.some((row) => row[key] == null);
|
|
542
|
+
return new Field(key, arrowType, nullable);
|
|
543
|
+
});
|
|
544
|
+
const inputSchema = new Schema(fields);
|
|
545
|
+
const children = inputSchema.fields.map((f) => {
|
|
546
|
+
const values = input.map((row) => row[f.name]);
|
|
547
|
+
return vectorFromArray2(values, f.type).data[0];
|
|
548
|
+
});
|
|
549
|
+
const structType = new Struct3(inputSchema.fields);
|
|
550
|
+
const data = makeData3({
|
|
551
|
+
type: structType,
|
|
552
|
+
length: input.length,
|
|
553
|
+
children,
|
|
554
|
+
nullCount: 0
|
|
555
|
+
});
|
|
556
|
+
const metadata = new Map;
|
|
557
|
+
metadata.set(STATE_KEY, this._stateToken);
|
|
558
|
+
const batch = new RecordBatch3(inputSchema, data, metadata);
|
|
559
|
+
return this._doExchange(inputSchema, [batch]);
|
|
368
560
|
}
|
|
369
|
-
|
|
370
|
-
|
|
561
|
+
async _doExchange(schema, batches) {
|
|
562
|
+
const body = serializeIpcStream(schema, batches);
|
|
563
|
+
const resp = await fetch(`${this._baseUrl}${this._prefix}/${this._method}/exchange`, {
|
|
564
|
+
method: "POST",
|
|
565
|
+
headers: this._buildHeaders(),
|
|
566
|
+
body: this._prepareBody(body)
|
|
567
|
+
});
|
|
568
|
+
const responseBody = await this._readResponse(resp);
|
|
569
|
+
const { batches: responseBatches } = await readResponseBatches(responseBody);
|
|
570
|
+
let resultRows = [];
|
|
571
|
+
for (const batch of responseBatches) {
|
|
572
|
+
if (batch.numRows === 0) {
|
|
573
|
+
dispatchLogOrError(batch, this._onLog);
|
|
574
|
+
const token2 = batch.metadata?.get(STATE_KEY);
|
|
575
|
+
if (token2) {
|
|
576
|
+
this._stateToken = token2;
|
|
577
|
+
}
|
|
578
|
+
continue;
|
|
579
|
+
}
|
|
580
|
+
const token = batch.metadata?.get(STATE_KEY);
|
|
581
|
+
if (token) {
|
|
582
|
+
this._stateToken = token;
|
|
583
|
+
}
|
|
584
|
+
resultRows = extractBatchRows(batch);
|
|
585
|
+
}
|
|
586
|
+
return resultRows;
|
|
371
587
|
}
|
|
372
|
-
|
|
373
|
-
|
|
374
|
-
|
|
375
|
-
|
|
376
|
-
|
|
377
|
-
|
|
378
|
-
|
|
379
|
-
const structType2 = new Struct(schema.fields);
|
|
380
|
-
const data2 = makeData({
|
|
381
|
-
type: structType2,
|
|
588
|
+
_buildEmptyBatch(schema) {
|
|
589
|
+
const children = schema.fields.map((f) => {
|
|
590
|
+
return makeData3({ type: f.type, length: 0, nullCount: 0 });
|
|
591
|
+
});
|
|
592
|
+
const structType = new Struct3(schema.fields);
|
|
593
|
+
const data = makeData3({
|
|
594
|
+
type: structType,
|
|
382
595
|
length: 0,
|
|
383
|
-
children
|
|
596
|
+
children,
|
|
384
597
|
nullCount: 0
|
|
385
598
|
});
|
|
386
|
-
return new
|
|
599
|
+
return new RecordBatch3(schema, data);
|
|
387
600
|
}
|
|
388
|
-
|
|
389
|
-
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
395
|
-
return new RecordBatch(schema, data, metadata);
|
|
396
|
-
}
|
|
397
|
-
|
|
398
|
-
// src/dispatch/describe.ts
|
|
399
|
-
import {
|
|
400
|
-
Schema as Schema2,
|
|
401
|
-
Field as Field2,
|
|
402
|
-
RecordBatch as RecordBatch2,
|
|
403
|
-
Utf8,
|
|
404
|
-
Bool,
|
|
405
|
-
Binary,
|
|
406
|
-
vectorFromArray as vectorFromArray2,
|
|
407
|
-
makeData as makeData2,
|
|
408
|
-
Struct as Struct2
|
|
409
|
-
} from "apache-arrow";
|
|
410
|
-
|
|
411
|
-
// src/util/schema.ts
|
|
412
|
-
import { RecordBatchStreamWriter as RecordBatchStreamWriter2 } from "apache-arrow";
|
|
413
|
-
function serializeSchema(schema) {
|
|
414
|
-
const writer = new RecordBatchStreamWriter2;
|
|
415
|
-
writer.reset(undefined, schema);
|
|
416
|
-
writer.close();
|
|
417
|
-
return writer.toUint8Array(true);
|
|
418
|
-
}
|
|
419
|
-
|
|
420
|
-
// src/dispatch/describe.ts
|
|
421
|
-
var DESCRIBE_SCHEMA = new Schema2([
|
|
422
|
-
new Field2("name", new Utf8, false),
|
|
423
|
-
new Field2("method_type", new Utf8, false),
|
|
424
|
-
new Field2("doc", new Utf8, true),
|
|
425
|
-
new Field2("has_return", new Bool, false),
|
|
426
|
-
new Field2("params_schema_ipc", new Binary, false),
|
|
427
|
-
new Field2("result_schema_ipc", new Binary, false),
|
|
428
|
-
new Field2("param_types_json", new Utf8, true),
|
|
429
|
-
new Field2("param_defaults_json", new Utf8, true),
|
|
430
|
-
new Field2("has_header", new Bool, false),
|
|
431
|
-
new Field2("header_schema_ipc", new Binary, true)
|
|
432
|
-
]);
|
|
433
|
-
function buildDescribeBatch(protocolName, methods, serverId) {
|
|
434
|
-
const sortedEntries = [...methods.entries()].sort(([a], [b]) => a.localeCompare(b));
|
|
435
|
-
const names = [];
|
|
436
|
-
const methodTypes = [];
|
|
437
|
-
const docs = [];
|
|
438
|
-
const hasReturns = [];
|
|
439
|
-
const paramsSchemas = [];
|
|
440
|
-
const resultSchemas = [];
|
|
441
|
-
const paramTypesJsons = [];
|
|
442
|
-
const paramDefaultsJsons = [];
|
|
443
|
-
const hasHeaders = [];
|
|
444
|
-
const headerSchemas = [];
|
|
445
|
-
for (const [name, method] of sortedEntries) {
|
|
446
|
-
names.push(name);
|
|
447
|
-
methodTypes.push(method.type);
|
|
448
|
-
docs.push(method.doc ?? null);
|
|
449
|
-
const hasReturn = method.type === "unary" && method.resultSchema.fields.length > 0;
|
|
450
|
-
hasReturns.push(hasReturn);
|
|
451
|
-
paramsSchemas.push(serializeSchema(method.paramsSchema));
|
|
452
|
-
resultSchemas.push(serializeSchema(method.resultSchema));
|
|
453
|
-
if (method.paramTypes && Object.keys(method.paramTypes).length > 0) {
|
|
454
|
-
paramTypesJsons.push(JSON.stringify(method.paramTypes));
|
|
455
|
-
} else {
|
|
456
|
-
paramTypesJsons.push(null);
|
|
601
|
+
async* [Symbol.asyncIterator]() {
|
|
602
|
+
for (const batch of this._pendingBatches) {
|
|
603
|
+
if (batch.numRows === 0) {
|
|
604
|
+
dispatchLogOrError(batch, this._onLog);
|
|
605
|
+
continue;
|
|
606
|
+
}
|
|
607
|
+
yield extractBatchRows(batch);
|
|
457
608
|
}
|
|
458
|
-
|
|
459
|
-
|
|
460
|
-
|
|
461
|
-
|
|
462
|
-
|
|
609
|
+
this._pendingBatches = [];
|
|
610
|
+
if (this._finished)
|
|
611
|
+
return;
|
|
612
|
+
if (this._stateToken === null)
|
|
613
|
+
return;
|
|
614
|
+
while (true) {
|
|
615
|
+
const responseBody = await this._sendContinuation(this._stateToken);
|
|
616
|
+
const { batches } = await readResponseBatches(responseBody);
|
|
617
|
+
let gotContinuation = false;
|
|
618
|
+
for (const batch of batches) {
|
|
619
|
+
if (batch.numRows === 0) {
|
|
620
|
+
const token = batch.metadata?.get(STATE_KEY);
|
|
621
|
+
if (token) {
|
|
622
|
+
this._stateToken = token;
|
|
623
|
+
gotContinuation = true;
|
|
624
|
+
continue;
|
|
625
|
+
}
|
|
626
|
+
dispatchLogOrError(batch, this._onLog);
|
|
627
|
+
continue;
|
|
463
628
|
}
|
|
629
|
+
yield extractBatchRows(batch);
|
|
464
630
|
}
|
|
465
|
-
|
|
466
|
-
|
|
467
|
-
paramDefaultsJsons.push(null);
|
|
631
|
+
if (!gotContinuation)
|
|
632
|
+
break;
|
|
468
633
|
}
|
|
469
|
-
hasHeaders.push(!!method.headerSchema);
|
|
470
|
-
headerSchemas.push(method.headerSchema ? serializeSchema(method.headerSchema) : null);
|
|
471
634
|
}
|
|
472
|
-
|
|
473
|
-
|
|
474
|
-
|
|
475
|
-
|
|
476
|
-
|
|
477
|
-
|
|
478
|
-
|
|
479
|
-
|
|
480
|
-
|
|
481
|
-
|
|
482
|
-
|
|
483
|
-
|
|
484
|
-
|
|
485
|
-
|
|
486
|
-
|
|
487
|
-
|
|
488
|
-
|
|
489
|
-
|
|
490
|
-
|
|
491
|
-
|
|
492
|
-
|
|
493
|
-
];
|
|
494
|
-
const structType = new Struct2(DESCRIBE_SCHEMA.fields);
|
|
495
|
-
const data = makeData2({
|
|
496
|
-
type: structType,
|
|
497
|
-
length: sortedEntries.length,
|
|
498
|
-
children,
|
|
499
|
-
nullCount: 0
|
|
500
|
-
});
|
|
501
|
-
const metadata = new Map;
|
|
502
|
-
metadata.set(PROTOCOL_NAME_KEY, protocolName);
|
|
503
|
-
metadata.set(REQUEST_VERSION_KEY, REQUEST_VERSION);
|
|
504
|
-
metadata.set(DESCRIBE_VERSION_KEY, DESCRIBE_VERSION);
|
|
505
|
-
metadata.set(SERVER_ID_KEY, serverId);
|
|
506
|
-
const batch = new RecordBatch2(DESCRIBE_SCHEMA, data, metadata);
|
|
507
|
-
return { batch, metadata };
|
|
635
|
+
async _sendContinuation(token) {
|
|
636
|
+
const emptySchema = new Schema([]);
|
|
637
|
+
const metadata = new Map;
|
|
638
|
+
metadata.set(STATE_KEY, token);
|
|
639
|
+
const structType = new Struct3(emptySchema.fields);
|
|
640
|
+
const data = makeData3({
|
|
641
|
+
type: structType,
|
|
642
|
+
length: 1,
|
|
643
|
+
children: [],
|
|
644
|
+
nullCount: 0
|
|
645
|
+
});
|
|
646
|
+
const batch = new RecordBatch3(emptySchema, data, metadata);
|
|
647
|
+
const body = serializeIpcStream(emptySchema, [batch]);
|
|
648
|
+
const resp = await fetch(`${this._baseUrl}${this._prefix}/${this._method}/exchange`, {
|
|
649
|
+
method: "POST",
|
|
650
|
+
headers: this._buildHeaders(),
|
|
651
|
+
body: this._prepareBody(body)
|
|
652
|
+
});
|
|
653
|
+
return this._readResponse(resp);
|
|
654
|
+
}
|
|
655
|
+
close() {}
|
|
508
656
|
}
|
|
509
657
|
|
|
510
|
-
// src/
|
|
511
|
-
|
|
512
|
-
|
|
513
|
-
|
|
514
|
-
|
|
515
|
-
|
|
516
|
-
|
|
517
|
-
|
|
518
|
-
|
|
519
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
523
|
-
|
|
524
|
-
|
|
525
|
-
|
|
526
|
-
|
|
527
|
-
|
|
528
|
-
this._producerMode = producerMode;
|
|
529
|
-
this._serverId = serverId;
|
|
530
|
-
this._requestId = requestId;
|
|
531
|
-
}
|
|
532
|
-
get outputSchema() {
|
|
533
|
-
return this._outputSchema;
|
|
534
|
-
}
|
|
535
|
-
get finished() {
|
|
536
|
-
return this._finished;
|
|
537
|
-
}
|
|
538
|
-
get batches() {
|
|
539
|
-
return this._batches;
|
|
540
|
-
}
|
|
541
|
-
emit(batchOrColumns, metadata) {
|
|
542
|
-
let batch;
|
|
543
|
-
if (batchOrColumns instanceof RecordBatch3) {
|
|
544
|
-
batch = batchOrColumns;
|
|
545
|
-
} else {
|
|
546
|
-
const coerced = coerceInt64(this._outputSchema, batchOrColumns);
|
|
547
|
-
batch = recordBatchFromArrays(coerced, this._outputSchema);
|
|
548
|
-
}
|
|
549
|
-
if (this._dataBatchIdx !== null) {
|
|
550
|
-
throw new Error("Only one data batch may be emitted per call");
|
|
551
|
-
}
|
|
552
|
-
this._dataBatchIdx = this._batches.length;
|
|
553
|
-
this._batches.push({ batch, metadata });
|
|
554
|
-
}
|
|
555
|
-
emitRow(values) {
|
|
556
|
-
const columns = {};
|
|
557
|
-
for (const [key, value] of Object.entries(values)) {
|
|
558
|
-
columns[key] = [value];
|
|
559
|
-
}
|
|
560
|
-
this.emit(columns);
|
|
658
|
+
// src/client/connect.ts
|
|
659
|
+
function httpConnect(baseUrl, options) {
|
|
660
|
+
const prefix = (options?.prefix ?? "/vgi").replace(/\/+$/, "");
|
|
661
|
+
const onLog = options?.onLog;
|
|
662
|
+
const compressionLevel = options?.compressionLevel;
|
|
663
|
+
let methodCache = null;
|
|
664
|
+
let compressFn;
|
|
665
|
+
let decompressFn;
|
|
666
|
+
let compressionLoaded = false;
|
|
667
|
+
async function ensureCompression() {
|
|
668
|
+
if (compressionLoaded || compressionLevel == null)
|
|
669
|
+
return;
|
|
670
|
+
compressionLoaded = true;
|
|
671
|
+
try {
|
|
672
|
+
const mod = await Promise.resolve().then(() => (init_zstd(), exports_zstd));
|
|
673
|
+
compressFn = mod.zstdCompress;
|
|
674
|
+
decompressFn = mod.zstdDecompress;
|
|
675
|
+
} catch {}
|
|
561
676
|
}
|
|
562
|
-
|
|
563
|
-
|
|
564
|
-
|
|
677
|
+
function buildHeaders() {
|
|
678
|
+
const headers = {
|
|
679
|
+
"Content-Type": ARROW_CONTENT_TYPE
|
|
680
|
+
};
|
|
681
|
+
if (compressionLevel != null) {
|
|
682
|
+
headers["Content-Encoding"] = "zstd";
|
|
683
|
+
headers["Accept-Encoding"] = "zstd";
|
|
565
684
|
}
|
|
566
|
-
|
|
567
|
-
}
|
|
568
|
-
clientLog(level, message, extra) {
|
|
569
|
-
const batch = buildLogBatch(this._outputSchema, level, message, extra, this._serverId, this._requestId);
|
|
570
|
-
this._batches.push({ batch });
|
|
571
|
-
}
|
|
572
|
-
}
|
|
573
|
-
|
|
574
|
-
// src/dispatch/unary.ts
|
|
575
|
-
async function dispatchUnary(method, params, writer, serverId, requestId) {
|
|
576
|
-
const schema = method.resultSchema;
|
|
577
|
-
const out = new OutputCollector(schema, true, serverId, requestId);
|
|
578
|
-
try {
|
|
579
|
-
const result = await method.handler(params, out);
|
|
580
|
-
const resultBatch = buildResultBatch(schema, result, serverId, requestId);
|
|
581
|
-
const batches = [...out.batches.map((b) => b.batch), resultBatch];
|
|
582
|
-
writer.writeStream(schema, batches);
|
|
583
|
-
} catch (error) {
|
|
584
|
-
const batch = buildErrorBatch(schema, error, serverId, requestId);
|
|
585
|
-
writer.writeStream(schema, [batch]);
|
|
685
|
+
return headers;
|
|
586
686
|
}
|
|
587
|
-
|
|
588
|
-
|
|
589
|
-
|
|
590
|
-
import { Schema as Schema3 } from "apache-arrow";
|
|
591
|
-
var EMPTY_SCHEMA = new Schema3([]);
|
|
592
|
-
async function dispatchStream(method, params, writer, reader, serverId, requestId) {
|
|
593
|
-
const isProducer = !method.inputSchema || method.inputSchema.fields.length === 0;
|
|
594
|
-
let state;
|
|
595
|
-
try {
|
|
596
|
-
if (isProducer) {
|
|
597
|
-
state = await method.producerInit(params);
|
|
598
|
-
} else {
|
|
599
|
-
state = await method.exchangeInit(params);
|
|
600
|
-
}
|
|
601
|
-
} catch (error) {
|
|
602
|
-
const errSchema = method.headerSchema ?? EMPTY_SCHEMA;
|
|
603
|
-
const errBatch = buildErrorBatch(errSchema, error, serverId, requestId);
|
|
604
|
-
writer.writeStream(errSchema, [errBatch]);
|
|
605
|
-
const inputSchema2 = await reader.openNextStream();
|
|
606
|
-
if (inputSchema2) {
|
|
607
|
-
while (await reader.readNextBatch() !== null) {}
|
|
687
|
+
function prepareBody(content) {
|
|
688
|
+
if (compressionLevel != null && compressFn) {
|
|
689
|
+
return compressFn(content, compressionLevel);
|
|
608
690
|
}
|
|
609
|
-
return;
|
|
691
|
+
return content;
|
|
610
692
|
}
|
|
611
|
-
|
|
612
|
-
|
|
613
|
-
|
|
614
|
-
|
|
615
|
-
const headerOut = new OutputCollector(method.headerSchema, true, serverId, requestId);
|
|
616
|
-
const headerValues = method.headerInit(params, state, headerOut);
|
|
617
|
-
const headerBatch = buildResultBatch(method.headerSchema, headerValues, serverId, requestId);
|
|
618
|
-
const headerBatches = [
|
|
619
|
-
...headerOut.batches.map((b) => b.batch),
|
|
620
|
-
headerBatch
|
|
621
|
-
];
|
|
622
|
-
writer.writeStream(method.headerSchema, headerBatches);
|
|
623
|
-
} catch (error) {
|
|
624
|
-
const errBatch = buildErrorBatch(method.headerSchema, error, serverId, requestId);
|
|
625
|
-
writer.writeStream(method.headerSchema, [errBatch]);
|
|
626
|
-
const inputSchema2 = await reader.openNextStream();
|
|
627
|
-
if (inputSchema2) {
|
|
628
|
-
while (await reader.readNextBatch() !== null) {}
|
|
629
|
-
}
|
|
630
|
-
return;
|
|
693
|
+
async function readResponse(resp) {
|
|
694
|
+
let body = new Uint8Array(await resp.arrayBuffer());
|
|
695
|
+
if (resp.headers.get("Content-Encoding") === "zstd" && decompressFn) {
|
|
696
|
+
body = new Uint8Array(decompressFn(body));
|
|
631
697
|
}
|
|
698
|
+
return body;
|
|
632
699
|
}
|
|
633
|
-
|
|
634
|
-
|
|
635
|
-
|
|
636
|
-
|
|
637
|
-
|
|
700
|
+
async function ensureMethodCache() {
|
|
701
|
+
if (methodCache)
|
|
702
|
+
return methodCache;
|
|
703
|
+
const desc = await httpIntrospect(baseUrl, { prefix });
|
|
704
|
+
methodCache = new Map(desc.methods.map((m) => [m.name, m]));
|
|
705
|
+
return methodCache;
|
|
638
706
|
}
|
|
639
|
-
|
|
640
|
-
|
|
641
|
-
|
|
642
|
-
const
|
|
643
|
-
|
|
644
|
-
|
|
645
|
-
|
|
646
|
-
if (isProducer) {
|
|
647
|
-
await method.producerFn(state, out);
|
|
648
|
-
} else {
|
|
649
|
-
await method.exchangeFn(state, inputBatch, out);
|
|
650
|
-
}
|
|
651
|
-
for (const emitted of out.batches) {
|
|
652
|
-
stream.write(emitted.batch);
|
|
653
|
-
}
|
|
654
|
-
if (out.finished) {
|
|
655
|
-
break;
|
|
707
|
+
return {
|
|
708
|
+
async call(method, params) {
|
|
709
|
+
await ensureCompression();
|
|
710
|
+
const methods = await ensureMethodCache();
|
|
711
|
+
const info = methods.get(method);
|
|
712
|
+
if (!info) {
|
|
713
|
+
throw new Error(`Unknown method: '${method}'`);
|
|
656
714
|
}
|
|
657
|
-
|
|
658
|
-
|
|
659
|
-
|
|
660
|
-
|
|
661
|
-
|
|
662
|
-
|
|
663
|
-
|
|
664
|
-
|
|
665
|
-
}
|
|
666
|
-
|
|
667
|
-
|
|
668
|
-
|
|
669
|
-
|
|
670
|
-
|
|
671
|
-
|
|
672
|
-
|
|
673
|
-
serverId;
|
|
674
|
-
describeBatch = null;
|
|
675
|
-
constructor(protocol, options) {
|
|
676
|
-
this.protocol = protocol;
|
|
677
|
-
this.enableDescribe = options?.enableDescribe ?? true;
|
|
678
|
-
this.serverId = options?.serverId ?? crypto.randomUUID().replace(/-/g, "").slice(0, 12);
|
|
679
|
-
if (this.enableDescribe) {
|
|
680
|
-
const { batch } = buildDescribeBatch(protocol.name, protocol.getMethods(), this.serverId);
|
|
681
|
-
this.describeBatch = batch;
|
|
682
|
-
}
|
|
683
|
-
}
|
|
684
|
-
async run() {
|
|
685
|
-
const stdin = process.stdin;
|
|
686
|
-
if (process.stdin.isTTY || process.stdout.isTTY) {
|
|
687
|
-
process.stderr.write("WARNING: This process communicates via Arrow IPC on stdin/stdout " + `and is not intended to be run interactively.
|
|
688
|
-
` + "It should be launched as a subprocess by an RPC client " + `(e.g. vgi_rpc.connect()).
|
|
689
|
-
`);
|
|
690
|
-
}
|
|
691
|
-
const reader = await IpcStreamReader.create(stdin);
|
|
692
|
-
const writer = new IpcStreamWriter;
|
|
693
|
-
try {
|
|
694
|
-
while (true) {
|
|
695
|
-
await this.serveOne(reader, writer);
|
|
715
|
+
const fullParams = { ...info.defaults ?? {}, ...params ?? {} };
|
|
716
|
+
const body = buildRequestIpc(info.paramsSchema, fullParams, method);
|
|
717
|
+
const resp = await fetch(`${baseUrl}${prefix}/${method}`, {
|
|
718
|
+
method: "POST",
|
|
719
|
+
headers: buildHeaders(),
|
|
720
|
+
body: prepareBody(body)
|
|
721
|
+
});
|
|
722
|
+
const responseBody = await readResponse(resp);
|
|
723
|
+
const { batches } = await readResponseBatches(responseBody);
|
|
724
|
+
let resultBatch = null;
|
|
725
|
+
for (const batch of batches) {
|
|
726
|
+
if (batch.numRows === 0) {
|
|
727
|
+
dispatchLogOrError(batch, onLog);
|
|
728
|
+
continue;
|
|
729
|
+
}
|
|
730
|
+
resultBatch = batch;
|
|
696
731
|
}
|
|
697
|
-
|
|
698
|
-
|
|
699
|
-
return;
|
|
732
|
+
if (!resultBatch) {
|
|
733
|
+
return null;
|
|
700
734
|
}
|
|
701
|
-
|
|
702
|
-
|
|
703
|
-
|
|
704
|
-
|
|
705
|
-
|
|
706
|
-
|
|
707
|
-
|
|
708
|
-
|
|
709
|
-
|
|
710
|
-
|
|
711
|
-
|
|
712
|
-
|
|
713
|
-
|
|
714
|
-
|
|
715
|
-
writer.writeStream(EMPTY_SCHEMA2, [errBatch]);
|
|
716
|
-
return;
|
|
717
|
-
}
|
|
718
|
-
const batch = batches[0];
|
|
719
|
-
let methodName;
|
|
720
|
-
let params;
|
|
721
|
-
let requestId;
|
|
722
|
-
try {
|
|
723
|
-
const parsed = parseRequest(schema, batch);
|
|
724
|
-
methodName = parsed.methodName;
|
|
725
|
-
params = parsed.params;
|
|
726
|
-
requestId = parsed.requestId;
|
|
727
|
-
} catch (e) {
|
|
728
|
-
const errBatch = buildErrorBatch(EMPTY_SCHEMA2, e, this.serverId, null);
|
|
729
|
-
writer.writeStream(EMPTY_SCHEMA2, [errBatch]);
|
|
730
|
-
if (e instanceof VersionError || e instanceof RpcError) {
|
|
731
|
-
return;
|
|
735
|
+
const rows = extractBatchRows(resultBatch);
|
|
736
|
+
if (rows.length === 0)
|
|
737
|
+
return null;
|
|
738
|
+
const result = rows[0];
|
|
739
|
+
if (info.resultSchema.fields.length === 0)
|
|
740
|
+
return null;
|
|
741
|
+
return result;
|
|
742
|
+
},
|
|
743
|
+
async stream(method, params) {
|
|
744
|
+
await ensureCompression();
|
|
745
|
+
const methods = await ensureMethodCache();
|
|
746
|
+
const info = methods.get(method);
|
|
747
|
+
if (!info) {
|
|
748
|
+
throw new Error(`Unknown method: '${method}'`);
|
|
732
749
|
}
|
|
733
|
-
|
|
734
|
-
|
|
735
|
-
|
|
736
|
-
|
|
737
|
-
|
|
738
|
-
|
|
739
|
-
|
|
740
|
-
|
|
741
|
-
|
|
742
|
-
|
|
743
|
-
const
|
|
744
|
-
|
|
745
|
-
|
|
746
|
-
|
|
747
|
-
|
|
748
|
-
|
|
749
|
-
|
|
750
|
-
|
|
751
|
-
|
|
752
|
-
|
|
753
|
-
|
|
754
|
-
}
|
|
755
|
-
|
|
756
|
-
|
|
757
|
-
|
|
758
|
-
|
|
759
|
-
|
|
760
|
-
|
|
761
|
-
|
|
762
|
-
|
|
763
|
-
|
|
764
|
-
|
|
765
|
-
|
|
766
|
-
|
|
767
|
-
|
|
768
|
-
|
|
769
|
-
|
|
770
|
-
|
|
771
|
-
|
|
772
|
-
|
|
773
|
-
|
|
774
|
-
|
|
775
|
-
|
|
776
|
-
|
|
777
|
-
|
|
778
|
-
|
|
779
|
-
|
|
780
|
-
|
|
781
|
-
|
|
782
|
-
|
|
783
|
-
|
|
784
|
-
|
|
785
|
-
|
|
786
|
-
|
|
787
|
-
|
|
788
|
-
|
|
789
|
-
|
|
790
|
-
|
|
791
|
-
|
|
792
|
-
|
|
793
|
-
}
|
|
794
|
-
|
|
795
|
-
|
|
796
|
-
|
|
797
|
-
|
|
798
|
-
|
|
799
|
-
|
|
800
|
-
|
|
801
|
-
|
|
802
|
-
|
|
803
|
-
|
|
804
|
-
|
|
805
|
-
|
|
806
|
-
|
|
807
|
-
|
|
808
|
-
|
|
809
|
-
|
|
810
|
-
|
|
811
|
-
|
|
812
|
-
|
|
813
|
-
|
|
814
|
-
|
|
815
|
-
|
|
816
|
-
|
|
817
|
-
|
|
818
|
-
|
|
819
|
-
|
|
820
|
-
|
|
821
|
-
|
|
822
|
-
|
|
823
|
-
|
|
824
|
-
|
|
825
|
-
|
|
826
|
-
|
|
827
|
-
|
|
828
|
-
|
|
829
|
-
|
|
830
|
-
|
|
831
|
-
|
|
832
|
-
|
|
833
|
-
|
|
834
|
-
|
|
835
|
-
|
|
836
|
-
|
|
837
|
-
|
|
838
|
-
|
|
839
|
-
|
|
840
|
-
|
|
841
|
-
|
|
842
|
-
|
|
843
|
-
|
|
844
|
-
|
|
845
|
-
|
|
846
|
-
|
|
847
|
-
|
|
848
|
-
|
|
849
|
-
|
|
850
|
-
|
|
851
|
-
|
|
852
|
-
|
|
853
|
-
resultSchema: EMPTY_SCHEMA3,
|
|
854
|
-
outputSchema: toSchema(config.outputSchema),
|
|
855
|
-
inputSchema: EMPTY_SCHEMA3,
|
|
856
|
-
producerInit: config.init,
|
|
857
|
-
producerFn: config.produce,
|
|
858
|
-
headerSchema: config.headerSchema ? toSchema(config.headerSchema) : undefined,
|
|
859
|
-
headerInit: config.headerInit,
|
|
860
|
-
doc: config.doc,
|
|
861
|
-
defaults: config.defaults,
|
|
862
|
-
paramTypes: config.paramTypes ?? inferParamTypes(params)
|
|
863
|
-
});
|
|
864
|
-
return this;
|
|
865
|
-
}
|
|
866
|
-
exchange(name, config) {
|
|
867
|
-
const params = toSchema(config.params);
|
|
868
|
-
this._methods.set(name, {
|
|
869
|
-
name,
|
|
870
|
-
type: "stream" /* STREAM */,
|
|
871
|
-
paramsSchema: params,
|
|
872
|
-
resultSchema: EMPTY_SCHEMA3,
|
|
873
|
-
inputSchema: toSchema(config.inputSchema),
|
|
874
|
-
outputSchema: toSchema(config.outputSchema),
|
|
875
|
-
exchangeInit: config.init,
|
|
876
|
-
exchangeFn: config.exchange,
|
|
877
|
-
headerSchema: config.headerSchema ? toSchema(config.headerSchema) : undefined,
|
|
878
|
-
headerInit: config.headerInit,
|
|
879
|
-
doc: config.doc,
|
|
880
|
-
defaults: config.defaults,
|
|
881
|
-
paramTypes: config.paramTypes ?? inferParamTypes(params)
|
|
882
|
-
});
|
|
883
|
-
return this;
|
|
884
|
-
}
|
|
885
|
-
getMethods() {
|
|
886
|
-
return new Map(this._methods);
|
|
887
|
-
}
|
|
888
|
-
}
|
|
889
|
-
// src/http/handler.ts
|
|
890
|
-
import { Schema as Schema9 } from "apache-arrow";
|
|
891
|
-
import { randomBytes } from "node:crypto";
|
|
892
|
-
|
|
893
|
-
// src/http/types.ts
|
|
894
|
-
var jsonStateSerializer = {
|
|
895
|
-
serialize(state) {
|
|
896
|
-
return new TextEncoder().encode(JSON.stringify(state, (_key, value) => typeof value === "bigint" ? `__bigint__:${value}` : value));
|
|
897
|
-
},
|
|
898
|
-
deserialize(bytes2) {
|
|
899
|
-
return JSON.parse(new TextDecoder().decode(bytes2), (_key, value) => typeof value === "string" && value.startsWith("__bigint__:") ? BigInt(value.slice(11)) : value);
|
|
900
|
-
}
|
|
901
|
-
};
|
|
902
|
-
|
|
903
|
-
// src/http/common.ts
|
|
904
|
-
import {
|
|
905
|
-
RecordBatchStreamWriter as RecordBatchStreamWriter3,
|
|
906
|
-
RecordBatchReader as RecordBatchReader2,
|
|
907
|
-
RecordBatch as RecordBatch4,
|
|
908
|
-
Struct as Struct3,
|
|
909
|
-
makeData as makeData3
|
|
910
|
-
} from "apache-arrow";
|
|
911
|
-
var ARROW_CONTENT_TYPE = "application/vnd.apache.arrow.stream";
|
|
912
|
-
|
|
913
|
-
class HttpRpcError extends Error {
|
|
914
|
-
statusCode;
|
|
915
|
-
constructor(message, statusCode) {
|
|
916
|
-
super(message);
|
|
917
|
-
this.statusCode = statusCode;
|
|
918
|
-
this.name = "HttpRpcError";
|
|
919
|
-
}
|
|
920
|
-
}
|
|
921
|
-
function conformBatchToSchema(batch, schema) {
|
|
922
|
-
if (batch.numRows === 0)
|
|
923
|
-
return batch;
|
|
924
|
-
const children = schema.fields.map((f, i) => batch.data.children[i].clone(f.type));
|
|
925
|
-
const structType = new Struct3(schema.fields);
|
|
926
|
-
const data = makeData3({
|
|
927
|
-
type: structType,
|
|
928
|
-
length: batch.numRows,
|
|
929
|
-
children,
|
|
930
|
-
nullCount: batch.data.nullCount,
|
|
931
|
-
nullBitmap: batch.data.nullBitmap
|
|
932
|
-
});
|
|
933
|
-
return new RecordBatch4(schema, data, batch.metadata);
|
|
934
|
-
}
|
|
935
|
-
function serializeIpcStream(schema, batches) {
|
|
936
|
-
const writer = new RecordBatchStreamWriter3;
|
|
937
|
-
writer.reset(undefined, schema);
|
|
938
|
-
for (const batch of batches) {
|
|
939
|
-
writer.write(conformBatchToSchema(batch, schema));
|
|
940
|
-
}
|
|
941
|
-
writer.close();
|
|
942
|
-
return writer.toUint8Array(true);
|
|
943
|
-
}
|
|
944
|
-
function arrowResponse(body, status = 200, extraHeaders) {
|
|
945
|
-
const headers = extraHeaders ?? new Headers;
|
|
946
|
-
headers.set("Content-Type", ARROW_CONTENT_TYPE);
|
|
947
|
-
return new Response(body, { status, headers });
|
|
948
|
-
}
|
|
949
|
-
async function readRequestFromBody(body) {
|
|
950
|
-
const reader = await RecordBatchReader2.from(body);
|
|
951
|
-
await reader.open();
|
|
952
|
-
const schema = reader.schema;
|
|
953
|
-
if (!schema) {
|
|
954
|
-
throw new HttpRpcError("Empty IPC stream: no schema", 400);
|
|
955
|
-
}
|
|
956
|
-
const batches = reader.readAll();
|
|
957
|
-
if (batches.length === 0) {
|
|
958
|
-
throw new HttpRpcError("IPC stream contains no batches", 400);
|
|
959
|
-
}
|
|
960
|
-
return { schema, batch: batches[0] };
|
|
961
|
-
}
|
|
962
|
-
|
|
963
|
-
// src/http/dispatch.ts
|
|
964
|
-
import { Schema as Schema8, RecordBatch as RecordBatch5 } from "apache-arrow";
|
|
965
|
-
|
|
966
|
-
// src/http/token.ts
|
|
967
|
-
import { createHmac, timingSafeEqual } from "node:crypto";
|
|
968
|
-
var TOKEN_VERSION = 2;
|
|
969
|
-
var HMAC_LEN = 32;
|
|
970
|
-
var MIN_TOKEN_LEN = 1 + 8 + 12 + HMAC_LEN;
|
|
971
|
-
function packStateToken(stateBytes, schemaBytes, inputSchemaBytes, signingKey, createdAt) {
|
|
972
|
-
const now = createdAt ?? Math.floor(Date.now() / 1000);
|
|
973
|
-
const payloadLen = 1 + 8 + 4 + stateBytes.length + 4 + schemaBytes.length + 4 + inputSchemaBytes.length;
|
|
974
|
-
const buf = Buffer.alloc(payloadLen);
|
|
975
|
-
let offset = 0;
|
|
976
|
-
buf.writeUInt8(TOKEN_VERSION, offset);
|
|
977
|
-
offset += 1;
|
|
978
|
-
buf.writeBigUInt64LE(BigInt(now), offset);
|
|
979
|
-
offset += 8;
|
|
980
|
-
buf.writeUInt32LE(stateBytes.length, offset);
|
|
981
|
-
offset += 4;
|
|
982
|
-
buf.set(stateBytes, offset);
|
|
983
|
-
offset += stateBytes.length;
|
|
984
|
-
buf.writeUInt32LE(schemaBytes.length, offset);
|
|
985
|
-
offset += 4;
|
|
986
|
-
buf.set(schemaBytes, offset);
|
|
987
|
-
offset += schemaBytes.length;
|
|
988
|
-
buf.writeUInt32LE(inputSchemaBytes.length, offset);
|
|
989
|
-
offset += 4;
|
|
990
|
-
buf.set(inputSchemaBytes, offset);
|
|
991
|
-
offset += inputSchemaBytes.length;
|
|
992
|
-
const mac = createHmac("sha256", signingKey).update(buf).digest();
|
|
993
|
-
const token = Buffer.concat([buf, mac]);
|
|
994
|
-
return token.toString("base64");
|
|
750
|
+
const fullParams = { ...info.defaults ?? {}, ...params ?? {} };
|
|
751
|
+
const body = buildRequestIpc(info.paramsSchema, fullParams, method);
|
|
752
|
+
const resp = await fetch(`${baseUrl}${prefix}/${method}/init`, {
|
|
753
|
+
method: "POST",
|
|
754
|
+
headers: buildHeaders(),
|
|
755
|
+
body: prepareBody(body)
|
|
756
|
+
});
|
|
757
|
+
const responseBody = await readResponse(resp);
|
|
758
|
+
let header = null;
|
|
759
|
+
let stateToken = null;
|
|
760
|
+
const pendingBatches = [];
|
|
761
|
+
let finished = false;
|
|
762
|
+
let streamSchema = null;
|
|
763
|
+
if (info.headerSchema) {
|
|
764
|
+
const reader = await readSequentialStreams(responseBody);
|
|
765
|
+
const headerStream = await reader.readStream();
|
|
766
|
+
if (headerStream) {
|
|
767
|
+
for (const batch of headerStream.batches) {
|
|
768
|
+
if (batch.numRows === 0) {
|
|
769
|
+
dispatchLogOrError(batch, onLog);
|
|
770
|
+
continue;
|
|
771
|
+
}
|
|
772
|
+
const rows = extractBatchRows(batch);
|
|
773
|
+
if (rows.length > 0) {
|
|
774
|
+
header = rows[0];
|
|
775
|
+
}
|
|
776
|
+
}
|
|
777
|
+
}
|
|
778
|
+
const dataStream = await reader.readStream();
|
|
779
|
+
if (dataStream) {
|
|
780
|
+
streamSchema = dataStream.schema;
|
|
781
|
+
}
|
|
782
|
+
const headerErrorBatches = [];
|
|
783
|
+
if (dataStream) {
|
|
784
|
+
for (const batch of dataStream.batches) {
|
|
785
|
+
if (batch.numRows === 0) {
|
|
786
|
+
const token = batch.metadata?.get(STATE_KEY);
|
|
787
|
+
if (token) {
|
|
788
|
+
stateToken = token;
|
|
789
|
+
continue;
|
|
790
|
+
}
|
|
791
|
+
const level = batch.metadata?.get(LOG_LEVEL_KEY);
|
|
792
|
+
if (level === "EXCEPTION") {
|
|
793
|
+
headerErrorBatches.push(batch);
|
|
794
|
+
continue;
|
|
795
|
+
}
|
|
796
|
+
dispatchLogOrError(batch, onLog);
|
|
797
|
+
continue;
|
|
798
|
+
}
|
|
799
|
+
pendingBatches.push(batch);
|
|
800
|
+
}
|
|
801
|
+
}
|
|
802
|
+
if (headerErrorBatches.length > 0) {
|
|
803
|
+
if (pendingBatches.length > 0 || stateToken !== null) {
|
|
804
|
+
pendingBatches.push(...headerErrorBatches);
|
|
805
|
+
} else {
|
|
806
|
+
for (const batch of headerErrorBatches) {
|
|
807
|
+
dispatchLogOrError(batch, onLog);
|
|
808
|
+
}
|
|
809
|
+
}
|
|
810
|
+
}
|
|
811
|
+
if (!dataStream && !stateToken) {
|
|
812
|
+
finished = true;
|
|
813
|
+
}
|
|
814
|
+
} else {
|
|
815
|
+
const { schema: responseSchema, batches } = await readResponseBatches(responseBody);
|
|
816
|
+
streamSchema = responseSchema;
|
|
817
|
+
const errorBatches = [];
|
|
818
|
+
for (const batch of batches) {
|
|
819
|
+
if (batch.numRows === 0) {
|
|
820
|
+
const token = batch.metadata?.get(STATE_KEY);
|
|
821
|
+
if (token) {
|
|
822
|
+
stateToken = token;
|
|
823
|
+
continue;
|
|
824
|
+
}
|
|
825
|
+
const level = batch.metadata?.get(LOG_LEVEL_KEY);
|
|
826
|
+
if (level === "EXCEPTION") {
|
|
827
|
+
errorBatches.push(batch);
|
|
828
|
+
continue;
|
|
829
|
+
}
|
|
830
|
+
dispatchLogOrError(batch, onLog);
|
|
831
|
+
continue;
|
|
832
|
+
}
|
|
833
|
+
pendingBatches.push(batch);
|
|
834
|
+
}
|
|
835
|
+
if (errorBatches.length > 0) {
|
|
836
|
+
if (pendingBatches.length > 0 || stateToken !== null) {
|
|
837
|
+
pendingBatches.push(...errorBatches);
|
|
838
|
+
} else {
|
|
839
|
+
for (const batch of errorBatches) {
|
|
840
|
+
dispatchLogOrError(batch, onLog);
|
|
841
|
+
}
|
|
842
|
+
}
|
|
843
|
+
}
|
|
844
|
+
}
|
|
845
|
+
if (pendingBatches.length === 0 && stateToken === null) {
|
|
846
|
+
finished = true;
|
|
847
|
+
}
|
|
848
|
+
const outputSchema = (streamSchema && streamSchema.fields.length > 0 ? streamSchema : null) ?? (pendingBatches.length > 0 ? pendingBatches[0].schema : null) ?? info.outputSchema ?? info.resultSchema;
|
|
849
|
+
return new HttpStreamSession({
|
|
850
|
+
baseUrl,
|
|
851
|
+
prefix,
|
|
852
|
+
method,
|
|
853
|
+
stateToken,
|
|
854
|
+
outputSchema,
|
|
855
|
+
inputSchema: info.inputSchema,
|
|
856
|
+
onLog,
|
|
857
|
+
pendingBatches,
|
|
858
|
+
finished,
|
|
859
|
+
header,
|
|
860
|
+
compressionLevel,
|
|
861
|
+
compressFn,
|
|
862
|
+
decompressFn
|
|
863
|
+
});
|
|
864
|
+
},
|
|
865
|
+
async describe() {
|
|
866
|
+
return httpIntrospect(baseUrl, { prefix });
|
|
867
|
+
},
|
|
868
|
+
close() {}
|
|
869
|
+
};
|
|
995
870
|
}
|
|
996
|
-
|
|
997
|
-
|
|
998
|
-
|
|
999
|
-
|
|
871
|
+
// src/client/pipe.ts
|
|
872
|
+
import {
|
|
873
|
+
Field as Field2,
|
|
874
|
+
makeData as makeData4,
|
|
875
|
+
RecordBatch as RecordBatch4,
|
|
876
|
+
RecordBatchStreamWriter as RecordBatchStreamWriter2,
|
|
877
|
+
Schema as Schema2,
|
|
878
|
+
Struct as Struct4,
|
|
879
|
+
vectorFromArray as vectorFromArray3
|
|
880
|
+
} from "@query-farm/apache-arrow";
|
|
881
|
+
class PipeIncrementalWriter {
|
|
882
|
+
writer;
|
|
883
|
+
writeFn;
|
|
884
|
+
closed = false;
|
|
885
|
+
constructor(writeFn, schema) {
|
|
886
|
+
this.writeFn = writeFn;
|
|
887
|
+
this.writer = new RecordBatchStreamWriter2;
|
|
888
|
+
this.writer.reset(undefined, schema);
|
|
889
|
+
this.drain();
|
|
1000
890
|
}
|
|
1001
|
-
|
|
1002
|
-
|
|
1003
|
-
|
|
1004
|
-
|
|
1005
|
-
|
|
891
|
+
write(batch) {
|
|
892
|
+
if (this.closed)
|
|
893
|
+
throw new Error("PipeIncrementalWriter already closed");
|
|
894
|
+
this.writer._writeRecordBatch(batch);
|
|
895
|
+
this.drain();
|
|
1006
896
|
}
|
|
1007
|
-
|
|
1008
|
-
|
|
1009
|
-
|
|
1010
|
-
|
|
1011
|
-
|
|
897
|
+
close() {
|
|
898
|
+
if (this.closed)
|
|
899
|
+
return;
|
|
900
|
+
this.closed = true;
|
|
901
|
+
const eos = new Uint8Array(new Int32Array([-1, 0]).buffer);
|
|
902
|
+
this.writeFn(eos);
|
|
1012
903
|
}
|
|
1013
|
-
|
|
1014
|
-
|
|
1015
|
-
|
|
1016
|
-
|
|
1017
|
-
if (now - createdAt > tokenTtl) {
|
|
1018
|
-
throw new Error("State token expired");
|
|
904
|
+
drain() {
|
|
905
|
+
const values = this.writer._sink._values;
|
|
906
|
+
for (const chunk of values) {
|
|
907
|
+
this.writeFn(chunk);
|
|
1019
908
|
}
|
|
909
|
+
values.length = 0;
|
|
1020
910
|
}
|
|
1021
|
-
const stateLen = payload.readUInt32LE(offset);
|
|
1022
|
-
offset += 4;
|
|
1023
|
-
if (offset + stateLen > payload.length) {
|
|
1024
|
-
throw new Error("State token truncated (state)");
|
|
1025
|
-
}
|
|
1026
|
-
const stateBytes = payload.slice(offset, offset + stateLen);
|
|
1027
|
-
offset += stateLen;
|
|
1028
|
-
const schemaLen = payload.readUInt32LE(offset);
|
|
1029
|
-
offset += 4;
|
|
1030
|
-
if (offset + schemaLen > payload.length) {
|
|
1031
|
-
throw new Error("State token truncated (schema)");
|
|
1032
|
-
}
|
|
1033
|
-
const schemaBytes = payload.slice(offset, offset + schemaLen);
|
|
1034
|
-
offset += schemaLen;
|
|
1035
|
-
const inputSchemaLen = payload.readUInt32LE(offset);
|
|
1036
|
-
offset += 4;
|
|
1037
|
-
if (offset + inputSchemaLen > payload.length) {
|
|
1038
|
-
throw new Error("State token truncated (input schema)");
|
|
1039
|
-
}
|
|
1040
|
-
const inputSchemaBytes = payload.slice(offset, offset + inputSchemaLen);
|
|
1041
|
-
return { stateBytes, schemaBytes, inputSchemaBytes, createdAt };
|
|
1042
911
|
}
|
|
1043
912
|
|
|
1044
|
-
|
|
1045
|
-
|
|
1046
|
-
|
|
1047
|
-
|
|
1048
|
-
|
|
1049
|
-
|
|
1050
|
-
|
|
1051
|
-
|
|
1052
|
-
|
|
1053
|
-
|
|
1054
|
-
|
|
1055
|
-
|
|
1056
|
-
|
|
1057
|
-
|
|
1058
|
-
|
|
1059
|
-
|
|
1060
|
-
|
|
1061
|
-
|
|
1062
|
-
|
|
1063
|
-
|
|
1064
|
-
} catch (error) {
|
|
1065
|
-
const errBatch = buildErrorBatch(schema, error, ctx.serverId, parsed.requestId);
|
|
1066
|
-
return arrowResponse(serializeIpcStream(schema, [errBatch]), 500);
|
|
1067
|
-
}
|
|
1068
|
-
}
|
|
1069
|
-
async function httpDispatchStreamInit(method, body, ctx) {
|
|
1070
|
-
const isProducer = !method.inputSchema || method.inputSchema.fields.length === 0;
|
|
1071
|
-
const outputSchema = method.outputSchema;
|
|
1072
|
-
const inputSchema = method.inputSchema ?? EMPTY_SCHEMA4;
|
|
1073
|
-
const { schema: reqSchema, batch: reqBatch } = await readRequestFromBody(body);
|
|
1074
|
-
const parsed = parseRequest(reqSchema, reqBatch);
|
|
1075
|
-
if (parsed.methodName !== method.name) {
|
|
1076
|
-
throw new HttpRpcError(`Method name in request '${parsed.methodName}' does not match URL '${method.name}'`, 400);
|
|
913
|
+
class PipeStreamSession {
|
|
914
|
+
_reader;
|
|
915
|
+
_writeFn;
|
|
916
|
+
_onLog;
|
|
917
|
+
_header;
|
|
918
|
+
_inputWriter = null;
|
|
919
|
+
_inputSchema = null;
|
|
920
|
+
_outputStreamOpened = false;
|
|
921
|
+
_closed = false;
|
|
922
|
+
_outputSchema;
|
|
923
|
+
_releaseBusy;
|
|
924
|
+
_setDrainPromise;
|
|
925
|
+
constructor(opts) {
|
|
926
|
+
this._reader = opts.reader;
|
|
927
|
+
this._writeFn = opts.writeFn;
|
|
928
|
+
this._onLog = opts.onLog;
|
|
929
|
+
this._header = opts.header;
|
|
930
|
+
this._outputSchema = opts.outputSchema;
|
|
931
|
+
this._releaseBusy = opts.releaseBusy;
|
|
932
|
+
this._setDrainPromise = opts.setDrainPromise;
|
|
1077
933
|
}
|
|
1078
|
-
|
|
1079
|
-
|
|
1080
|
-
if (isProducer) {
|
|
1081
|
-
state = await method.producerInit(parsed.params);
|
|
1082
|
-
} else {
|
|
1083
|
-
state = await method.exchangeInit(parsed.params);
|
|
1084
|
-
}
|
|
1085
|
-
} catch (error) {
|
|
1086
|
-
const errSchema = method.headerSchema ?? EMPTY_SCHEMA4;
|
|
1087
|
-
const errBatch = buildErrorBatch(errSchema, error, ctx.serverId, parsed.requestId);
|
|
1088
|
-
return arrowResponse(serializeIpcStream(errSchema, [errBatch]), 500);
|
|
934
|
+
get header() {
|
|
935
|
+
return this._header;
|
|
1089
936
|
}
|
|
1090
|
-
|
|
1091
|
-
|
|
1092
|
-
|
|
1093
|
-
|
|
1094
|
-
|
|
1095
|
-
|
|
1096
|
-
|
|
1097
|
-
|
|
1098
|
-
|
|
1099
|
-
|
|
1100
|
-
|
|
1101
|
-
];
|
|
1102
|
-
headerBytes = serializeIpcStream(method.headerSchema, headerBatches);
|
|
1103
|
-
} catch (error) {
|
|
1104
|
-
const errBatch = buildErrorBatch(method.headerSchema, error, ctx.serverId, parsed.requestId);
|
|
1105
|
-
return arrowResponse(serializeIpcStream(method.headerSchema, [errBatch]), 500);
|
|
937
|
+
async _readOutputBatch() {
|
|
938
|
+
while (true) {
|
|
939
|
+
const batch = await this._reader.readNextBatch();
|
|
940
|
+
if (batch === null)
|
|
941
|
+
return null;
|
|
942
|
+
if (batch.numRows === 0) {
|
|
943
|
+
if (dispatchLogOrError(batch, this._onLog)) {
|
|
944
|
+
continue;
|
|
945
|
+
}
|
|
946
|
+
}
|
|
947
|
+
return batch;
|
|
1106
948
|
}
|
|
1107
949
|
}
|
|
1108
|
-
|
|
1109
|
-
|
|
1110
|
-
|
|
1111
|
-
|
|
1112
|
-
const
|
|
1113
|
-
|
|
1114
|
-
|
|
1115
|
-
const tokenMeta = new Map;
|
|
1116
|
-
tokenMeta.set(STATE_KEY, token);
|
|
1117
|
-
const tokenBatch = buildEmptyBatch(resolvedOutputSchema, tokenMeta);
|
|
1118
|
-
const tokenStreamBytes = serializeIpcStream(resolvedOutputSchema, [tokenBatch]);
|
|
1119
|
-
let responseBody;
|
|
1120
|
-
if (headerBytes) {
|
|
1121
|
-
responseBody = concatBytes(headerBytes, tokenStreamBytes);
|
|
1122
|
-
} else {
|
|
1123
|
-
responseBody = tokenStreamBytes;
|
|
950
|
+
async _ensureOutputStream() {
|
|
951
|
+
if (this._outputStreamOpened)
|
|
952
|
+
return;
|
|
953
|
+
this._outputStreamOpened = true;
|
|
954
|
+
const schema = await this._reader.openNextStream();
|
|
955
|
+
if (!schema) {
|
|
956
|
+
throw new RpcError("ProtocolError", "Expected output stream but got EOF", "");
|
|
1124
957
|
}
|
|
1125
|
-
return arrowResponse(responseBody);
|
|
1126
|
-
}
|
|
1127
|
-
}
|
|
1128
|
-
async function httpDispatchStreamExchange(method, body, ctx) {
|
|
1129
|
-
const isProducer = !method.inputSchema || method.inputSchema.fields.length === 0;
|
|
1130
|
-
const { batch: reqBatch } = await readRequestFromBody(body);
|
|
1131
|
-
const tokenBase64 = reqBatch.metadata?.get(STATE_KEY);
|
|
1132
|
-
if (!tokenBase64) {
|
|
1133
|
-
throw new HttpRpcError("Missing state token in exchange request", 400);
|
|
1134
958
|
}
|
|
1135
|
-
|
|
1136
|
-
|
|
1137
|
-
|
|
1138
|
-
} catch (error) {
|
|
1139
|
-
throw new HttpRpcError(`Invalid state token: ${error.message}`, 400);
|
|
1140
|
-
}
|
|
1141
|
-
const state = ctx.stateSerializer.deserialize(unpacked.stateBytes);
|
|
1142
|
-
const outputSchema = state?.__outputSchema ?? method.outputSchema;
|
|
1143
|
-
const inputSchema = method.inputSchema ?? EMPTY_SCHEMA4;
|
|
1144
|
-
const effectiveProducer = state?.__isProducer ?? isProducer;
|
|
1145
|
-
if (effectiveProducer) {
|
|
1146
|
-
return produceStreamResponse(method, state, outputSchema, inputSchema, ctx, null, null);
|
|
1147
|
-
} else {
|
|
1148
|
-
const out = new OutputCollector(outputSchema, false, ctx.serverId, null);
|
|
1149
|
-
try {
|
|
1150
|
-
await method.exchangeFn(state, reqBatch, out);
|
|
1151
|
-
} catch (error) {
|
|
1152
|
-
const errBatch = buildErrorBatch(outputSchema, error, ctx.serverId, null);
|
|
1153
|
-
return arrowResponse(serializeIpcStream(outputSchema, [errBatch]), 500);
|
|
959
|
+
async exchange(input) {
|
|
960
|
+
if (this._closed) {
|
|
961
|
+
throw new RpcError("ProtocolError", "Stream session is closed", "");
|
|
1154
962
|
}
|
|
1155
|
-
|
|
1156
|
-
|
|
1157
|
-
|
|
1158
|
-
|
|
1159
|
-
|
|
1160
|
-
|
|
1161
|
-
|
|
1162
|
-
|
|
1163
|
-
|
|
1164
|
-
|
|
1165
|
-
|
|
963
|
+
let inputSchema;
|
|
964
|
+
let batch;
|
|
965
|
+
if (input.length === 0) {
|
|
966
|
+
inputSchema = this._inputSchema ?? this._outputSchema;
|
|
967
|
+
const children = inputSchema.fields.map((f) => {
|
|
968
|
+
return makeData4({ type: f.type, length: 0, nullCount: 0 });
|
|
969
|
+
});
|
|
970
|
+
const structType = new Struct4(inputSchema.fields);
|
|
971
|
+
const data = makeData4({
|
|
972
|
+
type: structType,
|
|
973
|
+
length: 0,
|
|
974
|
+
children,
|
|
975
|
+
nullCount: 0
|
|
976
|
+
});
|
|
977
|
+
batch = new RecordBatch4(inputSchema, data);
|
|
978
|
+
} else {
|
|
979
|
+
const keys = Object.keys(input[0]);
|
|
980
|
+
const fields = keys.map((key) => {
|
|
981
|
+
let sample;
|
|
982
|
+
for (const row of input) {
|
|
983
|
+
if (row[key] != null) {
|
|
984
|
+
sample = row[key];
|
|
985
|
+
break;
|
|
986
|
+
}
|
|
987
|
+
}
|
|
988
|
+
const arrowType = inferArrowType(sample);
|
|
989
|
+
return new Field2(key, arrowType, true);
|
|
990
|
+
});
|
|
991
|
+
inputSchema = new Schema2(fields);
|
|
992
|
+
if (this._inputSchema) {
|
|
993
|
+
const cached = this._inputSchema;
|
|
994
|
+
if (cached.fields.length !== inputSchema.fields.length || cached.fields.some((f, i) => f.name !== inputSchema.fields[i].name)) {
|
|
995
|
+
throw new RpcError("ProtocolError", `Exchange input schema changed: expected [${cached.fields.map((f) => f.name).join(", ")}] ` + `but got [${inputSchema.fields.map((f) => f.name).join(", ")}]`, "");
|
|
996
|
+
}
|
|
1166
997
|
} else {
|
|
1167
|
-
|
|
998
|
+
this._inputSchema = inputSchema;
|
|
1168
999
|
}
|
|
1000
|
+
const children = inputSchema.fields.map((f) => {
|
|
1001
|
+
const values = input.map((row) => row[f.name]);
|
|
1002
|
+
return vectorFromArray3(values, f.type).data[0];
|
|
1003
|
+
});
|
|
1004
|
+
const structType = new Struct4(inputSchema.fields);
|
|
1005
|
+
const data = makeData4({
|
|
1006
|
+
type: structType,
|
|
1007
|
+
length: input.length,
|
|
1008
|
+
children,
|
|
1009
|
+
nullCount: 0
|
|
1010
|
+
});
|
|
1011
|
+
batch = new RecordBatch4(inputSchema, data);
|
|
1169
1012
|
}
|
|
1170
|
-
|
|
1171
|
-
|
|
1172
|
-
}
|
|
1173
|
-
async function produceStreamResponse(method, state, outputSchema, inputSchema, ctx, requestId, headerBytes) {
|
|
1174
|
-
const allBatches = [];
|
|
1175
|
-
const maxBytes = ctx.maxStreamResponseBytes;
|
|
1176
|
-
let estimatedBytes = 0;
|
|
1177
|
-
while (true) {
|
|
1178
|
-
const out = new OutputCollector(outputSchema, true, ctx.serverId, requestId);
|
|
1179
|
-
try {
|
|
1180
|
-
await method.producerFn(state, out);
|
|
1181
|
-
} catch (error) {
|
|
1182
|
-
allBatches.push(buildErrorBatch(outputSchema, error, ctx.serverId, requestId));
|
|
1183
|
-
break;
|
|
1013
|
+
if (!this._inputWriter) {
|
|
1014
|
+
this._inputWriter = new PipeIncrementalWriter(this._writeFn, inputSchema);
|
|
1184
1015
|
}
|
|
1185
|
-
|
|
1186
|
-
|
|
1187
|
-
|
|
1188
|
-
|
|
1016
|
+
this._inputWriter.write(batch);
|
|
1017
|
+
await this._ensureOutputStream();
|
|
1018
|
+
try {
|
|
1019
|
+
const outputBatch = await this._readOutputBatch();
|
|
1020
|
+
if (outputBatch === null) {
|
|
1021
|
+
return [];
|
|
1189
1022
|
}
|
|
1023
|
+
return extractBatchRows(outputBatch);
|
|
1024
|
+
} catch (e) {
|
|
1025
|
+
await this._cleanup();
|
|
1026
|
+
throw e;
|
|
1190
1027
|
}
|
|
1191
|
-
if (out.finished) {
|
|
1192
|
-
break;
|
|
1193
|
-
}
|
|
1194
|
-
if (maxBytes != null && estimatedBytes >= maxBytes) {
|
|
1195
|
-
const stateBytes = ctx.stateSerializer.serialize(state);
|
|
1196
|
-
const schemaBytes = serializeSchema(outputSchema);
|
|
1197
|
-
const inputSchemaBytes = serializeSchema(inputSchema);
|
|
1198
|
-
const token = packStateToken(stateBytes, schemaBytes, inputSchemaBytes, ctx.signingKey);
|
|
1199
|
-
const tokenMeta = new Map;
|
|
1200
|
-
tokenMeta.set(STATE_KEY, token);
|
|
1201
|
-
allBatches.push(buildEmptyBatch(outputSchema, tokenMeta));
|
|
1202
|
-
break;
|
|
1203
|
-
}
|
|
1204
|
-
}
|
|
1205
|
-
const dataBytes = serializeIpcStream(outputSchema, allBatches);
|
|
1206
|
-
let responseBody;
|
|
1207
|
-
if (headerBytes) {
|
|
1208
|
-
responseBody = concatBytes(headerBytes, dataBytes);
|
|
1209
|
-
} else {
|
|
1210
|
-
responseBody = dataBytes;
|
|
1211
|
-
}
|
|
1212
|
-
return arrowResponse(responseBody);
|
|
1213
|
-
}
|
|
1214
|
-
function concatBytes(...arrays) {
|
|
1215
|
-
const totalLen = arrays.reduce((sum, a) => sum + a.length, 0);
|
|
1216
|
-
const result = new Uint8Array(totalLen);
|
|
1217
|
-
let offset = 0;
|
|
1218
|
-
for (const arr of arrays) {
|
|
1219
|
-
result.set(arr, offset);
|
|
1220
|
-
offset += arr.length;
|
|
1221
1028
|
}
|
|
1222
|
-
|
|
1223
|
-
|
|
1224
|
-
|
|
1225
|
-
|
|
1226
|
-
|
|
1227
|
-
|
|
1228
|
-
|
|
1229
|
-
const prefix = (options?.prefix ?? "/vgi").replace(/\/+$/, "");
|
|
1230
|
-
const signingKey = options?.signingKey ?? randomBytes(32);
|
|
1231
|
-
const tokenTtl = options?.tokenTtl ?? 3600;
|
|
1232
|
-
const corsOrigins = options?.corsOrigins;
|
|
1233
|
-
const maxRequestBytes = options?.maxRequestBytes;
|
|
1234
|
-
const maxStreamResponseBytes = options?.maxStreamResponseBytes;
|
|
1235
|
-
const serverId = options?.serverId ?? crypto.randomUUID().replace(/-/g, "").slice(0, 12);
|
|
1236
|
-
const methods = protocol.getMethods();
|
|
1237
|
-
const compressionLevel = options?.compressionLevel;
|
|
1238
|
-
const stateSerializer = options?.stateSerializer ?? jsonStateSerializer;
|
|
1239
|
-
const ctx = {
|
|
1240
|
-
signingKey,
|
|
1241
|
-
tokenTtl,
|
|
1242
|
-
serverId,
|
|
1243
|
-
maxStreamResponseBytes,
|
|
1244
|
-
stateSerializer
|
|
1245
|
-
};
|
|
1246
|
-
function addCorsHeaders(headers) {
|
|
1247
|
-
if (corsOrigins) {
|
|
1248
|
-
headers.set("Access-Control-Allow-Origin", corsOrigins);
|
|
1249
|
-
headers.set("Access-Control-Allow-Methods", "POST, OPTIONS");
|
|
1250
|
-
headers.set("Access-Control-Allow-Headers", "Content-Type");
|
|
1029
|
+
async _cleanup() {
|
|
1030
|
+
if (this._closed)
|
|
1031
|
+
return;
|
|
1032
|
+
this._closed = true;
|
|
1033
|
+
if (this._inputWriter) {
|
|
1034
|
+
this._inputWriter.close();
|
|
1035
|
+
this._inputWriter = null;
|
|
1251
1036
|
}
|
|
1037
|
+
try {
|
|
1038
|
+
if (this._outputStreamOpened) {
|
|
1039
|
+
while (await this._reader.readNextBatch() !== null) {}
|
|
1040
|
+
}
|
|
1041
|
+
} catch {}
|
|
1042
|
+
this._releaseBusy();
|
|
1252
1043
|
}
|
|
1253
|
-
async
|
|
1254
|
-
if (
|
|
1255
|
-
return
|
|
1256
|
-
|
|
1257
|
-
|
|
1258
|
-
|
|
1259
|
-
|
|
1260
|
-
|
|
1261
|
-
|
|
1262
|
-
|
|
1263
|
-
|
|
1264
|
-
|
|
1265
|
-
|
|
1266
|
-
|
|
1267
|
-
|
|
1268
|
-
|
|
1269
|
-
|
|
1270
|
-
|
|
1271
|
-
|
|
1272
|
-
|
|
1273
|
-
const url = new URL(request.url);
|
|
1274
|
-
const path = url.pathname;
|
|
1275
|
-
if (request.method === "OPTIONS") {
|
|
1276
|
-
if (path === `${prefix}/__capabilities__`) {
|
|
1277
|
-
const headers = new Headers;
|
|
1278
|
-
addCorsHeaders(headers);
|
|
1279
|
-
if (maxRequestBytes != null) {
|
|
1280
|
-
headers.set("VGI-Max-Request-Bytes", String(maxRequestBytes));
|
|
1044
|
+
async* [Symbol.asyncIterator]() {
|
|
1045
|
+
if (this._closed)
|
|
1046
|
+
return;
|
|
1047
|
+
try {
|
|
1048
|
+
const tickSchema = new Schema2([]);
|
|
1049
|
+
this._inputWriter = new PipeIncrementalWriter(this._writeFn, tickSchema);
|
|
1050
|
+
const structType = new Struct4(tickSchema.fields);
|
|
1051
|
+
const tickData = makeData4({
|
|
1052
|
+
type: structType,
|
|
1053
|
+
length: 0,
|
|
1054
|
+
children: [],
|
|
1055
|
+
nullCount: 0
|
|
1056
|
+
});
|
|
1057
|
+
const tickBatch = new RecordBatch4(tickSchema, tickData);
|
|
1058
|
+
while (true) {
|
|
1059
|
+
this._inputWriter.write(tickBatch);
|
|
1060
|
+
await this._ensureOutputStream();
|
|
1061
|
+
const outputBatch = await this._readOutputBatch();
|
|
1062
|
+
if (outputBatch === null) {
|
|
1063
|
+
break;
|
|
1281
1064
|
}
|
|
1282
|
-
|
|
1283
|
-
}
|
|
1284
|
-
if (corsOrigins) {
|
|
1285
|
-
const headers = new Headers;
|
|
1286
|
-
addCorsHeaders(headers);
|
|
1287
|
-
return new Response(null, { status: 204, headers });
|
|
1065
|
+
yield extractBatchRows(outputBatch);
|
|
1288
1066
|
}
|
|
1289
|
-
|
|
1290
|
-
|
|
1291
|
-
|
|
1292
|
-
|
|
1293
|
-
}
|
|
1294
|
-
const contentType = request.headers.get("Content-Type");
|
|
1295
|
-
if (!contentType || !contentType.includes(ARROW_CONTENT_TYPE)) {
|
|
1296
|
-
return new Response(`Unsupported Media Type: expected ${ARROW_CONTENT_TYPE}`, { status: 415 });
|
|
1297
|
-
}
|
|
1298
|
-
if (maxRequestBytes != null) {
|
|
1299
|
-
const contentLength = request.headers.get("Content-Length");
|
|
1300
|
-
if (contentLength && parseInt(contentLength) > maxRequestBytes) {
|
|
1301
|
-
return new Response("Request body too large", { status: 413 });
|
|
1067
|
+
} finally {
|
|
1068
|
+
if (this._inputWriter) {
|
|
1069
|
+
this._inputWriter.close();
|
|
1070
|
+
this._inputWriter = null;
|
|
1302
1071
|
}
|
|
1072
|
+
try {
|
|
1073
|
+
if (this._outputStreamOpened) {
|
|
1074
|
+
while (await this._reader.readNextBatch() !== null) {}
|
|
1075
|
+
}
|
|
1076
|
+
} catch {}
|
|
1077
|
+
this._closed = true;
|
|
1078
|
+
this._releaseBusy();
|
|
1303
1079
|
}
|
|
1304
|
-
|
|
1305
|
-
|
|
1306
|
-
|
|
1307
|
-
|
|
1308
|
-
|
|
1080
|
+
}
|
|
1081
|
+
close() {
|
|
1082
|
+
if (this._closed)
|
|
1083
|
+
return;
|
|
1084
|
+
this._closed = true;
|
|
1085
|
+
if (this._inputWriter) {
|
|
1086
|
+
this._inputWriter.close();
|
|
1087
|
+
this._inputWriter = null;
|
|
1088
|
+
} else {
|
|
1089
|
+
const emptySchema = new Schema2([]);
|
|
1090
|
+
const ipc = serializeIpcStream(emptySchema, []);
|
|
1091
|
+
this._writeFn(ipc);
|
|
1309
1092
|
}
|
|
1310
|
-
|
|
1093
|
+
const drainPromise = (async () => {
|
|
1311
1094
|
try {
|
|
1312
|
-
|
|
1313
|
-
|
|
1314
|
-
|
|
1315
|
-
|
|
1316
|
-
|
|
1095
|
+
if (!this._outputStreamOpened) {
|
|
1096
|
+
const schema = await this._reader.openNextStream();
|
|
1097
|
+
if (schema) {
|
|
1098
|
+
while (await this._reader.readNextBatch() !== null) {}
|
|
1099
|
+
}
|
|
1100
|
+
} else {
|
|
1101
|
+
while (await this._reader.readNextBatch() !== null) {}
|
|
1102
|
+
}
|
|
1103
|
+
} catch {} finally {
|
|
1104
|
+
this._releaseBusy();
|
|
1317
1105
|
}
|
|
1106
|
+
})();
|
|
1107
|
+
this._setDrainPromise(drainPromise);
|
|
1108
|
+
}
|
|
1109
|
+
}
|
|
1110
|
+
function pipeConnect(readable, writable, options) {
|
|
1111
|
+
const onLog = options?.onLog;
|
|
1112
|
+
let reader = null;
|
|
1113
|
+
let readerPromise = null;
|
|
1114
|
+
let methodCache = null;
|
|
1115
|
+
let protocolName = "";
|
|
1116
|
+
let _busy = false;
|
|
1117
|
+
let _drainPromise = null;
|
|
1118
|
+
let closed = false;
|
|
1119
|
+
const writeFn = (bytes) => {
|
|
1120
|
+
writable.write(bytes);
|
|
1121
|
+
writable.flush?.();
|
|
1122
|
+
};
|
|
1123
|
+
async function ensureReader() {
|
|
1124
|
+
if (reader)
|
|
1125
|
+
return reader;
|
|
1126
|
+
if (!readerPromise) {
|
|
1127
|
+
readerPromise = IpcStreamReader.create(readable);
|
|
1318
1128
|
}
|
|
1319
|
-
|
|
1320
|
-
|
|
1321
|
-
|
|
1322
|
-
|
|
1323
|
-
|
|
1324
|
-
|
|
1325
|
-
|
|
1326
|
-
methodName = subPath.slice(0, -5);
|
|
1327
|
-
action = "init";
|
|
1328
|
-
} else if (subPath.endsWith("/exchange")) {
|
|
1329
|
-
methodName = subPath.slice(0, -9);
|
|
1330
|
-
action = "exchange";
|
|
1331
|
-
} else {
|
|
1332
|
-
methodName = subPath;
|
|
1333
|
-
action = "call";
|
|
1129
|
+
reader = await readerPromise;
|
|
1130
|
+
return reader;
|
|
1131
|
+
}
|
|
1132
|
+
async function acquireBusy() {
|
|
1133
|
+
if (_drainPromise) {
|
|
1134
|
+
await _drainPromise;
|
|
1135
|
+
_drainPromise = null;
|
|
1334
1136
|
}
|
|
1335
|
-
|
|
1336
|
-
|
|
1337
|
-
const available = [...methods.keys()].sort();
|
|
1338
|
-
const err = new Error(`Unknown method: '${methodName}'. Available methods: [${available.join(", ")}]`);
|
|
1339
|
-
return compressIfAccepted(makeErrorResponse(err, 404), clientAcceptsZstd);
|
|
1137
|
+
if (_busy) {
|
|
1138
|
+
throw new Error("Pipe transport is busy — another call or stream is in progress. " + "Pipe connections are single-threaded; wait for the current operation to complete.");
|
|
1340
1139
|
}
|
|
1140
|
+
_busy = true;
|
|
1141
|
+
}
|
|
1142
|
+
function releaseBusy() {
|
|
1143
|
+
_busy = false;
|
|
1144
|
+
}
|
|
1145
|
+
function setDrainPromise(p) {
|
|
1146
|
+
_drainPromise = p;
|
|
1147
|
+
}
|
|
1148
|
+
async function ensureMethodCache() {
|
|
1149
|
+
if (methodCache)
|
|
1150
|
+
return methodCache;
|
|
1151
|
+
await acquireBusy();
|
|
1341
1152
|
try {
|
|
1342
|
-
|
|
1343
|
-
|
|
1344
|
-
|
|
1345
|
-
|
|
1153
|
+
const emptySchema = new Schema2([]);
|
|
1154
|
+
const body = buildRequestIpc(emptySchema, {}, DESCRIBE_METHOD_NAME);
|
|
1155
|
+
writeFn(body);
|
|
1156
|
+
const r = await ensureReader();
|
|
1157
|
+
const response = await r.readStream();
|
|
1158
|
+
if (!response) {
|
|
1159
|
+
throw new Error("EOF reading __describe__ response");
|
|
1160
|
+
}
|
|
1161
|
+
const desc = await parseDescribeResponse(response.batches, onLog);
|
|
1162
|
+
protocolName = desc.protocolName;
|
|
1163
|
+
methodCache = new Map(desc.methods.map((m) => [m.name, m]));
|
|
1164
|
+
return methodCache;
|
|
1165
|
+
} finally {
|
|
1166
|
+
releaseBusy();
|
|
1167
|
+
}
|
|
1168
|
+
}
|
|
1169
|
+
return {
|
|
1170
|
+
async call(method, params) {
|
|
1171
|
+
const methods = await ensureMethodCache();
|
|
1172
|
+
await acquireBusy();
|
|
1173
|
+
try {
|
|
1174
|
+
const info = methods.get(method);
|
|
1175
|
+
if (!info) {
|
|
1176
|
+
throw new Error(`Unknown method: '${method}'`);
|
|
1346
1177
|
}
|
|
1347
|
-
|
|
1348
|
-
|
|
1349
|
-
|
|
1350
|
-
|
|
1178
|
+
const r = await ensureReader();
|
|
1179
|
+
const fullParams = { ...info.defaults ?? {}, ...params ?? {} };
|
|
1180
|
+
const body = buildRequestIpc(info.paramsSchema, fullParams, method);
|
|
1181
|
+
writeFn(body);
|
|
1182
|
+
const response = await r.readStream();
|
|
1183
|
+
if (!response) {
|
|
1184
|
+
throw new Error("EOF reading response");
|
|
1351
1185
|
}
|
|
1352
|
-
|
|
1353
|
-
|
|
1354
|
-
|
|
1355
|
-
|
|
1186
|
+
let resultBatch = null;
|
|
1187
|
+
for (const batch of response.batches) {
|
|
1188
|
+
if (batch.numRows === 0) {
|
|
1189
|
+
dispatchLogOrError(batch, onLog);
|
|
1190
|
+
continue;
|
|
1191
|
+
}
|
|
1192
|
+
resultBatch = batch;
|
|
1356
1193
|
}
|
|
1357
|
-
|
|
1194
|
+
if (!resultBatch) {
|
|
1195
|
+
return null;
|
|
1196
|
+
}
|
|
1197
|
+
const rows = extractBatchRows(resultBatch);
|
|
1198
|
+
if (rows.length === 0)
|
|
1199
|
+
return null;
|
|
1200
|
+
if (info.resultSchema.fields.length === 0)
|
|
1201
|
+
return null;
|
|
1202
|
+
return rows[0];
|
|
1203
|
+
} finally {
|
|
1204
|
+
releaseBusy();
|
|
1358
1205
|
}
|
|
1359
|
-
|
|
1360
|
-
|
|
1361
|
-
|
|
1362
|
-
|
|
1363
|
-
|
|
1206
|
+
},
|
|
1207
|
+
async stream(method, params) {
|
|
1208
|
+
const methods = await ensureMethodCache();
|
|
1209
|
+
await acquireBusy();
|
|
1210
|
+
try {
|
|
1211
|
+
const info = methods.get(method);
|
|
1212
|
+
if (!info) {
|
|
1213
|
+
throw new Error(`Unknown method: '${method}'`);
|
|
1214
|
+
}
|
|
1215
|
+
const r = await ensureReader();
|
|
1216
|
+
const fullParams = { ...info.defaults ?? {}, ...params ?? {} };
|
|
1217
|
+
const body = buildRequestIpc(info.paramsSchema, fullParams, method);
|
|
1218
|
+
writeFn(body);
|
|
1219
|
+
let header = null;
|
|
1220
|
+
if (info.headerSchema) {
|
|
1221
|
+
const headerStream = await r.readStream();
|
|
1222
|
+
if (headerStream) {
|
|
1223
|
+
for (const batch of headerStream.batches) {
|
|
1224
|
+
if (batch.numRows === 0) {
|
|
1225
|
+
dispatchLogOrError(batch, onLog);
|
|
1226
|
+
continue;
|
|
1227
|
+
}
|
|
1228
|
+
const rows = extractBatchRows(batch);
|
|
1229
|
+
if (rows.length > 0) {
|
|
1230
|
+
header = rows[0];
|
|
1231
|
+
}
|
|
1232
|
+
}
|
|
1233
|
+
}
|
|
1234
|
+
}
|
|
1235
|
+
const outputSchema = info.outputSchema ?? info.resultSchema;
|
|
1236
|
+
return new PipeStreamSession({
|
|
1237
|
+
reader: r,
|
|
1238
|
+
writeFn,
|
|
1239
|
+
onLog,
|
|
1240
|
+
header,
|
|
1241
|
+
outputSchema,
|
|
1242
|
+
releaseBusy,
|
|
1243
|
+
setDrainPromise
|
|
1244
|
+
});
|
|
1245
|
+
} catch (e) {
|
|
1246
|
+
try {
|
|
1247
|
+
const r = await ensureReader();
|
|
1248
|
+
const emptySchema = new Schema2([]);
|
|
1249
|
+
const ipc = serializeIpcStream(emptySchema, []);
|
|
1250
|
+
writeFn(ipc);
|
|
1251
|
+
const outStream = await r.readStream();
|
|
1252
|
+
} catch {}
|
|
1253
|
+
releaseBusy();
|
|
1254
|
+
throw e;
|
|
1364
1255
|
}
|
|
1365
|
-
|
|
1256
|
+
},
|
|
1257
|
+
async describe() {
|
|
1258
|
+
const methods = await ensureMethodCache();
|
|
1259
|
+
return {
|
|
1260
|
+
protocolName,
|
|
1261
|
+
methods: [...methods.values()]
|
|
1262
|
+
};
|
|
1263
|
+
},
|
|
1264
|
+
close() {
|
|
1265
|
+
if (closed)
|
|
1266
|
+
return;
|
|
1267
|
+
closed = true;
|
|
1268
|
+
writable.end();
|
|
1366
1269
|
}
|
|
1367
1270
|
};
|
|
1368
1271
|
}
|
|
1369
|
-
|
|
1272
|
+
function subprocessConnect(cmd, options) {
|
|
1273
|
+
const proc = Bun.spawn(cmd, {
|
|
1274
|
+
stdin: "pipe",
|
|
1275
|
+
stdout: "pipe",
|
|
1276
|
+
stderr: options?.stderr ?? "ignore",
|
|
1277
|
+
cwd: options?.cwd,
|
|
1278
|
+
env: options?.env ? { ...process.env, ...options.env } : undefined
|
|
1279
|
+
});
|
|
1280
|
+
const stdout = proc.stdout;
|
|
1281
|
+
const writable = {
|
|
1282
|
+
write(data) {
|
|
1283
|
+
proc.stdin.write(data);
|
|
1284
|
+
},
|
|
1285
|
+
flush() {
|
|
1286
|
+
proc.stdin.flush();
|
|
1287
|
+
},
|
|
1288
|
+
end() {
|
|
1289
|
+
proc.stdin.end();
|
|
1290
|
+
}
|
|
1291
|
+
};
|
|
1292
|
+
const client = pipeConnect(stdout, writable, {
|
|
1293
|
+
onLog: options?.onLog
|
|
1294
|
+
});
|
|
1295
|
+
const originalClose = client.close;
|
|
1296
|
+
client.close = () => {
|
|
1297
|
+
originalClose.call(client);
|
|
1298
|
+
try {
|
|
1299
|
+
proc.kill();
|
|
1300
|
+
} catch {}
|
|
1301
|
+
};
|
|
1302
|
+
return client;
|
|
1303
|
+
}
|
|
1304
|
+
// src/http/handler.ts
|
|
1305
|
+
import { randomBytes } from "node:crypto";
|
|
1306
|
+
import { Schema as Schema5 } from "@query-farm/apache-arrow";
|
|
1307
|
+
|
|
1308
|
+
// src/types.ts
|
|
1309
|
+
import { RecordBatch as RecordBatch6, recordBatchFromArrays } from "@query-farm/apache-arrow";
|
|
1310
|
+
|
|
1311
|
+
// src/wire/response.ts
|
|
1370
1312
|
import {
|
|
1371
|
-
|
|
1372
|
-
|
|
1373
|
-
|
|
1374
|
-
|
|
1375
|
-
|
|
1376
|
-
|
|
1377
|
-
|
|
1378
|
-
|
|
1379
|
-
|
|
1380
|
-
|
|
1381
|
-
|
|
1382
|
-
|
|
1383
|
-
|
|
1384
|
-
|
|
1385
|
-
|
|
1386
|
-
|
|
1387
|
-
|
|
1388
|
-
|
|
1389
|
-
|
|
1390
|
-
|
|
1391
|
-
|
|
1392
|
-
|
|
1393
|
-
return new Binary3;
|
|
1394
|
-
return new Utf83;
|
|
1313
|
+
Data,
|
|
1314
|
+
DataType as DataType2,
|
|
1315
|
+
makeData as makeData5,
|
|
1316
|
+
RecordBatch as RecordBatch5,
|
|
1317
|
+
Struct as Struct5,
|
|
1318
|
+
vectorFromArray as vectorFromArray4
|
|
1319
|
+
} from "@query-farm/apache-arrow";
|
|
1320
|
+
function coerceInt64(schema, values) {
|
|
1321
|
+
const result = { ...values };
|
|
1322
|
+
for (const field of schema.fields) {
|
|
1323
|
+
const val = result[field.name];
|
|
1324
|
+
if (val === undefined)
|
|
1325
|
+
continue;
|
|
1326
|
+
if (!DataType2.isInt(field.type) || field.type.bitWidth !== 64)
|
|
1327
|
+
continue;
|
|
1328
|
+
if (Array.isArray(val)) {
|
|
1329
|
+
result[field.name] = val.map((v) => typeof v === "number" ? BigInt(v) : v);
|
|
1330
|
+
} else if (typeof val === "number") {
|
|
1331
|
+
result[field.name] = BigInt(val);
|
|
1332
|
+
}
|
|
1333
|
+
}
|
|
1334
|
+
return result;
|
|
1395
1335
|
}
|
|
1396
|
-
function
|
|
1397
|
-
|
|
1398
|
-
|
|
1399
|
-
if (
|
|
1400
|
-
|
|
1401
|
-
return BigInt(value);
|
|
1402
|
-
return value;
|
|
1336
|
+
function buildResultBatch(schema, values, serverId, requestId) {
|
|
1337
|
+
const metadata = new Map;
|
|
1338
|
+
metadata.set(SERVER_ID_KEY, serverId);
|
|
1339
|
+
if (requestId !== null) {
|
|
1340
|
+
metadata.set(REQUEST_ID_KEY, requestId);
|
|
1403
1341
|
}
|
|
1404
|
-
if (
|
|
1405
|
-
|
|
1406
|
-
|
|
1407
|
-
|
|
1408
|
-
|
|
1409
|
-
|
|
1410
|
-
|
|
1411
|
-
}
|
|
1412
|
-
return coerced;
|
|
1342
|
+
if (schema.fields.length === 0) {
|
|
1343
|
+
return buildEmptyBatch(schema, metadata);
|
|
1344
|
+
}
|
|
1345
|
+
for (const field of schema.fields) {
|
|
1346
|
+
if (values[field.name] === undefined && !field.nullable) {
|
|
1347
|
+
const got = Object.keys(values);
|
|
1348
|
+
throw new TypeError(`Handler result missing required field '${field.name}'. Got keys: [${got.join(", ")}]`);
|
|
1413
1349
|
}
|
|
1414
|
-
return value;
|
|
1415
1350
|
}
|
|
1416
|
-
|
|
1417
|
-
|
|
1418
|
-
|
|
1419
|
-
|
|
1351
|
+
const coerced = coerceInt64(schema, values);
|
|
1352
|
+
const children = schema.fields.map((f) => {
|
|
1353
|
+
const val = coerced[f.name];
|
|
1354
|
+
if (val instanceof Data) {
|
|
1355
|
+
return val;
|
|
1420
1356
|
}
|
|
1421
|
-
|
|
1357
|
+
const arr = vectorFromArray4([val], f.type);
|
|
1358
|
+
return arr.data[0];
|
|
1359
|
+
});
|
|
1360
|
+
const structType = new Struct5(schema.fields);
|
|
1361
|
+
const data = makeData5({
|
|
1362
|
+
type: structType,
|
|
1363
|
+
length: 1,
|
|
1364
|
+
children,
|
|
1365
|
+
nullCount: 0
|
|
1366
|
+
});
|
|
1367
|
+
return new RecordBatch5(schema, data, metadata);
|
|
1368
|
+
}
|
|
1369
|
+
function buildErrorBatch(schema, error, serverId, requestId) {
|
|
1370
|
+
const metadata = new Map;
|
|
1371
|
+
metadata.set(LOG_LEVEL_KEY, "EXCEPTION");
|
|
1372
|
+
metadata.set(LOG_MESSAGE_KEY, `${error.constructor.name}: ${error.message}`);
|
|
1373
|
+
const extra = {
|
|
1374
|
+
exception_type: error.constructor.name,
|
|
1375
|
+
exception_message: error.message,
|
|
1376
|
+
traceback: error.stack ?? ""
|
|
1377
|
+
};
|
|
1378
|
+
metadata.set(LOG_EXTRA_KEY, JSON.stringify(extra));
|
|
1379
|
+
metadata.set(SERVER_ID_KEY, serverId);
|
|
1380
|
+
if (requestId !== null) {
|
|
1381
|
+
metadata.set(REQUEST_ID_KEY, requestId);
|
|
1422
1382
|
}
|
|
1423
|
-
return
|
|
1383
|
+
return buildEmptyBatch(schema, metadata);
|
|
1424
1384
|
}
|
|
1425
|
-
function
|
|
1385
|
+
function buildLogBatch(schema, level, message, extra, serverId, requestId) {
|
|
1426
1386
|
const metadata = new Map;
|
|
1427
|
-
metadata.set(
|
|
1428
|
-
metadata.set(
|
|
1387
|
+
metadata.set(LOG_LEVEL_KEY, level);
|
|
1388
|
+
metadata.set(LOG_MESSAGE_KEY, message);
|
|
1389
|
+
if (extra) {
|
|
1390
|
+
metadata.set(LOG_EXTRA_KEY, JSON.stringify(extra));
|
|
1391
|
+
}
|
|
1392
|
+
if (serverId != null) {
|
|
1393
|
+
metadata.set(SERVER_ID_KEY, serverId);
|
|
1394
|
+
}
|
|
1395
|
+
if (requestId != null) {
|
|
1396
|
+
metadata.set(REQUEST_ID_KEY, requestId);
|
|
1397
|
+
}
|
|
1398
|
+
return buildEmptyBatch(schema, metadata);
|
|
1399
|
+
}
|
|
1400
|
+
function buildEmptyBatch(schema, metadata) {
|
|
1401
|
+
const children = schema.fields.map((f) => {
|
|
1402
|
+
return makeData5({ type: f.type, length: 0, nullCount: 0 });
|
|
1403
|
+
});
|
|
1429
1404
|
if (schema.fields.length === 0) {
|
|
1430
|
-
const structType2 = new
|
|
1431
|
-
const data2 =
|
|
1405
|
+
const structType2 = new Struct5(schema.fields);
|
|
1406
|
+
const data2 = makeData5({
|
|
1432
1407
|
type: structType2,
|
|
1433
|
-
length:
|
|
1408
|
+
length: 0,
|
|
1434
1409
|
children: [],
|
|
1435
1410
|
nullCount: 0
|
|
1436
1411
|
});
|
|
1437
|
-
|
|
1438
|
-
return serializeIpcStream(schema, [batch2]);
|
|
1412
|
+
return new RecordBatch5(schema, data2, metadata);
|
|
1439
1413
|
}
|
|
1440
|
-
const
|
|
1441
|
-
|
|
1442
|
-
return vectorFromArray3([val], f.type).data[0];
|
|
1443
|
-
});
|
|
1444
|
-
const structType = new Struct4(schema.fields);
|
|
1445
|
-
const data = makeData4({
|
|
1414
|
+
const structType = new Struct5(schema.fields);
|
|
1415
|
+
const data = makeData5({
|
|
1446
1416
|
type: structType,
|
|
1447
|
-
length:
|
|
1417
|
+
length: 0,
|
|
1448
1418
|
children,
|
|
1449
1419
|
nullCount: 0
|
|
1450
1420
|
});
|
|
1451
|
-
|
|
1452
|
-
return serializeIpcStream(schema, [batch]);
|
|
1421
|
+
return new RecordBatch5(schema, data, metadata);
|
|
1453
1422
|
}
|
|
1454
|
-
|
|
1455
|
-
|
|
1456
|
-
|
|
1457
|
-
|
|
1458
|
-
|
|
1459
|
-
|
|
1423
|
+
|
|
1424
|
+
// src/types.ts
|
|
1425
|
+
var MethodType;
|
|
1426
|
+
((MethodType2) => {
|
|
1427
|
+
MethodType2["UNARY"] = "unary";
|
|
1428
|
+
MethodType2["STREAM"] = "stream";
|
|
1429
|
+
})(MethodType ||= {});
|
|
1430
|
+
|
|
1431
|
+
class OutputCollector {
|
|
1432
|
+
_batches = [];
|
|
1433
|
+
_dataBatchIdx = null;
|
|
1434
|
+
_finished = false;
|
|
1435
|
+
_producerMode;
|
|
1436
|
+
_outputSchema;
|
|
1437
|
+
_serverId;
|
|
1438
|
+
_requestId;
|
|
1439
|
+
constructor(outputSchema, producerMode = true, serverId = "", requestId = null) {
|
|
1440
|
+
this._outputSchema = outputSchema;
|
|
1441
|
+
this._producerMode = producerMode;
|
|
1442
|
+
this._serverId = serverId;
|
|
1443
|
+
this._requestId = requestId;
|
|
1444
|
+
}
|
|
1445
|
+
get outputSchema() {
|
|
1446
|
+
return this._outputSchema;
|
|
1447
|
+
}
|
|
1448
|
+
get finished() {
|
|
1449
|
+
return this._finished;
|
|
1450
|
+
}
|
|
1451
|
+
get batches() {
|
|
1452
|
+
return this._batches;
|
|
1453
|
+
}
|
|
1454
|
+
emit(batchOrColumns, metadata) {
|
|
1455
|
+
let batch;
|
|
1456
|
+
if (batchOrColumns instanceof RecordBatch6) {
|
|
1457
|
+
batch = batchOrColumns;
|
|
1458
|
+
} else {
|
|
1459
|
+
const coerced = coerceInt64(this._outputSchema, batchOrColumns);
|
|
1460
|
+
batch = recordBatchFromArrays(coerced, this._outputSchema);
|
|
1461
|
+
}
|
|
1462
|
+
if (this._dataBatchIdx !== null) {
|
|
1463
|
+
throw new Error("Only one data batch may be emitted per call");
|
|
1464
|
+
}
|
|
1465
|
+
this._dataBatchIdx = this._batches.length;
|
|
1466
|
+
this._batches.push({ batch, metadata });
|
|
1467
|
+
}
|
|
1468
|
+
emitRow(values) {
|
|
1469
|
+
const columns = {};
|
|
1470
|
+
for (const [key, value] of Object.entries(values)) {
|
|
1471
|
+
columns[key] = [value];
|
|
1472
|
+
}
|
|
1473
|
+
this.emit(columns);
|
|
1474
|
+
}
|
|
1475
|
+
finish() {
|
|
1476
|
+
if (!this._producerMode) {
|
|
1477
|
+
throw new Error("finish() is not allowed on exchange streams; " + "exchange streams must emit exactly one data batch per call");
|
|
1478
|
+
}
|
|
1479
|
+
this._finished = true;
|
|
1480
|
+
}
|
|
1481
|
+
clientLog(level, message, extra) {
|
|
1482
|
+
const batch = buildLogBatch(this._outputSchema, level, message, extra, this._serverId, this._requestId);
|
|
1483
|
+
this._batches.push({ batch });
|
|
1460
1484
|
}
|
|
1461
|
-
const batches = reader.readAll();
|
|
1462
|
-
return { schema, batches };
|
|
1463
1485
|
}
|
|
1464
|
-
|
|
1465
|
-
|
|
1466
|
-
|
|
1467
|
-
|
|
1468
|
-
|
|
1469
|
-
|
|
1470
|
-
|
|
1471
|
-
|
|
1472
|
-
|
|
1473
|
-
|
|
1474
|
-
|
|
1475
|
-
|
|
1476
|
-
|
|
1477
|
-
|
|
1478
|
-
|
|
1479
|
-
|
|
1480
|
-
|
|
1481
|
-
|
|
1482
|
-
|
|
1483
|
-
|
|
1486
|
+
|
|
1487
|
+
// src/http/handler.ts
|
|
1488
|
+
init_zstd();
|
|
1489
|
+
|
|
1490
|
+
// src/http/dispatch.ts
|
|
1491
|
+
import { RecordBatch as RecordBatch8, RecordBatchReader as RecordBatchReader5, Schema as Schema4 } from "@query-farm/apache-arrow";
|
|
1492
|
+
|
|
1493
|
+
// src/dispatch/describe.ts
|
|
1494
|
+
import {
|
|
1495
|
+
Binary as Binary2,
|
|
1496
|
+
Bool as Bool2,
|
|
1497
|
+
Field as Field3,
|
|
1498
|
+
makeData as makeData6,
|
|
1499
|
+
RecordBatch as RecordBatch7,
|
|
1500
|
+
Schema as Schema3,
|
|
1501
|
+
Struct as Struct6,
|
|
1502
|
+
Utf8 as Utf82,
|
|
1503
|
+
vectorFromArray as vectorFromArray5
|
|
1504
|
+
} from "@query-farm/apache-arrow";
|
|
1505
|
+
|
|
1506
|
+
// src/util/schema.ts
|
|
1507
|
+
import { RecordBatchStreamWriter as RecordBatchStreamWriter3 } from "@query-farm/apache-arrow";
|
|
1508
|
+
function serializeSchema(schema) {
|
|
1509
|
+
const writer = new RecordBatchStreamWriter3;
|
|
1510
|
+
writer.reset(undefined, schema);
|
|
1511
|
+
writer.close();
|
|
1512
|
+
return writer.toUint8Array(true);
|
|
1513
|
+
}
|
|
1514
|
+
|
|
1515
|
+
// src/dispatch/describe.ts
|
|
1516
|
+
var DESCRIBE_SCHEMA = new Schema3([
|
|
1517
|
+
new Field3("name", new Utf82, false),
|
|
1518
|
+
new Field3("method_type", new Utf82, false),
|
|
1519
|
+
new Field3("doc", new Utf82, true),
|
|
1520
|
+
new Field3("has_return", new Bool2, false),
|
|
1521
|
+
new Field3("params_schema_ipc", new Binary2, false),
|
|
1522
|
+
new Field3("result_schema_ipc", new Binary2, false),
|
|
1523
|
+
new Field3("param_types_json", new Utf82, true),
|
|
1524
|
+
new Field3("param_defaults_json", new Utf82, true),
|
|
1525
|
+
new Field3("has_header", new Bool2, false),
|
|
1526
|
+
new Field3("header_schema_ipc", new Binary2, true)
|
|
1527
|
+
]);
|
|
1528
|
+
function buildDescribeBatch(protocolName, methods, serverId) {
|
|
1529
|
+
const sortedEntries = [...methods.entries()].sort(([a], [b]) => a.localeCompare(b));
|
|
1530
|
+
const names = [];
|
|
1531
|
+
const methodTypes = [];
|
|
1532
|
+
const docs = [];
|
|
1533
|
+
const hasReturns = [];
|
|
1534
|
+
const paramsSchemas = [];
|
|
1535
|
+
const resultSchemas = [];
|
|
1536
|
+
const paramTypesJsons = [];
|
|
1537
|
+
const paramDefaultsJsons = [];
|
|
1538
|
+
const hasHeaders = [];
|
|
1539
|
+
const headerSchemas = [];
|
|
1540
|
+
for (const [name, method] of sortedEntries) {
|
|
1541
|
+
names.push(name);
|
|
1542
|
+
methodTypes.push(method.type);
|
|
1543
|
+
docs.push(method.doc ?? null);
|
|
1544
|
+
const hasReturn = method.type === "unary" && method.resultSchema.fields.length > 0;
|
|
1545
|
+
hasReturns.push(hasReturn);
|
|
1546
|
+
paramsSchemas.push(serializeSchema(method.paramsSchema));
|
|
1547
|
+
resultSchemas.push(serializeSchema(method.resultSchema));
|
|
1548
|
+
if (method.paramTypes && Object.keys(method.paramTypes).length > 0) {
|
|
1549
|
+
paramTypesJsons.push(JSON.stringify(method.paramTypes));
|
|
1550
|
+
} else {
|
|
1551
|
+
paramTypesJsons.push(null);
|
|
1484
1552
|
}
|
|
1485
|
-
|
|
1553
|
+
if (method.defaults && Object.keys(method.defaults).length > 0) {
|
|
1554
|
+
const safe = {};
|
|
1555
|
+
for (const [k, v] of Object.entries(method.defaults)) {
|
|
1556
|
+
if (v === null || typeof v === "string" || typeof v === "number" || typeof v === "boolean") {
|
|
1557
|
+
safe[k] = v;
|
|
1558
|
+
}
|
|
1559
|
+
}
|
|
1560
|
+
paramDefaultsJsons.push(Object.keys(safe).length > 0 ? JSON.stringify(safe) : null);
|
|
1561
|
+
} else {
|
|
1562
|
+
paramDefaultsJsons.push(null);
|
|
1563
|
+
}
|
|
1564
|
+
hasHeaders.push(!!method.headerSchema);
|
|
1565
|
+
headerSchemas.push(method.headerSchema ? serializeSchema(method.headerSchema) : null);
|
|
1566
|
+
}
|
|
1567
|
+
const nameArr = vectorFromArray5(names, new Utf82);
|
|
1568
|
+
const methodTypeArr = vectorFromArray5(methodTypes, new Utf82);
|
|
1569
|
+
const docArr = vectorFromArray5(docs, new Utf82);
|
|
1570
|
+
const hasReturnArr = vectorFromArray5(hasReturns, new Bool2);
|
|
1571
|
+
const paramsSchemaArr = vectorFromArray5(paramsSchemas, new Binary2);
|
|
1572
|
+
const resultSchemaArr = vectorFromArray5(resultSchemas, new Binary2);
|
|
1573
|
+
const paramTypesArr = vectorFromArray5(paramTypesJsons, new Utf82);
|
|
1574
|
+
const paramDefaultsArr = vectorFromArray5(paramDefaultsJsons, new Utf82);
|
|
1575
|
+
const hasHeaderArr = vectorFromArray5(hasHeaders, new Bool2);
|
|
1576
|
+
const headerSchemaArr = vectorFromArray5(headerSchemas, new Binary2);
|
|
1577
|
+
const children = [
|
|
1578
|
+
nameArr.data[0],
|
|
1579
|
+
methodTypeArr.data[0],
|
|
1580
|
+
docArr.data[0],
|
|
1581
|
+
hasReturnArr.data[0],
|
|
1582
|
+
paramsSchemaArr.data[0],
|
|
1583
|
+
resultSchemaArr.data[0],
|
|
1584
|
+
paramTypesArr.data[0],
|
|
1585
|
+
paramDefaultsArr.data[0],
|
|
1586
|
+
hasHeaderArr.data[0],
|
|
1587
|
+
headerSchemaArr.data[0]
|
|
1588
|
+
];
|
|
1589
|
+
const structType = new Struct6(DESCRIBE_SCHEMA.fields);
|
|
1590
|
+
const data = makeData6({
|
|
1591
|
+
type: structType,
|
|
1592
|
+
length: sortedEntries.length,
|
|
1593
|
+
children,
|
|
1594
|
+
nullCount: 0
|
|
1595
|
+
});
|
|
1596
|
+
const metadata = new Map;
|
|
1597
|
+
metadata.set(PROTOCOL_NAME_KEY, protocolName);
|
|
1598
|
+
metadata.set(REQUEST_VERSION_KEY, REQUEST_VERSION);
|
|
1599
|
+
metadata.set(DESCRIBE_VERSION_KEY, DESCRIBE_VERSION);
|
|
1600
|
+
metadata.set(SERVER_ID_KEY, serverId);
|
|
1601
|
+
const batch = new RecordBatch7(DESCRIBE_SCHEMA, data, metadata);
|
|
1602
|
+
return { batch, metadata };
|
|
1603
|
+
}
|
|
1604
|
+
|
|
1605
|
+
// src/wire/request.ts
|
|
1606
|
+
import { DataType as DataType3 } from "@query-farm/apache-arrow";
|
|
1607
|
+
function parseRequest(schema, batch) {
|
|
1608
|
+
const metadata = batch.metadata ?? new Map;
|
|
1609
|
+
const methodName = metadata.get(RPC_METHOD_KEY);
|
|
1610
|
+
if (methodName === undefined) {
|
|
1611
|
+
throw new RpcError("ProtocolError", "Missing 'vgi_rpc.method' in request batch custom_metadata. " + "Each request batch must carry a 'vgi_rpc.method' key in its Arrow IPC custom_metadata " + "with the method name as a UTF-8 string.", "");
|
|
1486
1612
|
}
|
|
1487
|
-
|
|
1488
|
-
|
|
1489
|
-
|
|
1490
|
-
if (extraStr) {
|
|
1491
|
-
try {
|
|
1492
|
-
extra = JSON.parse(extraStr);
|
|
1493
|
-
} catch {}
|
|
1494
|
-
}
|
|
1495
|
-
onLog({ level, message, extra });
|
|
1613
|
+
const version = metadata.get(REQUEST_VERSION_KEY);
|
|
1614
|
+
if (version === undefined) {
|
|
1615
|
+
throw new VersionError("Missing 'vgi_rpc.request_version' in request batch custom_metadata. " + `Set the 'vgi_rpc.request_version' custom_metadata value to '${REQUEST_VERSION}'.`);
|
|
1496
1616
|
}
|
|
1497
|
-
|
|
1498
|
-
}
|
|
1499
|
-
|
|
1500
|
-
const
|
|
1501
|
-
|
|
1502
|
-
|
|
1503
|
-
|
|
1504
|
-
|
|
1505
|
-
|
|
1506
|
-
|
|
1507
|
-
|
|
1508
|
-
|
|
1509
|
-
|
|
1617
|
+
if (version !== REQUEST_VERSION) {
|
|
1618
|
+
throw new VersionError(`Unsupported request version '${version}', expected '${REQUEST_VERSION}'. ` + `Set the 'vgi_rpc.request_version' custom_metadata value to '${REQUEST_VERSION}'.`);
|
|
1619
|
+
}
|
|
1620
|
+
const requestId = metadata.get(REQUEST_ID_KEY) ?? null;
|
|
1621
|
+
const params = {};
|
|
1622
|
+
if (schema.fields.length > 0 && batch.numRows !== 1) {
|
|
1623
|
+
throw new RpcError("ProtocolError", `Expected 1 row in request batch, got ${batch.numRows}. ` + "Each parameter is a column (not a row). The batch should have exactly 1 row.", "");
|
|
1624
|
+
}
|
|
1625
|
+
for (let i = 0;i < schema.fields.length; i++) {
|
|
1626
|
+
const field = schema.fields[i];
|
|
1627
|
+
if (DataType3.isMap(field.type)) {
|
|
1628
|
+
params[field.name] = batch.getChildAt(i).data[0];
|
|
1629
|
+
continue;
|
|
1630
|
+
}
|
|
1631
|
+
let value = batch.getChildAt(i)?.get(0);
|
|
1632
|
+
if (typeof value === "bigint") {
|
|
1633
|
+
if (value >= BigInt(Number.MIN_SAFE_INTEGER) && value <= BigInt(Number.MAX_SAFE_INTEGER)) {
|
|
1634
|
+
value = Number(value);
|
|
1510
1635
|
}
|
|
1511
|
-
row[field.name] = value;
|
|
1512
1636
|
}
|
|
1513
|
-
|
|
1637
|
+
params[field.name] = value;
|
|
1514
1638
|
}
|
|
1515
|
-
return
|
|
1516
|
-
|
|
1517
|
-
|
|
1518
|
-
|
|
1519
|
-
|
|
1520
|
-
|
|
1521
|
-
|
|
1522
|
-
|
|
1523
|
-
});
|
|
1524
|
-
return IpcStreamReader.create(stream);
|
|
1639
|
+
return {
|
|
1640
|
+
methodName,
|
|
1641
|
+
requestVersion: version,
|
|
1642
|
+
requestId,
|
|
1643
|
+
schema,
|
|
1644
|
+
params,
|
|
1645
|
+
rawMetadata: metadata
|
|
1646
|
+
};
|
|
1525
1647
|
}
|
|
1526
1648
|
|
|
1527
|
-
// src/
|
|
1528
|
-
import {
|
|
1529
|
-
|
|
1530
|
-
|
|
1531
|
-
|
|
1532
|
-
|
|
1533
|
-
|
|
1649
|
+
// src/http/token.ts
|
|
1650
|
+
import { createHmac, timingSafeEqual } from "node:crypto";
|
|
1651
|
+
var TOKEN_VERSION = 2;
|
|
1652
|
+
var HMAC_LEN = 32;
|
|
1653
|
+
var MIN_TOKEN_LEN = 1 + 8 + 12 + HMAC_LEN;
|
|
1654
|
+
function packStateToken(stateBytes, schemaBytes, inputSchemaBytes, signingKey, createdAt) {
|
|
1655
|
+
const now = createdAt ?? Math.floor(Date.now() / 1000);
|
|
1656
|
+
const payloadLen = 1 + 8 + 4 + stateBytes.length + 4 + schemaBytes.length + 4 + inputSchemaBytes.length;
|
|
1657
|
+
const buf = Buffer.alloc(payloadLen);
|
|
1658
|
+
let offset = 0;
|
|
1659
|
+
buf.writeUInt8(TOKEN_VERSION, offset);
|
|
1660
|
+
offset += 1;
|
|
1661
|
+
buf.writeBigUInt64LE(BigInt(now), offset);
|
|
1662
|
+
offset += 8;
|
|
1663
|
+
buf.writeUInt32LE(stateBytes.length, offset);
|
|
1664
|
+
offset += 4;
|
|
1665
|
+
buf.set(stateBytes, offset);
|
|
1666
|
+
offset += stateBytes.length;
|
|
1667
|
+
buf.writeUInt32LE(schemaBytes.length, offset);
|
|
1668
|
+
offset += 4;
|
|
1669
|
+
buf.set(schemaBytes, offset);
|
|
1670
|
+
offset += schemaBytes.length;
|
|
1671
|
+
buf.writeUInt32LE(inputSchemaBytes.length, offset);
|
|
1672
|
+
offset += 4;
|
|
1673
|
+
buf.set(inputSchemaBytes, offset);
|
|
1674
|
+
offset += inputSchemaBytes.length;
|
|
1675
|
+
const mac = createHmac("sha256", signingKey).update(buf).digest();
|
|
1676
|
+
const token = Buffer.concat([buf, mac]);
|
|
1677
|
+
return token.toString("base64");
|
|
1534
1678
|
}
|
|
1535
|
-
|
|
1536
|
-
|
|
1537
|
-
|
|
1538
|
-
|
|
1539
|
-
dispatchLogOrError(batch, onLog);
|
|
1540
|
-
continue;
|
|
1541
|
-
}
|
|
1542
|
-
dataBatch = batch;
|
|
1679
|
+
function unpackStateToken(tokenBase64, signingKey, tokenTtl) {
|
|
1680
|
+
const token = Buffer.from(tokenBase64, "base64");
|
|
1681
|
+
if (token.length < MIN_TOKEN_LEN) {
|
|
1682
|
+
throw new Error("State token too short");
|
|
1543
1683
|
}
|
|
1544
|
-
|
|
1545
|
-
|
|
1684
|
+
const payload = token.subarray(0, token.length - HMAC_LEN);
|
|
1685
|
+
const receivedMac = token.subarray(token.length - HMAC_LEN);
|
|
1686
|
+
const expectedMac = createHmac("sha256", signingKey).update(payload).digest();
|
|
1687
|
+
if (!timingSafeEqual(receivedMac, expectedMac)) {
|
|
1688
|
+
throw new Error("State token HMAC verification failed");
|
|
1546
1689
|
}
|
|
1547
|
-
|
|
1548
|
-
const
|
|
1549
|
-
|
|
1550
|
-
|
|
1551
|
-
|
|
1552
|
-
|
|
1553
|
-
|
|
1554
|
-
|
|
1555
|
-
|
|
1556
|
-
const
|
|
1557
|
-
|
|
1558
|
-
|
|
1559
|
-
const hasHeader = dataBatch.getChildAt(8).get(i);
|
|
1560
|
-
const headerIpc = dataBatch.getChildAt(9)?.get(i);
|
|
1561
|
-
const paramsSchema = await deserializeSchema(paramsIpc);
|
|
1562
|
-
const resultSchema = await deserializeSchema(resultIpc);
|
|
1563
|
-
let paramTypes;
|
|
1564
|
-
if (paramTypesJson) {
|
|
1565
|
-
try {
|
|
1566
|
-
paramTypes = JSON.parse(paramTypesJson);
|
|
1567
|
-
} catch {}
|
|
1568
|
-
}
|
|
1569
|
-
let defaults;
|
|
1570
|
-
if (paramDefaultsJson) {
|
|
1571
|
-
try {
|
|
1572
|
-
defaults = JSON.parse(paramDefaultsJson);
|
|
1573
|
-
} catch {}
|
|
1574
|
-
}
|
|
1575
|
-
const info = {
|
|
1576
|
-
name,
|
|
1577
|
-
type: methodType,
|
|
1578
|
-
paramsSchema,
|
|
1579
|
-
resultSchema,
|
|
1580
|
-
doc: doc ?? undefined,
|
|
1581
|
-
paramTypes,
|
|
1582
|
-
defaults
|
|
1583
|
-
};
|
|
1584
|
-
if (methodType === "stream") {
|
|
1585
|
-
info.outputSchema = resultSchema;
|
|
1586
|
-
}
|
|
1587
|
-
if (hasHeader && headerIpc) {
|
|
1588
|
-
info.headerSchema = await deserializeSchema(headerIpc);
|
|
1690
|
+
let offset = 0;
|
|
1691
|
+
const version = payload.readUInt8(offset);
|
|
1692
|
+
offset += 1;
|
|
1693
|
+
if (version !== TOKEN_VERSION) {
|
|
1694
|
+
throw new Error(`Unsupported state token version: ${version}`);
|
|
1695
|
+
}
|
|
1696
|
+
const createdAt = Number(payload.readBigUInt64LE(offset));
|
|
1697
|
+
offset += 8;
|
|
1698
|
+
if (tokenTtl > 0) {
|
|
1699
|
+
const now = Math.floor(Date.now() / 1000);
|
|
1700
|
+
if (now - createdAt > tokenTtl) {
|
|
1701
|
+
throw new Error("State token expired");
|
|
1589
1702
|
}
|
|
1590
|
-
methods.push(info);
|
|
1591
1703
|
}
|
|
1592
|
-
|
|
1593
|
-
|
|
1594
|
-
|
|
1595
|
-
|
|
1596
|
-
|
|
1597
|
-
const
|
|
1598
|
-
|
|
1599
|
-
|
|
1600
|
-
|
|
1601
|
-
|
|
1602
|
-
|
|
1603
|
-
|
|
1604
|
-
const
|
|
1605
|
-
|
|
1704
|
+
const stateLen = payload.readUInt32LE(offset);
|
|
1705
|
+
offset += 4;
|
|
1706
|
+
if (offset + stateLen > payload.length) {
|
|
1707
|
+
throw new Error("State token truncated (state)");
|
|
1708
|
+
}
|
|
1709
|
+
const stateBytes = payload.slice(offset, offset + stateLen);
|
|
1710
|
+
offset += stateLen;
|
|
1711
|
+
const schemaLen = payload.readUInt32LE(offset);
|
|
1712
|
+
offset += 4;
|
|
1713
|
+
if (offset + schemaLen > payload.length) {
|
|
1714
|
+
throw new Error("State token truncated (schema)");
|
|
1715
|
+
}
|
|
1716
|
+
const schemaBytes = payload.slice(offset, offset + schemaLen);
|
|
1717
|
+
offset += schemaLen;
|
|
1718
|
+
const inputSchemaLen = payload.readUInt32LE(offset);
|
|
1719
|
+
offset += 4;
|
|
1720
|
+
if (offset + inputSchemaLen > payload.length) {
|
|
1721
|
+
throw new Error("State token truncated (input schema)");
|
|
1722
|
+
}
|
|
1723
|
+
const inputSchemaBytes = payload.slice(offset, offset + inputSchemaLen);
|
|
1724
|
+
return { stateBytes, schemaBytes, inputSchemaBytes, createdAt };
|
|
1606
1725
|
}
|
|
1607
1726
|
|
|
1608
|
-
// src/
|
|
1609
|
-
|
|
1610
|
-
|
|
1611
|
-
|
|
1612
|
-
|
|
1613
|
-
|
|
1614
|
-
|
|
1615
|
-
|
|
1616
|
-
}
|
|
1617
|
-
|
|
1618
|
-
|
|
1619
|
-
|
|
1620
|
-
|
|
1621
|
-
|
|
1622
|
-
|
|
1623
|
-
|
|
1624
|
-
|
|
1625
|
-
|
|
1626
|
-
_finished;
|
|
1627
|
-
_header;
|
|
1628
|
-
_compressionLevel;
|
|
1629
|
-
_compressFn;
|
|
1630
|
-
_decompressFn;
|
|
1631
|
-
constructor(opts) {
|
|
1632
|
-
this._baseUrl = opts.baseUrl;
|
|
1633
|
-
this._prefix = opts.prefix;
|
|
1634
|
-
this._method = opts.method;
|
|
1635
|
-
this._stateToken = opts.stateToken;
|
|
1636
|
-
this._outputSchema = opts.outputSchema;
|
|
1637
|
-
this._inputSchema = opts.inputSchema;
|
|
1638
|
-
this._onLog = opts.onLog;
|
|
1639
|
-
this._pendingBatches = opts.pendingBatches;
|
|
1640
|
-
this._finished = opts.finished;
|
|
1641
|
-
this._header = opts.header;
|
|
1642
|
-
this._compressionLevel = opts.compressionLevel;
|
|
1643
|
-
this._compressFn = opts.compressFn;
|
|
1644
|
-
this._decompressFn = opts.decompressFn;
|
|
1727
|
+
// src/http/dispatch.ts
|
|
1728
|
+
async function deserializeSchema2(bytes) {
|
|
1729
|
+
const reader = await RecordBatchReader5.from(bytes);
|
|
1730
|
+
await reader.open();
|
|
1731
|
+
return reader.schema;
|
|
1732
|
+
}
|
|
1733
|
+
var EMPTY_SCHEMA = new Schema4([]);
|
|
1734
|
+
function httpDispatchDescribe(protocolName, methods, serverId) {
|
|
1735
|
+
const { batch } = buildDescribeBatch(protocolName, methods, serverId);
|
|
1736
|
+
const body = serializeIpcStream(DESCRIBE_SCHEMA, [batch]);
|
|
1737
|
+
return arrowResponse(body);
|
|
1738
|
+
}
|
|
1739
|
+
async function httpDispatchUnary(method, body, ctx) {
|
|
1740
|
+
const schema = method.resultSchema;
|
|
1741
|
+
const { schema: reqSchema, batch: reqBatch } = await readRequestFromBody(body);
|
|
1742
|
+
const parsed = parseRequest(reqSchema, reqBatch);
|
|
1743
|
+
if (parsed.methodName !== method.name) {
|
|
1744
|
+
throw new HttpRpcError(`Method name in request '${parsed.methodName}' does not match URL '${method.name}'`, 400);
|
|
1645
1745
|
}
|
|
1646
|
-
|
|
1647
|
-
|
|
1746
|
+
const out = new OutputCollector(schema, true, ctx.serverId, parsed.requestId);
|
|
1747
|
+
try {
|
|
1748
|
+
const result = await method.handler(parsed.params, out);
|
|
1749
|
+
const resultBatch = buildResultBatch(schema, result, ctx.serverId, parsed.requestId);
|
|
1750
|
+
const batches = [...out.batches.map((b) => b.batch), resultBatch];
|
|
1751
|
+
return arrowResponse(serializeIpcStream(schema, batches));
|
|
1752
|
+
} catch (error) {
|
|
1753
|
+
const errBatch = buildErrorBatch(schema, error, ctx.serverId, parsed.requestId);
|
|
1754
|
+
return arrowResponse(serializeIpcStream(schema, [errBatch]), 500);
|
|
1648
1755
|
}
|
|
1649
|
-
|
|
1650
|
-
|
|
1651
|
-
|
|
1652
|
-
|
|
1653
|
-
|
|
1654
|
-
|
|
1655
|
-
|
|
1656
|
-
|
|
1657
|
-
|
|
1756
|
+
}
|
|
1757
|
+
async function httpDispatchStreamInit(method, body, ctx) {
|
|
1758
|
+
const isProducer = !!method.producerFn;
|
|
1759
|
+
const outputSchema = method.outputSchema;
|
|
1760
|
+
const inputSchema = method.inputSchema ?? EMPTY_SCHEMA;
|
|
1761
|
+
const { schema: reqSchema, batch: reqBatch } = await readRequestFromBody(body);
|
|
1762
|
+
const parsed = parseRequest(reqSchema, reqBatch);
|
|
1763
|
+
if (parsed.methodName !== method.name) {
|
|
1764
|
+
throw new HttpRpcError(`Method name in request '${parsed.methodName}' does not match URL '${method.name}'`, 400);
|
|
1658
1765
|
}
|
|
1659
|
-
|
|
1660
|
-
|
|
1661
|
-
|
|
1766
|
+
let state;
|
|
1767
|
+
try {
|
|
1768
|
+
if (isProducer) {
|
|
1769
|
+
state = await method.producerInit(parsed.params);
|
|
1770
|
+
} else {
|
|
1771
|
+
state = await method.exchangeInit(parsed.params);
|
|
1662
1772
|
}
|
|
1663
|
-
|
|
1773
|
+
} catch (error) {
|
|
1774
|
+
const errSchema = method.headerSchema ?? EMPTY_SCHEMA;
|
|
1775
|
+
const errBatch = buildErrorBatch(errSchema, error, ctx.serverId, parsed.requestId);
|
|
1776
|
+
return arrowResponse(serializeIpcStream(errSchema, [errBatch]), 500);
|
|
1664
1777
|
}
|
|
1665
|
-
|
|
1666
|
-
|
|
1667
|
-
|
|
1668
|
-
|
|
1778
|
+
const resolvedOutputSchema = state?.__outputSchema ?? outputSchema;
|
|
1779
|
+
const effectiveProducer = state?.__isProducer ?? isProducer;
|
|
1780
|
+
let headerBytes = null;
|
|
1781
|
+
if (method.headerSchema && method.headerInit) {
|
|
1782
|
+
try {
|
|
1783
|
+
const headerOut = new OutputCollector(method.headerSchema, true, ctx.serverId, parsed.requestId);
|
|
1784
|
+
const headerValues = method.headerInit(parsed.params, state, headerOut);
|
|
1785
|
+
const headerBatch = buildResultBatch(method.headerSchema, headerValues, ctx.serverId, parsed.requestId);
|
|
1786
|
+
const headerBatches = [...headerOut.batches.map((b) => b.batch), headerBatch];
|
|
1787
|
+
headerBytes = serializeIpcStream(method.headerSchema, headerBatches);
|
|
1788
|
+
} catch (error) {
|
|
1789
|
+
const errBatch = buildErrorBatch(method.headerSchema, error, ctx.serverId, parsed.requestId);
|
|
1790
|
+
return arrowResponse(serializeIpcStream(method.headerSchema, [errBatch]), 500);
|
|
1669
1791
|
}
|
|
1670
|
-
return body;
|
|
1671
1792
|
}
|
|
1672
|
-
|
|
1673
|
-
|
|
1674
|
-
|
|
1793
|
+
if (effectiveProducer) {
|
|
1794
|
+
return produceStreamResponse(method, state, resolvedOutputSchema, inputSchema, ctx, parsed.requestId, headerBytes);
|
|
1795
|
+
} else {
|
|
1796
|
+
const stateBytes = ctx.stateSerializer.serialize(state);
|
|
1797
|
+
const schemaBytes = serializeSchema(resolvedOutputSchema);
|
|
1798
|
+
const inputSchemaBytes = serializeSchema(inputSchema);
|
|
1799
|
+
const token = packStateToken(stateBytes, schemaBytes, inputSchemaBytes, ctx.signingKey);
|
|
1800
|
+
const tokenMeta = new Map;
|
|
1801
|
+
tokenMeta.set(STATE_KEY, token);
|
|
1802
|
+
const tokenBatch = buildEmptyBatch(resolvedOutputSchema, tokenMeta);
|
|
1803
|
+
const tokenStreamBytes = serializeIpcStream(resolvedOutputSchema, [tokenBatch]);
|
|
1804
|
+
let responseBody;
|
|
1805
|
+
if (headerBytes) {
|
|
1806
|
+
responseBody = concatBytes(headerBytes, tokenStreamBytes);
|
|
1807
|
+
} else {
|
|
1808
|
+
responseBody = tokenStreamBytes;
|
|
1675
1809
|
}
|
|
1676
|
-
|
|
1677
|
-
|
|
1678
|
-
|
|
1679
|
-
|
|
1680
|
-
|
|
1681
|
-
|
|
1682
|
-
|
|
1810
|
+
return arrowResponse(responseBody);
|
|
1811
|
+
}
|
|
1812
|
+
}
|
|
1813
|
+
async function httpDispatchStreamExchange(method, body, ctx) {
|
|
1814
|
+
const isProducer = !!method.producerFn;
|
|
1815
|
+
const { batch: reqBatch } = await readRequestFromBody(body);
|
|
1816
|
+
const tokenBase64 = reqBatch.metadata?.get(STATE_KEY);
|
|
1817
|
+
if (!tokenBase64) {
|
|
1818
|
+
throw new HttpRpcError("Missing state token in exchange request", 400);
|
|
1819
|
+
}
|
|
1820
|
+
let unpacked;
|
|
1821
|
+
try {
|
|
1822
|
+
unpacked = unpackStateToken(tokenBase64, ctx.signingKey, ctx.tokenTtl);
|
|
1823
|
+
} catch (error) {
|
|
1824
|
+
throw new HttpRpcError(`Invalid state token: ${error.message}`, 400);
|
|
1825
|
+
}
|
|
1826
|
+
let state;
|
|
1827
|
+
try {
|
|
1828
|
+
state = ctx.stateSerializer.deserialize(unpacked.stateBytes);
|
|
1829
|
+
} catch (error) {
|
|
1830
|
+
console.error(`[httpDispatchStreamExchange] state deserialize error:`, error.message);
|
|
1831
|
+
throw new HttpRpcError(`State deserialization failed: ${error.message}`, 500);
|
|
1832
|
+
}
|
|
1833
|
+
let outputSchema;
|
|
1834
|
+
if (unpacked.schemaBytes.length > 0) {
|
|
1835
|
+
outputSchema = await deserializeSchema2(unpacked.schemaBytes);
|
|
1836
|
+
} else {
|
|
1837
|
+
outputSchema = state?.__outputSchema ?? method.outputSchema;
|
|
1838
|
+
}
|
|
1839
|
+
let inputSchema;
|
|
1840
|
+
if (unpacked.inputSchemaBytes.length > 0) {
|
|
1841
|
+
inputSchema = await deserializeSchema2(unpacked.inputSchemaBytes);
|
|
1842
|
+
} else {
|
|
1843
|
+
inputSchema = method.inputSchema ?? EMPTY_SCHEMA;
|
|
1844
|
+
}
|
|
1845
|
+
const effectiveProducer = state?.__isProducer ?? isProducer;
|
|
1846
|
+
if (process.env.VGI_DISPATCH_DEBUG)
|
|
1847
|
+
console.error(`[httpDispatchStreamExchange] method=${method.name} effectiveProducer=${effectiveProducer} stateKeys=${Object.keys(state || {})}`);
|
|
1848
|
+
if (effectiveProducer) {
|
|
1849
|
+
return produceStreamResponse(method, state, outputSchema, inputSchema, ctx, null, null);
|
|
1850
|
+
} else {
|
|
1851
|
+
const out = new OutputCollector(outputSchema, effectiveProducer, ctx.serverId, null);
|
|
1852
|
+
try {
|
|
1853
|
+
if (method.exchangeFn) {
|
|
1854
|
+
await method.exchangeFn(state, reqBatch, out);
|
|
1855
|
+
} else {
|
|
1856
|
+
await method.producerFn(state, out);
|
|
1857
|
+
}
|
|
1858
|
+
} catch (error) {
|
|
1859
|
+
if (process.env.VGI_DISPATCH_DEBUG)
|
|
1860
|
+
console.error(`[httpDispatchStreamExchange] exchange handler error:`, error.message, error.stack?.split(`
|
|
1861
|
+
`).slice(0, 5).join(`
|
|
1862
|
+
`));
|
|
1863
|
+
const errBatch = buildErrorBatch(outputSchema, error, ctx.serverId, null);
|
|
1864
|
+
return arrowResponse(serializeIpcStream(outputSchema, [errBatch]), 500);
|
|
1683
1865
|
}
|
|
1684
|
-
const
|
|
1685
|
-
|
|
1686
|
-
|
|
1687
|
-
|
|
1688
|
-
if (row[key] != null) {
|
|
1689
|
-
sample = row[key];
|
|
1690
|
-
break;
|
|
1691
|
-
}
|
|
1866
|
+
const batches = [];
|
|
1867
|
+
if (out.finished) {
|
|
1868
|
+
for (const emitted of out.batches) {
|
|
1869
|
+
batches.push(emitted.batch);
|
|
1692
1870
|
}
|
|
1693
|
-
|
|
1694
|
-
const
|
|
1695
|
-
|
|
1696
|
-
|
|
1697
|
-
|
|
1698
|
-
|
|
1699
|
-
|
|
1700
|
-
|
|
1701
|
-
|
|
1702
|
-
|
|
1703
|
-
|
|
1704
|
-
|
|
1705
|
-
|
|
1706
|
-
children,
|
|
1707
|
-
nullCount: 0
|
|
1708
|
-
});
|
|
1709
|
-
const metadata = new Map;
|
|
1710
|
-
metadata.set(STATE_KEY, this._stateToken);
|
|
1711
|
-
const batch = new RecordBatch8(inputSchema, data, metadata);
|
|
1712
|
-
return this._doExchange(inputSchema, [batch]);
|
|
1713
|
-
}
|
|
1714
|
-
async _doExchange(schema, batches) {
|
|
1715
|
-
const body = serializeIpcStream(schema, batches);
|
|
1716
|
-
const resp = await fetch(`${this._baseUrl}${this._prefix}/${this._method}/exchange`, {
|
|
1717
|
-
method: "POST",
|
|
1718
|
-
headers: this._buildHeaders(),
|
|
1719
|
-
body: this._prepareBody(body)
|
|
1720
|
-
});
|
|
1721
|
-
const responseBody = await this._readResponse(resp);
|
|
1722
|
-
const { batches: responseBatches } = await readResponseBatches(responseBody);
|
|
1723
|
-
let resultRows = [];
|
|
1724
|
-
for (const batch of responseBatches) {
|
|
1725
|
-
if (batch.numRows === 0) {
|
|
1726
|
-
dispatchLogOrError(batch, this._onLog);
|
|
1727
|
-
const token2 = batch.metadata?.get(STATE_KEY);
|
|
1728
|
-
if (token2) {
|
|
1729
|
-
this._stateToken = token2;
|
|
1871
|
+
} else {
|
|
1872
|
+
const stateBytes = ctx.stateSerializer.serialize(state);
|
|
1873
|
+
const schemaBytes = serializeSchema(outputSchema);
|
|
1874
|
+
const inputSchemaBytes = serializeSchema(inputSchema);
|
|
1875
|
+
const token = packStateToken(stateBytes, schemaBytes, inputSchemaBytes, ctx.signingKey);
|
|
1876
|
+
for (const emitted of out.batches) {
|
|
1877
|
+
const batch = emitted.batch;
|
|
1878
|
+
if (batch.numRows > 0) {
|
|
1879
|
+
const mergedMeta = new Map(batch.metadata ?? []);
|
|
1880
|
+
mergedMeta.set(STATE_KEY, token);
|
|
1881
|
+
batches.push(new RecordBatch8(batch.schema, batch.data, mergedMeta));
|
|
1882
|
+
} else {
|
|
1883
|
+
batches.push(batch);
|
|
1730
1884
|
}
|
|
1731
|
-
continue;
|
|
1732
1885
|
}
|
|
1733
|
-
|
|
1734
|
-
|
|
1735
|
-
|
|
1886
|
+
if (!batches.some((b) => b.metadata?.get(STATE_KEY))) {
|
|
1887
|
+
const tokenMeta = new Map;
|
|
1888
|
+
tokenMeta.set(STATE_KEY, token);
|
|
1889
|
+
batches.push(buildEmptyBatch(outputSchema, tokenMeta));
|
|
1736
1890
|
}
|
|
1737
|
-
resultRows = extractBatchRows(batch);
|
|
1738
1891
|
}
|
|
1739
|
-
return
|
|
1740
|
-
}
|
|
1741
|
-
_buildEmptyBatch(schema) {
|
|
1742
|
-
const children = schema.fields.map((f) => {
|
|
1743
|
-
return makeData5({ type: f.type, length: 0, nullCount: 0 });
|
|
1744
|
-
});
|
|
1745
|
-
const structType = new Struct5(schema.fields);
|
|
1746
|
-
const data = makeData5({
|
|
1747
|
-
type: structType,
|
|
1748
|
-
length: 0,
|
|
1749
|
-
children,
|
|
1750
|
-
nullCount: 0
|
|
1751
|
-
});
|
|
1752
|
-
return new RecordBatch8(schema, data);
|
|
1892
|
+
return arrowResponse(serializeIpcStream(outputSchema, batches));
|
|
1753
1893
|
}
|
|
1754
|
-
|
|
1755
|
-
|
|
1756
|
-
|
|
1757
|
-
|
|
1758
|
-
|
|
1894
|
+
}
|
|
1895
|
+
async function produceStreamResponse(method, state, outputSchema, inputSchema, ctx, requestId, headerBytes) {
|
|
1896
|
+
const allBatches = [];
|
|
1897
|
+
const maxBytes = ctx.maxStreamResponseBytes;
|
|
1898
|
+
let estimatedBytes = 0;
|
|
1899
|
+
while (true) {
|
|
1900
|
+
const out = new OutputCollector(outputSchema, true, ctx.serverId, requestId);
|
|
1901
|
+
try {
|
|
1902
|
+
if (method.producerFn) {
|
|
1903
|
+
await method.producerFn(state, out);
|
|
1904
|
+
} else {
|
|
1905
|
+
const tickBatch = buildEmptyBatch(inputSchema);
|
|
1906
|
+
await method.exchangeFn(state, tickBatch, out);
|
|
1759
1907
|
}
|
|
1760
|
-
|
|
1908
|
+
} catch (error) {
|
|
1909
|
+
if (process.env.VGI_DISPATCH_DEBUG)
|
|
1910
|
+
console.error(`[produceStreamResponse] error:`, error.message, error.stack?.split(`
|
|
1911
|
+
`).slice(0, 3).join(`
|
|
1912
|
+
`));
|
|
1913
|
+
allBatches.push(buildErrorBatch(outputSchema, error, ctx.serverId, requestId));
|
|
1914
|
+
break;
|
|
1761
1915
|
}
|
|
1762
|
-
|
|
1763
|
-
|
|
1764
|
-
|
|
1765
|
-
|
|
1766
|
-
return;
|
|
1767
|
-
while (true) {
|
|
1768
|
-
const responseBody = await this._sendContinuation(this._stateToken);
|
|
1769
|
-
const { batches } = await readResponseBatches(responseBody);
|
|
1770
|
-
let gotContinuation = false;
|
|
1771
|
-
for (const batch of batches) {
|
|
1772
|
-
if (batch.numRows === 0) {
|
|
1773
|
-
const token = batch.metadata?.get(STATE_KEY);
|
|
1774
|
-
if (token) {
|
|
1775
|
-
this._stateToken = token;
|
|
1776
|
-
gotContinuation = true;
|
|
1777
|
-
continue;
|
|
1778
|
-
}
|
|
1779
|
-
dispatchLogOrError(batch, this._onLog);
|
|
1780
|
-
continue;
|
|
1781
|
-
}
|
|
1782
|
-
yield extractBatchRows(batch);
|
|
1916
|
+
for (const emitted of out.batches) {
|
|
1917
|
+
allBatches.push(emitted.batch);
|
|
1918
|
+
if (maxBytes != null) {
|
|
1919
|
+
estimatedBytes += emitted.batch.data.byteLength;
|
|
1783
1920
|
}
|
|
1784
|
-
|
|
1785
|
-
|
|
1921
|
+
}
|
|
1922
|
+
if (out.finished) {
|
|
1923
|
+
break;
|
|
1924
|
+
}
|
|
1925
|
+
if (maxBytes != null && estimatedBytes >= maxBytes) {
|
|
1926
|
+
const stateBytes = ctx.stateSerializer.serialize(state);
|
|
1927
|
+
const schemaBytes = serializeSchema(outputSchema);
|
|
1928
|
+
const inputSchemaBytes = serializeSchema(inputSchema);
|
|
1929
|
+
const token = packStateToken(stateBytes, schemaBytes, inputSchemaBytes, ctx.signingKey);
|
|
1930
|
+
const tokenMeta = new Map;
|
|
1931
|
+
tokenMeta.set(STATE_KEY, token);
|
|
1932
|
+
allBatches.push(buildEmptyBatch(outputSchema, tokenMeta));
|
|
1933
|
+
break;
|
|
1786
1934
|
}
|
|
1787
1935
|
}
|
|
1788
|
-
|
|
1789
|
-
|
|
1790
|
-
|
|
1791
|
-
|
|
1792
|
-
|
|
1793
|
-
|
|
1794
|
-
type: structType,
|
|
1795
|
-
length: 1,
|
|
1796
|
-
children: [],
|
|
1797
|
-
nullCount: 0
|
|
1798
|
-
});
|
|
1799
|
-
const batch = new RecordBatch8(emptySchema, data, metadata);
|
|
1800
|
-
const body = serializeIpcStream(emptySchema, [batch]);
|
|
1801
|
-
const resp = await fetch(`${this._baseUrl}${this._prefix}/${this._method}/exchange`, {
|
|
1802
|
-
method: "POST",
|
|
1803
|
-
headers: this._buildHeaders(),
|
|
1804
|
-
body: this._prepareBody(body)
|
|
1805
|
-
});
|
|
1806
|
-
return this._readResponse(resp);
|
|
1936
|
+
const dataBytes = serializeIpcStream(outputSchema, allBatches);
|
|
1937
|
+
let responseBody;
|
|
1938
|
+
if (headerBytes) {
|
|
1939
|
+
responseBody = concatBytes(headerBytes, dataBytes);
|
|
1940
|
+
} else {
|
|
1941
|
+
responseBody = dataBytes;
|
|
1807
1942
|
}
|
|
1808
|
-
|
|
1943
|
+
return arrowResponse(responseBody);
|
|
1944
|
+
}
|
|
1945
|
+
function concatBytes(...arrays) {
|
|
1946
|
+
const totalLen = arrays.reduce((sum, a) => sum + a.length, 0);
|
|
1947
|
+
const result = new Uint8Array(totalLen);
|
|
1948
|
+
let offset = 0;
|
|
1949
|
+
for (const arr of arrays) {
|
|
1950
|
+
result.set(arr, offset);
|
|
1951
|
+
offset += arr.length;
|
|
1952
|
+
}
|
|
1953
|
+
return result;
|
|
1809
1954
|
}
|
|
1810
1955
|
|
|
1811
|
-
// src/
|
|
1812
|
-
|
|
1956
|
+
// src/http/types.ts
|
|
1957
|
+
var jsonStateSerializer = {
|
|
1958
|
+
serialize(state) {
|
|
1959
|
+
return new TextEncoder().encode(JSON.stringify(state, (_key, value) => typeof value === "bigint" ? `__bigint__:${value}` : value));
|
|
1960
|
+
},
|
|
1961
|
+
deserialize(bytes) {
|
|
1962
|
+
return JSON.parse(new TextDecoder().decode(bytes), (_key, value) => typeof value === "string" && value.startsWith("__bigint__:") ? BigInt(value.slice(11)) : value);
|
|
1963
|
+
}
|
|
1964
|
+
};
|
|
1965
|
+
|
|
1966
|
+
// src/http/handler.ts
|
|
1967
|
+
var EMPTY_SCHEMA2 = new Schema5([]);
|
|
1968
|
+
function createHttpHandler(protocol, options) {
|
|
1813
1969
|
const prefix = (options?.prefix ?? "/vgi").replace(/\/+$/, "");
|
|
1814
|
-
const
|
|
1970
|
+
const signingKey = options?.signingKey ?? randomBytes(32);
|
|
1971
|
+
const tokenTtl = options?.tokenTtl ?? 3600;
|
|
1972
|
+
const corsOrigins = options?.corsOrigins;
|
|
1973
|
+
const maxRequestBytes = options?.maxRequestBytes;
|
|
1974
|
+
const maxStreamResponseBytes = options?.maxStreamResponseBytes;
|
|
1975
|
+
const serverId = options?.serverId ?? crypto.randomUUID().replace(/-/g, "").slice(0, 12);
|
|
1976
|
+
const methods = protocol.getMethods();
|
|
1815
1977
|
const compressionLevel = options?.compressionLevel;
|
|
1816
|
-
|
|
1817
|
-
|
|
1818
|
-
|
|
1819
|
-
|
|
1820
|
-
|
|
1821
|
-
|
|
1822
|
-
|
|
1823
|
-
|
|
1824
|
-
|
|
1825
|
-
|
|
1826
|
-
|
|
1827
|
-
|
|
1828
|
-
|
|
1829
|
-
}
|
|
1830
|
-
function buildHeaders() {
|
|
1831
|
-
const headers = {
|
|
1832
|
-
"Content-Type": ARROW_CONTENT_TYPE
|
|
1833
|
-
};
|
|
1834
|
-
if (compressionLevel != null) {
|
|
1835
|
-
headers["Content-Encoding"] = "zstd";
|
|
1836
|
-
headers["Accept-Encoding"] = "zstd";
|
|
1837
|
-
}
|
|
1838
|
-
return headers;
|
|
1839
|
-
}
|
|
1840
|
-
function prepareBody(content) {
|
|
1841
|
-
if (compressionLevel != null && compressFn) {
|
|
1842
|
-
return compressFn(content, compressionLevel);
|
|
1978
|
+
const stateSerializer = options?.stateSerializer ?? jsonStateSerializer;
|
|
1979
|
+
const ctx = {
|
|
1980
|
+
signingKey,
|
|
1981
|
+
tokenTtl,
|
|
1982
|
+
serverId,
|
|
1983
|
+
maxStreamResponseBytes,
|
|
1984
|
+
stateSerializer
|
|
1985
|
+
};
|
|
1986
|
+
function addCorsHeaders(headers) {
|
|
1987
|
+
if (corsOrigins) {
|
|
1988
|
+
headers.set("Access-Control-Allow-Origin", corsOrigins);
|
|
1989
|
+
headers.set("Access-Control-Allow-Methods", "POST, OPTIONS");
|
|
1990
|
+
headers.set("Access-Control-Allow-Headers", "Content-Type");
|
|
1843
1991
|
}
|
|
1844
|
-
return content;
|
|
1845
1992
|
}
|
|
1846
|
-
async function
|
|
1847
|
-
|
|
1848
|
-
|
|
1849
|
-
|
|
1850
|
-
|
|
1851
|
-
|
|
1993
|
+
async function compressIfAccepted(response, clientAcceptsZstd) {
|
|
1994
|
+
if (compressionLevel == null || !clientAcceptsZstd)
|
|
1995
|
+
return response;
|
|
1996
|
+
const responseBody = new Uint8Array(await response.arrayBuffer());
|
|
1997
|
+
const compressed = zstdCompress(responseBody, compressionLevel);
|
|
1998
|
+
const headers = new Headers(response.headers);
|
|
1999
|
+
headers.set("Content-Encoding", "zstd");
|
|
2000
|
+
return new Response(compressed, {
|
|
2001
|
+
status: response.status,
|
|
2002
|
+
headers
|
|
2003
|
+
});
|
|
1852
2004
|
}
|
|
1853
|
-
|
|
1854
|
-
|
|
1855
|
-
|
|
1856
|
-
const
|
|
1857
|
-
|
|
1858
|
-
return
|
|
2005
|
+
function makeErrorResponse(error, statusCode, schema = EMPTY_SCHEMA2) {
|
|
2006
|
+
const errBatch = buildErrorBatch(schema, error, serverId, null);
|
|
2007
|
+
const body = serializeIpcStream(schema, [errBatch]);
|
|
2008
|
+
const resp = arrowResponse(body, statusCode);
|
|
2009
|
+
addCorsHeaders(resp.headers);
|
|
2010
|
+
return resp;
|
|
1859
2011
|
}
|
|
1860
|
-
return {
|
|
1861
|
-
|
|
1862
|
-
|
|
1863
|
-
|
|
1864
|
-
|
|
1865
|
-
|
|
1866
|
-
|
|
1867
|
-
|
|
1868
|
-
|
|
1869
|
-
const body = buildRequestIpc(info.paramsSchema, fullParams, method);
|
|
1870
|
-
const resp = await fetch(`${baseUrl}${prefix}/${method}`, {
|
|
1871
|
-
method: "POST",
|
|
1872
|
-
headers: buildHeaders(),
|
|
1873
|
-
body: prepareBody(body)
|
|
1874
|
-
});
|
|
1875
|
-
const responseBody = await readResponse(resp);
|
|
1876
|
-
const { batches } = await readResponseBatches(responseBody);
|
|
1877
|
-
let resultBatch = null;
|
|
1878
|
-
for (const batch of batches) {
|
|
1879
|
-
if (batch.numRows === 0) {
|
|
1880
|
-
dispatchLogOrError(batch, onLog);
|
|
1881
|
-
continue;
|
|
2012
|
+
return async function handler(request) {
|
|
2013
|
+
const url = new URL(request.url);
|
|
2014
|
+
const path = url.pathname;
|
|
2015
|
+
if (request.method === "OPTIONS") {
|
|
2016
|
+
if (path === `${prefix}/__capabilities__`) {
|
|
2017
|
+
const headers = new Headers;
|
|
2018
|
+
addCorsHeaders(headers);
|
|
2019
|
+
if (maxRequestBytes != null) {
|
|
2020
|
+
headers.set("VGI-Max-Request-Bytes", String(maxRequestBytes));
|
|
1882
2021
|
}
|
|
1883
|
-
|
|
2022
|
+
return new Response(null, { status: 204, headers });
|
|
1884
2023
|
}
|
|
1885
|
-
if (
|
|
1886
|
-
|
|
2024
|
+
if (corsOrigins) {
|
|
2025
|
+
const headers = new Headers;
|
|
2026
|
+
addCorsHeaders(headers);
|
|
2027
|
+
return new Response(null, { status: 204, headers });
|
|
1887
2028
|
}
|
|
1888
|
-
|
|
1889
|
-
|
|
1890
|
-
|
|
1891
|
-
|
|
1892
|
-
|
|
1893
|
-
|
|
1894
|
-
|
|
1895
|
-
|
|
1896
|
-
|
|
1897
|
-
|
|
1898
|
-
const
|
|
1899
|
-
|
|
1900
|
-
|
|
1901
|
-
throw new Error(`Unknown method: '${method}'`);
|
|
2029
|
+
return new Response(null, { status: 405 });
|
|
2030
|
+
}
|
|
2031
|
+
if (request.method !== "POST") {
|
|
2032
|
+
return new Response("Method Not Allowed", { status: 405 });
|
|
2033
|
+
}
|
|
2034
|
+
const contentType = request.headers.get("Content-Type");
|
|
2035
|
+
if (!contentType || !contentType.includes(ARROW_CONTENT_TYPE)) {
|
|
2036
|
+
return new Response(`Unsupported Media Type: expected ${ARROW_CONTENT_TYPE}`, { status: 415 });
|
|
2037
|
+
}
|
|
2038
|
+
if (maxRequestBytes != null) {
|
|
2039
|
+
const contentLength = request.headers.get("Content-Length");
|
|
2040
|
+
if (contentLength && parseInt(contentLength, 10) > maxRequestBytes) {
|
|
2041
|
+
return new Response("Request body too large", { status: 413 });
|
|
1902
2042
|
}
|
|
1903
|
-
|
|
1904
|
-
|
|
1905
|
-
|
|
1906
|
-
|
|
1907
|
-
|
|
1908
|
-
|
|
1909
|
-
|
|
1910
|
-
|
|
1911
|
-
|
|
1912
|
-
|
|
1913
|
-
|
|
1914
|
-
|
|
1915
|
-
|
|
1916
|
-
|
|
1917
|
-
|
|
1918
|
-
|
|
1919
|
-
|
|
1920
|
-
|
|
1921
|
-
|
|
1922
|
-
|
|
1923
|
-
|
|
1924
|
-
|
|
1925
|
-
|
|
1926
|
-
|
|
1927
|
-
|
|
1928
|
-
|
|
1929
|
-
|
|
1930
|
-
|
|
1931
|
-
|
|
1932
|
-
|
|
1933
|
-
|
|
1934
|
-
|
|
1935
|
-
|
|
1936
|
-
|
|
1937
|
-
|
|
1938
|
-
|
|
1939
|
-
|
|
1940
|
-
|
|
1941
|
-
|
|
1942
|
-
|
|
1943
|
-
|
|
1944
|
-
|
|
1945
|
-
|
|
1946
|
-
headerErrorBatches.push(batch);
|
|
1947
|
-
continue;
|
|
1948
|
-
}
|
|
1949
|
-
dispatchLogOrError(batch, onLog);
|
|
1950
|
-
continue;
|
|
1951
|
-
}
|
|
1952
|
-
pendingBatches.push(batch);
|
|
1953
|
-
}
|
|
1954
|
-
}
|
|
1955
|
-
if (headerErrorBatches.length > 0) {
|
|
1956
|
-
if (pendingBatches.length > 0 || stateToken !== null) {
|
|
1957
|
-
pendingBatches.push(...headerErrorBatches);
|
|
1958
|
-
} else {
|
|
1959
|
-
for (const batch of headerErrorBatches) {
|
|
1960
|
-
dispatchLogOrError(batch, onLog);
|
|
1961
|
-
}
|
|
1962
|
-
}
|
|
2043
|
+
}
|
|
2044
|
+
const clientAcceptsZstd = (request.headers.get("Accept-Encoding") ?? "").includes("zstd");
|
|
2045
|
+
let body = new Uint8Array(await request.arrayBuffer());
|
|
2046
|
+
const contentEncoding = request.headers.get("Content-Encoding");
|
|
2047
|
+
if (contentEncoding === "zstd") {
|
|
2048
|
+
body = zstdDecompress(body);
|
|
2049
|
+
}
|
|
2050
|
+
if (path === `${prefix}/${DESCRIBE_METHOD_NAME}`) {
|
|
2051
|
+
try {
|
|
2052
|
+
const response = httpDispatchDescribe(protocol.name, methods, serverId);
|
|
2053
|
+
addCorsHeaders(response.headers);
|
|
2054
|
+
return compressIfAccepted(response, clientAcceptsZstd);
|
|
2055
|
+
} catch (error) {
|
|
2056
|
+
return compressIfAccepted(makeErrorResponse(error, 500), clientAcceptsZstd);
|
|
2057
|
+
}
|
|
2058
|
+
}
|
|
2059
|
+
if (!path.startsWith(`${prefix}/`)) {
|
|
2060
|
+
return new Response("Not Found", { status: 404 });
|
|
2061
|
+
}
|
|
2062
|
+
const subPath = path.slice(prefix.length + 1);
|
|
2063
|
+
let methodName;
|
|
2064
|
+
let action;
|
|
2065
|
+
if (subPath.endsWith("/init")) {
|
|
2066
|
+
methodName = subPath.slice(0, -5);
|
|
2067
|
+
action = "init";
|
|
2068
|
+
} else if (subPath.endsWith("/exchange")) {
|
|
2069
|
+
methodName = subPath.slice(0, -9);
|
|
2070
|
+
action = "exchange";
|
|
2071
|
+
} else {
|
|
2072
|
+
methodName = subPath;
|
|
2073
|
+
action = "call";
|
|
2074
|
+
}
|
|
2075
|
+
const method = methods.get(methodName);
|
|
2076
|
+
if (!method) {
|
|
2077
|
+
const available = [...methods.keys()].sort();
|
|
2078
|
+
const err = new Error(`Unknown method: '${methodName}'. Available methods: [${available.join(", ")}]`);
|
|
2079
|
+
return compressIfAccepted(makeErrorResponse(err, 404), clientAcceptsZstd);
|
|
2080
|
+
}
|
|
2081
|
+
try {
|
|
2082
|
+
let response;
|
|
2083
|
+
if (action === "call") {
|
|
2084
|
+
if (method.type !== "unary" /* UNARY */) {
|
|
2085
|
+
throw new HttpRpcError(`Method '${methodName}' is a stream method. Use /init and /exchange endpoints.`, 400);
|
|
1963
2086
|
}
|
|
1964
|
-
|
|
1965
|
-
|
|
2087
|
+
response = await httpDispatchUnary(method, body, ctx);
|
|
2088
|
+
} else if (action === "init") {
|
|
2089
|
+
if (method.type !== "stream" /* STREAM */) {
|
|
2090
|
+
throw new HttpRpcError(`Method '${methodName}' is a unary method. Use POST ${prefix}/${methodName} instead.`, 400);
|
|
1966
2091
|
}
|
|
2092
|
+
response = await httpDispatchStreamInit(method, body, ctx);
|
|
1967
2093
|
} else {
|
|
1968
|
-
|
|
1969
|
-
|
|
1970
|
-
const errorBatches = [];
|
|
1971
|
-
for (const batch of batches) {
|
|
1972
|
-
if (batch.numRows === 0) {
|
|
1973
|
-
const token = batch.metadata?.get(STATE_KEY);
|
|
1974
|
-
if (token) {
|
|
1975
|
-
stateToken = token;
|
|
1976
|
-
continue;
|
|
1977
|
-
}
|
|
1978
|
-
const level = batch.metadata?.get(LOG_LEVEL_KEY);
|
|
1979
|
-
if (level === "EXCEPTION") {
|
|
1980
|
-
errorBatches.push(batch);
|
|
1981
|
-
continue;
|
|
1982
|
-
}
|
|
1983
|
-
dispatchLogOrError(batch, onLog);
|
|
1984
|
-
continue;
|
|
1985
|
-
}
|
|
1986
|
-
pendingBatches.push(batch);
|
|
1987
|
-
}
|
|
1988
|
-
if (errorBatches.length > 0) {
|
|
1989
|
-
if (pendingBatches.length > 0 || stateToken !== null) {
|
|
1990
|
-
pendingBatches.push(...errorBatches);
|
|
1991
|
-
} else {
|
|
1992
|
-
for (const batch of errorBatches) {
|
|
1993
|
-
dispatchLogOrError(batch, onLog);
|
|
1994
|
-
}
|
|
1995
|
-
}
|
|
2094
|
+
if (method.type !== "stream" /* STREAM */) {
|
|
2095
|
+
throw new HttpRpcError(`Method '${methodName}' is a unary method. Use POST ${prefix}/${methodName} instead.`, 400);
|
|
1996
2096
|
}
|
|
2097
|
+
response = await httpDispatchStreamExchange(method, body, ctx);
|
|
1997
2098
|
}
|
|
1998
|
-
|
|
1999
|
-
|
|
2099
|
+
addCorsHeaders(response.headers);
|
|
2100
|
+
return compressIfAccepted(response, clientAcceptsZstd);
|
|
2101
|
+
} catch (error) {
|
|
2102
|
+
if (error instanceof HttpRpcError) {
|
|
2103
|
+
return compressIfAccepted(makeErrorResponse(error, error.statusCode), clientAcceptsZstd);
|
|
2000
2104
|
}
|
|
2001
|
-
|
|
2002
|
-
|
|
2003
|
-
baseUrl,
|
|
2004
|
-
prefix,
|
|
2005
|
-
method,
|
|
2006
|
-
stateToken,
|
|
2007
|
-
outputSchema,
|
|
2008
|
-
inputSchema: info.inputSchema,
|
|
2009
|
-
onLog,
|
|
2010
|
-
pendingBatches,
|
|
2011
|
-
finished,
|
|
2012
|
-
header,
|
|
2013
|
-
compressionLevel,
|
|
2014
|
-
compressFn,
|
|
2015
|
-
decompressFn
|
|
2016
|
-
});
|
|
2017
|
-
},
|
|
2018
|
-
async describe() {
|
|
2019
|
-
return httpIntrospect(baseUrl, { prefix });
|
|
2020
|
-
},
|
|
2021
|
-
close() {}
|
|
2105
|
+
return compressIfAccepted(makeErrorResponse(error, 500), clientAcceptsZstd);
|
|
2106
|
+
}
|
|
2022
2107
|
};
|
|
2023
2108
|
}
|
|
2024
|
-
// src/
|
|
2109
|
+
// src/protocol.ts
|
|
2110
|
+
import { Schema as Schema7 } from "@query-farm/apache-arrow";
|
|
2111
|
+
|
|
2112
|
+
// src/schema.ts
|
|
2025
2113
|
import {
|
|
2026
|
-
|
|
2027
|
-
|
|
2028
|
-
|
|
2029
|
-
Field as
|
|
2030
|
-
|
|
2031
|
-
|
|
2032
|
-
|
|
2033
|
-
|
|
2034
|
-
|
|
2035
|
-
|
|
2036
|
-
|
|
2037
|
-
|
|
2038
|
-
|
|
2039
|
-
|
|
2040
|
-
|
|
2041
|
-
|
|
2042
|
-
|
|
2043
|
-
|
|
2044
|
-
|
|
2045
|
-
|
|
2046
|
-
|
|
2047
|
-
|
|
2048
|
-
|
|
2049
|
-
|
|
2050
|
-
|
|
2051
|
-
|
|
2052
|
-
|
|
2053
|
-
|
|
2054
|
-
|
|
2055
|
-
|
|
2114
|
+
Binary as Binary3,
|
|
2115
|
+
Bool as Bool3,
|
|
2116
|
+
DataType as DataType4,
|
|
2117
|
+
Field as Field4,
|
|
2118
|
+
Float32,
|
|
2119
|
+
Float64 as Float642,
|
|
2120
|
+
Int16,
|
|
2121
|
+
Int32,
|
|
2122
|
+
Int64 as Int642,
|
|
2123
|
+
Schema as Schema6,
|
|
2124
|
+
Utf8 as Utf83
|
|
2125
|
+
} from "@query-farm/apache-arrow";
|
|
2126
|
+
var str = new Utf83;
|
|
2127
|
+
var bytes = new Binary3;
|
|
2128
|
+
var int = new Int642;
|
|
2129
|
+
var int32 = new Int32;
|
|
2130
|
+
var float = new Float642;
|
|
2131
|
+
var float32 = new Float32;
|
|
2132
|
+
var bool = new Bool3;
|
|
2133
|
+
function toSchema(spec) {
|
|
2134
|
+
if (spec instanceof Schema6)
|
|
2135
|
+
return spec;
|
|
2136
|
+
const fields = [];
|
|
2137
|
+
for (const [name, value] of Object.entries(spec)) {
|
|
2138
|
+
if (value instanceof Field4) {
|
|
2139
|
+
fields.push(value);
|
|
2140
|
+
} else if (value instanceof DataType4) {
|
|
2141
|
+
fields.push(new Field4(name, value, false));
|
|
2142
|
+
} else {
|
|
2143
|
+
throw new TypeError(`Invalid schema value for "${name}": expected DataType or Field, got ${typeof value}`);
|
|
2144
|
+
}
|
|
2056
2145
|
}
|
|
2057
|
-
|
|
2058
|
-
|
|
2059
|
-
|
|
2060
|
-
|
|
2146
|
+
return new Schema6(fields);
|
|
2147
|
+
}
|
|
2148
|
+
var TYPE_MAP = [
|
|
2149
|
+
[Utf83, "str"],
|
|
2150
|
+
[Binary3, "bytes"],
|
|
2151
|
+
[Bool3, "bool"],
|
|
2152
|
+
[Float642, "float"],
|
|
2153
|
+
[Float32, "float"],
|
|
2154
|
+
[Int642, "int"],
|
|
2155
|
+
[Int32, "int"],
|
|
2156
|
+
[Int16, "int"]
|
|
2157
|
+
];
|
|
2158
|
+
function inferParamTypes(spec) {
|
|
2159
|
+
const schema = toSchema(spec);
|
|
2160
|
+
if (schema.fields.length === 0)
|
|
2161
|
+
return;
|
|
2162
|
+
const result = {};
|
|
2163
|
+
for (const field of schema.fields) {
|
|
2164
|
+
let mapped;
|
|
2165
|
+
for (const [ctor, name] of TYPE_MAP) {
|
|
2166
|
+
if (field.type instanceof ctor) {
|
|
2167
|
+
mapped = name;
|
|
2168
|
+
break;
|
|
2169
|
+
}
|
|
2061
2170
|
}
|
|
2062
|
-
|
|
2171
|
+
if (!mapped)
|
|
2172
|
+
return;
|
|
2173
|
+
result[field.name] = mapped;
|
|
2063
2174
|
}
|
|
2175
|
+
return result;
|
|
2064
2176
|
}
|
|
2065
2177
|
|
|
2066
|
-
|
|
2067
|
-
|
|
2068
|
-
|
|
2069
|
-
|
|
2070
|
-
|
|
2071
|
-
|
|
2072
|
-
|
|
2073
|
-
|
|
2074
|
-
_closed = false;
|
|
2075
|
-
_outputSchema;
|
|
2076
|
-
_releaseBusy;
|
|
2077
|
-
_setDrainPromise;
|
|
2078
|
-
constructor(opts) {
|
|
2079
|
-
this._reader = opts.reader;
|
|
2080
|
-
this._writeFn = opts.writeFn;
|
|
2081
|
-
this._onLog = opts.onLog;
|
|
2082
|
-
this._header = opts.header;
|
|
2083
|
-
this._outputSchema = opts.outputSchema;
|
|
2084
|
-
this._releaseBusy = opts.releaseBusy;
|
|
2085
|
-
this._setDrainPromise = opts.setDrainPromise;
|
|
2178
|
+
// src/protocol.ts
|
|
2179
|
+
var EMPTY_SCHEMA3 = new Schema7([]);
|
|
2180
|
+
|
|
2181
|
+
class Protocol {
|
|
2182
|
+
name;
|
|
2183
|
+
_methods = new Map;
|
|
2184
|
+
constructor(name) {
|
|
2185
|
+
this.name = name;
|
|
2086
2186
|
}
|
|
2087
|
-
|
|
2088
|
-
|
|
2187
|
+
unary(name, config) {
|
|
2188
|
+
const params = toSchema(config.params);
|
|
2189
|
+
this._methods.set(name, {
|
|
2190
|
+
name,
|
|
2191
|
+
type: "unary" /* UNARY */,
|
|
2192
|
+
paramsSchema: params,
|
|
2193
|
+
resultSchema: toSchema(config.result),
|
|
2194
|
+
handler: config.handler,
|
|
2195
|
+
doc: config.doc,
|
|
2196
|
+
defaults: config.defaults,
|
|
2197
|
+
paramTypes: config.paramTypes ?? inferParamTypes(params)
|
|
2198
|
+
});
|
|
2199
|
+
return this;
|
|
2089
2200
|
}
|
|
2090
|
-
|
|
2091
|
-
|
|
2092
|
-
|
|
2093
|
-
|
|
2094
|
-
|
|
2095
|
-
|
|
2096
|
-
|
|
2097
|
-
|
|
2098
|
-
|
|
2099
|
-
|
|
2100
|
-
|
|
2201
|
+
producer(name, config) {
|
|
2202
|
+
const params = toSchema(config.params);
|
|
2203
|
+
this._methods.set(name, {
|
|
2204
|
+
name,
|
|
2205
|
+
type: "stream" /* STREAM */,
|
|
2206
|
+
paramsSchema: params,
|
|
2207
|
+
resultSchema: EMPTY_SCHEMA3,
|
|
2208
|
+
outputSchema: toSchema(config.outputSchema),
|
|
2209
|
+
inputSchema: EMPTY_SCHEMA3,
|
|
2210
|
+
producerInit: config.init,
|
|
2211
|
+
producerFn: config.produce,
|
|
2212
|
+
headerSchema: config.headerSchema ? toSchema(config.headerSchema) : undefined,
|
|
2213
|
+
headerInit: config.headerInit,
|
|
2214
|
+
doc: config.doc,
|
|
2215
|
+
defaults: config.defaults,
|
|
2216
|
+
paramTypes: config.paramTypes ?? inferParamTypes(params)
|
|
2217
|
+
});
|
|
2218
|
+
return this;
|
|
2219
|
+
}
|
|
2220
|
+
exchange(name, config) {
|
|
2221
|
+
const params = toSchema(config.params);
|
|
2222
|
+
this._methods.set(name, {
|
|
2223
|
+
name,
|
|
2224
|
+
type: "stream" /* STREAM */,
|
|
2225
|
+
paramsSchema: params,
|
|
2226
|
+
resultSchema: EMPTY_SCHEMA3,
|
|
2227
|
+
inputSchema: toSchema(config.inputSchema),
|
|
2228
|
+
outputSchema: toSchema(config.outputSchema),
|
|
2229
|
+
exchangeInit: config.init,
|
|
2230
|
+
exchangeFn: config.exchange,
|
|
2231
|
+
headerSchema: config.headerSchema ? toSchema(config.headerSchema) : undefined,
|
|
2232
|
+
headerInit: config.headerInit,
|
|
2233
|
+
doc: config.doc,
|
|
2234
|
+
defaults: config.defaults,
|
|
2235
|
+
paramTypes: config.paramTypes ?? inferParamTypes(params)
|
|
2236
|
+
});
|
|
2237
|
+
return this;
|
|
2238
|
+
}
|
|
2239
|
+
getMethods() {
|
|
2240
|
+
return new Map(this._methods);
|
|
2241
|
+
}
|
|
2242
|
+
}
|
|
2243
|
+
// src/server.ts
|
|
2244
|
+
import { Schema as Schema9 } from "@query-farm/apache-arrow";
|
|
2245
|
+
|
|
2246
|
+
// src/dispatch/stream.ts
|
|
2247
|
+
import { Schema as Schema8 } from "@query-farm/apache-arrow";
|
|
2248
|
+
var EMPTY_SCHEMA4 = new Schema8([]);
|
|
2249
|
+
async function dispatchStream(method, params, writer, reader, serverId, requestId) {
|
|
2250
|
+
const isProducer = !!method.producerFn;
|
|
2251
|
+
let state;
|
|
2252
|
+
try {
|
|
2253
|
+
if (isProducer) {
|
|
2254
|
+
state = await method.producerInit(params);
|
|
2255
|
+
} else {
|
|
2256
|
+
state = await method.exchangeInit(params);
|
|
2257
|
+
}
|
|
2258
|
+
} catch (error) {
|
|
2259
|
+
const errSchema = method.headerSchema ?? EMPTY_SCHEMA4;
|
|
2260
|
+
const errBatch = buildErrorBatch(errSchema, error, serverId, requestId);
|
|
2261
|
+
writer.writeStream(errSchema, [errBatch]);
|
|
2262
|
+
const inputSchema2 = await reader.openNextStream();
|
|
2263
|
+
if (inputSchema2) {
|
|
2264
|
+
while (await reader.readNextBatch() !== null) {}
|
|
2101
2265
|
}
|
|
2266
|
+
return;
|
|
2102
2267
|
}
|
|
2103
|
-
|
|
2104
|
-
|
|
2268
|
+
const outputSchema = state?.__outputSchema ?? method.outputSchema;
|
|
2269
|
+
const effectiveProducer = state?.__isProducer ?? isProducer;
|
|
2270
|
+
if (method.headerSchema && method.headerInit) {
|
|
2271
|
+
try {
|
|
2272
|
+
const headerOut = new OutputCollector(method.headerSchema, true, serverId, requestId);
|
|
2273
|
+
const headerValues = method.headerInit(params, state, headerOut);
|
|
2274
|
+
const headerBatch = buildResultBatch(method.headerSchema, headerValues, serverId, requestId);
|
|
2275
|
+
const headerBatches = [...headerOut.batches.map((b) => b.batch), headerBatch];
|
|
2276
|
+
writer.writeStream(method.headerSchema, headerBatches);
|
|
2277
|
+
} catch (error) {
|
|
2278
|
+
const errBatch = buildErrorBatch(method.headerSchema, error, serverId, requestId);
|
|
2279
|
+
writer.writeStream(method.headerSchema, [errBatch]);
|
|
2280
|
+
const inputSchema2 = await reader.openNextStream();
|
|
2281
|
+
if (inputSchema2) {
|
|
2282
|
+
while (await reader.readNextBatch() !== null) {}
|
|
2283
|
+
}
|
|
2105
2284
|
return;
|
|
2106
|
-
this._outputStreamOpened = true;
|
|
2107
|
-
const schema = await this._reader.openNextStream();
|
|
2108
|
-
if (!schema) {
|
|
2109
|
-
throw new RpcError("ProtocolError", "Expected output stream but got EOF", "");
|
|
2110
2285
|
}
|
|
2111
2286
|
}
|
|
2112
|
-
|
|
2113
|
-
|
|
2114
|
-
|
|
2115
|
-
|
|
2116
|
-
|
|
2117
|
-
|
|
2118
|
-
|
|
2119
|
-
|
|
2120
|
-
|
|
2121
|
-
|
|
2122
|
-
|
|
2123
|
-
|
|
2124
|
-
const
|
|
2125
|
-
|
|
2126
|
-
|
|
2127
|
-
children,
|
|
2128
|
-
nullCount: 0
|
|
2129
|
-
});
|
|
2130
|
-
batch = new RecordBatch9(inputSchema, data);
|
|
2131
|
-
} else {
|
|
2132
|
-
const keys = Object.keys(input[0]);
|
|
2133
|
-
const fields = keys.map((key) => {
|
|
2134
|
-
let sample = undefined;
|
|
2135
|
-
for (const row of input) {
|
|
2136
|
-
if (row[key] != null) {
|
|
2137
|
-
sample = row[key];
|
|
2138
|
-
break;
|
|
2139
|
-
}
|
|
2140
|
-
}
|
|
2141
|
-
const arrowType = inferArrowType(sample);
|
|
2142
|
-
return new Field5(key, arrowType, true);
|
|
2143
|
-
});
|
|
2144
|
-
inputSchema = new Schema12(fields);
|
|
2145
|
-
if (this._inputSchema) {
|
|
2146
|
-
const cached = this._inputSchema;
|
|
2147
|
-
if (cached.fields.length !== inputSchema.fields.length || cached.fields.some((f, i) => f.name !== inputSchema.fields[i].name)) {
|
|
2148
|
-
throw new RpcError("ProtocolError", `Exchange input schema changed: expected [${cached.fields.map((f) => f.name).join(", ")}] ` + `but got [${inputSchema.fields.map((f) => f.name).join(", ")}]`, "");
|
|
2149
|
-
}
|
|
2287
|
+
const inputSchema = await reader.openNextStream();
|
|
2288
|
+
if (!inputSchema) {
|
|
2289
|
+
const errBatch = buildErrorBatch(outputSchema, new Error("Expected input stream but got EOF"), serverId, requestId);
|
|
2290
|
+
writer.writeStream(outputSchema, [errBatch]);
|
|
2291
|
+
return;
|
|
2292
|
+
}
|
|
2293
|
+
const stream = writer.openStream(outputSchema);
|
|
2294
|
+
try {
|
|
2295
|
+
while (true) {
|
|
2296
|
+
const inputBatch = await reader.readNextBatch();
|
|
2297
|
+
if (!inputBatch)
|
|
2298
|
+
break;
|
|
2299
|
+
const out = new OutputCollector(outputSchema, effectiveProducer, serverId, requestId);
|
|
2300
|
+
if (isProducer) {
|
|
2301
|
+
await method.producerFn(state, out);
|
|
2150
2302
|
} else {
|
|
2151
|
-
|
|
2303
|
+
await method.exchangeFn(state, inputBatch, out);
|
|
2304
|
+
}
|
|
2305
|
+
for (const emitted of out.batches) {
|
|
2306
|
+
stream.write(emitted.batch);
|
|
2307
|
+
}
|
|
2308
|
+
if (out.finished) {
|
|
2309
|
+
break;
|
|
2152
2310
|
}
|
|
2153
|
-
const children = inputSchema.fields.map((f) => {
|
|
2154
|
-
const values = input.map((row) => row[f.name]);
|
|
2155
|
-
return vectorFromArray5(values, f.type).data[0];
|
|
2156
|
-
});
|
|
2157
|
-
const structType = new Struct6(inputSchema.fields);
|
|
2158
|
-
const data = makeData6({
|
|
2159
|
-
type: structType,
|
|
2160
|
-
length: input.length,
|
|
2161
|
-
children,
|
|
2162
|
-
nullCount: 0
|
|
2163
|
-
});
|
|
2164
|
-
batch = new RecordBatch9(inputSchema, data);
|
|
2165
|
-
}
|
|
2166
|
-
if (!this._inputWriter) {
|
|
2167
|
-
this._inputWriter = new PipeIncrementalWriter(this._writeFn, inputSchema);
|
|
2168
2311
|
}
|
|
2169
|
-
|
|
2170
|
-
|
|
2312
|
+
} catch (error) {
|
|
2313
|
+
stream.write(buildErrorBatch(outputSchema, error, serverId, requestId));
|
|
2314
|
+
}
|
|
2315
|
+
stream.close();
|
|
2316
|
+
try {
|
|
2317
|
+
while (await reader.readNextBatch() !== null) {}
|
|
2318
|
+
} catch {}
|
|
2319
|
+
}
|
|
2320
|
+
|
|
2321
|
+
// src/dispatch/unary.ts
|
|
2322
|
+
async function dispatchUnary(method, params, writer, serverId, requestId) {
|
|
2323
|
+
const schema = method.resultSchema;
|
|
2324
|
+
const out = new OutputCollector(schema, true, serverId, requestId);
|
|
2325
|
+
try {
|
|
2326
|
+
const result = await method.handler(params, out);
|
|
2327
|
+
const resultBatch = buildResultBatch(schema, result, serverId, requestId);
|
|
2328
|
+
const batches = [...out.batches.map((b) => b.batch), resultBatch];
|
|
2329
|
+
writer.writeStream(schema, batches);
|
|
2330
|
+
} catch (error) {
|
|
2331
|
+
const batch = buildErrorBatch(schema, error, serverId, requestId);
|
|
2332
|
+
writer.writeStream(schema, [batch]);
|
|
2333
|
+
}
|
|
2334
|
+
}
|
|
2335
|
+
|
|
2336
|
+
// src/wire/writer.ts
|
|
2337
|
+
import { writeSync } from "node:fs";
|
|
2338
|
+
import { RecordBatchStreamWriter as RecordBatchStreamWriter4 } from "@query-farm/apache-arrow";
|
|
2339
|
+
var STDOUT_FD = 1;
|
|
2340
|
+
function writeAll(fd, data) {
|
|
2341
|
+
let offset = 0;
|
|
2342
|
+
while (offset < data.length) {
|
|
2171
2343
|
try {
|
|
2172
|
-
const
|
|
2173
|
-
if (
|
|
2174
|
-
|
|
2175
|
-
|
|
2176
|
-
return extractBatchRows(outputBatch);
|
|
2344
|
+
const written = writeSync(fd, data, offset, data.length - offset);
|
|
2345
|
+
if (written <= 0)
|
|
2346
|
+
throw new Error(`writeSync returned ${written}`);
|
|
2347
|
+
offset += written;
|
|
2177
2348
|
} catch (e) {
|
|
2178
|
-
|
|
2349
|
+
if (e.code === "EAGAIN") {
|
|
2350
|
+
Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 1);
|
|
2351
|
+
continue;
|
|
2352
|
+
}
|
|
2179
2353
|
throw e;
|
|
2180
2354
|
}
|
|
2181
2355
|
}
|
|
2182
|
-
|
|
2183
|
-
|
|
2184
|
-
|
|
2185
|
-
|
|
2186
|
-
|
|
2187
|
-
|
|
2188
|
-
this._inputWriter = null;
|
|
2189
|
-
}
|
|
2190
|
-
try {
|
|
2191
|
-
if (this._outputStreamOpened) {
|
|
2192
|
-
while (await this._reader.readNextBatch() !== null) {}
|
|
2193
|
-
}
|
|
2194
|
-
} catch {}
|
|
2195
|
-
this._releaseBusy();
|
|
2356
|
+
}
|
|
2357
|
+
|
|
2358
|
+
class IpcStreamWriter {
|
|
2359
|
+
fd;
|
|
2360
|
+
constructor(fd = STDOUT_FD) {
|
|
2361
|
+
this.fd = fd;
|
|
2196
2362
|
}
|
|
2197
|
-
|
|
2198
|
-
|
|
2199
|
-
|
|
2200
|
-
|
|
2201
|
-
|
|
2202
|
-
this._inputWriter = new PipeIncrementalWriter(this._writeFn, tickSchema);
|
|
2203
|
-
const structType = new Struct6(tickSchema.fields);
|
|
2204
|
-
const tickData = makeData6({
|
|
2205
|
-
type: structType,
|
|
2206
|
-
length: 0,
|
|
2207
|
-
children: [],
|
|
2208
|
-
nullCount: 0
|
|
2209
|
-
});
|
|
2210
|
-
const tickBatch = new RecordBatch9(tickSchema, tickData);
|
|
2211
|
-
while (true) {
|
|
2212
|
-
this._inputWriter.write(tickBatch);
|
|
2213
|
-
await this._ensureOutputStream();
|
|
2214
|
-
const outputBatch = await this._readOutputBatch();
|
|
2215
|
-
if (outputBatch === null) {
|
|
2216
|
-
break;
|
|
2217
|
-
}
|
|
2218
|
-
yield extractBatchRows(outputBatch);
|
|
2219
|
-
}
|
|
2220
|
-
} finally {
|
|
2221
|
-
if (this._inputWriter) {
|
|
2222
|
-
this._inputWriter.close();
|
|
2223
|
-
this._inputWriter = null;
|
|
2224
|
-
}
|
|
2225
|
-
try {
|
|
2226
|
-
if (this._outputStreamOpened) {
|
|
2227
|
-
while (await this._reader.readNextBatch() !== null) {}
|
|
2228
|
-
}
|
|
2229
|
-
} catch {}
|
|
2230
|
-
this._closed = true;
|
|
2231
|
-
this._releaseBusy();
|
|
2363
|
+
writeStream(schema, batches) {
|
|
2364
|
+
const writer = new RecordBatchStreamWriter4;
|
|
2365
|
+
writer.reset(undefined, schema);
|
|
2366
|
+
for (const batch of batches) {
|
|
2367
|
+
writer._writeRecordBatch(batch);
|
|
2232
2368
|
}
|
|
2369
|
+
writer.close();
|
|
2370
|
+
const bytes2 = writer.toUint8Array(true);
|
|
2371
|
+
writeAll(this.fd, bytes2);
|
|
2372
|
+
}
|
|
2373
|
+
openStream(schema) {
|
|
2374
|
+
return new IncrementalStream(this.fd, schema);
|
|
2375
|
+
}
|
|
2376
|
+
}
|
|
2377
|
+
|
|
2378
|
+
class IncrementalStream {
|
|
2379
|
+
writer;
|
|
2380
|
+
fd;
|
|
2381
|
+
closed = false;
|
|
2382
|
+
constructor(fd, schema) {
|
|
2383
|
+
this.fd = fd;
|
|
2384
|
+
this.writer = new RecordBatchStreamWriter4;
|
|
2385
|
+
this.writer.reset(undefined, schema);
|
|
2386
|
+
this.drain();
|
|
2387
|
+
}
|
|
2388
|
+
write(batch) {
|
|
2389
|
+
if (this.closed)
|
|
2390
|
+
throw new Error("Stream already closed");
|
|
2391
|
+
this.writer._writeRecordBatch(batch);
|
|
2392
|
+
this.drain();
|
|
2233
2393
|
}
|
|
2234
2394
|
close() {
|
|
2235
|
-
if (this.
|
|
2395
|
+
if (this.closed)
|
|
2236
2396
|
return;
|
|
2237
|
-
this.
|
|
2238
|
-
|
|
2239
|
-
|
|
2240
|
-
this._inputWriter = null;
|
|
2241
|
-
} else {
|
|
2242
|
-
const emptySchema = new Schema12([]);
|
|
2243
|
-
const ipc = serializeIpcStream(emptySchema, []);
|
|
2244
|
-
this._writeFn(ipc);
|
|
2245
|
-
}
|
|
2246
|
-
const drainPromise = (async () => {
|
|
2247
|
-
try {
|
|
2248
|
-
if (!this._outputStreamOpened) {
|
|
2249
|
-
const schema = await this._reader.openNextStream();
|
|
2250
|
-
if (schema) {
|
|
2251
|
-
while (await this._reader.readNextBatch() !== null) {}
|
|
2252
|
-
}
|
|
2253
|
-
} else {
|
|
2254
|
-
while (await this._reader.readNextBatch() !== null) {}
|
|
2255
|
-
}
|
|
2256
|
-
} catch {} finally {
|
|
2257
|
-
this._releaseBusy();
|
|
2258
|
-
}
|
|
2259
|
-
})();
|
|
2260
|
-
this._setDrainPromise(drainPromise);
|
|
2397
|
+
this.closed = true;
|
|
2398
|
+
const eos = new Uint8Array(new Int32Array([-1, 0]).buffer);
|
|
2399
|
+
writeAll(this.fd, eos);
|
|
2261
2400
|
}
|
|
2262
|
-
|
|
2263
|
-
|
|
2264
|
-
|
|
2265
|
-
|
|
2266
|
-
let readerPromise = null;
|
|
2267
|
-
let methodCache = null;
|
|
2268
|
-
let protocolName = "";
|
|
2269
|
-
let _busy = false;
|
|
2270
|
-
let _drainPromise = null;
|
|
2271
|
-
let closed = false;
|
|
2272
|
-
const writeFn = (bytes2) => {
|
|
2273
|
-
writable.write(bytes2);
|
|
2274
|
-
writable.flush?.();
|
|
2275
|
-
};
|
|
2276
|
-
async function ensureReader() {
|
|
2277
|
-
if (reader)
|
|
2278
|
-
return reader;
|
|
2279
|
-
if (!readerPromise) {
|
|
2280
|
-
readerPromise = IpcStreamReader.create(readable);
|
|
2401
|
+
drain() {
|
|
2402
|
+
const values = this.writer._sink._values;
|
|
2403
|
+
for (const chunk of values) {
|
|
2404
|
+
writeAll(this.fd, chunk);
|
|
2281
2405
|
}
|
|
2282
|
-
|
|
2283
|
-
return reader;
|
|
2406
|
+
values.length = 0;
|
|
2284
2407
|
}
|
|
2285
|
-
|
|
2286
|
-
|
|
2287
|
-
|
|
2288
|
-
|
|
2289
|
-
|
|
2290
|
-
|
|
2291
|
-
|
|
2408
|
+
}
|
|
2409
|
+
|
|
2410
|
+
// src/server.ts
|
|
2411
|
+
var EMPTY_SCHEMA5 = new Schema9([]);
|
|
2412
|
+
|
|
2413
|
+
class VgiRpcServer {
|
|
2414
|
+
protocol;
|
|
2415
|
+
enableDescribe;
|
|
2416
|
+
serverId;
|
|
2417
|
+
describeBatch = null;
|
|
2418
|
+
constructor(protocol, options) {
|
|
2419
|
+
this.protocol = protocol;
|
|
2420
|
+
this.enableDescribe = options?.enableDescribe ?? true;
|
|
2421
|
+
this.serverId = options?.serverId ?? crypto.randomUUID().replace(/-/g, "").slice(0, 12);
|
|
2422
|
+
if (this.enableDescribe) {
|
|
2423
|
+
const { batch } = buildDescribeBatch(protocol.name, protocol.getMethods(), this.serverId);
|
|
2424
|
+
this.describeBatch = batch;
|
|
2292
2425
|
}
|
|
2293
|
-
_busy = true;
|
|
2294
|
-
}
|
|
2295
|
-
function releaseBusy() {
|
|
2296
|
-
_busy = false;
|
|
2297
2426
|
}
|
|
2298
|
-
|
|
2299
|
-
|
|
2300
|
-
|
|
2301
|
-
|
|
2302
|
-
|
|
2303
|
-
|
|
2304
|
-
|
|
2427
|
+
async run() {
|
|
2428
|
+
const stdin = process.stdin;
|
|
2429
|
+
if (process.stdin.isTTY || process.stdout.isTTY) {
|
|
2430
|
+
process.stderr.write("WARNING: This process communicates via Arrow IPC on stdin/stdout " + `and is not intended to be run interactively.
|
|
2431
|
+
` + "It should be launched as a subprocess by an RPC client " + `(e.g. vgi_rpc.connect()).
|
|
2432
|
+
`);
|
|
2433
|
+
}
|
|
2434
|
+
const reader = await IpcStreamReader.create(stdin);
|
|
2435
|
+
const writer = new IpcStreamWriter;
|
|
2305
2436
|
try {
|
|
2306
|
-
|
|
2307
|
-
|
|
2308
|
-
writeFn(body);
|
|
2309
|
-
const r = await ensureReader();
|
|
2310
|
-
const response = await r.readStream();
|
|
2311
|
-
if (!response) {
|
|
2312
|
-
throw new Error("EOF reading __describe__ response");
|
|
2437
|
+
while (true) {
|
|
2438
|
+
await this.serveOne(reader, writer);
|
|
2313
2439
|
}
|
|
2314
|
-
|
|
2315
|
-
|
|
2316
|
-
|
|
2317
|
-
|
|
2440
|
+
} catch (e) {
|
|
2441
|
+
if (e.message?.includes("closed") || e.message?.includes("Expected Schema Message") || e.message?.includes("null or length 0") || e.code === "EPIPE" || e.code === "ERR_STREAM_PREMATURE_CLOSE" || e.code === "ERR_STREAM_DESTROYED" || e instanceof Error && e.message.includes("EOF")) {
|
|
2442
|
+
return;
|
|
2443
|
+
}
|
|
2444
|
+
throw e;
|
|
2318
2445
|
} finally {
|
|
2319
|
-
|
|
2446
|
+
await reader.cancel();
|
|
2320
2447
|
}
|
|
2321
2448
|
}
|
|
2322
|
-
|
|
2323
|
-
|
|
2324
|
-
|
|
2325
|
-
|
|
2326
|
-
try {
|
|
2327
|
-
const info = methods.get(method);
|
|
2328
|
-
if (!info) {
|
|
2329
|
-
throw new Error(`Unknown method: '${method}'`);
|
|
2330
|
-
}
|
|
2331
|
-
const r = await ensureReader();
|
|
2332
|
-
const fullParams = { ...info.defaults ?? {}, ...params ?? {} };
|
|
2333
|
-
const body = buildRequestIpc(info.paramsSchema, fullParams, method);
|
|
2334
|
-
writeFn(body);
|
|
2335
|
-
const response = await r.readStream();
|
|
2336
|
-
if (!response) {
|
|
2337
|
-
throw new Error("EOF reading response");
|
|
2338
|
-
}
|
|
2339
|
-
let resultBatch = null;
|
|
2340
|
-
for (const batch of response.batches) {
|
|
2341
|
-
if (batch.numRows === 0) {
|
|
2342
|
-
dispatchLogOrError(batch, onLog);
|
|
2343
|
-
continue;
|
|
2344
|
-
}
|
|
2345
|
-
resultBatch = batch;
|
|
2346
|
-
}
|
|
2347
|
-
if (!resultBatch) {
|
|
2348
|
-
return null;
|
|
2349
|
-
}
|
|
2350
|
-
const rows = extractBatchRows(resultBatch);
|
|
2351
|
-
if (rows.length === 0)
|
|
2352
|
-
return null;
|
|
2353
|
-
if (info.resultSchema.fields.length === 0)
|
|
2354
|
-
return null;
|
|
2355
|
-
return rows[0];
|
|
2356
|
-
} finally {
|
|
2357
|
-
releaseBusy();
|
|
2358
|
-
}
|
|
2359
|
-
},
|
|
2360
|
-
async stream(method, params) {
|
|
2361
|
-
const methods = await ensureMethodCache();
|
|
2362
|
-
await acquireBusy();
|
|
2363
|
-
try {
|
|
2364
|
-
const info = methods.get(method);
|
|
2365
|
-
if (!info) {
|
|
2366
|
-
throw new Error(`Unknown method: '${method}'`);
|
|
2367
|
-
}
|
|
2368
|
-
const r = await ensureReader();
|
|
2369
|
-
const fullParams = { ...info.defaults ?? {}, ...params ?? {} };
|
|
2370
|
-
const body = buildRequestIpc(info.paramsSchema, fullParams, method);
|
|
2371
|
-
writeFn(body);
|
|
2372
|
-
let header = null;
|
|
2373
|
-
if (info.headerSchema) {
|
|
2374
|
-
const headerStream = await r.readStream();
|
|
2375
|
-
if (headerStream) {
|
|
2376
|
-
for (const batch of headerStream.batches) {
|
|
2377
|
-
if (batch.numRows === 0) {
|
|
2378
|
-
dispatchLogOrError(batch, onLog);
|
|
2379
|
-
continue;
|
|
2380
|
-
}
|
|
2381
|
-
const rows = extractBatchRows(batch);
|
|
2382
|
-
if (rows.length > 0) {
|
|
2383
|
-
header = rows[0];
|
|
2384
|
-
}
|
|
2385
|
-
}
|
|
2386
|
-
}
|
|
2387
|
-
}
|
|
2388
|
-
const outputSchema = info.outputSchema ?? info.resultSchema;
|
|
2389
|
-
return new PipeStreamSession({
|
|
2390
|
-
reader: r,
|
|
2391
|
-
writeFn,
|
|
2392
|
-
onLog,
|
|
2393
|
-
header,
|
|
2394
|
-
outputSchema,
|
|
2395
|
-
releaseBusy,
|
|
2396
|
-
setDrainPromise
|
|
2397
|
-
});
|
|
2398
|
-
} catch (e) {
|
|
2399
|
-
try {
|
|
2400
|
-
const r = await ensureReader();
|
|
2401
|
-
const emptySchema = new Schema12([]);
|
|
2402
|
-
const ipc = serializeIpcStream(emptySchema, []);
|
|
2403
|
-
writeFn(ipc);
|
|
2404
|
-
const outStream = await r.readStream();
|
|
2405
|
-
} catch {}
|
|
2406
|
-
releaseBusy();
|
|
2407
|
-
throw e;
|
|
2408
|
-
}
|
|
2409
|
-
},
|
|
2410
|
-
async describe() {
|
|
2411
|
-
const methods = await ensureMethodCache();
|
|
2412
|
-
return {
|
|
2413
|
-
protocolName,
|
|
2414
|
-
methods: [...methods.values()]
|
|
2415
|
-
};
|
|
2416
|
-
},
|
|
2417
|
-
close() {
|
|
2418
|
-
if (closed)
|
|
2419
|
-
return;
|
|
2420
|
-
closed = true;
|
|
2421
|
-
writable.end();
|
|
2449
|
+
async serveOne(reader, writer) {
|
|
2450
|
+
const stream = await reader.readStream();
|
|
2451
|
+
if (!stream) {
|
|
2452
|
+
throw new Error("EOF");
|
|
2422
2453
|
}
|
|
2423
|
-
|
|
2424
|
-
|
|
2425
|
-
|
|
2426
|
-
|
|
2427
|
-
|
|
2428
|
-
|
|
2429
|
-
stderr: options?.stderr ?? "ignore",
|
|
2430
|
-
cwd: options?.cwd,
|
|
2431
|
-
env: options?.env ? { ...process.env, ...options.env } : undefined
|
|
2432
|
-
});
|
|
2433
|
-
const stdout = proc.stdout;
|
|
2434
|
-
const writable = {
|
|
2435
|
-
write(data) {
|
|
2436
|
-
proc.stdin.write(data);
|
|
2437
|
-
},
|
|
2438
|
-
flush() {
|
|
2439
|
-
proc.stdin.flush();
|
|
2440
|
-
},
|
|
2441
|
-
end() {
|
|
2442
|
-
proc.stdin.end();
|
|
2454
|
+
const { schema, batches } = stream;
|
|
2455
|
+
if (batches.length === 0) {
|
|
2456
|
+
const err = new RpcError("ProtocolError", "Request stream contains no batches", "");
|
|
2457
|
+
const errBatch = buildErrorBatch(EMPTY_SCHEMA5, err, this.serverId, null);
|
|
2458
|
+
writer.writeStream(EMPTY_SCHEMA5, [errBatch]);
|
|
2459
|
+
return;
|
|
2443
2460
|
}
|
|
2444
|
-
|
|
2445
|
-
|
|
2446
|
-
|
|
2447
|
-
|
|
2448
|
-
const originalClose = client.close;
|
|
2449
|
-
client.close = () => {
|
|
2450
|
-
originalClose.call(client);
|
|
2461
|
+
const batch = batches[0];
|
|
2462
|
+
let methodName;
|
|
2463
|
+
let params;
|
|
2464
|
+
let requestId;
|
|
2451
2465
|
try {
|
|
2452
|
-
|
|
2453
|
-
|
|
2454
|
-
|
|
2455
|
-
|
|
2466
|
+
const parsed = parseRequest(schema, batch);
|
|
2467
|
+
methodName = parsed.methodName;
|
|
2468
|
+
params = parsed.params;
|
|
2469
|
+
requestId = parsed.requestId;
|
|
2470
|
+
} catch (e) {
|
|
2471
|
+
const errBatch = buildErrorBatch(EMPTY_SCHEMA5, e, this.serverId, null);
|
|
2472
|
+
writer.writeStream(EMPTY_SCHEMA5, [errBatch]);
|
|
2473
|
+
if (e instanceof VersionError || e instanceof RpcError) {
|
|
2474
|
+
return;
|
|
2475
|
+
}
|
|
2476
|
+
throw e;
|
|
2477
|
+
}
|
|
2478
|
+
if (methodName === DESCRIBE_METHOD_NAME && this.describeBatch) {
|
|
2479
|
+
writer.writeStream(this.describeBatch.schema, [this.describeBatch]);
|
|
2480
|
+
return;
|
|
2481
|
+
}
|
|
2482
|
+
const methods = this.protocol.getMethods();
|
|
2483
|
+
const method = methods.get(methodName);
|
|
2484
|
+
if (!method) {
|
|
2485
|
+
const available = [...methods.keys()].sort();
|
|
2486
|
+
const err = new Error(`Unknown method: '${methodName}'. Available methods: [${available.join(", ")}]`);
|
|
2487
|
+
const errBatch = buildErrorBatch(EMPTY_SCHEMA5, err, this.serverId, requestId);
|
|
2488
|
+
writer.writeStream(EMPTY_SCHEMA5, [errBatch]);
|
|
2489
|
+
return;
|
|
2490
|
+
}
|
|
2491
|
+
if (method.type === "unary" /* UNARY */) {
|
|
2492
|
+
await dispatchUnary(method, params, writer, this.serverId, requestId);
|
|
2493
|
+
} else {
|
|
2494
|
+
await dispatchStream(method, params, writer, reader, this.serverId, requestId);
|
|
2495
|
+
}
|
|
2496
|
+
}
|
|
2456
2497
|
}
|
|
2457
2498
|
export {
|
|
2499
|
+
unpackStateToken,
|
|
2458
2500
|
toSchema,
|
|
2459
2501
|
subprocessConnect,
|
|
2460
2502
|
str,
|
|
2461
2503
|
pipeConnect,
|
|
2462
2504
|
parseDescribeResponse,
|
|
2505
|
+
jsonStateSerializer,
|
|
2463
2506
|
int32,
|
|
2464
2507
|
int,
|
|
2465
2508
|
inferParamTypes,
|
|
@@ -2494,4 +2537,4 @@ export {
|
|
|
2494
2537
|
ARROW_CONTENT_TYPE
|
|
2495
2538
|
};
|
|
2496
2539
|
|
|
2497
|
-
//# debugId=
|
|
2540
|
+
//# debugId=E5E571F45BA75F8D64756E2164756E21
|