From 88db0a4ea134593ba5bf2c7209f94eae466519bd Mon Sep 17 00:00:00 2001
From: Matthew Wear <matthew.wear@gmail.com>
Date: Fri, 14 Jan 2022 01:38:20 -0800
Subject: [PATCH] fix: allow mysql streamable query with values array (#790)

Co-authored-by: Rauno Viskus <Rauno56@users.noreply.github.com>
---
 .../src/instrumentation.ts                    | 20 +++---
 .../test/index.test.ts                        | 69 +++++++++++++++++++
 2 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts
index 8a0652d503..00968d7dce 100644
--- a/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-mysql/src/instrumentation.ts
@@ -279,7 +279,11 @@ export class MySQLInstrumentation extends InstrumentationBase<
           getDbStatement(query, format, values)
         );
 
-        if (arguments.length === 1) {
+        const cbIndex = Array.from(arguments).findIndex(
+          arg => typeof arg === 'function'
+        );
+
+        if (cbIndex === -1) {
           const streamableQuery: mysqlTypes.Query = originalQuery.apply(
             connection,
             arguments
@@ -295,15 +299,15 @@ export class MySQLInstrumentation extends InstrumentationBase<
             .on('end', () => {
               span.end();
             });
-        }
+        } else {
+          thisPlugin._wrap(
+            arguments,
+            cbIndex,
+            thisPlugin._patchCallbackQuery(span)
+          );
 
-        if (typeof arguments[1] === 'function') {
-          thisPlugin._wrap(arguments, 1, thisPlugin._patchCallbackQuery(span));
-        } else if (typeof arguments[2] === 'function') {
-          thisPlugin._wrap(arguments, 2, thisPlugin._patchCallbackQuery(span));
+          return originalQuery.apply(connection, arguments);
         }
-
-        return originalQuery.apply(connection, arguments);
       };
     };
   }
diff --git a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts b/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts
index 746e4b501b..438f875499 100644
--- a/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-mysql/test/index.test.ts
@@ -205,6 +205,28 @@ describe('mysql@2.x', () => {
       });
     });
 
+    it('should intercept connection.query(text: options, values: [])', done => {
+      const span = provider.getTracer('default').startSpan('test span');
+      context.with(trace.setSpan(context.active(), span), () => {
+        const sql = 'SELECT 1+? as solution';
+        let rows = 0;
+        const query = connection.query(sql, [1]);
+
+        query.on('result', row => {
+          assert.strictEqual(row.solution, 2);
+          rows += 1;
+        });
+
+        query.on('end', () => {
+          assert.strictEqual(rows, 1);
+          const spans = memoryExporter.getFinishedSpans();
+          assert.strictEqual(spans.length, 1);
+          assertSpan(spans[0], sql, [1]);
+          done();
+        });
+      });
+    });
+
     it('should intercept connection.query(text: options, values: [], callback)', done => {
       const span = provider.getTracer('default').startSpan('test span');
       context.with(trace.setSpan(context.active(), span), () => {
@@ -365,6 +387,28 @@ describe('mysql@2.x', () => {
       });
     });
 
+    it('should intercept pool.query(text: options, values: [])', done => {
+      const span = provider.getTracer('default').startSpan('test span');
+      context.with(trace.setSpan(context.active(), span), () => {
+        const sql = 'SELECT 1+? as solution';
+        let rows = 0;
+        const query = pool.query(sql, [1]);
+
+        query.on('result', row => {
+          assert.strictEqual(row.solution, 2);
+          rows += 1;
+        });
+
+        query.on('end', () => {
+          assert.strictEqual(rows, 1);
+          const spans = memoryExporter.getFinishedSpans();
+          assert.strictEqual(spans.length, 1);
+          assertSpan(spans[0], sql, [1]);
+          done();
+        });
+      });
+    });
+
     it('should intercept pool.query(text: options, values: [], callback)', done => {
       const span = provider.getTracer('default').startSpan('test span');
       context.with(trace.setSpan(context.active(), span), () => {
@@ -511,6 +555,31 @@ describe('mysql@2.x', () => {
       });
     });
 
+    it('should intercept poolClusterConnection.query(text: options, values: [])', done => {
+      poolCluster.getConnection((err, poolClusterConnection) => {
+        assert.ifError(err);
+        const span = provider.getTracer('default').startSpan('test span');
+        context.with(trace.setSpan(context.active(), span), () => {
+          const sql = 'SELECT 1+? as solution';
+          const query = poolClusterConnection.query(sql, [1]);
+          let rows = 0;
+
+          query.on('result', row => {
+            assert.strictEqual(row.solution, 2);
+            rows += 1;
+          });
+
+          query.on('end', () => {
+            assert.strictEqual(rows, 1);
+            const spans = memoryExporter.getFinishedSpans();
+            assert.strictEqual(spans.length, 1);
+            assertSpan(spans[0], sql, [1]);
+            done();
+          });
+        });
+      });
+    });
+
     it('should intercept poolClusterConnection.query(text: options, values: [], callback)', done => {
       poolCluster.getConnection((err, poolClusterConnection) => {
         assert.ifError(err);