diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts index b7cd5d3406..56029e568e 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts +++ b/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 './semconv'; +function extractModuleExports(module: any) { + return module[Symbol.toStringTag] === 'Module' + ? module.default // ESM + : module; // CommonJS +} + export class PgInstrumentation extends InstrumentationBase { private _operationDuration!: Histogram; private _connectionsCount!: UpDownCounter; @@ -125,45 +132,38 @@ export class PgInstrumentation extends InstrumentationBase=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'); - } + const SUPPORTED_PG_VERSIONS = ['>=8.0.3 <9']; - if (isWrapped(moduleExports.Client.prototype.connect)) { - this._unwrap(moduleExports.Client.prototype, 'connect'); - } + 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, - 'query', - this._getClientQueryPatch() as any - ); + const modulePgClient = new InstrumentationNodeModuleFile( + 'pg/lib/client.js', + SUPPORTED_PG_VERSIONS, + this._patchPgClient.bind(this), + this._unpatchPgClient.bind(this) + ); - this._wrap( - moduleExports.Client.prototype, - 'connect', - this._getClientConnectPatch() as any - ); + const modulePG = new InstrumentationNodeModuleDefinition( + 'pg', + SUPPORTED_PG_VERSIONS, + (module: any) => { + const moduleExports = extractModuleExports(module); + this._patchPgClient(moduleExports.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'); - } - } + const moduleExports = extractModuleExports(module); + + this._unpatchPgClient(moduleExports.Client); + return module; + }, + [modulePgClient, modulePgNativeClient] ); const modulePGPool = new InstrumentationNodeModuleDefinition( @@ -190,6 +190,50 @@ export class PgInstrumentation extends InstrumentationBase { diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/fixtures/use-pg.mjs b/plugins/node/opentelemetry-instrumentation-pg/test/fixtures/use-pg.mjs new file mode 100644 index 0000000000..a9c8c5a6a4 --- /dev/null +++ b/plugins/node/opentelemetry-instrumentation-pg/test/fixtures/use-pg.mjs @@ -0,0 +1,57 @@ +/* + * 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 use-pg.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 || 'postgres', + host: process.env.POSTGRES_HOST || 'localhost', + port: process.env.POSTGRES_PORT + ? parseInt(process.env.POSTGRES_PORT, 10) + : 54320, +}; + +const sdk = createTestNodeSdk({ + serviceName: 'use-pg', + instrumentations: [new PgInstrumentation()], +}); +sdk.start(); + +import pg from 'pg'; +const client = new pg.Client(CONFIG); + +await client.connect(); + +const tracer = trace.getTracer(); + +await tracer.startActiveSpan('test-span', async span => { + const res = await client.query('SELECT NOW()'); + + assert.ok(res); + span.end(); + + await client.end(); + await sdk.shutdown(); +}); diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/pg-pool.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/pg-pool.test.ts index 98d9a1b7c8..0e4f71d42c 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/pg-pool.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/pg-pool.test.ts @@ -114,6 +114,10 @@ describe('pg-pool', () => { function create(config: PgInstrumentationConfig = {}) { instrumentation.setConfig(config); instrumentation.enable(); + + // Disable and enable the instrumentation to visit unwrap calls + instrumentation.disable(); + instrumentation.enable(); } let pool: pgPool; diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts index 8fda8a8fb8..8dd33aa55f 100644 --- a/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts +++ b/plugins/node/opentelemetry-instrumentation-pg/test/pg.test.ts @@ -108,6 +108,10 @@ describe('pg', () => { function create(config: PgInstrumentationConfig = {}) { instrumentation.setConfig(config); instrumentation.enable(); + + // Disable and enable the instrumentation to visit unwrap calls + instrumentation.disable(); + instrumentation.enable(); } let postgres: typeof pg; @@ -152,6 +156,7 @@ describe('pg', () => { postgres = require('pg'); client = new postgres.Client(CONFIG); + await client.connect(); }); @@ -159,6 +164,7 @@ describe('pg', () => { if (testPostgresLocally) { testUtils.cleanUpDocker('postgres'); } + await client.end(); }); @@ -1087,3 +1093,32 @@ describe('pg', () => { }); }); }); + +describe('pg (ESM)', () => { + 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, 3); + + assert.strictEqual(spans[0].name, 'pg.connect'); + assert.strictEqual(spans[0].kind, 3); + assert.strictEqual(spans[1].name, 'test-span'); + assert.strictEqual(spans[1].kind, 1); + assert.strictEqual(spans[2].name, 'pg.query:SELECT otel_pg_database'); + assert.strictEqual(spans[2].kind, 3); + }, + }); + }); +});