From 9c3d872be75341d7c2e2eeb1a6e577afdefc3493 Mon Sep 17 00:00:00 2001
From: Mike Lang <mike.lang1010@gmail.com>
Date: Tue, 10 Oct 2017 16:04:16 -0500
Subject: [PATCH 1/2] Add support for named prepared statements

---
 lib/instrumentation/pg.js  | 12 ++++++++++++
 lib/shim/datastore-shim.js | 20 +++++++++++++++++++-
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/lib/instrumentation/pg.js b/lib/instrumentation/pg.js
index 3bbb2f853c..cabeb4ec46 100644
--- a/lib/instrumentation/pg.js
+++ b/lib/instrumentation/pg.js
@@ -14,6 +14,16 @@ function getQuery(shim, original, name, args) {
   return statement
 }
 
+function getQueryName(shim, original, name, args) {
+  var config = args[0]
+
+  if (config && config.name) {
+    return config.name
+  }
+
+  return null
+}
+
 module.exports = function initialize(agent, pgsql, moduleName, shim) {
   shim.setDatastore(shim.POSTGRES)
   // allows for native wrapping to not happen if not necessary
@@ -44,6 +54,7 @@ module.exports = function initialize(agent, pgsql, moduleName, shim) {
     shim.recordQuery(this, 'query', {
       callback: shim.LAST,
       query: getQuery,
+      queryName: getQueryName,
       stream: 'row',
       parameters: getInstanceParameters(shim, this),
       internal: false
@@ -82,6 +93,7 @@ module.exports = function initialize(agent, pgsql, moduleName, shim) {
       return {
         callback: shim.LAST,
         query: getQuery,
+        queryName: getQueryName,
         stream: 'row',
         parameters: getInstanceParameters(shim, this),
         internal: false
diff --git a/lib/shim/datastore-shim.js b/lib/shim/datastore-shim.js
index 4c7a655280..94d36b4608 100644
--- a/lib/shim/datastore-shim.js
+++ b/lib/shim/datastore-shim.js
@@ -721,9 +721,15 @@ function _recordQuery(suffix, nodule, properties, querySpec) {
     }
     shim.logger.trace('Found query %j', queryStr)
 
+    var queryName = _extractQueryName.call(shim, fn, fnName, queryDesc, this, args)
+
     // Parse the query and assemble the name.
     var parsed = shim.parseQuery(queryStr, this)
-    var name = (parsed.collection || 'other') + '/' + parsed.operation + suffix
+    // If the intstrumentation provides a name use it, otherwise construct a name
+    // from the parsed query string
+    var name = queryName ? 
+      queryName : 
+      (parsed.collection || 'other') + '/' + parsed.operation + suffix
 
     // Return the segment descriptor.
     return {
@@ -827,6 +833,18 @@ function _extractQueryStr(fn, fnName, spec, ctx, args) {
   return queryStr
 }
 
+function _extractQueryName(fn, fnName, spec, ctx, args) {
+  var queryName = spec.queryName
+  // Adds support for an instrumentation to specify
+  // a function to calculate a name that should be used
+  // for generating the transaction segment name for the query
+  if (!queryName || !this.isFunction(queryName)) {
+    return null
+  }
+
+  return queryName.call(ctx, this, fn, fnName, args)
+}
+
 /**
  * Normalizes segment parameter values.
  *

From aedb20962adc55018bd4e79e3f1a762e8abd05dd Mon Sep 17 00:00:00 2001
From: Mike Lang <mike.lang1010@gmail.com>
Date: Tue, 10 Oct 2017 17:01:10 -0500
Subject: [PATCH 2/2] Try making a named query look like a select

---
 lib/shim/datastore-shim.js | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/shim/datastore-shim.js b/lib/shim/datastore-shim.js
index 94d36b4608..508a87d765 100644
--- a/lib/shim/datastore-shim.js
+++ b/lib/shim/datastore-shim.js
@@ -728,7 +728,7 @@ function _recordQuery(suffix, nodule, properties, querySpec) {
     // If the intstrumentation provides a name use it, otherwise construct a name
     // from the parsed query string
     var name = queryName ? 
-      queryName : 
+      queryName + '/' + 'select' + suffix : 
       (parsed.collection || 'other') + '/' + parsed.operation + suffix
 
     // Return the segment descriptor.