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