From 73df1c79c3e80e8e2a8e1d49e367facea5b75ca9 Mon Sep 17 00:00:00 2001 From: Scott Trinh Date: Wed, 4 Oct 2023 17:43:49 -0400 Subject: [PATCH] Support multirange in the driver --- packages/driver/src/codecs/array.ts | 15 +- packages/driver/src/codecs/ifaces.ts | 3 +- packages/driver/src/codecs/range.ts | 194 +++++++++++++----- packages/driver/src/codecs/registry.ts | 15 +- packages/driver/src/datatypes/range.ts | 24 ++- packages/driver/src/index.shared.ts | 4 +- .../driver/src/reflection/analyzeQuery.ts | 10 +- packages/driver/test/client.test.ts | 65 +++--- 8 files changed, 243 insertions(+), 87 deletions(-) diff --git a/packages/driver/src/codecs/array.ts b/packages/driver/src/codecs/array.ts index c497b6d99..cc8ac0879 100644 --- a/packages/driver/src/codecs/array.ts +++ b/packages/driver/src/codecs/array.ts @@ -16,10 +16,11 @@ * limitations under the License. */ -import { ICodec, Codec, ScalarCodec, uuid, CodecKind } from "./ifaces"; +import type { ICodec, uuid, CodecKind } from "./ifaces"; +import { Codec, ScalarCodec } from "./ifaces"; import { WriteBuffer, ReadBuffer } from "../primitives/buffer"; import { TupleCodec } from "./tuple"; -import { RangeCodec } from "./range"; +import { MultiRangeCodec, RangeCodec } from "./range"; import { InvalidArgumentError, ProtocolError } from "../errors"; import { NamedTupleCodec } from "./namedtuple"; @@ -39,7 +40,8 @@ export class ArrayCodec extends Codec implements ICodec { this.subCodec instanceof ScalarCodec || this.subCodec instanceof TupleCodec || this.subCodec instanceof NamedTupleCodec || - this.subCodec instanceof RangeCodec + this.subCodec instanceof RangeCodec || + this.subCodec instanceof MultiRangeCodec ) ) { throw new InvalidArgumentError( @@ -48,7 +50,9 @@ export class ArrayCodec extends Codec implements ICodec { } if (!Array.isArray(obj) && !isTypedArray(obj)) { - throw new InvalidArgumentError("an array was expected"); + throw new InvalidArgumentError( + `an array or multirange was expected (got type ${obj.constructor.name})` + ); } const subCodec = this.subCodec; @@ -60,8 +64,7 @@ export class ArrayCodec extends Codec implements ICodec { throw new InvalidArgumentError("too many elements in array"); } - for (let i = 0; i < objLen; i++) { - const item = obj[i]; + for (const item of obj) { if (item == null) { elemData.writeInt32(-1); } else { diff --git a/packages/driver/src/codecs/ifaces.ts b/packages/driver/src/codecs/ifaces.ts index a9b9ff2fa..bde43efc1 100644 --- a/packages/driver/src/codecs/ifaces.ts +++ b/packages/driver/src/codecs/ifaces.ts @@ -29,7 +29,8 @@ export type CodecKind = | "set" | "scalar" | "sparse_object" - | "range"; + | "range" + | "multirange"; export interface ICodec { readonly tid: uuid; diff --git a/packages/driver/src/codecs/range.ts b/packages/driver/src/codecs/range.ts index e067996e2..8a96a080e 100644 --- a/packages/driver/src/codecs/range.ts +++ b/packages/driver/src/codecs/range.ts @@ -16,10 +16,11 @@ * limitations under the License. */ -import { ICodec, Codec, uuid, CodecKind } from "./ifaces"; +import type { ICodec, uuid, CodecKind } from "./ifaces"; +import { Codec } from "./ifaces"; import { WriteBuffer, ReadBuffer } from "../primitives/buffer"; -import { Range } from "../datatypes/range"; -import { InvalidArgumentError } from "../errors"; +import { MultiRange, Range } from "../datatypes/range"; +import { InvalidArgumentError, ProtocolError } from "../errors"; enum RangeFlags { EMPTY = 1 << 0, @@ -29,6 +30,68 @@ enum RangeFlags { EMPTY_UPPER = 1 << 4, } +const MAXINT32 = 0x7fffffff; + +function encodeRange(buf: WriteBuffer, obj: any, subCodec: ICodec): void { + if (!(obj instanceof Range)) { + throw new InvalidArgumentError("a Range was expected"); + } + + const elemData = new WriteBuffer(); + + if (obj.lower !== null) { + subCodec.encode(elemData, obj.lower); + } + if (obj.upper !== null) { + subCodec.encode(elemData, obj.upper); + } + + const elemBuf = elemData.unwrap(); + + buf.writeInt32(1 + elemBuf.length); + buf.writeUInt8( + obj.isEmpty + ? RangeFlags.EMPTY + : (obj.incLower ? RangeFlags.INC_LOWER : 0) | + (obj.incUpper ? RangeFlags.INC_UPPER : 0) | + (obj.lower === null ? RangeFlags.EMPTY_LOWER : 0) | + (obj.upper === null ? RangeFlags.EMPTY_UPPER : 0) + ); + buf.writeBuffer(elemBuf); +} + +function decodeRange(buf: ReadBuffer, subCodec: ICodec): any { + const flags = buf.readUInt8(); + + if (flags & RangeFlags.EMPTY) { + return Range.empty(); + } + + const elemBuf = ReadBuffer.alloc(); + + let lower: any = null; + let upper: any = null; + + if (!(flags & RangeFlags.EMPTY_LOWER)) { + buf.sliceInto(elemBuf, buf.readInt32()); + lower = subCodec.decode(elemBuf); + elemBuf.finish(); + } + + if (!(flags & RangeFlags.EMPTY_UPPER)) { + buf.sliceInto(elemBuf, buf.readInt32()); + upper = subCodec.decode(elemBuf); + elemBuf.finish(); + } + + return new Range( + lower, + upper, + !!(flags & RangeFlags.INC_LOWER), + !!(flags & RangeFlags.INC_UPPER) + ); +} + export class RangeCodec extends Codec implements ICodec { private subCodec: ICodec; @@ -37,66 +100,101 @@ export class RangeCodec extends Codec implements ICodec { this.subCodec = subCodec; } + encode(buf: WriteBuffer, obj: any) { + return encodeRange(buf, obj, this.subCodec); + } + + decode(buf: ReadBuffer): any { + return decodeRange(buf, this.subCodec); + } + + getSubcodecs(): ICodec[] { + return [this.subCodec]; + } + + getKind(): CodecKind { + return "range"; + } +} + +export class MultiRangeCodec extends Codec implements ICodec { + private subCodec: ICodec; + + constructor(tid: uuid, subCodec: ICodec) { + super(tid); + this.subCodec = subCodec; + } + encode(buf: WriteBuffer, obj: any): void { - if (!(obj instanceof Range)) { - throw new InvalidArgumentError("a Range was expected"); + if (!(obj instanceof MultiRange)) { + throw new TypeError( + `a MultiRange expected (got type ${obj.constructor.name})` + ); } - const subCodec = this.subCodec; - const elemData = new WriteBuffer(); + const objLen = obj.length; + if (objLen > MAXINT32) { + throw new InvalidArgumentError("too many elements in array"); + } - if (obj.lower !== null) { - subCodec.encode(elemData, obj.lower); + const elemData = new WriteBuffer(); + for (const item of obj) { + try { + encodeRange(elemData, item, this.subCodec); + } catch (e) { + if (e instanceof InvalidArgumentError) { + throw new InvalidArgumentError( + `invalid multirange element: ${e.message}` + ); + } else { + throw e; + } + } } - if (obj.upper !== null) { - subCodec.encode(elemData, obj.upper); + + const elemBuf = elemData.unwrap() + const elemDataLen = elemBuf.length; + if (elemDataLen > MAXINT32 - 4) { + throw new InvalidArgumentError( + `size of encoded multirange datum exceeds the maximum allowed ${ + MAXINT32 - 4 + } bytes` + ); } - const elemBuf = elemData.unwrap(); - - buf.writeInt32(1 + elemBuf.length); - buf.writeUInt8( - obj.isEmpty - ? RangeFlags.EMPTY - : (obj.incLower ? RangeFlags.INC_LOWER : 0) | - (obj.incUpper ? RangeFlags.INC_UPPER : 0) | - (obj.lower === null ? RangeFlags.EMPTY_LOWER : 0) | - (obj.upper === null ? RangeFlags.EMPTY_UPPER : 0) - ); + // Datum length + buf.writeInt32(4 + elemDataLen); + + // Number of elements in multirange + buf.writeInt32(objLen); buf.writeBuffer(elemBuf); } decode(buf: ReadBuffer): any { - const flags = buf.readUInt8(); - - if (flags & RangeFlags.EMPTY) { - return Range.empty(); - } - + const elemCount = buf.readInt32(); + const result = new Array(elemCount); const elemBuf = ReadBuffer.alloc(); const subCodec = this.subCodec; - let lower: any = null; - let upper: any = null; - - if (!(flags & RangeFlags.EMPTY_LOWER)) { - buf.sliceInto(elemBuf, buf.readInt32()); - lower = subCodec.decode(elemBuf); - elemBuf.finish(); + for (let i = 0; i < elemCount; i++) { + const elemLen = buf.readInt32(); + if (elemLen === -1) { + throw new ProtocolError("unexpected NULL element in multirange value"); + } else { + buf.sliceInto(elemBuf, elemLen); + const elem = decodeRange(elemBuf, subCodec); + if (elemBuf.length) { + throw new ProtocolError( + `unexpected trailing data in buffer after multirange element decoding: ${elemBuf.length}` + ); + } + + result[i] = elem; + elemBuf.finish(); + } } - if (!(flags & RangeFlags.EMPTY_UPPER)) { - buf.sliceInto(elemBuf, buf.readInt32()); - upper = subCodec.decode(elemBuf); - elemBuf.finish(); - } - - return new Range( - lower, - upper, - !!(flags & RangeFlags.INC_LOWER), - !!(flags & RangeFlags.INC_UPPER) - ); + return new MultiRange(result); } getSubcodecs(): ICodec[] { @@ -104,6 +202,6 @@ export class RangeCodec extends Codec implements ICodec { } getKind(): CodecKind { - return "range"; + return "multirange"; } } diff --git a/packages/driver/src/codecs/registry.ts b/packages/driver/src/codecs/registry.ts index d5f861552..b951f31e0 100644 --- a/packages/driver/src/codecs/registry.ts +++ b/packages/driver/src/codecs/registry.ts @@ -30,7 +30,7 @@ import { NamedTupleCodec } from "./namedtuple"; import { EnumCodec } from "./enum"; import { ObjectCodec } from "./object"; import { SetCodec } from "./set"; -import { RangeCodec } from "./range"; +import { MultiRangeCodec, RangeCodec } from "./range"; import { ProtocolVersion } from "../ifaces"; import { versionGreaterThanOrEqual } from "../utils"; import { SparseObjectCodec } from "./sparseObject"; @@ -49,6 +49,7 @@ const CTYPE_ARRAY = 6; const CTYPE_ENUM = 7; const CTYPE_INPUT_SHAPE = 8; const CTYPE_RANGE = 9; +const CTYPE_MULTIRANGE = 12; export interface CustomCodecSpec { int64_bigint?: boolean; @@ -446,6 +447,18 @@ export class CodecsRegistry { res = new RangeCodec(tid, subCodec); break; } + + case CTYPE_MULTIRANGE: { + const pos = frb.readUInt16(); + const subCodec = cl[pos]; + if (subCodec == null) { + throw new ProtocolError( + "could not build range codec: missing subcodec" + ); + } + res = new MultiRangeCodec(tid, subCodec); + break; + } } if (res == null) { diff --git a/packages/driver/src/datatypes/range.ts b/packages/driver/src/datatypes/range.ts index 7593ca83e..d881ccbf0 100644 --- a/packages/driver/src/datatypes/range.ts +++ b/packages/driver/src/datatypes/range.ts @@ -16,7 +16,7 @@ * limitations under the License. */ -import { Duration, LocalDate, LocalDateTime } from "./datetime"; +import type { Duration, LocalDate, LocalDateTime } from "./datetime"; export class Range< T extends number | Date | LocalDate | LocalDateTime | Duration @@ -63,3 +63,25 @@ export class Range< }; } } + +export class MultiRange< + T extends number | Date | LocalDate | LocalDateTime | Duration +> { + constructor(private readonly _ranges: Range[] = []) {} + + get length() { + return this._ranges.length; + } + + *[Symbol.iterator]() { + for (const range of this._ranges) { + yield range; + } + } + + toJSON() { + return { + ranges: this._ranges, + }; + } +} diff --git a/packages/driver/src/index.shared.ts b/packages/driver/src/index.shared.ts index c411b92fd..a12976798 100644 --- a/packages/driver/src/index.shared.ts +++ b/packages/driver/src/index.shared.ts @@ -27,14 +27,14 @@ export { DateDuration, } from "./datatypes/datetime"; export { ConfigMemory } from "./datatypes/memory"; -export { Range } from "./datatypes/range"; +export { Range, MultiRange } from "./datatypes/range"; export type { Executor } from "./ifaces"; export * from "./errors"; /* Private APIs */ -import * as codecs from "./codecs/ifaces"; +import type * as codecs from "./codecs/ifaces"; import * as reg from "./codecs/registry"; import * as buf from "./primitives/buffer"; export const _CodecsRegistry = reg.CodecsRegistry; diff --git a/packages/driver/src/reflection/analyzeQuery.ts b/packages/driver/src/reflection/analyzeQuery.ts index 63cd7a8e1..aa72d75e4 100644 --- a/packages/driver/src/reflection/analyzeQuery.ts +++ b/packages/driver/src/reflection/analyzeQuery.ts @@ -5,7 +5,7 @@ import { EnumCodec } from "../codecs/enum"; import { ICodec, ScalarCodec } from "../codecs/ifaces"; import { NamedTupleCodec } from "../codecs/namedtuple"; import { ObjectCodec } from "../codecs/object"; -import { RangeCodec } from "../codecs/range"; +import { MultiRangeCodec, RangeCodec } from "../codecs/range"; import { NullCodec } from "../codecs/codecs"; import { SetCodec } from "../codecs/set"; import { TupleCodec } from "../codecs/tuple"; @@ -144,5 +144,13 @@ function walkCodec( ctx.imports.add("Range"); return `Range<${subCodec.tsType}>`; } + if (codec instanceof MultiRangeCodec) { + const subCodec = codec.getSubcodecs()[0]; + if (!(subCodec instanceof ScalarCodec)) { + throw Error("expected range subtype to be scalar type"); + } + ctx.imports.add("MultiRange"); + return `MultiRange<${subCodec.tsType}>`; + } throw Error(`Unexpected codec kind: ${codec.getKind()}`); } diff --git a/packages/driver/test/client.test.ts b/packages/driver/test/client.test.ts index 56fc36d7f..f473dd179 100644 --- a/packages/driver/test/client.test.ts +++ b/packages/driver/test/client.test.ts @@ -18,15 +18,15 @@ import fc from "fast-check"; import { parseConnectArguments } from "../src/conUtils.server"; +import type { Client, Executor, _ICodec } from "../src/index.node"; import { - Client, DivisionByZeroError, Duration, EdgeDBError, - Executor, LocalDate, LocalDateTime, Range, + MultiRange, MissingRequiredError, NoDataError, RelativeDuration, @@ -34,7 +34,6 @@ import { QueryArgumentError, _CodecsRegistry, _ReadBuffer, - _ICodec, Session, AuthenticationError, InvalidReferenceError, @@ -43,7 +42,7 @@ import { import { retryingConnect } from "../src/retry"; import { RawConnection } from "../src/rawConn"; import { AdminUIFetchConnection } from "../src/fetchConn"; -import { CustomCodecSpec } from "../src/codecs/registry"; +import type { CustomCodecSpec } from "../src/codecs/registry"; import { getAvailableFeatures, getClient, @@ -1224,32 +1223,32 @@ test("fetch: ConfigMemory", async () => { } }); -if (getEdgeDBVersion().major >= 2) { - function expandRangeEQL(lower: string, upper: string) { - return [ - [false, false], - [true, false], - [false, true], - [true, true], - ] - .map( - ([incl, incu]) => - `range(${lower}, ${upper}, inc_lower := ${incl}, inc_upper := ${incu})` - ) - .join(",\n"); - } +function expandRangeEQL(lower: string, upper: string) { + return [ + [false, false], + [true, false], + [false, true], + [true, true], + ] + .map( + ([incl, incu]) => + `range(${lower}, ${upper}, inc_lower := ${incl}, inc_upper := ${incu})` + ) + .join(",\n"); +} - function expandRangeJS(lower: any, upper: any) { - return [ - new Range(lower, upper, false, false), - new Range(lower, upper, true, false), - new Range(lower, upper, false, true), - new Range(lower, upper, true, true), - ]; - } +function expandRangeJS(lower: any, upper: any) { + return [ + new Range(lower, upper, false, false), + new Range(lower, upper, true, false), + new Range(lower, upper, false, true), + new Range(lower, upper, true, true), + ]; +} +if (getEdgeDBVersion().major >= 2) { test("fetch: ranges", async () => { - const client = await getClient(); + const client = getClient(); try { const res = await client.querySingle(` @@ -1360,6 +1359,18 @@ if (getEdgeDBVersion().major >= 2) { }); } +if (getEdgeDBVersion().major >= 4) { + test("fetch: multirange", async () => { + const client = getClient(); + const multiRangeRes = await client.query( + "select >$mr;", + { mr: new MultiRange([new Range(1, 2)])} + ); + + expect(multiRangeRes).toEqual([new MultiRange([new Range(1, 2)])]); + }); +} + test("fetch: tuple", async () => { const con = getClient(); let res: any;