Add PgPlugin - PosgreSQL

README.md

@@ -57,7 +57,7 @@ Environment Variable | Description | Default
 | `SW_AGENT_LOGGING_LEVEL` | The logging level, could be one of `CRITICAL`, `FATAL`, `ERROR`, `WARN`(`WARNING`), `INFO`, `DEBUG` | `INFO` |
 | `SW_IGNORE_SUFFIX` | The suffices of endpoints that will be ignored (not traced), comma separated | `.jpg,.jpeg,.js,.css,.png,.bmp,.gif,.ico,.mp3,.mp4,.html,.svg` |
 | `SW_TRACE_IGNORE_PATH` | The paths of endpoints that will be ignored (not traced), comma separated | `` |
-| `SW_MYSQL_SQL_PARAMETERS_MAX_LENGTH` | The maximum string length of MySQL parameters to log | `512` |
+| `SW_SQL_PARAMETERS_MAX_LENGTH` | The maximum string length of SQL parameters to log | `512` |
 | `SW_AGENT_MAX_BUFFER_SIZE` | The maximum buffer size before sending the segment data to backend | `'1000'` |
 
 ## Supported Libraries
@@ -70,6 +70,7 @@ Library | Plugin Name
 | [`express`](https://expressjs.com) | `express` |
 | [`axios`](https://github.com/axios/axios) | `axios` |
 | [`mysql`](https://github.com/mysqljs/mysql) | `mysql` |
+| [`pg`](https://github.com/brianc/node-postgres) | `pg` |
 
 ### Compatible Libraries
 

package.json

@@ -46,9 +46,11 @@
     "@types/uuid": "^8.0.0",
     "axios": "^0.21.0",
     "express": "^4.17.1",
-    "grpc-tools": "^1.10.0",
     "grpc_tools_node_protoc_ts": "^4.0.0",
+    "grpc-tools": "^1.10.0",
     "jest": "^26.6.3",
+    "mysql": "^2.18.1",
+    "pg": "^8.5.1",
     "prettier": "^2.0.5",
     "testcontainers": "^6.2.0",
     "ts-jest": "^26.4.4",

src/config/AgentConfig.ts

@@ -27,7 +27,7 @@ export type AgentConfig = {
   maxBufferSize?: number;
   ignoreSuffix?: string;
   traceIgnorePath?: string;
-  mysql_sql_parameters_max_length?: number;
+  sql_parameters_max_length?: number;
   // the following is internal state computed from config values
   reIgnoreOperation?: RegExp;
 };
@@ -60,6 +60,6 @@ export default {
     Number.parseInt(process.env.SW_AGENT_MAX_BUFFER_SIZE as string, 10) : 1000,
   ignoreSuffix: process.env.SW_IGNORE_SUFFIX ?? '.jpg,.jpeg,.js,.css,.png,.bmp,.gif,.ico,.mp3,.mp4,.html,.svg',
   traceIgnorePath: process.env.SW_TRACE_IGNORE_PATH || '',
-  mysql_sql_parameters_max_length: Math.trunc(Math.max(0, Number(process.env.SW_MYSQL_SQL_PARAMETERS_MAX_LENGTH))) || 512,
+  sql_parameters_max_length: Math.trunc(Math.max(0, Number(process.env.SW_MYSQL_SQL_PARAMETERS_MAX_LENGTH))) || 512,
   reIgnoreOperation: RegExp(''),  // temporary placeholder so Typescript doesn't throw a fit
 };

src/plugins/MySQLPlugin.ts

@@ -43,6 +43,8 @@ class MySQLPlugin implements SwPlugin {
     Connection.prototype.query = function(sql: any, values: any, cb: any) {
       const wrapCallback = (_cb: any) => {
         return function(this: any, error: any, results: any, fields: any) {
+          span.resync();
+
           if (error)
             span.error(error);
 
@@ -52,10 +54,19 @@ class MySQLPlugin implements SwPlugin {
         }
       };
 
+      let query: any;
+
       const host = `${this.config.host}:${this.config.port}`;
       const span = ContextManager.current.newExitSpan('mysql/query', host).start();
 
       try {
+        span.component = Component.MYSQL;
+        span.layer = SpanLayer.DATABASE;
+        span.peer = host;
+
+        span.tag(Tag.dbType('Mysql'));
+        span.tag(Tag.dbInstance(`${this.config.database}`));
+
         let _sql: any;
         let _values: any;
         let streaming: any;
@@ -103,38 +114,41 @@ class MySQLPlugin implements SwPlugin {
           }
         }
 
-        span.component = Component.MYSQL;
-        span.layer = SpanLayer.DATABASE;
-        span.peer = host;
-
-        span.tag(Tag.dbType('mysql'));
-        span.tag(Tag.dbInstance(this.config.database || ''));
-        span.tag(Tag.dbStatement(_sql || ''));
+        span.tag(Tag.dbStatement(`${_sql}`));
 
         if (_values) {
-          let vals = _values.map((v: any) => `${v}`).join(', ');
+          let vals = _values.map((v: any) => v === undefined ? 'undefined' : JSON.stringify(v)).join(', ');
 
-          if (vals.length > config.mysql_sql_parameters_max_length)
-            vals = vals.splice(0, config.mysql_sql_parameters_max_length);
+          if (vals.length > config.sql_parameters_max_length)
+            vals = vals.splice(0, config.sql_parameters_max_length);
 
-            span.tag(Tag.dbSqlParameters(`[${vals}]`));
+          span.tag(Tag.dbSqlParameters(`[${vals}]`));
         }
 
-        const query = _query.call(this, sql, values, cb);
+        query = _query.call(this, sql, values, cb);
 
         if (streaming) {
-          query.on('error', (e: any) => span.error(e));
-          query.on('end', () => span.stop());
+          query.on('error', (e: any) => {
+            span.resync();
+            span.error(e);
+          });
+
+          query.on('end', () => {
+            span.resync();  // may have already been done in 'error' but safe to do multiple times
+            span.stop()
+          });
         }
 
-        return query;
-
       } catch (e) {
         span.error(e);
         span.stop();
 
         throw e;
       }
+
+      span.async();
+
+      return query;
     };
   }
 }

src/plugins/PgPlugin.ts

@@ -0,0 +1,138 @@
+/*!
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ *     http://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.
+ *
+ */
+
+import SwPlugin from '../core/SwPlugin';
+import ContextManager from '../trace/context/ContextManager';
+import { Component } from '../trace/Component';
+import Tag from '../Tag';
+import { SpanLayer } from '../proto/language-agent/Tracing_pb';
+import { createLogger } from '../logging';
+import PluginInstaller from '../core/PluginInstaller';
+import agentConfig from '../config/AgentConfig';
+
+const logger = createLogger(__filename);
+
+class MySQLPlugin implements SwPlugin {
+  readonly module = 'pg';
+  readonly versions = '*';
+
+  install(installer: PluginInstaller): void {
+    if (logger.isDebugEnabled()) {
+      logger.debug('installing pg plugin');
+    }
+
+    const Client = installer.require('pg/lib/client');
+    const _query = Client.prototype.query;
+
+    Client.prototype.query = function(config: any, values: any, callback: any) {
+      const wrapCallback = (_cb: any) => {
+        return function(this: any, err: any, res: any) {
+          span.resync();
+
+          if (err)
+            span.error(err);
+
+          span.stop();
+
+          return _cb.call(this, err, res);
+        }
+      };
+
+      let query: any;
+
+      const host = `${this.host}:${this.port}`;
+      const span = ContextManager.current.newExitSpan('pg/query', host).start();
+
+      try {
+        span.component = Component.POSTGRESQL;
+        span.layer = SpanLayer.DATABASE;
+        span.peer = host;
+
+        span.tag(Tag.dbType('PostgreSQL'));
+        span.tag(Tag.dbInstance(`${this.connectionParameters.database}`));
+
+        let _sql: any;
+        let _values: any;
+
+        if (typeof config === 'string')
+          _sql = config;
+
+        else if (config !== null && config !== undefined) {
+          _sql    = config.text;
+          _values = config.values;
+
+          if (typeof config.callback === 'function')
+            config.callback = wrapCallback(config.callback);
+        }
+
+        if (typeof values === 'function')
+          values = wrapCallback(values);
+        else
+          _values = values;
+
+        if (typeof callback === 'function')
+          callback = wrapCallback(callback);
+
+        span.tag(Tag.dbStatement(`${_sql}`));
+
+        if (_values) {
+          let vals = _values.map((v: any) => v === undefined ? 'undefined' : JSON.stringify(v)).join(', ');
+
+          if (vals.length > agentConfig.sql_parameters_max_length)
+            vals = vals.splice(0, agentConfig.sql_parameters_max_length);
+
+            span.tag(Tag.dbSqlParameters(`[${vals}]`));
+        }
+
+        query = _query.call(this, config, values, callback);
+
+        if (query && typeof query.then === 'function' && typeof query.catch === 'function')  // generic Promise check
+          query = query.then(
+            (res: any) => {
+              span.resync();
+              span.stop();
+
+              return res;
+            },
+
+            (err: any) => {
+              span.resync();
+              span.error(err);
+              span.stop();
+
+              return Promise.reject(err);
+            }
+          );
+
+      } catch (e) {
+        span.error(e);
+        span.stop();
+
+        throw e;
+      }
+
+      span.async();
+
+      return query;
+    };
+  }
+}
+
+// noinspection JSUnusedGlobalSymbols
+export default new MySQLPlugin();

src/trace/Component.ts

@@ -22,6 +22,7 @@ export class Component {
   static readonly HTTP = new Component(2);
   static readonly MYSQL = new Component(5);
   static readonly MONGODB = new Component(9);
+  static readonly POSTGRESQL = new Component(22);
   static readonly HTTP_SERVER = new Component(49);
   static readonly EXPRESS = new Component(4002);
   static readonly AXIOS = new Component(4005);

tests/plugins/mysql/client.ts

@@ -0,0 +1,38 @@
+/*!
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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
+ *
+ *     http://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.
+ *
+ */
+
+import * as http from 'http';
+import agent from '../../../src';
+
+agent.start({
+  serviceName: 'client',
+  maxBufferSize: 1000,
+})
+
+const server = http.createServer((req, res) => {
+  http
+    .request(`http://${process.env.SERVER || 'localhost:5000'}${req.url}`, (r) => {
+      let data = '';
+      r.on('data', (chunk) => (data += chunk));
+      r.on('end', () => res.end(data));
+    })
+    .end();
+});
+
+server.listen(5001, () => console.info('Listening on port 5001...'));