fix: properly create/end otel spans

Author: andykais

deno.json

@@ -3,8 +3,8 @@
   "version": "1.9.6",
 
   "tasks": {
-    "test":       "deno test --check --import-map=test/import_map.json --allow-read --allow-write --allow-ffi --allow-env=HOME,DENO_DIR,XDG_CACHE_HOME,DENO_SQLITE_PATH,DENO_SQLITE_LOCAL --allow-net test",
-    "test:watch": "deno test --check --import-map=test/import_map.json --allow-read --allow-write --allow-ffi --allow-env=HOME,DENO_DIR,XDG_CACHE_HOME,DENO_SQLITE_PATH,DENO_SQLITE_LOCAL --allow-net --watch"
+    "test":       "deno test --check --import-map=test/import_map.json --allow-read --allow-write test",
+    "test:watch": "deno test --check --import-map=test/import_map.json --allow-read --allow-write --watch"
   },
 
   "lint": {

src/statement.ts

@@ -89,9 +89,9 @@ abstract class StatementBase<DriverStatement, Params extends SchemaGeneric, Resu
     else throw new Error('A statement cannot be used until init() is called')
   }
 
-  protected encode_params = (params: Params | undefined): EncodedParams => {
-    // SchemaGeneric is not assignable to an attribute, but I think we can trust the opentelemetry serialization. If not, we could switch over to the driver encoded params, which we know are definitely serialization friendly
-    const span = telemetry.start_span('statement.encode_params', {params: params as any})
+  @telemetry.instrument()
+  protected encode_params(params: Params | undefined): EncodedParams {
+    telemetry.attribute('params', params as any) // SchemaGeneric is not assignable to an attribute, but I think we can trust the opentelemetry serialization. If not, we could switch over to the driver encoded params, which we know are definitely serialization friendly
     const encoded_params: {[field: string]: any} = {}
     for (const field of Object.values(this.params)) {
       const val = params ? params[field.field_name] : undefined
@@ -102,18 +102,15 @@ abstract class StatementBase<DriverStatement, Params extends SchemaGeneric, Resu
           const message = e.format()._errors.join(',')
           throw new TypeError(`${message} on column ${field.table_name}.${field.field_name}:\n${Deno.inspect(params, { colors: true })}`)
         } else {
-          span.recordException(e as Error)
-          span.end()
           throw e
         }
       }
     }
-    span.end()
     return encoded_params
   }
 
-  protected decode_result = (result: RawRowData): Result => {
-    const span = telemetry.start_span('statement.decode_result')
+  @telemetry.instrument()
+  protected decode_result(result: RawRowData): Result {
     const decoded_result: {[field: string]: any} = {}
     for (const [field_name, val] of Object.entries(result)) {
       const field = this.get_result_field(field_name)
@@ -126,8 +123,6 @@ abstract class StatementBase<DriverStatement, Params extends SchemaGeneric, Resu
         //   decoded_result[field_name] = val
         // }
       } catch (e) {
-        span.recordException(e as Error)
-        span.end()
         if (e instanceof z.ZodError) {
           const message = e.format()._errors.join(',')
           throw new TypeError(`${message} on column ${field.table_name}.${field.field_name}:\n${Deno.inspect(result, { colors: true })}`)
@@ -137,67 +132,45 @@ abstract class StatementBase<DriverStatement, Params extends SchemaGeneric, Resu
       }
     }
     // I dont know how to convert these partial types into full types, so we just cast here
-    span.end()
     return decoded_result as Result
   }
 
-  public one = (...params: OptionalOnEmpty<Params>): Result | undefined => {
-    const span = telemetry.start_span('statement.one', {sql: this.sql})
-    try {
-      return this.one_impl(...params)
-    } catch (e) {
-      span.recordException(e as Error)
-      throw e
-    } finally {
-      span.end()
-    }
+  @telemetry.instrument()
+  public one(...params: OptionalOnEmpty<Params>): Result | undefined {
+    telemetry.attribute('sql', this.sql)
+    return this.one_impl(...params)
   }
   protected abstract one_impl(...params: OptionalOnEmpty<Params>): Result | undefined
 
-  public all = (...params: OptionalOnEmpty<Params>): Result[] => {
-    const span = telemetry.start_span('statement.all', {sql: this.sql})
-    span.setAttribute('params', params as any) // SchemaGeneric is not assignable to an attribute, but I think we can trust the opentelemetry serialization. If not, we could switch over to the driver encoded params, which we know are definitely serialization friendly
-    try {
-      return this.all_impl(...params)
-    } catch (e) {
-      span.recordException(e as Error)
-      throw e
-    } finally {
-      span.end()
-    }
+  @telemetry.instrument()
+  public all(...params: OptionalOnEmpty<Params>): Result[] {
+    telemetry.attribute('sql', this.sql)
+    return this.all_impl(...params)
   }
   protected abstract all_impl(...params: OptionalOnEmpty<Params>): Result[]
 
-  public exec = (...params: OptionalOnEmpty<Params>): ExecInfo => {
-    const span = telemetry.start_span('statement.exec', {sql: this.sql})
-    span.setAttribute('params', params as any) // SchemaGeneric is not assignable to an attribute, but I think we can trust the opentelemetry serialization. If not, we could switch over to the driver encoded params, which we know are definitely serialization friendly
-    try {
-      return this.exec_impl(...params)
-    } catch (e) {
-      span.recordException(e as Error)
-      throw e
-    } finally {
-      span.end()
-    }
+  @telemetry.instrument()
+  public exec(...params: OptionalOnEmpty<Params>): ExecInfo {
+    telemetry.attribute('sql', this.sql)
+    telemetry.attribute('params', params as any) // SchemaGeneric is not assignable to an attribute, but I think we can trust the opentelemetry serialization. If not, we could switch over to the driver encoded params, which we know are definitely serialization friendly
+    return this.exec_impl(...params)
   }
   protected abstract exec_impl(...params: OptionalOnEmpty<Params>): ExecInfo
 
-  public prepare = (driver: Driver): void => {
-    const span = telemetry.start_span('statement.exec', {sql: this.sql})
+  @telemetry.instrument()
+  public prepare(driver: Driver): void {
+    telemetry.attribute('sql', this.sql)
     this._driver = driver
     try {
       this._stmt = this.prepare_impl(this.sql)
     } catch (e) {
-      span.recordException(e as Error)
       if (e instanceof Error === false) throw e
       throw new Error(`${e.message}
 ${'```'}
 ${this.sql}
 ${'```'}`, {
   cause: e
 })
-    } finally {
-      span.end()
     }
   }
 

src/telemetry.ts

@@ -2,9 +2,46 @@ import * as otel from '@opentelemetry/api'
 
 const TRACE_NAME = '@torm/sqlite'
 
-const trace = otel.trace.getTracer(TRACE_NAME)
+const tracer = otel.trace.getTracer(TRACE_NAME)
 
-export function start_span(name: string, attributes?: otel.Attributes) {
-  const span = trace.startSpan(name, {attributes})
-  return span
+export function attribute(name: string, value: any) {
+  otel.trace.getActiveSpan()?.setAttribute(name, value)
+}
+
+export function instrument(span_name?: string) {
+  return function(original_method: any, context: ClassMethodDecoratorContext) {
+    const method_name = String(context.name);
+
+    function replacement_method(...args: any[]) {
+      // @ts-ignore
+      const thisClass = this as any
+      const final_span_name = span_name || `${thisClass.constructor.name}.${method_name}`;
+
+      return tracer.startActiveSpan(final_span_name, (span: otel.Span) => {
+        try {
+          span.setAttributes({
+            'class.name': thisClass.constructor.name,
+            'method.name': method_name,
+          });
+
+          // execute the original method
+          // note that currentl we only support synchronous methods. In the future, we can easily add promise support here, but since our only driver is sqlite, this is fine for now
+          return original_method.call(thisClass, ...args);
+        } catch (error) {
+          span.recordException(error as Error);
+          throw error;
+        } finally {
+          span.end()
+        }
+      });
+    }
+
+    context.addInitializer(function (this) {
+      const thisClass = this as any
+      // also bind any method we set up this way
+      thisClass[method_name] = thisClass[method_name].bind(thisClass)
+    });
+
+    return replacement_method;
+  };
 }