From ca84959c281b3d0fa04a5a3d77a9add7cf62b399 Mon Sep 17 00:00:00 2001
From: Matthew Loring <matthewloring@users.noreply.github.com>
Date: Thu, 12 Jan 2017 13:57:02 -0800
Subject: [PATCH] Ensure agent is always set for mysql/grpc (#343)

If either of these modules switch to lazily importing any of our patched
files, we would fail to set the agent in some cases.

PR-URL: https://github.com/GoogleCloudPlatform/cloud-trace-nodejs/pull/343
---
 src/hooks/userspace/hook-grpc.js  | 7 ++++++-
 src/hooks/userspace/hook-mysql.js | 7 ++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/hooks/userspace/hook-grpc.js b/src/hooks/userspace/hook-grpc.js
index 98f9d6395..699dfb09a 100644
--- a/src/hooks/userspace/hook-grpc.js
+++ b/src/hooks/userspace/hook-grpc.js
@@ -431,7 +431,9 @@ module.exports = function(version_, agent_) {
   return {
     'src/node/src/client.js': {
       patch: function(client) {
-        agent = agent_;
+        if (!agent) {
+          agent = agent_;
+        }
         shimmer.wrap(client, 'makeClientConstructor',
             makeClientConstructorWrap);
       },
@@ -445,6 +447,9 @@ module.exports = function(version_, agent_) {
     },
     'src/node/src/server.js': {
       patch: function(server) {
+        if (!agent) {
+          agent = agent_;
+        }
         shimmer.wrap(server.Server.prototype, 'register', serverRegisterWrap);
       },
       unpatch: function(server) {
diff --git a/src/hooks/userspace/hook-mysql.js b/src/hooks/userspace/hook-mysql.js
index 85859c365..e52164989 100644
--- a/src/hooks/userspace/hook-mysql.js
+++ b/src/hooks/userspace/hook-mysql.js
@@ -86,7 +86,9 @@ module.exports = function(version_, agent_) {
   return {
     'lib/Connection.js': {
       patch: function(Connection) {
-        agent = agent_;
+        if (!agent) {
+          agent = agent_;
+        }
         shimmer.wrap(Connection, 'createQuery', createQueryWrap);
       },
       unpatch: function(Connection) {
@@ -96,6 +98,9 @@ module.exports = function(version_, agent_) {
     },
     'lib/Pool.js': {
       patch: function(Pool) {
+        if (!agent) {
+          agent = agent_;
+        }
         shimmer.wrap(Pool.prototype, 'getConnection', wrapGetConnection);
       },
       unpatch: function(Pool) {