add neo4j plugin

.github/workflows/plugins.yml

@@ -740,6 +740,29 @@ jobs:
         uses: ./.github/actions/testagent/logs
       - uses: codecov/codecov-action@v2
 
+  neo4j:
+    runs-on: ubuntu-latest
+    services:
+      neo4j:
+        image: neo4j:4.2.3
+        env:
+          NEO4J_AUTH: 'neo4j/test'
+        ports:
+          - 7474:7474
+          - 11011:7687
+    env:
+      PLUGINS: neo4j
+      SERVICES: neo4j
+    steps:
+      - uses: actions/checkout@v2
+      - uses: ./.github/actions/node/setup
+      - run: yarn install
+      - uses: ./.github/actions/node/oldest
+      - run: yarn test:plugins:ci
+      - uses: ./.github/actions/node/latest
+      - run: yarn test:plugins:ci
+      - uses: codecov/codecov-action@v2
+
   net:
     runs-on: ubuntu-latest
     env:

docker-compose.yml

@@ -82,6 +82,13 @@ services:
     image: ridedott/pubsub-emulator
     ports:
       - "127.0.0.1:8081:8081"
+  neo4j:
+    image: neo4j:4.2.3
+    environment:
+      - NEO4J_AUTH=neo4j/test
+    ports:
+      - "7474:7474"
+      - "11011:7687"
   localstack:
     # TODO: Figure out why SNS doesn't work in >=1.2
     # https://github.com/localstack/localstack/issues/7479

docs/API.md

@@ -76,6 +76,9 @@ tracer.use('pg', {
 <h5 id="mysql2"></h5>
 <h5 id="mysql2-tags"></h5>
 <h5 id="mysql2-config"></h5>
+<h5 id="neo4j"></h5>
+<h5 id="neo4j-tags"></h5>
+<h5 id="neo4j-config"></h5>
 <h5 id="net"></h5>
 <h5 id="next"></h5>
 <h5 id="opensearch"></h5>
@@ -131,6 +134,7 @@ tracer.use('pg', {
 * [mongodb-core](./interfaces/plugins.mongodb_core.html)
 * [mysql](./interfaces/plugins.mysql.html)
 * [mysql2](./interfaces/plugins.mysql2.html)
+* [neo4j](./interfaces/plugins.neo4j.html)
 * [net](./interfaces/plugins.net.html)
 * [next](./interfaces/plugins.next.html)
 * [opensearch](./interfaces/plugins.opensearch.html)

docs/test.ts

@@ -315,6 +315,7 @@ tracer.use('mysql');
 tracer.use('mysql', { service: () => `my-custom-mysql` });
 tracer.use('mysql2');
 tracer.use('mysql2', { service: () => `my-custom-mysql2` });
+tracer.use('neo4j');
 tracer.use('net');
 tracer.use('next');
 tracer.use('next', nextOptions);

index.d.ts

@@ -827,6 +827,7 @@ interface Plugins {
   "mongoose": plugins.mongoose;
   "mysql": plugins.mysql;
   "mysql2": plugins.mysql2;
+  "neo4j": plugins.neo4j;
   "net": plugins.net;
   "next": plugins.next;
   "opensearch": plugins.opensearch;
@@ -1483,6 +1484,11 @@ declare namespace plugins {
    * [mysql2](https://github.com/sidorares/node-mysql2) module.
    */
   interface mysql2 extends mysql {}
+  /**
+   * This plugin automatically instruments the
+   * [neo4j](https://github.com/neo4j/neo4j-javascript-driver) module.
+   */
+  interface neo4j extends Instrumentation {}
 
   /**
    * This plugin automatically instruments the

packages/datadog-instrumentations/src/neo4j.js

@@ -0,0 +1,103 @@
+'use strict'
+
+const {
+  channel,
+  addHook,
+  AsyncResource
+} = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const startCh = channel('apm:neo4j:query:start')
+const finishCh = channel('apm:neo4j:query:finish')
+const errorCh = channel('apm:neo4j:query:error')
+
+addHook({ name: 'neo4j-driver-core', file: 'lib/session.js', versions: ['>=4.3.0'] }, exports => {
+  const Session = exports.default
+  shimmer.wrap(Session.prototype, 'run', wrapRun)
+  return Session
+})
+
+addHook({ name: 'neo4j-driver-core', file: 'lib/transaction.js', versions: ['>=4.3.0'] }, exports => {
+  const Transaction = exports.default
+  shimmer.wrap(Transaction.prototype, 'run', wrapRun)
+  return Transaction
+})
+
+addHook({ name: 'neo4j-driver', file: 'lib/session.js', versions: ['<4.3.0', '>=4.0.0'] }, exports => {
+  const Session = exports.default
+  shimmer.wrap(Session.prototype, 'run', wrapRun)
+  return Session
+})
+
+addHook({ name: 'neo4j-driver', file: 'lib/transaction.js', versions: ['<4.3.0', '>=4.0.0'] }, exports => {
+  const Transaction = exports.default
+  shimmer.wrap(Transaction.prototype, 'run', wrapRun)
+  return Transaction
+})
+
+function wrapRun (run) {
+  return function (statement) {
+    if (!startCh.hasSubscribers) {
+      return run.apply(this, arguments)
+    }
+
+    if (!statement) return run.apply(this, arguments)
+
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    const attributes = getAttributesFromNeo4jSession(this)
+
+    return asyncResource.runInAsyncScope(() => {
+      startCh.publish({ attributes, statement })
+
+      try {
+        const promise = run.apply(this, arguments)
+        if (promise && typeof promise.then === 'function') {
+          const onResolve = asyncResource.bind(() => finish())
+          const onReject = asyncResource.bind(e => finish(e))
+
+          promise.then(onResolve, onReject)
+        } else {
+          finish()
+        }
+        return promise
+      } catch (err) {
+        err.stack // trigger getting the stack at the original throwing point
+        errorCh.publish(err)
+
+        throw err
+      }
+    })
+  }
+}
+
+function finish (error) {
+  if (error) {
+    errorCh.publish(error)
+  }
+  finishCh.publish()
+}
+
+function getAttributesFromNeo4jSession (session) {
+  const connectionHolder =
+    (session._mode === 'WRITE' ? session._writeConnectionHolder : session._readConnectionHolder) ||
+    session._connectionHolder ||
+    {}
+  const connectionProvider = connectionHolder._connectionProvider || {}
+
+  // seedRouter is used when connecting to a url that starts with "neo4j", usually aura
+  const address = connectionProvider._address || connectionProvider._seedRouter
+  const auth = connectionProvider._authToken || {}
+
+  const attributes = {
+    // "neo4j" is the default database name. When used, "session._database" is an empty string
+    dbName: session._database ? session._database : 'neo4j'
+  }
+  if (address) {
+    attributes.host = address._host
+    attributes.port = address._port
+  }
+  if (auth.principal) {
+    attributes.dbUser = auth.principal
+  }
+  return attributes
+}

packages/datadog-plugin-neo4j/src/index.js

@@ -0,0 +1,48 @@
+'use strict'
+
+const Plugin = require('../../dd-trace/src/plugins/plugin')
+const { storage } = require('../../datadog-core')
+const analyticsSampler = require('../../dd-trace/src/analytics_sampler')
+
+class Neo4jPlugin extends Plugin {
+  static get name () {
+    return 'neo4j'
+  }
+
+  constructor (...args) {
+    super(...args)
+
+    this.addSub('apm:neo4j:query:start', ({ attributes, statement }) => {
+      const store = storage.getStore()
+      const childOf = store ? store.span : store
+      const span = this.tracer.startSpan('neo4j.query', {
+        childOf,
+        tags: {
+          'db.name': attributes.dbName,
+          'db.type': 'neo4j',
+          'db.user': attributes.dbUser,
+          'out.host': attributes.host,
+          'out.port': attributes.port,
+          'resource.name': statement,
+          'service.name': this.config.service || `${this.tracer._service}-neo4j`,
+          'span.kind': 'client',
+          'span.type': 'cypher'
+        }
+      })
+      analyticsSampler.sample(span, this.config.measured)
+      this.enter(span, store)
+    })
+
+    this.addSub('apm:neo4j:query:error', err => {
+      const span = storage.getStore().span
+      span.setTag('error', err)
+    })
+
+    this.addSub('apm:neo4j:query:finish', () => {
+      const span = storage.getStore().span
+      span.finish()
+    })
+  }
+}
+
+module.exports = Neo4jPlugin