Move getCurrentTraceFromAgent out of logging package

googleapis · Sep 1, 2017 · cc052c7 · cc052c7
1 parent 0e13f5a
commit cc052c7
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 38 deletions.
diff --git a/packages/logging-bunyan/src/index.js b/packages/logging-bunyan/src/index.js
@@ -186,6 +186,30 @@ LoggingBunyan.prototype.formatEntry_ = function(record) {
   return this.log_.entry(entryMetadata, record);
 };
 
+/**
+ * Gets the current fully qualified trace ID when available from the
+ * @google-cloud/trace-agent library in the LogEntry.trace field format of:
+ * "projects/[PROJECT-ID]/traces/[TRACE-ID]".
+ */
+function getCurrentTraceFromAgent() {
+  var agent = global._google_trace_agent;
+  if (!agent || !agent.getCurrentContextId || !agent.getWriterProjectId) {
+    return null;
+  }
+
+  var traceId = agent.getCurrentContextId();
+  if (!traceId) {
+    return null;
+  }
+
+  var traceProjectId = agent.getWriterProjectId();
+  if (!traceProjectId) {
+    return null;
+  }
+
+  return `projects/${traceProjectId}/traces/${traceId}`;
+}
+
 /**
  * Intercept log entries as they are written so we can attempt to add the trace
  * ID in the same continuation as the function that wrote the log, because the
@@ -197,7 +221,7 @@ LoggingBunyan.prototype.formatEntry_ = function(record) {
 LoggingBunyan.prototype.write = function(record, encoding, callback) {
   record = extend({}, record);
   if (!record[LOGGING_TRACE_KEY]) {
-    var trace = logging.getCurrentTraceFromAgent();
+    var trace = getCurrentTraceFromAgent();
     if (trace) {
       record[LOGGING_TRACE_KEY] = trace;
     }

diff --git a/packages/logging-bunyan/test/index.js b/packages/logging-bunyan/test/index.js
@@ -298,6 +298,7 @@ describe('logging-bunyan', function() {
     });
 
     it('should leave prefixed trace property as is if set', function(done) {
+      var oldTraceAgent = global._google_trace_agent;
       global._google_trace_agent = {
         getCurrentContextId: function() { return 'trace-from-agent'; },
         getWriterProjectId: function() { return 'project1'; }
@@ -315,9 +316,38 @@ describe('logging-bunyan', function() {
       };
 
       loggingBunyan.write(recordWithTraceAlreadySet, '', assert.ifError);
+
+      global._google_trace_agent = oldTraceAgent;
     });
   });
 
+  it('should not set prefixed trace property if trace unavailable', function() {
+    FakeWritable.prototype.write = function(record, encoding, callback) {
+      assert.deepStrictEqual(record, RECORD);
+      assert.strictEqual(encoding, '');
+      assert.strictEqual(callback, assert.ifError);
+      assert.strictEqual(this, loggingBunyan);
+    };
+    var oldTraceAgent = global._google_trace_agent;
+
+    global._google_trace_agent = {};
+    loggingBunyan.write(RECORD, '', assert.ifError);
+
+    global._google_trace_agent = {
+      getCurrentContextId: function() { return null; },
+      getWriterProjectId: function() { return null; }
+    };
+    loggingBunyan.write(RECORD, '', assert.ifError);
+
+    global._google_trace_agent = {
+      getCurrentContextId: function() { return null; },
+      getWriterProjectId: function() { return 'project1'; }
+    };
+    loggingBunyan.write(RECORD, '', assert.ifError);
+
+    global._google_trace_agent = oldTraceAgent;
+  });
+
   describe('_write', function() {
     beforeEach(function() {
       fakeLogInstance.entry = function() {};

diff --git a/packages/logging-winston/src/index.js b/packages/logging-winston/src/index.js
@@ -130,6 +130,30 @@ function LoggingWinston(options) {
 winston.transports.StackdriverLogging = LoggingWinston;
 util.inherits(LoggingWinston, winston.Transport);
 
+/**
+ * Gets the current fully qualified trace ID when available from the
+ * @google-cloud/trace-agent library in the LogEntry.trace field format of:
+ * "projects/[PROJECT-ID]/traces/[TRACE-ID]".
+ */
+function getCurrentTraceFromAgent() {
+  var agent = global._google_trace_agent;
+  if (!agent || !agent.getCurrentContextId || !agent.getWriterProjectId) {
+    return null;
+  }
+
+  var traceId = agent.getCurrentContextId();
+  if (!traceId) {
+    return null;
+  }
+
+  var traceProjectId = agent.getWriterProjectId();
+  if (!traceProjectId) {
+    return null;
+  }
+
+  return `projects/${traceProjectId}/traces/${traceId}`;
+}
+
 /**
  * Relay a log entry to the logging agent. This is normally called by winston.
  *
@@ -207,7 +231,7 @@ LoggingWinston.prototype.log = function(levelName, msg, metadata, callback) {
     entryMetadata.trace = metadata[LOGGING_TRACE_KEY];
     delete data.metadata[LOGGING_TRACE_KEY];
   } else {
-    var trace = logging.getCurrentTraceFromAgent();
+    var trace = getCurrentTraceFromAgent();
     if (trace) {
       entryMetadata.trace = trace;
     }

diff --git a/packages/logging-winston/test/index.js b/packages/logging-winston/test/index.js
@@ -350,6 +350,33 @@ describe('logging-winston', function() {
       global._google_trace_agent = oldTraceAgent;
     });
 
+    it('should leave out trace metadata if trace unavailable', function() {
+      loggingWinston.log_.entry = function(entryMetadata, data) {
+        assert.deepStrictEqual(entryMetadata, {
+          resource: loggingWinston.resource_,
+        });
+      };
+
+      var oldTraceAgent = global._google_trace_agent;
+
+      global._google_trace_agent = {};
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = {
+        getCurrentContextId: function() { return null; },
+        getWriterProjectId: function() { return null; }
+      };
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = {
+        getCurrentContextId: function() { return null; },
+        getWriterProjectId: function() { return 'project1'; }
+      };
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = oldTraceAgent;
+    });
+
     it('should write to the log', function(done) {
       var entry = {};
 

diff --git a/packages/logging/src/index.js b/packages/logging/src/index.js
@@ -773,41 +773,6 @@ Logging.prototype.setAclForTopic_ = function(name, config, callback) {
   });
 };
 
-/**
- * Attempts to get the current fully qualified trace ID from the
- * @google-cloud/trace-agent (if installed) in the format
- * "projects/[PROJECT-ID]/traces/[TRACE-ID]".
- *
- * That is the specified format to set the LogEntry.trace field so that the
- * Cloud Console logs and trace viewer can correlate log entries and associate
- * them with traces.
- *
- * This returns null if the trace agent is not installed, there is no available
- * trace context or no trace writer project ID.
- */
-function getCurrentTraceFromAgent() {
-  var traceAgent = global._google_trace_agent;
-  if (!traceAgent || !traceAgent.getCurrentContextId ||
-      !traceAgent.getWriterProjectId) {
-    // Trace agent is not installed or is an older version
-    return null;
-  }
-
-  var traceId = traceAgent.getCurrentContextId();
-  if (!traceId) {
-    // No trace context because request not sampled or context lost
-    return null;
-  }
-
-  var traceProjectId = traceAgent.getWriterProjectId();
-  if (!traceProjectId) {
-    // No project, possibly because async project auto-discovery not yet done
-    return null;
-  }
-
-  return `projects/${traceProjectId}/traces/${traceId}`;
-}
-
 /*! Developer Documentation
  *
  * These methods can be auto-paginated.
@@ -832,7 +797,6 @@ Logging.Entry = Entry;
 Logging.Log = Log;
 Logging.Logging = Logging;
 Logging.Sink = Sink;
-Logging.getCurrentTraceFromAgent = getCurrentTraceFromAgent;
 
 module.exports = Logging;
 module.exports.v2 = v2;