Inject tenant in all* logging records (#7033)

edb/common/log.py

@@ -0,0 +1,57 @@
+#
+# This source file is part of the EdgeDB open source project.
+#
+# Copyright 2024-present MagicStack Inc. and the EdgeDB authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+from __future__ import annotations
+
+# DON'T IMPORT asyncio or any package that creates their own logger here,
+# or the "tenant" value cannot be injected.
+import contextvars
+import logging
+
+
+current_tenant = contextvars.ContextVar("current_tenant", default="-")
+
+
+class EdgeDBLogger(logging.Logger):
+
+    def makeRecord(
+        self,
+        name,
+        level,
+        fn,
+        lno,
+        msg,
+        args,
+        exc_info,
+        func=None,
+        extra=None,
+        sinfo=None,
+    ):
+        # Unlike the standard Logger class, we allow overwriting
+        # all attributes of the log record with stuff from *extra*.
+        factory = logging.getLogRecordFactory()
+        rv = factory(name, level, fn, lno, msg, args, exc_info, func, sinfo)
+        rv.__dict__["tenant"] = current_tenant.get()
+        if extra is not None:
+            rv.__dict__.update(extra)
+        return rv
+
+
+def early_setup():
+    logging.setLoggerClass(EdgeDBLogger)

edb/server/logsetup.py

@@ -124,46 +124,18 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         fmt = EdgeDBLogFormatter(
-            '{levelname} {process} {asctime} {name}: {message}',
+            '{levelname} {process} {tenant} {asctime} {name}: {message}',
             style='{')
 
         self.setFormatter(fmt)
 
 
-class EdgeDBLogger(logging.Logger):
-
-    def makeRecord(
-        self,
-        name,
-        level,
-        fn,
-        lno,
-        msg,
-        args,
-        exc_info,
-        func=None,
-        extra=None,
-        sinfo=None,
-    ):
-        # Unlike the standard Logger class, we allow overwriting
-        # all attributes of the log record with stuff from *extra*.
-        factory = logging.getLogRecordFactory()
-        rv = factory(name, level, fn, lno, msg, args, exc_info, func, sinfo)
-        if extra is not None:
-            rv.__dict__.update(extra)
-        return rv
-
-
 IGNORE_DEPRECATIONS_IN = {
     'graphql',
     'promise',
 }
 
 
-def early_setup():
-    logging.setLoggerClass(EdgeDBLogger)
-
-
 def setup_logging(log_level, log_destination):
     log_level = log_level.upper()
     try:
@@ -179,7 +151,7 @@ def setup_logging(log_level, log_destination):
 
     if log_destination == 'syslog':
         fmt = logging.Formatter(
-            '{processName}[{process}]: {name}: {message}',
+            '{processName}[{process}]: {tenant}: {name}: {message}',
             style='{')
         handler = logging.handlers.SysLogHandler(
             '/dev/log',
@@ -191,7 +163,7 @@ def setup_logging(log_level, log_destination):
 
     else:
         fmt = logging.Formatter(
-            '{levelname} {process} {asctime} {name}: {message}',
+            '{levelname} {process} {tenant} {asctime} {name}: {message}',
             style='{')
         handler = logging.FileHandler(log_destination)
         handler.setFormatter(fmt)

edb/server/main.py

@@ -30,6 +30,10 @@
     TYPE_CHECKING,
 )
 
+from edb.common.log import early_setup
+# ruff: noqa: E402
+early_setup()
+
 import asyncio
 import contextlib
 import dataclasses
@@ -48,10 +52,6 @@
 import setproctitle
 import uvloop
 
-from . import logsetup
-# ruff: noqa: E402
-logsetup.early_setup()
-
 from edb import buildmeta
 from edb.ir import statypes
 from edb.common import exceptions
@@ -64,6 +64,7 @@
 from . import compiler as edbcompiler
 from . import daemon
 from . import defines
+from . import logsetup
 from . import pgconnparams
 from . import pgcluster
 from . import service_manager

edb/server/multitenant.py

@@ -35,6 +35,7 @@
 from edb import errors
 from edb.common import retryloop
 from edb.common import signalctl
+from edb.common.log import current_tenant
 from edb.pgsql import params as pgparams
 from edb.server import compiler as edbcompiler
 
@@ -278,6 +279,7 @@ def _warn(e):
             )
 
         async def _add_tenant():
+            current_tenant.set(conf["instance-name"])
             rloop = retryloop.RetryLoop(
                 backoff=retryloop.exp_backoff(),
                 timeout=300,
@@ -313,6 +315,7 @@ async def _remove_tenant(self, serial: int, sni: str):
                 if serial > self._tenants_serial.get(sni, 0):
                     if sni in self._tenants:
                         tenant = self._tenants.pop(sni)
+                        current_tenant.set(tenant.get_instance_name())
                         await self._destroy_tenant(tenant)
                         logger.info("Removed Tenant %s", sni)
                     self._tenants_serial[sni] = serial
@@ -324,6 +327,7 @@ async def _reload_tenant(self, serial: int, sni: str, conf: TenantConfig):
             async with self._tenants_lock[sni]:
                 if serial > self._tenants_serial.get(sni, 0):
                     if tenant := self._tenants.get(sni):
+                        current_tenant.set(tenant.get_instance_name())
                         tenant.set_reloadable_files(
                             readiness_state_file=conf.get(
                                 "readiness-state-file"),

edb/server/protocol/pg_ext.pyx

@@ -36,6 +36,7 @@ from libc.stdint cimport int32_t, int16_t, uint32_t
 
 from edb import errors
 from edb.common import debug
+from edb.common.log import current_tenant
 from edb.pgsql.parser import exceptions as parser_errors
 from edb.server import args as srvargs
 from edb.server import defines, metrics
@@ -493,6 +494,8 @@ cdef class PgConnection(frontend.FrontendConnection):
                     self.tenant = self.server.retrieve_tenant(
                         self._transport.get_extra_info("ssl_object")
                     )
+                    if self.tenant is not None:
+                        current_tenant.set(self.tenant.get_instance_name())
                     self.is_tls = True
 
                 elif proto_ver_minor == 5680:  # GSSENCRequest

edb/server/protocol/protocol.pyx

@@ -34,6 +34,7 @@ import httptools
 from edb import errors
 from edb.common import debug
 from edb.common import markup
+from edb.common.log import current_tenant
 
 from edb.graphql import extension as graphql_ext
 
@@ -405,6 +406,8 @@ cdef class HttpProtocol:
         )
         sslobj = self.transport.get_extra_info('ssl_object')
         self.tenant = self.server.retrieve_tenant(sslobj)
+        if self.tenant is not None:
+            current_tenant.set(self.get_tenant_label())
         if sslobj.selected_alpn_protocol() == 'edgedb-binary':
             self._switch_to_binary_protocol()
         else:

edb/server/server.py

@@ -55,6 +55,7 @@
 from edb.common import lru
 from edb.common import secretkey
 from edb.common import windowedsum
+from edb.common.log import current_tenant
 
 from edb.schema import reflection as s_refl
 from edb.schema import schema as s_schema
@@ -462,9 +463,9 @@ def _idle_gc_collector(self):
             for conn in self._binary_conns:
                 try:
                     if conn.is_idle(expiry_time):
-                        metrics.idle_client_connections.inc(
-                            1.0, conn.get_tenant_label()
-                        )
+                        label = conn.get_tenant_label()
+                        metrics.idle_client_connections.inc(1.0, label)
+                        current_tenant.set(label)
                         conn.close_for_idling()
                     elif conn.is_alive():
                         # We are sorting connections in

edb/server/tenant.py

@@ -48,6 +48,7 @@
 
 from edb import errors
 from edb.common import retryloop
+from edb.common.log import current_tenant
 
 from . import args as srvargs
 from . import config
@@ -441,6 +442,7 @@ def create_task(
         # Therefore, it is an error trying to create a task while the server is
         # not expecting one, so always couple the call with an additional check
         if self._accept_new_tasks and self._task_group is not None:
+            current_tenant.set(self.get_instance_name())
             if interruptable:
                 rv = self.__loop.create_task(coro)
             else: