feat(pg): Patch client inside lib and lib/pg-native

Author: onurtemizkan

plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts

@@ -18,6 +18,7 @@ import {
   InstrumentationBase,
   InstrumentationNodeModuleDefinition,
   safeExecuteInTheMiddle,
+  InstrumentationNodeModuleFile,
 } from '@opentelemetry/instrumentation';
 import {
   context,
@@ -67,6 +68,12 @@ import {
   ATTR_DB_OPERATION_NAME,
 } from '@opentelemetry/semantic-conventions/incubating';
 
+function extractModuleExports(module: any) {
+  return module[Symbol.toStringTag] === 'Module'
+    ? module.default // ESM
+    : module; // CommonJS
+}
+
 export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConfig> {
   private _operationDuration!: Histogram;
   private _connectionsCount!: UpDownCounter;
@@ -125,45 +132,33 @@ export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConf
   }
 
   protected init() {
-    const modulePG = new InstrumentationNodeModuleDefinition(
-      'pg',
-      ['>=8.0.3 <9'],
-      (module: any) => {
-        const moduleExports: typeof pgTypes =
-          module[Symbol.toStringTag] === 'Module'
-            ? module.default // ESM
-            : module; // CommonJS
-        if (isWrapped(moduleExports.Client.prototype.query)) {
-          this._unwrap(moduleExports.Client.prototype, 'query');
-        }
-
-        if (isWrapped(moduleExports.Client.prototype.connect)) {
-          this._unwrap(moduleExports.Client.prototype, 'connect');
-        }
+    const SUPPORTED_PG_VERSIONS = ['>=8.0.3 <9'];
 
-        this._wrap(
-          moduleExports.Client.prototype,
-          'query',
-          this._getClientQueryPatch() as any
-        );
+    const modulePgNativeClient = new InstrumentationNodeModuleFile(
+      'pg/lib/native/client.js',
+      SUPPORTED_PG_VERSIONS,
+      this._patchPgClient.bind(this),
+      this._unpatchPgClient.bind(this)
+    );
 
-        this._wrap(
-          moduleExports.Client.prototype,
-          'connect',
-          this._getClientConnectPatch() as any
-        );
+    const modulePgClient = new InstrumentationNodeModuleFile(
+      'pg/lib/client.js',
+      SUPPORTED_PG_VERSIONS,
+      this._patchPgClient.bind(this),
+      this._unpatchPgClient.bind(this)
+    );
 
+    const modulePG = new InstrumentationNodeModuleDefinition(
+      'pg',
+      SUPPORTED_PG_VERSIONS,
+      (module: any) => {
+        this._patchPgClient(module.Client);
         return module;
       },
       (module: any) => {
-        const moduleExports: typeof pgTypes =
-          module[Symbol.toStringTag] === 'Module'
-            ? module.default // ESM
-            : module; // CommonJS
-        if (isWrapped(moduleExports.Client.prototype.query)) {
-          this._unwrap(moduleExports.Client.prototype, 'query');
-        }
-      }
+        this._unpatchPgClient(module.Client);
+        return module;
+      }, [modulePgClient, modulePgNativeClient]
     );
 
     const modulePGPool = new InstrumentationNodeModuleDefinition(
@@ -190,6 +185,51 @@ export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConf
     return [modulePG, modulePGPool];
   }
 
+  private _patchPgClient(
+    module: any,
+  ) {
+    const moduleExports = extractModuleExports(module);
+
+    if (isWrapped(moduleExports.prototype.query)) {
+      this._unwrap(moduleExports.prototype, 'query');
+    }
+
+    if (isWrapped(moduleExports.prototype.connect)) {
+      this._unwrap(moduleExports.prototype, 'connect');
+    }
+
+    this._wrap(
+      moduleExports.prototype,
+      'query',
+      this._getClientQueryPatch() as any
+    );
+
+    this._wrap(
+      moduleExports.prototype,
+      'connect',
+      this._getClientConnectPatch() as any
+    );
+
+    return module;
+  }
+
+  private _unpatchPgClient(
+    module: any,
+  ) {
+    const moduleExports = extractModuleExports(module);
+
+    if (isWrapped(moduleExports.prototype.query)) {
+      this._unwrap(moduleExports.prototype, 'query');
+    }
+
+    if (isWrapped(moduleExports.prototype.connect)) {
+      this._unwrap(moduleExports.prototype, 'connect');
+    }
+
+    return module;
+  }
+
+
   private _getClientConnectPatch() {
     const plugin = this;
     return (original: PgClientConnect) => {
@@ -272,12 +312,12 @@ export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConf
         // to properly narrow arg0, but TS 4.3.5 does not.
         const queryConfig = firstArgIsString
           ? {
-              text: arg0 as string,
-              values: Array.isArray(args[1]) ? args[1] : undefined,
-            }
+            text: arg0 as string,
+            values: Array.isArray(args[1]) ? args[1] : undefined,
+          }
           : firstArgIsQueryObjectWithText
-          ? (arg0 as utils.ObjectWithText)
-          : undefined;
+            ? (arg0 as utils.ObjectWithText)
+            : undefined;
 
         const attributes: Attributes = {
           [SEMATTRS_DB_SYSTEM]: DBSYSTEMVALUES_POSTGRESQL,

plugins/node/opentelemetry-instrumentation-pg/test/fixtures/use-pg.mjs

@@ -0,0 +1,44 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Use postgres from an ES module:
+//    node --experimental-loader=@opentelemetry/instrumentation/hook.mjs pg-esm.mjs
+
+import { createTestNodeSdk } from '@opentelemetry/contrib-test-utils';
+
+import { PgInstrumentation } from '../../build/src/index.js';
+
+const sdk = createTestNodeSdk({
+  serviceName: 'use-pg',
+  instrumentations: [
+    new PgInstrumentation()
+  ]
+})
+sdk.start();
+
+import pg from 'pg';
+
+const client = new pg.Client();
+
+client.connect();
+
+client.query('SELECT NOW()', (err, res) => {
+  console.log(err, res);
+  client.end();
+});
+
+
+

plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts

@@ -1086,4 +1086,24 @@ describe('pg', () => {
       });
     });
   });
+
+  it('should work with ESM usage', async () => {
+    await testUtils.runTestFixture({
+      cwd: __dirname,
+      argv: ['fixtures/use-pg.mjs'],
+      env: {
+        NODE_OPTIONS:
+          '--experimental-loader=@opentelemetry/instrumentation/hook.mjs',
+        NODE_NO_WARNINGS: '1',
+      },
+      checkResult: (err, stdout, stderr) => {
+        assert.ifError(err);
+      },
+      checkCollector: (collector: testUtils.TestCollector) => {
+        const spans = collector.sortedSpans;
+
+        assert.strictEqual(spans.length, 2);
+      },
+    });
+  });
 });