From aeadca8da626164828852489ab749dfd0aa1d981 Mon Sep 17 00:00:00 2001
From: GiviMAD <GiviMAD@users.noreply.github.com>
Date: Tue, 12 Oct 2021 10:13:02 +0200
Subject: [PATCH] fix(opentelemetry-instrumentation-mongodb): fix span
 attributes with unified topology (#663)

Co-authored-by: Valentin Marchaud <contact@vmarchaud.fr>
Co-authored-by: Daniel Dyla <dyladan@users.noreply.github.com>
---
 .../src/instrumentation.ts                    | 26 +++++++---
 .../src/types.ts                              |  4 ++
 .../test/mongodb.test.ts                      | 50 +++++++++++++++++++
 .../test/utils.ts                             |  5 +-
 4 files changed, 75 insertions(+), 10 deletions(-)

diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts
index d7444c7645..eeb01820ea 100644
--- a/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts
+++ b/plugins/node/opentelemetry-instrumentation-mongodb/src/instrumentation.ts
@@ -382,14 +382,24 @@ export class MongoDBInstrumentation extends InstrumentationBase<
   ) {
     // add network attributes to determine the remote server
     if (topology && topology.s) {
-      span.setAttributes({
-        [SemanticAttributes.NET_HOST_NAME]: `${
-          topology.s.options?.host ?? topology.s.host
-        }`,
-        [SemanticAttributes.NET_HOST_PORT]: `${
-          topology.s.options?.port ?? topology.s.port
-        }`,
-      });
+      let host = topology.s.options?.host ?? topology.s.host;
+      let port: string | undefined = (
+        topology.s.options?.port ?? topology.s.port
+      )?.toString();
+      if (host == null || port == null) {
+        const address = topology.description?.address;
+        if (address) {
+          const addressSegments = address.split(':');
+          host = addressSegments[0];
+          port = addressSegments[1];
+        }
+      }
+      if (host?.length && port?.length) {
+        span.setAttributes({
+          [SemanticAttributes.NET_HOST_NAME]: host,
+          [SemanticAttributes.NET_HOST_PORT]: port,
+        });
+      }
     }
 
     // The namespace is a combination of the database name and the name of the
diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts b/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts
index fd54361449..0292ecb62c 100644
--- a/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts
+++ b/plugins/node/opentelemetry-instrumentation-mongodb/src/types.ts
@@ -148,6 +148,10 @@ export type MongoInternalTopology = {
     host?: string;
     port?: number;
   };
+  // mongodb@3 with useUnifiedTopology option
+  description?: {
+    address?: string;
+  };
 };
 
 export enum MongodbCommandType {
diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts b/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts
index d86c8490be..d202515925 100644
--- a/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts
+++ b/plugins/node/opentelemetry-instrumentation-mongodb/test/mongodb.test.ts
@@ -25,6 +25,7 @@ import {
   getTestSpans,
   resetMemoryExporter,
 } from '@opentelemetry/contrib-test-utils';
+import { lookup } from 'dns';
 
 const instrumentation = registerInstrumentationTesting(
   new MongoDBInstrumentation()
@@ -404,6 +405,55 @@ describe('MongoDBInstrumentation', () => {
     });
   });
 
+  describe('MongoDb useUnifiedTopology enabled', () => {
+    let client: mongodb.MongoClient;
+    let collection: mongodb.Collection;
+    before(done => {
+      accessCollection(URL, DB_NAME, COLLECTION_NAME, {
+        useUnifiedTopology: true,
+      })
+        .then(result => {
+          client = result.client;
+          collection = result.collection;
+          done();
+        })
+        .catch((err: Error) => {
+          console.log(
+            'Skipping test-mongodb. Could not connect. Run MongoDB to test'
+          );
+          shouldTest = false;
+          done();
+        });
+    });
+    after(() => {
+      if (client) {
+        client.close();
+      }
+    });
+    it('should generate correct span attributes', done => {
+      const span = trace.getTracer('default').startSpan('findRootSpan');
+      context.with(trace.setSpan(context.active(), span), () => {
+        collection.find({ a: 1 }).toArray((err, results) => {
+          span.end();
+          const [mongoSpan] = getTestSpans();
+          assert.ifError(err);
+          lookup(process.env.MONGODB_HOST || 'localhost', (err, address) => {
+            if (err) return done(err);
+            assert.strictEqual(
+              mongoSpan.attributes[SemanticAttributes.NET_HOST_NAME],
+              address
+            );
+            assert.strictEqual(
+              mongoSpan.attributes[SemanticAttributes.NET_HOST_PORT],
+              process.env.MONGODB_PORT || '27017'
+            );
+            done();
+          });
+        });
+      });
+    });
+  });
+
   /** Should intercept command */
   describe('Removing Instrumentation', () => {
     it('should unpatch plugin', () => {
diff --git a/plugins/node/opentelemetry-instrumentation-mongodb/test/utils.ts b/plugins/node/opentelemetry-instrumentation-mongodb/test/utils.ts
index cf83e6e7ba..7d734277fe 100644
--- a/plugins/node/opentelemetry-instrumentation-mongodb/test/utils.ts
+++ b/plugins/node/opentelemetry-instrumentation-mongodb/test/utils.ts
@@ -34,10 +34,11 @@ export interface MongoDBAccess {
 export function accessCollection(
   url: string,
   dbName: string,
-  collectionName: string
+  collectionName: string,
+  options: mongodb.MongoClientOptions = {}
 ): Promise<MongoDBAccess> {
   return new Promise((resolve, reject) => {
-    mongodb.MongoClient.connect(url, (err, client) => {
+    mongodb.MongoClient.connect(url, options, (err, client) => {
       if (err) {
         reject(err);
         return;