fix(pg): fix instrumentation of ESM-imported pg-pool

packages/opentelemetry-test-utils/src/test-utils.ts

@@ -45,7 +45,7 @@ const dockerRunCmds = {
   oracledb:
     'docker run --rm -d --name otel-oracledb -p 1521:1521 -e ORACLE_PASSWORD=oracle -e APP_USER=otel -e APP_USER_PASSWORD=secret gvenzl/oracle-free:slim',
   postgres:
-    'docker run --rm -d --name otel-postgres -p 54320:5432 -e POSTGRES_PASSWORD=postgres postgres:16-alpine',
+    'docker run --rm -d --name otel-postgres -p 54320:5432 -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=otel_pg_database postgres:16-alpine',
   redis: 'docker run --rm -d --name otel-redis -p 63790:6379 redis:alpine',
 };
 

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

@@ -133,6 +133,7 @@ export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConf
 
   protected init() {
     const SUPPORTED_PG_VERSIONS = ['>=8.0.3 <9'];
+    const SUPPORTED_PG_POOL_VERSIONS = ['>=2.0.0 <4'];
 
     const modulePgNativeClient = new InstrumentationNodeModuleFile(
       'pg/lib/native/client.js',
@@ -168,8 +169,9 @@ export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConf
 
     const modulePGPool = new InstrumentationNodeModuleDefinition(
       'pg-pool',
-      ['>=2.0.0 <4'],
-      (moduleExports: typeof pgPoolTypes) => {
+      SUPPORTED_PG_POOL_VERSIONS,
+      (module: any) => {
+        const moduleExports = extractModuleExports(module);
         if (isWrapped(moduleExports.prototype.connect)) {
           this._unwrap(moduleExports.prototype, 'connect');
         }
@@ -180,7 +182,8 @@ export class PgInstrumentation extends InstrumentationBase<PgInstrumentationConf
         );
         return moduleExports;
       },
-      (moduleExports: typeof pgPoolTypes) => {
+      (module: any) => {
+        const moduleExports = extractModuleExports(module);
         if (isWrapped(moduleExports.prototype.connect)) {
           this._unwrap(moduleExports.prototype, 'connect');
         }

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

@@ -0,0 +1,59 @@
+/*
+ * 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 pg-pool from an ES module:
+//    node --experimental-loader=@opentelemetry/instrumentation/hook.mjs use-pg-pool.mjs
+
+import { trace } from '@opentelemetry/api';
+import { createTestNodeSdk } from '@opentelemetry/contrib-test-utils';
+import assert from 'assert';
+
+import { PgInstrumentation } from '../../build/src/index.js';
+
+const CONFIG = {
+  user: process.env.POSTGRES_USER || 'postgres',
+  password: process.env.POSTGRES_PASSWORD || 'postgres',
+  database: process.env.POSTGRES_DB || 'otel_pg_database',
+  host: process.env.POSTGRES_HOST || 'localhost',
+  port: process.env.POSTGRES_PORT
+    ? parseInt(process.env.POSTGRES_PORT, 10)
+    : 54320,
+};
+
+const sdk = createTestNodeSdk({
+  serviceName: 'use-pg-pool',
+  instrumentations: [new PgInstrumentation()],
+});
+sdk.start();
+
+import Pool from 'pg-pool';
+const pgPool = new Pool(CONFIG);
+
+const tracer = trace.getTracer();
+
+await tracer.startActiveSpan('test-span', async span => {
+  const client = await pgPool.connect();
+  try {
+    const res = await pgPool.query('SELECT NOW()');
+    assert.ok(res);
+    console.log('rows:', res.rows);
+  } finally {
+    client.release();
+    pgPool.end();
+    span.end();
+    sdk.shutdown();
+  }
+});

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

@@ -26,7 +26,7 @@ import { PgInstrumentation } from '../../build/src/index.js';
 const CONFIG = {
   user: process.env.POSTGRES_USER || 'postgres',
   password: process.env.POSTGRES_PASSWORD || 'postgres',
-  database: process.env.POSTGRES_DB || 'postgres',
+  database: process.env.POSTGRES_DB || 'otel_pg_database',
   host: process.env.POSTGRES_HOST || 'localhost',
   port: process.env.POSTGRES_PORT
     ? parseInt(process.env.POSTGRES_PORT, 10)

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

@@ -650,6 +650,12 @@ describe('pg-pool', () => {
 
     it('should not add duplicate event listeners to PgPool events', done => {
       const poolAux: pgPool<pg.Client> = new pgPool(CONFIG);
+
+      const finish = () => {
+        poolAux.end();
+        done();
+      };
+
       let completed = 0;
       poolAux.connect((err, client, release) => {
         if (err) {
@@ -687,7 +693,7 @@ describe('pg-pool', () => {
 
         completed++;
         if (completed >= 2) {
-          done();
+          finish();
         }
       });
 
@@ -727,7 +733,7 @@ describe('pg-pool', () => {
 
         completed++;
         if (completed >= 2) {
-          done();
+          finish();
         }
       });
     });
@@ -808,6 +814,8 @@ describe('pg-pool', () => {
             1,
             'expected to have 1 used connection'
           );
+
+          poolAux.end();
           done();
         });
       });
@@ -817,6 +825,12 @@ describe('pg-pool', () => {
       const pool1: pgPool<pg.Client> = new pgPool(CONFIG);
       const pool2: pgPool<pg.Client> = new pgPool(CONFIG);
 
+      const finish = () => {
+        pool1.end();
+        pool2.end();
+        done();
+      };
+
       let completed = 0;
       pool1.connect((err, client, release) => {
         if (err) {
@@ -862,7 +876,7 @@ describe('pg-pool', () => {
 
         completed++;
         if (completed >= 2) {
-          done();
+          finish();
         }
       });
 
@@ -910,9 +924,48 @@ describe('pg-pool', () => {
 
         completed++;
         if (completed >= 2) {
-          done();
+          finish();
         }
       });
     });
   });
 });
+
+describe('pg-pool (ESM)', () => {
+  it('should work with ESM usage', async () => {
+    await testUtils.runTestFixture({
+      cwd: __dirname,
+      argv: ['fixtures/use-pg-pool.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, 6);
+
+        let span = spans.shift()!;
+        assert.strictEqual(span.name, 'test-span');
+        assert.strictEqual(span.kind, 1 /* OtlpSpanKind.INTERNAL */);
+        const expectedRemainingSpanNames = [
+          // I believe two sets of `*.connect` spans because pg-pool opens
+          // two connections to start.
+          'pg-pool.connect',
+          'pg.connect',
+          'pg-pool.connect',
+          'pg.connect',
+          'pg.query:SELECT otel_pg_database',
+        ];
+        for (const expectedName of expectedRemainingSpanNames) {
+          span = spans.shift()!;
+          assert.strictEqual(span.name, expectedName);
+        }
+      },
+    });
+  });
+});