Break down analyzeQuery for more generation flexibility (#885)

edgedb · Mar 4, 2024 · 20fe059 · 20fe059
1 parent ebf6351
commit 20fe059
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 29 deletions.
diff --git a/packages/driver/src/baseConn.ts b/packages/driver/src/baseConn.ts
@@ -97,12 +97,12 @@ const OLD_ERROR_CODES = new Map([
 ]);
 
 export type ParseResult = [
-  Cardinality,
-  ICodec,
-  ICodec,
-  number,
-  Uint8Array | null,
-  Uint8Array | null
+  cardinality: Cardinality,
+  inCodec: ICodec,
+  outCodec: ICodec,
+  capabilities: number,
+  inCodecBuffer: Uint8Array | null,
+  outCodecBuffer: Uint8Array | null
 ];
 
 export type connConstructor = new (

diff --git a/packages/driver/src/reflection/analyzeQuery.ts b/packages/driver/src/reflection/analyzeQuery.ts
@@ -1,8 +1,8 @@
-import type { ParseResult } from "../baseConn";
 import { ArrayCodec } from "../codecs/array";
 import { AT_LEAST_ONE, AT_MOST_ONE, MANY, ONE } from "../codecs/consts";
 import { EnumCodec } from "../codecs/enum";
-import { ICodec, ScalarCodec } from "../codecs/ifaces";
+import type { ICodec } from "../codecs/ifaces";
+import { ScalarCodec } from "../codecs/ifaces";
 import { NamedTupleCodec } from "../codecs/namedtuple";
 import { ObjectCodec } from "../codecs/object";
 import { MultiRangeCodec, RangeCodec } from "../codecs/range";
@@ -25,33 +25,16 @@ export async function analyzeQuery(
   client: Client,
   query: string
 ): Promise<QueryType> {
-  let parseResult: ParseResult;
-  const pool: BaseClientPool = (client as any).pool;
+  const [cardinality, inCodec, outCodec] = await parseQuery(client, query);
 
-  const holder = await pool.acquireHolder(Options.defaults());
-  try {
-    const cxn = await holder._getConnection();
-    parseResult = await cxn._parse(
-      query,
-      OutputFormat.BINARY,
-      Cardinality.MANY,
-      Session.defaults()
-    );
-  } finally {
-    await holder.release();
-  }
-
-  const cardinality = parseResult[0];
-  const inCodec = parseResult[1];
-  const outCodec = parseResult[2];
   const imports = new Set<string>();
   const args = walkCodec(inCodec, {
     indent: "",
     optionalNulls: true,
     imports,
   });
 
-  const result = generateSetType(
+  const result = applyCardinalityToTsType(
     walkCodec(outCodec, {
       indent: "",
       optionalNulls: false,
@@ -69,7 +52,27 @@ export async function analyzeQuery(
   };
 }
 
-function generateSetType(type: string, cardinality: Cardinality): string {
+export async function parseQuery(client: Client, query: string) {
+  const pool: BaseClientPool = (client as any).pool;
+
+  const holder = await pool.acquireHolder(Options.defaults());
+  try {
+    const cxn = await holder._getConnection();
+    return await cxn._parse(
+      query,
+      OutputFormat.BINARY,
+      Cardinality.MANY,
+      Session.defaults()
+    );
+  } finally {
+    await holder.release();
+  }
+}
+
+export function applyCardinalityToTsType(
+  type: string,
+  cardinality: Cardinality
+): string {
   switch (cardinality) {
     case Cardinality.MANY:
       return `${type}[]`;
@@ -85,6 +88,7 @@ function generateSetType(type: string, cardinality: Cardinality): string {
 
 // type AtLeastOne<T> = [T, ...T[]];
 
+export { walkCodec as walkCodecToTsType };
 function walkCodec(
   codec: ICodec,
   ctx: { indent: string; optionalNulls: boolean; imports: Set<string> }
@@ -120,7 +124,7 @@ function walkCodec(
         }
         return `${ctx.indent}  ${JSON.stringify(field.name)}${
           ctx.optionalNulls && field.cardinality === AT_MOST_ONE ? "?" : ""
-        }: ${generateSetType(
+        }: ${applyCardinalityToTsType(
           walkCodec(subCodec, { ...ctx, indent: ctx.indent + "  " }),
           field.cardinality
         )};`;