From 36b58f4dc08886e4d4ab2a72a3231197ff496739 Mon Sep 17 00:00:00 2001
From: Steven Bal <steven@maykinmedia.nl>
Date: Mon, 14 Jul 2025 16:45:46 +0200
Subject: [PATCH] :bug: Workaround to fix APM for psycopg pooling

---
 src/objects/conf/base.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/objects/conf/base.py b/src/objects/conf/base.py
index c0114829..eb060307 100644
--- a/src/objects/conf/base.py
+++ b/src/objects/conf/base.py
@@ -128,6 +128,19 @@
 
 
 if DB_POOL_ENABLED:
+    # FIXME Workaround for https://github.com/elastic/apm-agent-python/issues/2094
+    # apm-agent-python does not instrument ConnectionPool yet
+    import psycopg
+
+    class WrapperConnectionClass(psycopg.Connection):
+        @classmethod
+        def connect(
+            cls,
+            conninfo: str = "",
+            **kwargs,
+        ) -> "psycopg.Connection[psycopg.rows.TupleRow]":
+            return psycopg.connect(conninfo, **kwargs)
+
     DATABASES["default"]["OPTIONS"] = {
         "pool": {
             "min_size": DB_POOL_MIN_SIZE,
@@ -138,6 +151,7 @@
             "max_idle": DB_POOL_MAX_IDLE,
             "reconnect_timeout": DB_POOL_RECONNECT_TIMEOUT,
             "num_workers": DB_POOL_NUM_WORKERS,
+            "connection_class": WrapperConnectionClass,
         }
     }