Skip to content

Commit ed1786d

Browse files
1st11st1
committed
Implement withCodecs/encode support
1 parent aebeb4c commit ed1786d

24 files changed

+300
-315
lines changed

‎packages/driver/src/baseConn.ts‎

Lines changed: 30 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,8 @@
1818

1919
import { INVALID_CODEC, NullCodec, NULL_CODEC } from "./codecs/codecs";
2020
import type { ICodec } from "./codecs/ifaces";
21-
import { NamedTupleCodec } from "./codecs/namedtuple";
2221
import { ObjectCodec } from "./codecs/object";
2322
import type { CodecsRegistry } from "./codecs/registry";
24-
import { EmptyTupleCodec, EMPTY_TUPLE_CODEC, TupleCodec } from "./codecs/tuple";
2523
import { versionGreaterThan, versionGreaterThanOrEqual } from "./utils";
2624
import * as errors from "./errors";
2725
import { resolveErrorCode, errorFromJSON } from "./errors/resolve";
@@ -360,14 +358,13 @@ export class BaseRawConnection {
360358
private _parseDataMessages(
361359
codec: ICodec,
362360
result: any[] | WriteBuffer,
363-
state: Options,
361+
ctx: CodecContext,
364362
): void {
365363
const frb = ReadBuffer.alloc();
366364
const $D = chars.$D;
367365
const buffer = this.buffer;
368366

369367
if (Array.isArray(result)) {
370-
const ctx = new CodecContext(state.codecs);
371368
while (buffer.takeMessageType($D)) {
372369
buffer.consumeMessageInto(frb);
373370
frb.discard(6);
@@ -460,42 +457,27 @@ export class BaseRawConnection {
460457
}
461458
}
462459

463-
private _encodeArgs(args: QueryArgs, inCodec: ICodec): Uint8Array {
464-
if (versionGreaterThanOrEqual(this.protocolVersion, [0, 12])) {
465-
if (inCodec === NULL_CODEC) {
466-
if (args != null) {
467-
throw new errors.QueryArgumentError(
468-
`This query does not contain any query parameters, ` +
469-
`but query arguments were provided to the 'query*()' method`,
470-
);
471-
}
472-
return NullCodec.BUFFER;
473-
}
474-
475-
if (inCodec instanceof ObjectCodec) {
476-
return inCodec.encodeArgs(args);
477-
}
478-
479-
// Shouldn't ever happen.
480-
throw new errors.ProtocolError("invalid input codec");
481-
} else {
482-
if (inCodec === EMPTY_TUPLE_CODEC) {
483-
if (args != null) {
484-
throw new errors.QueryArgumentError(
485-
`This query does not contain any query parameters, ` +
486-
`but query arguments were provided to the 'query*()' method`,
487-
);
488-
}
489-
return EmptyTupleCodec.BUFFER;
490-
}
491-
492-
if (inCodec instanceof NamedTupleCodec || inCodec instanceof TupleCodec) {
493-
return inCodec.encodeArgs(args);
460+
private _encodeArgs(
461+
args: QueryArgs,
462+
inCodec: ICodec,
463+
ctx: CodecContext
464+
): Uint8Array {
465+
if (inCodec === NULL_CODEC) {
466+
if (args != null) {
467+
throw new errors.QueryArgumentError(
468+
`This query does not contain any query parameters, ` +
469+
`but query arguments were provided to the 'query*()' method`,
470+
);
494471
}
472+
return NullCodec.BUFFER;
473+
}
495474

496-
// Shouldn't ever happen.
497-
throw new errors.ProtocolError("invalid input codec");
475+
if (inCodec instanceof ObjectCodec) {
476+
return inCodec.encodeArgs(args, ctx);
498477
}
478+
479+
// Shouldn't ever happen.
480+
throw new errors.ProtocolError("invalid input codec");
499481
}
500482

501483
private _encodeParseParams(
@@ -555,7 +537,7 @@ export class BaseRawConnection {
555537
} else {
556538
if (!this.stateCache.has(state)) {
557539
const buf = new WriteBuffer();
558-
this.stateCodec.encode(buf, state._serialise());
540+
this.stateCodec.encode(buf, state._serialise(), NOOP_CODEC_CONTEXT);
559541
this.stateCache.set(state, buf.unwrap());
560542
}
561543
wb.writeBuffer(this.stateCache.get(state)!);
@@ -698,6 +680,7 @@ export class BaseRawConnection {
698680
capabilitiesFlags: number = RESTRICTED_CAPABILITIES,
699681
options?: QueryOptions,
700682
): Promise<errors.EdgeDBError[]> {
683+
let ctx: CodecContext | null = null;
701684
const wb = new WriteMessageBuffer();
702685
wb.beginMessage(chars.$O);
703686

@@ -716,7 +699,10 @@ export class BaseRawConnection {
716699
wb.writeBuffer(outCodec.tidBuffer);
717700

718701
if (inCodec) {
719-
wb.writeBuffer(this._encodeArgs(args, inCodec));
702+
if (ctx == null) {
703+
ctx = new CodecContext(state.codecs);
704+
}
705+
wb.writeBuffer(this._encodeArgs(args, inCodec, ctx));
720706
} else {
721707
wb.writeInt32(0);
722708
}
@@ -730,6 +716,10 @@ export class BaseRawConnection {
730716
let parsing = true;
731717
let warnings: errors.EdgeDBError[] = [];
732718

719+
if (ctx == null) {
720+
ctx = new CodecContext(state.codecs);
721+
}
722+
733723
while (parsing) {
734724
if (!this.buffer.takeMessage()) {
735725
await this._waitForMessage();
@@ -741,7 +731,7 @@ export class BaseRawConnection {
741731
case chars.$D: {
742732
if (error == null) {
743733
try {
744-
this._parseDataMessages(outCodec!, result, state);
734+
this._parseDataMessages(outCodec!, result, ctx);
745735
} catch (e: any) {
746736
error = e;
747737
this.buffer.finishMessage();

‎packages/driver/src/codecs/array.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ export class ArrayCodec extends Codec implements ICodec {
3737
this.typeName = typeName;
3838
}
3939

40-
encode(buf: WriteBuffer, obj: any): void {
40+
encode(buf: WriteBuffer, obj: any, ctx: CodecContext): void {
4141
if (
4242
!(
4343
this.subCodec instanceof ScalarCodec ||
@@ -72,7 +72,7 @@ export class ArrayCodec extends Codec implements ICodec {
7272
if (item == null) {
7373
elemData.writeInt32(-1);
7474
} else {
75-
subCodec.encode(elemData, item);
75+
subCodec.encode(elemData, item, ctx);
7676
}
7777
}
7878
const elemBuf = elemData.unwrap();

‎packages/driver/src/codecs/boolean.ts‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,16 @@ import type { CodecContext } from "./context";
2525
export class BoolCodec extends ScalarCodec implements ICodec {
2626
override tsType = "boolean";
2727

28-
encode(buf: WriteBuffer, object: any): void {
29-
const typeOf = typeof object;
28+
encode(buf: WriteBuffer, object: any, ctx: CodecContext): void {
29+
const val = ctx.preEncode<Codecs.BoolCodec>(this, object);
30+
const typeOf = typeof val;
3031
if (typeOf !== "boolean" && typeOf !== "number") {
3132
throw new InvalidArgumentError(
32-
`a boolean or a number was expected, got "${object}"`,
33+
`a boolean or a number was expected, got "${val}"`,
3334
);
3435
}
3536
buf.writeInt32(1);
36-
buf.writeChar(object ? 1 : 0);
37+
buf.writeChar(val ? 1 : 0);
3738
}
3839

3940
decode(buf: ReadBuffer, ctx: CodecContext): any {

‎packages/driver/src/codecs/bytes.ts‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,16 @@ import type { CodecContext } from "./context";
2525
export class BytesCodec extends ScalarCodec implements ICodec {
2626
override tsType = "Uint8Array";
2727

28-
encode(buf: WriteBuffer, object: any): void {
29-
if (!(object instanceof Uint8Array)) {
28+
encode(buf: WriteBuffer, object: any, ctx: CodecContext): void {
29+
const val = ctx.preEncode<Codecs.BytesCodec>(this, object);
30+
if (!(val instanceof Uint8Array)) {
3031
throw new InvalidArgumentError(
31-
`a Uint8Array or Buffer was expected, got "${object}"`,
32+
`a Uint8Array or Buffer was expected, got "${val}"`,
3233
);
3334
}
3435

35-
buf.writeInt32(object.length);
36-
buf.writeBuffer(object);
36+
buf.writeInt32(val.length);
37+
buf.writeBuffer(val);
3738
}
3839

3940
decode(buf: ReadBuffer, ctx: CodecContext): any {

‎packages/driver/src/codecs/codecs.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ export namespace Codecs {
8888
export type MemoryCodec = Codec<bigint>;
8989
export type PgVectorCodec = Codec<Float32Array>;
9090
export type PGVectorSparseCodec = Codec<
91-
[dimensions: number, indexes: Uint32Array, data: Float32Array]
91+
[dimensions: number, indexes: Uint32Array, values: Float32Array]
9292
>;
9393
export type StrCodec = Codec<string>;
9494
export type UUIDCodec = Codec<Uint8Array>;

‎packages/driver/src/codecs/context.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import type { Codecs } from "./codecs";
44
export type CodecMap = Map<string, Codecs.AnyCodec>;
55
export type ReadonlyCodecMap = ReadonlyMap<string, Codecs.AnyCodec>;
66

7-
export type CodecValueType<S> =
7+
type CodecValueType<S> =
88
S extends Codecs.KnownCodecs[keyof Codecs.KnownCodecs]
99
? S extends Codecs.Codec<infer T>
1010
? T

‎packages/driver/src/codecs/datetime.ts‎

Lines changed: 92 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,12 +40,27 @@ import type { Codecs } from "./codecs";
4040
* which is specified in the below constant (in milliseconds.)
4141
*/
4242
const TIMESHIFT = 946684800000;
43+
const BI_TIMESHIFT = BigInt(TIMESHIFT);
44+
const BI_TIMESHIFT_US = BigInt(TIMESHIFT) * 1000n;
45+
4346
const DATESHIFT_ORD = ymd2ord(2000, 1, 1);
4447

4548
export class DateTimeCodec extends ScalarCodec implements ICodec {
4649
override tsType = "Date";
4750

48-
encode(buf: WriteBuffer, object: unknown): void {
51+
encode(buf: WriteBuffer, object: unknown, ctx: CodecContext): void {
52+
if (ctx.hasOverload(this)) {
53+
const val = ctx.preEncode<Codecs.DateTimeCodec>(this, object);
54+
if (typeof val != 'bigint') {
55+
throw new InvalidArgumentError(
56+
`a bigint was expected out of a custom std::datetime codec`
57+
);
58+
}
59+
buf.writeInt32(8);
60+
buf.writeBigInt64((val - BI_TIMESHIFT) * 1000n);
61+
return;
62+
}
63+
4964
if (!(object instanceof Date)) {
5065
throw new InvalidArgumentError(
5166
`a Date instance was expected, got "${object}"`,
@@ -79,7 +94,20 @@ export class LocalDateTimeCodec extends ScalarCodec implements ICodec {
7994
override tsType = "LocalDateTime";
8095
override tsModule = "edgedb";
8196

82-
encode(buf: WriteBuffer, object: unknown): void {
97+
encode(buf: WriteBuffer, object: unknown, ctx: CodecContext): void {
98+
if (ctx.hasOverload(this)) {
99+
let us = ctx.preEncode<Codecs.LocalDateTimeCodec>(this, object);
100+
if (typeof us != 'bigint') {
101+
throw new InvalidArgumentError(
102+
`a bigint was expected out of a custom cal_local_datetime codec`
103+
);
104+
}
105+
us -= BI_TIMESHIFT_US;
106+
buf.writeInt32(8);
107+
buf.writeBigInt64(us);
108+
return;
109+
}
110+
83111
if (!(object instanceof LocalDateTime)) {
84112
throw new InvalidArgumentError(
85113
`a LocalDateTime instance was expected, got "${object}"`,
@@ -105,7 +133,7 @@ export class LocalDateTimeCodec extends ScalarCodec implements ICodec {
105133
}
106134

107135
buf.writeInt32(8);
108-
buf.writeBigInt64(us as bigint);
136+
buf.writeBigInt64(us);
109137
}
110138

111139
decode(buf: ReadBuffer, ctx: CodecContext): LocalDateTime {
@@ -142,14 +170,24 @@ export class LocalDateTimeCodec extends ScalarCodec implements ICodec {
142170
export class LocalDateCodec extends ScalarCodec implements ICodec {
143171
override tsType = "LocalDate";
144172
override tsModule = "edgedb";
145-
encode(buf: WriteBuffer, object: unknown): void {
173+
encode(buf: WriteBuffer, object: unknown, ctx: CodecContext): void {
174+
if (ctx.hasOverload(this)) {
175+
const ret = ctx.preEncode<Codecs.LocalDateCodec>(this, object);
176+
const ord = ymd2ord(...ret);
177+
buf.writeInt32(4);
178+
buf.writeInt32(ord - DATESHIFT_ORD);
179+
return;
180+
}
181+
146182
if (!(object instanceof LocalDate)) {
147183
throw new InvalidArgumentError(
148184
`a LocalDate instance was expected, got "${object}"`,
149185
);
150186
}
187+
188+
const ord = LocalDateToOrdinal(object);
151189
buf.writeInt32(4);
152-
buf.writeInt32(LocalDateToOrdinal(object) - DATESHIFT_ORD);
190+
buf.writeInt32(ord - DATESHIFT_ORD);
153191
}
154192

155193
decode(buf: ReadBuffer, ctx: CodecContext): LocalDate {
@@ -166,7 +204,19 @@ export class LocalDateCodec extends ScalarCodec implements ICodec {
166204
export class LocalTimeCodec extends ScalarCodec implements ICodec {
167205
override tsType = "LocalTime";
168206
override tsModule = "edgedb";
169-
encode(buf: WriteBuffer, object: unknown): void {
207+
encode(buf: WriteBuffer, object: unknown, ctx: CodecContext): void {
208+
if (ctx.hasOverload(this)) {
209+
const us = ctx.preEncode<Codecs.LocalTimeCodec>(this, object);
210+
if (typeof us != 'bigint') {
211+
throw new InvalidArgumentError(
212+
`a bigint was expected out of a custom cal::local_time codec`
213+
);
214+
}
215+
buf.writeInt32(8);
216+
buf.writeBigInt64(us);
217+
return;
218+
}
219+
170220
if (!(object instanceof LocalTime)) {
171221
throw new InvalidArgumentError(
172222
`a LocalTime instance was expected, got "${object}"`,
@@ -230,7 +280,22 @@ export function checkValidEdgeDBDuration(duration: Duration): null | string {
230280
export class DurationCodec extends ScalarCodec implements ICodec {
231281
override tsType = "Duration";
232282
override tsModule = "edgedb";
233-
encode(buf: WriteBuffer, object: unknown): void {
283+
encode(buf: WriteBuffer, object: unknown, ctx: CodecContext): void {
284+
if (ctx.hasOverload(this)) {
285+
const us = ctx.preEncode<Codecs.DurationCodec>(this, object);
286+
if (typeof us != 'bigint') {
287+
throw new InvalidArgumentError(
288+
`a bigint was expected out of a custom std::duration codec`
289+
);
290+
}
291+
292+
buf.writeInt32(16);
293+
buf.writeBigInt64(us);
294+
buf.writeInt32(0);
295+
buf.writeInt32(0);
296+
return;
297+
}
298+
234299
if (!(object instanceof Duration)) {
235300
throw new InvalidArgumentError(
236301
`a Duration instance was expected, got "${object}"`,
@@ -318,7 +383,16 @@ export class DurationCodec extends ScalarCodec implements ICodec {
318383
export class RelativeDurationCodec extends ScalarCodec implements ICodec {
319384
override tsType = "RelativeDuration";
320385
override tsModule = "edgedb";
321-
encode(buf: WriteBuffer, object: unknown): void {
386+
encode(buf: WriteBuffer, object: unknown, ctx: CodecContext): void {
387+
if (ctx.hasOverload(this)) {
388+
const ret = ctx.preEncode<Codecs.RelativeDurationCodec>(this, object);
389+
buf.writeInt32(16);
390+
buf.writeBigInt64(ret[2]);
391+
buf.writeInt32(ret[1]);
392+
buf.writeInt32(ret[0]);
393+
return;
394+
}
395+
322396
if (!(object instanceof RelativeDuration)) {
323397
throw new InvalidArgumentError(`
324398
a RelativeDuration instance was expected, got "${object}"
@@ -393,7 +467,16 @@ export class RelativeDurationCodec extends ScalarCodec implements ICodec {
393467
export class DateDurationCodec extends ScalarCodec implements ICodec {
394468
override tsType = "DateDuration";
395469
override tsModule = "edgedb";
396-
encode(buf: WriteBuffer, object: unknown): void {
470+
encode(buf: WriteBuffer, object: unknown, ctx: CodecContext): void {
471+
if (ctx.hasOverload(this)) {
472+
const ret = ctx.preEncode<Codecs.DateDurationCodec>(this, object);
473+
buf.writeInt32(16);
474+
buf.writeInt64(0);
475+
buf.writeInt32(ret[1]);
476+
buf.writeInt32(ret[0]);
477+
return;
478+
}
479+
397480
if (!(object instanceof DateDuration)) {
398481
throw new InvalidArgumentError(`
399482
a DateDuration instance was expected, got "${object}"

0 commit comments

Comments
 (0)