Add tracing to EdgeDB queries

Author: CarsonF

src/core/edgedb/edgedb.service.ts

@@ -3,6 +3,7 @@ import { Injectable, Optional } from '@nestjs/common';
 import { $, ConstraintViolationError, EdgeDBError, Executor } from 'edgedb';
 import { QueryArgs } from 'edgedb/dist/ifaces';
 import { retry, RetryOptions } from '~/common/retry';
+import { TracingService } from '~/core/tracing';
 import { jestSkipFileInExceptionSource } from '../exception';
 import { TypedEdgeQL } from './edgeql';
 import { enhanceConstraintError } from './errors';
@@ -17,6 +18,7 @@ export class EdgeDB {
     private readonly client: Client,
     private readonly transactionContext: TransactionContext,
     private readonly optionsContext: OptionsContext,
+    private readonly tracing: TracingService,
     @Optional() private readonly childOptions: ApplyOptions[] = [],
     @Optional() private childExecutor?: Executor,
   ) {}
@@ -26,6 +28,7 @@ export class EdgeDB {
       this.client,
       this.transactionContext,
       this.optionsContext,
+      this.tracing,
       [...this.childOptions],
       this.childExecutor,
     );
@@ -104,55 +107,62 @@ export class EdgeDB {
 
   async run(query: any, args?: any) {
     const executor = this.childExecutor ?? this.transactionContext.current;
-    try {
-      if (query instanceof TypedEdgeQL) {
-        const found = InlineQueryRuntimeMap.get(query.query);
-        if (!found) {
-          throw new Error(`Query was not found from inline query generation`);
+    return await this.tracing.capture('EdgeDB Query', async (segment) => {
+      // Show this segment separately in the service map
+      segment.namespace = 'remote';
+      // Help ID the segment as being for a database
+      segment.sql = {};
+
+      try {
+        if (query instanceof TypedEdgeQL) {
+          const found = InlineQueryRuntimeMap.get(query.query);
+          if (!found) {
+            throw new Error(`Query was not found from inline query generation`);
+          }
+          const exeMethod = cardinalityToExecutorMethod[found.cardinality];
+
+          // eslint-disable-next-line @typescript-eslint/return-await
+          return await executor[exeMethod](found.query, args);
         }
-        const exeMethod = cardinalityToExecutorMethod[found.cardinality];
 
-        // eslint-disable-next-line @typescript-eslint/return-await
-        return await executor[exeMethod](found.query, args);
-      }
-
-      if (query.run) {
-        // eslint-disable-next-line @typescript-eslint/return-await
-        return await query.run(executor, args);
-      }
+        if (query.run) {
+          // eslint-disable-next-line @typescript-eslint/return-await
+          return await query.run(executor, args);
+        }
 
-      if (typeof query === 'function') {
-        // eslint-disable-next-line @typescript-eslint/return-await
-        return await query(executor, args);
-      }
+        if (typeof query === 'function') {
+          // eslint-disable-next-line @typescript-eslint/return-await
+          return await query(executor, args);
+        }
 
-      // For REPL, as this is untyped and assumes many/empty cardinality
-      if (typeof query === 'string') {
-        return await executor.query(query, args);
-      }
-    } catch (e) {
-      // Ignore this call in stack trace. This puts the actual query as the first.
-      e.stack = e.stack!.replace(/^\s+at(?: async)? EdgeDB\.run.+$\n/m, '');
-
-      // Don't present abstract repositories as the src block in jest reports
-      // for DB execution errors.
-      // There shouldn't be anything specific to there to be helpful.
-      // This is a bit of a broad assumption though, so only do for jest and
-      // keep the frame for actual use from users/devs.
-      if (e instanceof EdgeDBError) {
-        jestSkipFileInExceptionSource(
-          e,
-          /^\s+at .+src[/|\\]core[/|\\]edgedb[/|\\].+\.repository\..+$\n/gm,
-        );
-      }
+        // For REPL, as this is untyped and assumes many/empty cardinality
+        if (typeof query === 'string') {
+          return await executor.query(query, args);
+        }
+      } catch (e) {
+        // Ignore this call in stack trace. This puts the actual query as the first.
+        e.stack = e.stack!.replace(/^\s+at(?: async)? EdgeDB\.run.+$\n/m, '');
+
+        // Don't present abstract repositories as the src block in jest reports
+        // for DB execution errors.
+        // There shouldn't be anything specific to there to be helpful.
+        // This is a bit of a broad assumption though, so only do for jest and
+        // keep the frame for actual use from users/devs.
+        if (e instanceof EdgeDBError) {
+          jestSkipFileInExceptionSource(
+            e,
+            /^\s+at .+src[/|\\]core[/|\\]edgedb[/|\\].+\.repository\..+$\n/gm,
+          );
+        }
 
-      if (e instanceof ConstraintViolationError) {
-        throw enhanceConstraintError(e);
+        if (e instanceof ConstraintViolationError) {
+          throw enhanceConstraintError(e);
+        }
+        throw e;
       }
-      throw e;
-    }
 
-    throw new Error('Could not figure out how to run given query');
+      throw new Error('Could not figure out how to run given query');
+    });
   }
 }