From aa90a59c9a3e56321ed789d6c9120c1c455b85ae Mon Sep 17 00:00:00 2001
From: Holly Gong <39108850+hogo6002@users.noreply.github.com>
Date: Mon, 28 Oct 2024 15:51:58 +1100
Subject: [PATCH] feat(api): log queries for test instance (#2796)

We're seeing some 502 and 504 errors reported from API queries. Without
the details of each query, it's hard to reproduce and find the root
causes.
Adds logs for the test instance. It does not record actual user queries,
only queries from our load tests.
---
 gcp/api/server.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/gcp/api/server.py b/gcp/api/server.py
index d816dace371..61f0577026c 100644
--- a/gcp/api/server.py
+++ b/gcp/api/server.py
@@ -97,6 +97,8 @@
 # Prefix for the
 _TAG_PREFIX = "refs/tags/"
 
+_TEST_INSTANCE = 'oss-vdb-test'
+
 # ----
 # Type Aliases:
 
@@ -200,6 +202,10 @@ def QueryAffected(self, request, context: grpc.ServicerContext):
     else:
       logging.info('QueryAffected for %s', qtype)
 
+    # Log queries for test instance.
+    # This is for debugging purposes. Production queries will not be recorded.
+    if get_gcp_project() == _TEST_INSTANCE:
+      logging.info('Query: %s', request.query)
     try:
       page_token = QueryCursor.from_page_token(request.query.page_token)
     except ValueError as e:
@@ -284,6 +290,10 @@ def QueryAffectedBatch(self, request, context: grpc.ServicerContext):
         extra={'json_fields': {
             'details': query_details
         }})
+    # Log queries for test instance.
+    # This is for debugging purposes. Production queries will not be recorded.
+    if get_gcp_project() == _TEST_INSTANCE:
+      logging.info('Batch query: %s', request.query)
 
     if len(request.query.queries) > _MAX_BATCH_QUERY:
       context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'Too many queries.')