From 38eaafcb2c0e60184633c5df1cb12e365b0c85ae Mon Sep 17 00:00:00 2001
From: Ryan Barrett <git@ryanb.org>
Date: Sat, 6 Mar 2021 11:31:36 -0800
Subject: [PATCH] appengine_config: logging: manually find and inject Google
 Cloud trace id

makes log lines attach to their HTTP request and show grouped in cloud console log browser. google-cloud-logging used to do this automatically for webapp2, based on the X-Cloud-Trace-Context header, but it dropped webapp2 as of google-cloud-logging 2.0.

background:
https://github.com/googleapis/python-logging/issues/110#issuecomment-745534629
https://github.com/googleapis/python-logging/issues/149#issuecomment-782693201

also note that AppEngineHandler is evidently deprecated, so I may need to port that whole class into webutil eventually. :(
https://github.com/googleapis/python-logging/issues/202
---
 appengine_config.py | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/appengine_config.py b/appengine_config.py
index 0d093aa..56c03a5 100644
--- a/appengine_config.py
+++ b/appengine_config.py
@@ -1,7 +1,7 @@
 """App Engine config. dev_appserver vs prod, logging, Google API clients, etc."""
 import os
 
-from .appengine_info import DEBUG
+from .appengine_info import APP_ID, DEBUG
 
 # Use lxml for BeautifulSoup explicitly.
 from . import util
@@ -70,11 +70,32 @@
 try:
   import google.cloud.logging
   logging_client = google.cloud.logging.Client()
+
   if not DEBUG:
-    # https://stackoverflow.com/a/58296028/186123
-    # https://googleapis.dev/python/logging/latest/usage.html#cloud-logging-handler
     from google.cloud.logging.handlers import AppEngineHandler, setup_logging
-    setup_logging(AppEngineHandler(logging_client, name='stdout'),
+
+    class Webapp2TraceHandler(AppEngineHandler):
+      """Log handler that adds trace id based on webapp2 request header.
+
+      https://github.com/googleapis/python-logging/issues/110#issuecomment-745534629
+      https://github.com/googleapis/python-logging/issues/149#issuecomment-782693201
+
+      Also note that AppEngineHandler is evidently deprecated, so I may need to
+      port that whole class into webutil eventually. :(
+      https://github.com/googleapis/python-logging/issues/202
+      """
+      def emit(self, record):
+        try:
+          import webapp2
+          trace = webapp2.get_request().headers.get('X-Cloud-Trace-Context')
+          if trace:
+            trace_id = trace.split('/', 1)[0]
+            record.trace = f'projects/{APP_ID}/traces/{trace_id}'
+        except (ImportError, AssertionError):
+          pass
+        return super().emit(record)
+
+    setup_logging(Webapp2TraceHandler(logging_client, name='stdout'),
                   log_level=logging.DEBUG)
     # this currently occasionally hits the 256KB stackdriver logging limit and
     # crashes in the background service. i've tried batch_size=1 and