Merge pull request #2930 from SeedCompany/edgedb/inline-edgeql

src/core/edgedb/edgedb.module.ts

@@ -3,7 +3,8 @@ import { createClient } from 'edgedb';
 import { KNOWN_TYPENAMES } from 'edgedb/dist/codecs/consts.js';
 import { ScalarCodec } from 'edgedb/dist/codecs/ifaces.js';
 import { Class } from 'type-fest';
-import { Client, EdgeDB } from './reexports';
+import { EdgeDB } from './edgedb.service';
+import { Client } from './reexports';
 import { LuxonCalendarDateCodec, LuxonDateTimeCodec } from './temporal.codecs';
 
 @Module({
@@ -21,10 +22,7 @@ import { LuxonCalendarDateCodec, LuxonDateTimeCodec } from './temporal.codecs';
         return client;
       },
     },
-    {
-      provide: EdgeDB,
-      useExisting: Client,
-    },
+    EdgeDB,
   ],
   exports: [EdgeDB, Client],
 })

src/core/edgedb/edgedb.service.ts

@@ -0,0 +1,68 @@
+/* eslint-disable @typescript-eslint/unified-signatures */
+import { Injectable } from '@nestjs/common';
+import { $, Executor } from 'edgedb';
+import { TypedEdgeQL } from './edgeql';
+import { InlineQueryCardinalityMap } from './generated-client/inline-queries';
+import { Client } from './reexports';
+
+@Injectable()
+export class EdgeDB {
+  constructor(private readonly client: Client) {}
+
+  /** Run a query from an edgeql string */
+  run<R>(query: TypedEdgeQL<null, R>): Promise<R>;
+  /** Run a query from an edgeql string */
+  run<Args extends Record<string, any>, R>(
+    query: TypedEdgeQL<Args, R>,
+    args: Args,
+  ): Promise<R>;
+
+  /** Run a query from a edgeql file */
+  run<Args, R>(
+    query: (client: Executor, args: Args) => Promise<R>,
+    args: Args,
+  ): Promise<R>;
+  /** Run a query from a edgeql file */
+  run<R>(query: (client: Executor) => Promise<R>): Promise<R>;
+
+  /** Run a query from the query builder */
+  run<R>(query: { run: (client: Executor) => Promise<R> }): Promise<R>;
+
+  async run(query: any, args?: any) {
+    if (query instanceof TypedEdgeQL) {
+      const cardinality = InlineQueryCardinalityMap.get(query.query);
+      if (!cardinality) {
+        throw new Error(`Query was not found from inline query generation`);
+      }
+      const exeMethod = cardinalityToExecutorMethod[cardinality];
+
+      // eslint-disable-next-line @typescript-eslint/return-await
+      return await this.client[exeMethod](query.query, args);
+    }
+
+    if (query.run) {
+      // eslint-disable-next-line @typescript-eslint/return-await
+      return await query.run(this.client);
+    }
+
+    if (typeof query === 'function') {
+      // eslint-disable-next-line @typescript-eslint/return-await
+      return await query(this.client, args);
+    }
+
+    // For REPL, as this is untyped and assumes many/empty cardinality
+    if (typeof query === 'string') {
+      return await this.client.query(query, args);
+    }
+
+    throw new Error('Could not figure out how to run given query');
+  }
+}
+
+const cardinalityToExecutorMethod = {
+  One: 'queryRequiredSingle',
+  AtMostOne: 'querySingle',
+  Many: 'query',
+  AtLeastOne: 'query',
+  Empty: 'query',
+} satisfies Record<`${$.Cardinality}`, keyof Executor>;

src/core/edgedb/edgeql.ts

@@ -0,0 +1,25 @@
+import { stripIndent } from 'common-tags';
+import { InlineQueryMap as QueryMap } from './generated-client/inline-queries';
+
+export const edgeql = <const Query extends string>(
+  query: Query,
+): Query extends keyof QueryMap ? QueryMap[Query] : unknown => {
+  return new TypedEdgeQL(stripIndent(query)) as any;
+};
+
+export type EdgeQLArgsOf<T extends TypedEdgeQL<any, any>> =
+  T extends TypedEdgeQL<infer Args, any> ? Args : never;
+
+export type EdgeQLReturnOf<T extends TypedEdgeQL<any, any>> =
+  T extends TypedEdgeQL<any, infer Return> ? Return : never;
+
+// Internal symbol to mark a query as typed
+const edgeqlTS = Symbol('edgeqlTS');
+
+export class TypedEdgeQL<Args, Return> {
+  constructor(readonly query: string) {}
+  [edgeqlTS]?: {
+    args: Args;
+    return: Return;
+  };
+}

src/core/edgedb/generate.ts

@@ -10,15 +10,20 @@ import {
   stringifyImports,
 } from '@edgedb/generate/dist/queries.js';
 import { groupBy, mapKeys } from '@seedcompany/common';
+import { stripIndent } from 'common-tags';
 import { $, adapter, Client, createClient } from 'edgedb';
 import { SCALAR_CODECS } from 'edgedb/dist/codecs/codecs.js';
 import { KNOWN_TYPENAMES } from 'edgedb/dist/codecs/consts.js';
+import { Cardinality } from 'edgedb/dist/ifaces.js';
+import { $ as $$ } from 'execa';
 import {
   IndentationText,
+  Node,
   Project,
   QuoteKind,
   SourceFile,
   SyntaxKind,
+  VariableDeclarationKind,
 } from 'ts-morph';
 
 const customScalarList: readonly CustomScalar[] = [
@@ -115,19 +120,152 @@ async function generateAll({
   const schemaFile = project.addSourceFileAtPath(generatedSchemaFile);
   addCustomScalarImports(schemaFile, customScalars.values());
 
-  // Quick/naive fix for the type generated from
+  // Fix the type generated from
   // Project::Resource extending default::Resource
-  schemaFile.replaceWithText(
-    schemaFile
-      .getFullText()
-      .replace(
-        'interface Resource extends Resource',
-        'interface Resource extends DefaultResource',
-      ) + '\ntype DefaultResource = Resource;\n',
-  );
+  schemaFile.addTypeAlias({ name: 'DefaultResource', type: 'Resource' });
+  schemaFile
+    .getModule('Project')!
+    .getInterface('Resource')!
+    .removeExtends(0)
+    .addExtends('DefaultResource');
 
   await generateQueryFiles({ client, root });
   fixCustomScalarImportsInGeneratedEdgeqlFiles(project);
+
+  await generateInlineQueries({ client, project, root });
+}
+
+async function generateInlineQueries({
+  client,
+  project,
+  root,
+}: {
+  client: Client;
+  project: Project;
+  root: string;
+}) {
+  console.log('Generating queries for edgeql() calls...');
+
+  const grepForShortList =
+    await $$`grep -lR edgeql src --exclude-dir=src/core/edgedb`;
+  const shortList = project.addSourceFilesAtPaths(
+    grepForShortList.stdout.split('\n'),
+  );
+
+  const queries =
+    shortList.flatMap((file) =>
+      file.getDescendantsOfKind(SyntaxKind.CallExpression).flatMap((call) => {
+        if (call.getExpression().getText() !== 'edgeql') {
+          return [];
+        }
+        const args = call.getArguments();
+
+        // // 1000x slower to confirm edgeql import
+        // const defs = call
+        //   .getExpressionIfKindOrThrow(SyntaxKind.Identifier)
+        //   .getDefinitionNodes();
+        // if (
+        //   !defs[0].getSourceFile().getFilePath().endsWith('edgedb/edgeql.ts')
+        // ) {
+        //   return [];
+        // }
+
+        if (
+          args.length > 1 ||
+          (!Node.isStringLiteral(args[0]) &&
+            !Node.isNoSubstitutionTemplateLiteral(args[0]))
+        ) {
+          return [];
+        }
+
+        const query = args[0].getText().slice(1, -1);
+        return { query, call };
+      }),
+    ) ?? [];
+
+  const inlineQueriesFile = project.createSourceFile(
+    'src/core/edgedb/generated-client/inline-queries.ts',
+    `import type { TypedEdgeQL } from '../edgeql';`,
+    { overwrite: true },
+  );
+  const queryMap = inlineQueriesFile.addInterface({
+    name: 'InlineQueryMap',
+    isExported: true,
+  });
+
+  const imports = new Set<string>();
+  const seen = new Set<string>();
+  const cardinalityMap = new Map<string, $.Cardinality>();
+  for (const { query, call } of queries) {
+    // Prevent duplicate keys in QueryMap in the off chance that two queries are identical
+    if (seen.has(query)) {
+      continue;
+    }
+    seen.add(query);
+
+    const path = adapter.path.posix.relative(
+      root,
+      call.getSourceFile().getFilePath(),
+    );
+    const lineNumber = call.getStartLineNumber();
+    const source = `./${path}:${lineNumber}`;
+
+    let types;
+    let error;
+    try {
+      types = await $.analyzeQuery(client, query);
+      console.log(`   ${source}`);
+    } catch (err) {
+      error = err as Error;
+      console.log(`Error in query '${source}': ${String(err)}`);
+    }
+
+    if (types) {
+      // Save cardinality for use at runtime.
+      cardinalityMap.set(
+        stripIndent(query),
+        cardinalityMapping[types.cardinality],
+      );
+      // Add imports to the used imports list
+      [...types.imports].forEach((i) => imports.add(i));
+    }
+
+    queryMap.addProperty({
+      name: `[\`${query}\`]`,
+      type: types
+        ? `TypedEdgeQL<${types.args}, ${types.result}>`
+        : error
+        ? `{ ${error.name}: \`${error.message.trim()}\` }`
+        : 'unknown',
+      leadingTrivia:
+        (queryMap.getProperties().length > 0 ? '\n' : '') +
+        `/** {@link import('${path}')} L${lineNumber} */\n`,
+    });
+  }
+
+  addCustomScalarImports(
+    inlineQueriesFile,
+    [...imports].flatMap((i) => customScalars.get(i) ?? []),
+    0,
+  );
+  const builtIn = ['$', ...[...imports].filter((i) => !customScalars.has(i))];
+  inlineQueriesFile.insertImportDeclaration(0, {
+    isTypeOnly: true,
+    namedImports: builtIn,
+    moduleSpecifier: 'edgedb',
+  });
+
+  const cardinalitiesAsStr = JSON.stringify([...cardinalityMap], null, 2);
+  inlineQueriesFile.addVariableStatement({
+    isExported: true,
+    declarationKind: VariableDeclarationKind.Const,
+    declarations: [
+      {
+        name: 'InlineQueryCardinalityMap',
+        initializer: `new Map<string, \`\${$.Cardinality}\`>(${cardinalitiesAsStr})`,
+      },
+    ],
+  });
 }
 
 async function generateQueryFiles({
@@ -151,7 +289,7 @@ async function generateQueryFiles({
         path,
         types,
       });
-      const prettyPath = './' + adapter.path.posix.relative(srcDir, path);
+      const prettyPath = './' + adapter.path.posix.relative(root, path);
       console.log(`   ${prettyPath}`);
       await adapter.fs.writeFile(
         path + '.ts',
@@ -250,9 +388,10 @@ function createTsMorphProject() {
 function addCustomScalarImports(
   file: SourceFile,
   scalars: Iterable<CustomScalar>,
+  index = 2,
 ) {
-  file.insertImportDeclarations(
-    2,
+  return file.insertImportDeclarations(
+    index,
     [...scalars].map((scalar, i) => ({
       isTypeOnly: true,
       namedImports: [scalar.ts],
@@ -261,3 +400,11 @@ function addCustomScalarImports(
     })),
   );
 }
+
+const cardinalityMapping = {
+  [Cardinality.NO_RESULT]: $.Cardinality.Empty,
+  [Cardinality.AT_MOST_ONE]: $.Cardinality.AtMostOne,
+  [Cardinality.ONE]: $.Cardinality.One,
+  [Cardinality.MANY]: $.Cardinality.Many,
+  [Cardinality.AT_LEAST_ONE]: $.Cardinality.AtLeastOne,
+};

src/core/edgedb/index.ts

@@ -1,2 +1,4 @@
 export * from './reexports';
+export { edgeql, EdgeQLArgsOf, EdgeQLReturnOf } from './edgeql';
+export * from './edgedb.service';
 export { default as e } from './generated-client/index.mjs';

src/core/edgedb/reexports.ts

@@ -1,12 +1,3 @@
-import { Transaction } from 'edgedb/dist/transaction.js';
-
 // Only use this if you need the extra methods.
 // Otherwise, for querying, use EdgeDb from below.
 export { Client } from 'edgedb/dist/baseClient.js';
-
-// Using this as it's a runtime symbol for the Executor TS shape
-// which makes it perfect for injection.
-// @ts-expect-error private constructor; doesn't matter that's not how we are using it.
-// We could just export the Transaction as an aliased named, but this also
-// allows REPL to reference it with the correct name.
-export abstract class EdgeDB extends Transaction {}