From 6e8989de4d60c32bc9b3e3ef760e92c68ae7f491 Mon Sep 17 00:00:00 2001
From: Ryan Eberhardt <reberhardt7@gmail.com>
Date: Tue, 23 Jul 2024 07:25:58 -0700
Subject: [PATCH] feat(instr-knex): implement requireParentSpan config flag
 (#2288)

Co-authored-by: Amir Blum <amirgiraffe@gmail.com>
---
 .../README.md                                 |  1 +
 .../src/instrumentation.ts                    | 15 +++-
 .../src/types.ts                              |  2 +
 .../test/index.test.ts                        | 74 ++++++++++++++++++-
 4 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-knex/README.md b/plugins/node/opentelemetry-instrumentation-knex/README.md
index 61404aa594..e467172648 100644
--- a/plugins/node/opentelemetry-instrumentation-knex/README.md
+++ b/plugins/node/opentelemetry-instrumentation-knex/README.md
@@ -47,6 +47,7 @@ registerInstrumentations({
 | Options | Type | Example | Description |
 | ------- | ---- | ------- | ----------- |
 | `maxQueryLength` | `number` | `100` | Truncate `db.statement` attribute to a maximum length. If the statement is truncated `'..'` is added to it's end. Default `1022`. `-1` leaves `db.statement` untouched. |
+| `requireParentSpan` | `boolean` | `false` | Don't create spans unless they are part of an existing trace. Default is `false`. |
 
 ## Semantic Conventions
 
diff --git a/plugins/node/opentelemetry-instrumentation-knex/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-knex/src/instrumentation.ts
index 5dbe4d3035..191bf1522d 100644
--- a/plugins/node/opentelemetry-instrumentation-knex/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-knex/src/instrumentation.ts
@@ -40,6 +40,7 @@ import { KnexInstrumentationConfig } from './types';
 const contextSymbol = Symbol('opentelemetry.instrumentation-knex.context');
 const DEFAULT_CONFIG: KnexInstrumentationConfig = {
   maxQueryLength: 1022,
+  requireParentSpan: false,
 };
 
 export class KnexInstrumentation extends InstrumentationBase<KnexInstrumentationConfig> {
@@ -120,7 +121,7 @@ export class KnexInstrumentation extends InstrumentationBase<KnexInstrumentation
 
   private createQueryWrapper(moduleVersion?: string) {
     const instrumentation = this;
-    return function wrapQuery(original: () => any) {
+    return function wrapQuery(original: (...args: any[]) => any) {
       return function wrapped_logging_method(this: any, query: any) {
         const config = this.client.config;
 
@@ -152,14 +153,22 @@ export class KnexInstrumentation extends InstrumentationBase<KnexInstrumentation
           );
         }
 
-        const parent = this.builder[contextSymbol];
+        const parentContext =
+          this.builder[contextSymbol] || api.context.active();
+        const parentSpan = api.trace.getSpan(parentContext);
+        const hasActiveParent =
+          parentSpan && api.trace.isSpanContextValid(parentSpan.spanContext());
+        if (instrumentation._config.requireParentSpan && !hasActiveParent) {
+          return original.bind(this)(...arguments);
+        }
+
         const span = instrumentation.tracer.startSpan(
           utils.getName(name, operation, table),
           {
             kind: api.SpanKind.CLIENT,
             attributes,
           },
-          parent
+          parentContext
         );
         const spanContext = api.trace.setSpan(api.context.active(), span);
 
diff --git a/plugins/node/opentelemetry-instrumentation-knex/src/types.ts b/plugins/node/opentelemetry-instrumentation-knex/src/types.ts
index a952a66e3d..4d0e8495f7 100644
--- a/plugins/node/opentelemetry-instrumentation-knex/src/types.ts
+++ b/plugins/node/opentelemetry-instrumentation-knex/src/types.ts
@@ -18,4 +18,6 @@ import { InstrumentationConfig } from '@opentelemetry/instrumentation';
 export interface KnexInstrumentationConfig extends InstrumentationConfig {
   /** max query length in db.statement attribute ".." is added to the end when query is truncated  */
   maxQueryLength?: number;
+  /** only create spans if part of an existing trace */
+  requireParentSpan?: boolean;
 }
diff --git a/plugins/node/opentelemetry-instrumentation-knex/test/index.test.ts b/plugins/node/opentelemetry-instrumentation-knex/test/index.test.ts
index 344fb94400..e84dcbc398 100644
--- a/plugins/node/opentelemetry-instrumentation-knex/test/index.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-knex/test/index.test.ts
@@ -14,7 +14,12 @@
  * limitations under the License.
  */
 
-import { SpanKind, context, trace } from '@opentelemetry/api';
+import {
+  INVALID_SPAN_CONTEXT,
+  SpanKind,
+  context,
+  trace,
+} from '@opentelemetry/api';
 import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
 import { AsyncHooksContextManager } from '@opentelemetry/context-async-hooks';
 import {
@@ -452,6 +457,73 @@ describe('Knex instrumentation', () => {
       );
     });
   });
+
+  describe('Setting requireParentSpan=true', () => {
+    beforeEach(() => {
+      plugin.disable();
+      plugin.setConfig({ requireParentSpan: true });
+      plugin.enable();
+    });
+
+    it('should not create new spans when there is no parent', async () => {
+      await client.schema.createTable('testTable1', (table: any) => {
+        table.string('title');
+      });
+      assert.deepEqual(await client('testTable1').select('*'), []);
+      assert.deepEqual(await client.raw('select 1 as result'), [{ result: 1 }]);
+      assert.strictEqual(memoryExporter.getFinishedSpans().length, 0);
+    });
+
+    it('should not create new spans when there is an INVALID_SPAN_CONTEXT parent', async () => {
+      const parentSpan = trace.wrapSpanContext(INVALID_SPAN_CONTEXT);
+      await context.with(
+        trace.setSpan(context.active(), parentSpan),
+        async () => {
+          await client.schema.createTable('testTable1', (table: any) => {
+            table.string('title');
+          });
+          assert.deepEqual(await client('testTable1').select('*'), []);
+          assert.deepEqual(await client.raw('select 1 as result'), [
+            { result: 1 },
+          ]);
+        }
+      );
+      assert.strictEqual(memoryExporter.getFinishedSpans().length, 0);
+    });
+
+    it('should create new spans when there is a parent', async () => {
+      await tracer.startActiveSpan('parentSpan', async parentSpan => {
+        await client.schema.createTable('testTable1', (table: any) => {
+          table.string('title');
+        });
+        assert.deepEqual(await client('testTable1').select('*'), []);
+        assert.deepEqual(await client.raw('select 1 as result'), [
+          { result: 1 },
+        ]);
+        parentSpan.end();
+      });
+      assert.strictEqual(memoryExporter.getFinishedSpans().length, 4);
+      const instrumentationSpans = memoryExporter.getFinishedSpans();
+      const last = instrumentationSpans.pop() as any;
+      assertSpans(instrumentationSpans, [
+        {
+          statement: 'create table `testTable1` (`title` varchar(255))',
+          parentSpan: last,
+        },
+        {
+          op: 'select',
+          table: 'testTable1',
+          statement: 'select * from `testTable1`',
+          parentSpan: last,
+        },
+        {
+          op: 'raw',
+          statement: 'select 1 as result',
+          parentSpan: last,
+        },
+      ]);
+    });
+  });
 });
 
 const assertSpans = (actualSpans: any[], expectedSpans: any[]) => {