Move rawParse/Execute methods to AdminFetchConn + add ability to prov…

…ide AbortSignal
edgedb · Sep 14, 2023 · e2ae37b · e2ae37b
1 parent 95eeee7
commit e2ae37b
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 117 deletions.
diff --git a/packages/driver/src/baseConn.ts b/packages/driver/src/baseConn.ts
@@ -78,19 +78,13 @@ new DataView(NO_TRANSACTION_CAPABILITIES_BYTES.buffer).setUint32(
   NO_TRANSACTION_CAPABILITIES
 );
 
-const RESTRICTED_CAPABILITIES =
+export const RESTRICTED_CAPABILITIES =
   (Capabilities.ALL &
     ~Capabilities.TRANSACTION &
     ~Capabilities.SESSION_CONFIG &
     ~Capabilities.SET_GLOBAL) >>>
   0;
 
-const STUDIO_CAPABILITIES =
-  (RESTRICTED_CAPABILITIES |
-    Capabilities.SESSION_CONFIG |
-    Capabilities.SET_GLOBAL) >>>
-  0;
-
 enum CompilationFlag {
   INJECT_OUTPUT_TYPE_IDS = 1 << 0,
   INJECT_OUTPUT_TYPE_NAMES = 1 << 1,
@@ -1038,7 +1032,7 @@ export class BaseRawConnection {
     ];
   }
 
-  private async _executeFlow(
+  protected async _executeFlow(
     query: string,
     args: QueryArgs,
     outputFormat: OutputFormat,
@@ -1447,54 +1441,4 @@ export class BaseRawConnection {
   async close(): Promise<void> {
     this._abort();
   }
-
-  // These methods are exposed for use by EdgeDB Studio
-  public async rawParse(
-    query: string,
-    state: Session,
-    options?: QueryOptions
-  ): Promise<
-    [ICodec, ICodec, Uint8Array, Uint8Array, ProtocolVersion, number]
-  > {
-    const result = (await this._parse(
-      query,
-      OutputFormat.BINARY,
-      Cardinality.MANY,
-      state,
-      STUDIO_CAPABILITIES,
-      options
-    ))!;
-    return [
-      result[1],
-      result[2],
-      result[4]!,
-      result[5]!,
-      this.protocolVersion,
-      result[3],
-    ];
-  }
-
-  public async rawExecute(
-    query: string,
-    state: Session,
-    outCodec?: ICodec,
-    options?: QueryOptions,
-    inCodec?: ICodec,
-    args: QueryArgs = null
-  ): Promise<Uint8Array> {
-    const result = new WriteBuffer();
-    await this._executeFlow(
-      query,
-      args,
-      outCodec ? OutputFormat.BINARY : OutputFormat.NONE,
-      Cardinality.MANY,
-      state,
-      inCodec ?? NULL_CODEC,
-      outCodec ?? NULL_CODEC,
-      result,
-      STUDIO_CAPABILITIES,
-      options
-    );
-    return result.unwrap();
-  }
 }
diff --git a/packages/driver/src/fetchConn.ts b/packages/driver/src/fetchConn.ts
@@ -16,13 +16,29 @@
  * limitations under the License.
  */
 
+import {
+  BaseRawConnection,
+  Capabilities,
+  PROTO_VER,
+  RESTRICTED_CAPABILITIES,
+} from "./baseConn";
+import { NULL_CODEC } from "./codecs/codecs";
+import { ICodec } from "./codecs/ifaces";
 import { CodecsRegistry } from "./codecs/registry";
 import { Address, NormalizedConnectConfig } from "./conUtils";
-import { PROTO_VER, BaseRawConnection } from "./baseConn";
-import Event from "./primitives/event";
-import * as chars from "./primitives/chars";
 import { InternalClientError, ProtocolError } from "./errors";
 import type { HttpSCRAMAuth } from "./httpScram";
+import {
+  Cardinality,
+  OutputFormat,
+  ProtocolVersion,
+  QueryArgs,
+  QueryOptions,
+} from "./ifaces";
+import { Session } from "./options";
+import { WriteBuffer } from "./primitives/buffer";
+import * as chars from "./primitives/chars";
+import Event from "./primitives/event";
 
 interface FetchConfig {
   address: Address | string;
@@ -34,9 +50,16 @@ interface FetchConfig {
 
 const PROTO_MIME = `application/x.edgedb.v_${PROTO_VER[0]}_${PROTO_VER[1]}.binary'`;
 
+const STUDIO_CAPABILITIES =
+  (RESTRICTED_CAPABILITIES |
+    Capabilities.SESSION_CONFIG |
+    Capabilities.SET_GLOBAL) >>>
+  0;
+
 class BaseFetchConnection extends BaseRawConnection {
   protected config: FetchConfig;
   protected addr: string;
+  protected abortSignal: AbortSignal | null = null;
 
   constructor(config: FetchConfig, registry: CodecsRegistry) {
     super(registry);
@@ -96,6 +119,7 @@ class BaseFetchConnection extends BaseRawConnection {
         method: "post",
         body: data,
         headers,
+        signal: this.abortSignal,
       });
 
       if (!resp.ok) {
@@ -129,16 +153,17 @@ class BaseFetchConnection extends BaseRawConnection {
     this.__sendData(data);
   }
 
-  static create(
+  static create<T extends typeof BaseFetchConnection>(
+    this: T,
     config: FetchConfig,
     registry: CodecsRegistry
-  ): BaseFetchConnection {
+  ): InstanceType<T> {
     const conn = new this(config, registry);
 
     conn.connected = true;
     conn.connWaiter.set();
 
-    return conn;
+    return conn as InstanceType<T>;
   }
 }
 
@@ -158,6 +183,62 @@ export class AdminUIFetchConnection extends BaseFetchConnection {
     const baseUrl = `${protocol}://${address[0]}:${address[1]}`;
     return `${baseUrl}/db/${database}`;
   }
+
+  // These methods are exposed for use by EdgeDB Studio
+  public async rawParse(
+    query: string,
+    state: Session,
+    options?: QueryOptions,
+    abortSignal?: AbortSignal | null
+  ): Promise<
+    [ICodec, ICodec, Uint8Array, Uint8Array, ProtocolVersion, number]
+  > {
+    this.abortSignal = abortSignal ?? null;
+
+    const result = (await this._parse(
+      query,
+      OutputFormat.BINARY,
+      Cardinality.MANY,
+      state,
+      STUDIO_CAPABILITIES,
+      options
+    ))!;
+    return [
+      result[1],
+      result[2],
+      result[4]!,
+      result[5]!,
+      this.protocolVersion,
+      result[3],
+    ];
+  }
+
+  public async rawExecute(
+    query: string,
+    state: Session,
+    outCodec?: ICodec,
+    options?: QueryOptions,
+    inCodec?: ICodec,
+    args: QueryArgs = null,
+    abortSignal?: AbortSignal | null
+  ): Promise<Uint8Array> {
+    this.abortSignal = abortSignal ?? null;
+
+    const result = new WriteBuffer();
+    await this._executeFlow(
+      query,
+      args,
+      outCodec ? OutputFormat.BINARY : OutputFormat.NONE,
+      Cardinality.MANY,
+      state,
+      inCodec ?? NULL_CODEC,
+      outCodec ?? NULL_CODEC,
+      result,
+      STUDIO_CAPABILITIES,
+      options
+    );
+    return result.unwrap();
+  }
 }
 
 const _tokens = new WeakMap<NormalizedConnectConfig, string>();

diff --git a/packages/driver/test/client.test.ts b/packages/driver/test/client.test.ts
@@ -40,8 +40,6 @@ import {
   InvalidReferenceError,
 } from "../src/index.node";
 
-import { retryingConnect } from "../src/retry";
-import { RawConnection } from "../src/rawConn";
 import { AdminUIFetchConnection } from "../src/fetchConn";
 import { CustomCodecSpec } from "../src/codecs/registry";
 import {
@@ -52,6 +50,7 @@ import {
   isDeno,
 } from "./testbase";
 import { PG_VECTOR_MAX_DIM } from "../src/codecs/pgvector";
+import { HTTPSCRAMAuth } from "../src/httpScram";
 
 function setCustomCodecs(codecs: (keyof CustomCodecSpec)[], client: Client) {
   // @ts-ignore
@@ -2041,53 +2040,24 @@ function _decodeResultBuffer(outCodec: _ICodec, resultData: Uint8Array) {
   return result;
 }
 
-// 'raw' methods are only used by edgedb studio, so only need to work with
-// EdgeDB 2.0 or greater
-if (getEdgeDBVersion().major >= 2) {
-  test("'implicit*' headers", async () => {
-    const config = await parseConnectArguments(getConnectOptions());
-    const registry = new _CodecsRegistry();
-    const con = await retryingConnect(
-      RawConnection.connectWithTimeout.bind(RawConnection),
-      config,
-      registry
-    );
-    try {
-      const query = `SELECT Function {
-        name
-      }`;
-      const state = new Session({ module: "schema" });
-      const options = {
-        injectTypenames: true,
-        implicitLimit: BigInt(5),
-      } as const;
-      const [_, outCodec] = await con.rawParse(query, state, options);
-      const resultData = await con.rawExecute(query, state, outCodec, options);
-
-      const result = _decodeResultBuffer(outCodec, resultData);
-
-      expect(result).toHaveLength(5);
-      expect(result[0]["__tname__"]).toBe("schema::Function");
-    } finally {
-      await con.close();
-    }
-  });
-}
-
 if (!isDeno && getAvailableFeatures().has("binary-over-http")) {
-  // @ts-ignore // make deno ignore skip
-  test.skip("binary protocol over http", async () => {
-    //@ts-ignore
-    const tokenFile = require("path").join(__dirname, "keys", "jwt");
-    //@ts-ignore
-    const token = require("fs").readFileSync(tokenFile, "utf8").trim();
+  test("binary protocol over http", async () => {
     const codecsRegistry = new _CodecsRegistry();
     const config = await parseConnectArguments(getConnectOptions());
+
+    const { address, user, password } = config.connectionParams;
+    const token = await HTTPSCRAMAuth(
+      `http://${address[0]}:${address[1]}`,
+      user,
+      password!
+    );
+
     const fetchConn = AdminUIFetchConnection.create(
       {
         address: config.connectionParams.address,
         database: config.connectionParams.database,
         user: config.connectionParams.user,
+        tlsSecurity: "insecure",
         token: token,
       },
       codecsRegistry

diff --git a/packages/driver/test/keys/jwt b/packages/driver/test/keys/jwt
diff --git a/packages/driver/test/keys/private.pem b/packages/driver/test/keys/private.pem
diff --git a/packages/driver/test/keys/public.pem b/packages/driver/test/keys/public.pem
diff --git a/packages/driver/test/testUtil.ts b/packages/driver/test/testUtil.ts
@@ -113,8 +113,6 @@ export const getServerCommand = (
 
   if (help.includes("--admin-ui")) {
     args.push("--http-endpoint-security=optional");
-    // args.push("--jws-key-file", path.join(__dirname, "keys", "public.pem"));
-    // args.push("--jwe-key-file", path.join(__dirname, "keys", "private.pem"));
     args.push("--jose-key-mode=generate");
 
     availableFeatures.push("binary-over-http");