feat(client): Tracing Connect and Disconnect (#14527)

* feat: add Lib tracing for connect and disconnect

* refactor: correct headers and testing for lib connect and discconect

* refactor: export tracing from engine-core

* fix: add disconnect race back

* feat: add connect and disconnect to binary engine tracing

* refactor: temp remove --additional-headers flag

* fix: disconnect race

packages/client/src/runtime/core/request/PrismaPromise.ts

@@ -1,4 +1,4 @@
-import { TransactionTracer } from '../../../utils/TransactionTracer'
+import { TransactionTracer } from '@prisma/engine-core'
 
 /**
  * Prisma's `Promise` that is backwards-compatible. All additions on top of the

packages/client/src/runtime/core/request/createPrismaPromise.ts

@@ -1,4 +1,5 @@
-import { TransactionTracer } from '../../../utils/TransactionTracer'
+import { TransactionTracer } from '@prisma/engine-core'
+
 import type { PrismaPromise } from './PrismaPromise'
 
 /**

packages/client/src/runtime/getPrismaClient.ts

@@ -1,7 +1,17 @@
 import { SpanOptions } from '@opentelemetry/api'
 import Debug from '@prisma/debug'
 import type { DatasourceOverwrite, Engine, EngineConfig, EngineEventType, Options } from '@prisma/engine-core'
-import { BinaryEngine, DataProxyEngine, getTracingConfig, LibraryEngine, TransactionHeaders } from '@prisma/engine-core'
+import {
+  BinaryEngine,
+  DataProxyEngine,
+  getTraceParent,
+  getTracingConfig,
+  LibraryEngine,
+  runInActiveSpan,
+  runInSpan,
+  TransactionHeaders,
+  TransactionTracer,
+} from '@prisma/engine-core'
 import type { DataSource, GeneratorConfig } from '@prisma/generator-helper'
 import {
   ClientEngineType,
@@ -19,7 +29,6 @@ import * as sqlTemplateTag from 'sql-template-tag'
 
 import { getPrismaClientDMMF } from '../generation/getDMMF'
 import type { InlineDatasources } from '../generation/utils/buildInlineDatasources'
-import { TransactionTracer } from '../utils/TransactionTracer'
 import { PrismaClientValidationError } from '.'
 import { MetricsClient } from './core/metrics/MetricsClient'
 import { applyModels } from './core/model/applyModels'
@@ -39,8 +48,6 @@ import { clientVersion } from './utils/clientVersion'
 import { getOutputTypeName } from './utils/common'
 import { deserializeRawResults } from './utils/deserializeRawResults'
 import { mssqlPreparedStatement } from './utils/mssqlPreparedStatement'
-import { getTraceParent } from './utils/otel/getTraceParent'
-import { runInActiveSpan, runInSpan } from './utils/otel/runInSpan'
 import { printJsonWithErrors } from './utils/printJsonErrors'
 import type { InstanceRejectOnNotFound, RejectOnNotFound } from './utils/rejectOnNotFound'
 import { getRejectOnNotFound } from './utils/rejectOnNotFound'

packages/client/tests/functional/tracing/tests.ts

@@ -68,6 +68,10 @@ afterAll(() => {
 })
 
 testMatrix.setupTestSuite(({ provider }) => {
+  beforeEach(async () => {
+    await prisma.$connect()
+  })
+
   beforeEach(() => {
     inMemorySpanExporter.reset()
   })
@@ -562,8 +566,10 @@ testMatrix.setupTestSuite(({ provider }) => {
     // @ts-ignore
     let _prisma: PrismaClient
 
-    beforeAll(() => {
+    beforeAll(async () => {
       _prisma = newPrismaClient()
+
+      await _prisma.$connect()
     })
 
     test('tracing with middleware', async () => {
@@ -640,4 +646,78 @@ testMatrix.setupTestSuite(({ provider }) => {
       expect(dbQuery4.span.attributes['db.statement']).toContain('COMMIT')
     })
   })
+
+  describe('Tracing connect', () => {
+    // @ts-ignore
+    let _prisma: PrismaClient
+
+    beforeAll(() => {
+      _prisma = newPrismaClient()
+    })
+
+    afterAll(async () => {
+      await _prisma.$disconnect()
+    })
+
+    test('should trace the implict $connect call', async () => {
+      const email = faker.internet.email()
+
+      await _prisma.user.findMany({
+        where: {
+          email: email,
+        },
+      })
+
+      const tree = await waitForSpanTree()
+
+      expect(tree.span.name).toEqual('prisma')
+      expect(tree.span.attributes['method']).toEqual('findMany')
+      expect(tree.span.attributes['model']).toEqual('User')
+
+      expect(tree.children).toHaveLength(2)
+
+      const connect = (tree?.children || [])[0] as unknown as Tree
+      expect(connect.span.name).toEqual('prisma:connect')
+
+      const engine = (tree?.children || [])[1] as unknown as Tree
+      expect(engine.span.name).toEqual('prisma:query_builder')
+
+      const getConnection = (engine.children || [])[0]
+      expect(getConnection.span.name).toEqual('prisma:connection')
+
+      if (provider === 'mongodb') {
+        expect(engine.children).toHaveLength(2)
+
+        const dbQuery1 = (engine.children || [])[1]
+        expect(dbQuery1.span.name).toEqual('prisma:db_query')
+        expect(dbQuery1.span.attributes['db.statement']).toContain('db.User.findMany(*)')
+
+        return
+      }
+
+      expect(engine.children).toHaveLength(2)
+
+      const select = (engine.children || [])[1]
+      expect(select.span.name).toEqual('prisma:db_query')
+      expect(select.span.attributes['db.statement']).toContain('SELECT')
+    })
+  })
+
+  describe('Tracing disconnect', () => {
+    // @ts-ignore
+    let _prisma: PrismaClient
+
+    beforeAll(async () => {
+      _prisma = newPrismaClient()
+      await _prisma.$connect()
+    })
+
+    test('should trace $disconnect', async () => {
+      await _prisma.$disconnect()
+
+      const tree = await waitForSpanTree()
+
+      expect(tree.span.name).toEqual('prisma:disconnect')
+    })
+  })
 })

packages/engine-core/package.json

@@ -31,7 +31,6 @@
   "dependencies": {
     "@opentelemetry/sdk-trace-base": "^1.4.0",
     "@opentelemetry/api": "^1.1.0",
-    "@opentelemetry/core": "^1.4.0",
     "@prisma/debug": "workspace:*",
     "@prisma/engines": "workspace:*",
     "@prisma/generator-helper": "workspace:*",

packages/engine-core/src/binary/BinaryEngine.ts

@@ -31,12 +31,12 @@ import { prismaGraphQLToJSError } from '../common/errors/utils/prismaGraphQLToJS
 import { EngineMetricsOptions, Metrics, MetricsOptionsJson, MetricsOptionsPrometheus } from '../common/types/Metrics'
 import type { EngineSpanEvent, QueryEngineRequestHeaders, QueryEngineResult } from '../common/types/QueryEngine'
 import type * as Tx from '../common/types/Transaction'
-import { createSpan } from '../common/utils/createSpan'
 import { getTracingConfig } from '../common/utils/getTracingConfig'
 import { printGeneratorConfig } from '../common/utils/printGeneratorConfig'
 import { fixBinaryTargets, plusX } from '../common/utils/util'
 import byline from '../tools/byline'
 import { omit } from '../tools/omit'
+import { createSpan, getTraceParent, runInActiveSpan } from '../tracing'
 import type { Result } from './Connection'
 import { Connection } from './Connection'
 
@@ -469,21 +469,28 @@ ${chalk.dim("In case we're mistaken, please report this to us 🙏.")}`)
       await this.stopPromise
     }
 
-    if (!this.startPromise) {
-      this.startCount++
-      this.startPromise = this.internalStart()
-    }
+    const startFn = async () => {
+      if (!this.startPromise) {
+        this.startCount++
+        this.startPromise = this.internalStart()
+      }
 
-    await this.startPromise
+      await this.startPromise
 
-    if (!this.child && !this.engineEndpoint) {
-      throw new PrismaClientUnknownRequestError(
-        `Can't perform request, as the Engine has already been stopped`,
-        this.clientVersion!,
-      )
+      if (!this.child && !this.engineEndpoint) {
+        throw new PrismaClientUnknownRequestError(
+          `Can't perform request, as the Engine has already been stopped`,
+          this.clientVersion!,
+        )
+      }
     }
 
-    return this.startPromise
+    const tracingConfig = getTracingConfig(this)
+    if (tracingConfig.enabled && !this.startPromise) {
+      return runInActiveSpan({ name: 'prisma:connect', callback: () => startFn() })
+    } else {
+      return startFn()
+    }
   }
 
   private getEngineEnvVars() {
@@ -562,6 +569,16 @@ ${chalk.dim("In case we're mistaken, please report this to us 🙏.")}`)
         this.port = await this.getFreePort()
         flags.push('--port', String(this.port))
 
+        // TODO - This should be uncommended(and tested) when this PR is merged: https://github.com/prisma/prisma-engines/pull/3087
+        // const additionalHeaders: { traceparent?: string } = {}
+
+        // const tracingConfig = getTracingConfig(this)
+        // if (tracingConfig.enabled) {
+        //   additionalHeaders.traceparent = getTraceParent()
+        // }
+
+        // flags.push('--additional-headers', JSON.stringify(additionalHeaders))
+
         debug({ flags })
 
         const env = this.getEngineEnvVars()
@@ -768,11 +785,20 @@ You very likely have the wrong "binaryTarget" defined in the schema.prisma file.
   }
 
   async stop(): Promise<void> {
-    if (!this.stopPromise) {
-      this.stopPromise = this._stop()
+    const stopFn = async () => {
+      if (!this.stopPromise) {
+        this.stopPromise = this._stop()
+      }
+
+      return this.stopPromise
     }
 
-    return this.stopPromise
+    const tracingConfig = getTracingConfig(this)
+    if (tracingConfig.enabled) {
+      return runInActiveSpan({ name: 'prisma:disconnect', callback: () => stopFn() })
+    } else {
+      return stopFn()
+    }
   }
 
   /**

packages/engine-core/src/index.ts

@@ -17,3 +17,4 @@ export { plusX } from './common/utils/util'
 export { DataProxyEngine } from './data-proxy/DataProxyEngine'
 export { LibraryEngine } from './library/LibraryEngine'
 export * as NodeAPILibraryTypes from './library/types/Library'
+export * from './tracing'