Merge pull request #12 from aiokitchen/fix/hasql-metrics-acquire-time

Fix acquire_time metrics

hasql/aiopg_sa.py

@@ -32,7 +32,9 @@ def _driver_metrics(self) -> Sequence[DriverMetrics]:
                 idle=p.freesize,
                 used=p.size - p.freesize,
                 host=parse_dsn(p.dsn).get("host", ""),
-            ) for p in self.pools
+            )
+            for p in self.pools
+            if p
         ]
 
 

hasql/asyncpg.py

@@ -70,7 +70,9 @@ def _driver_metrics(self) -> Sequence[DriverMetrics]:
                 idle=self.get_pool_freesize(p),
                 used=p._maxsize - self.get_pool_freesize(p),
                 host=self.host(p),
-            ) for p in self.pools
+            )
+            for p in self.pools
+            if p
         ]
 
 

hasql/asyncsqlalchemy.py

@@ -62,7 +62,9 @@ def _driver_metrics(self) -> Sequence[DriverMetrics]:
                 idle=p.sync_engine.pool.checkedin(),
                 used=p.sync_engine.pool.checkedout(),
                 host=p.sync_engine.url.host,
-            ) for p in self.pools
+            )
+            for p in self.pools
+            if p
         ]
 
 

hasql/base.py

@@ -73,7 +73,7 @@ async def execute():
                     master_as_replica_weight=self.master_as_replica_weight,
                 )
 
-            with self.metrics.with_acquire():
+            with self.metrics.with_acquire(self.pool_manager.host(self.pool)):
                 return await self.pool_manager.acquire_from_pool(
                     self.pool, **self.kwargs,
                 )
@@ -91,7 +91,7 @@ async def go():
                     fallback_master=self.fallback_master,
                     master_as_replica_weight=self.master_as_replica_weight,
                 )
-            with self.metrics.with_acquire():
+            with self.metrics.with_acquire(self.pool_manager.host(self.pool)):
                 self.context = self.pool_manager.acquire_from_pool(
                     self.pool,
                     **self.kwargs,

hasql/metrics.py

@@ -1,4 +1,5 @@
 import time
+from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Dict, Sequence
@@ -17,18 +18,20 @@ class DriverMetrics:
 class HasqlMetrics:
     pool: int
     pool_time: float
-    acquire: int
-    acquire_time: float
+    acquire: Dict[str, int]
+    acquire_time: Dict[str, float]
     add_connections: Dict[str, int]
     remove_connections: Dict[str, int]
 
 
 @dataclass
 class CalculateMetrics:
     _pool: int = 0
-    _pool_time: int = 0
-    _acquire: int = 0
-    _acquire_time: int = 0
+    _pool_time: float = 0.
+    _acquire: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
+    _acquire_time: Dict[str, float] = field(
+        default_factory=lambda: defaultdict(int)
+    )
     _add_connections: Dict[str, int] = field(default_factory=dict)
     _remove_connections: Dict[str, int] = field(default_factory=dict)
 
@@ -37,7 +40,7 @@ def metrics(self) -> HasqlMetrics:
             pool=self._pool,
             pool_time=self._pool_time,
             acquire=self._acquire,
-            acquire_time=self._acquire,
+            acquire_time=self._acquire_time,
             add_connections=self._add_connections,
             remove_connections=self._remove_connections,
         )
@@ -50,11 +53,11 @@ def with_get_pool(self):
         self._pool_time += time.monotonic() - tt
 
     @contextmanager
-    def with_acquire(self):
-        self._acquire += 1
+    def with_acquire(self, pool: str):
+        self._acquire[pool] += 1
         tt = time.monotonic()
         yield
-        self._acquire_time += time.monotonic() - tt
+        self._acquire_time[pool] += time.monotonic() - tt
 
     def add_connection(self, dsn: str):
         self._add_connections[dsn] = (

hasql/psycopg3.py

@@ -97,7 +97,9 @@ def _driver_metrics(self) -> Sequence[DriverMetrics]:
             {
                 **p.get_stats(),
                 "host": self.host(p)
-            } for p in self.pools
+            }
+            for p in self.pools
+            if p
         ]
         return [
             DriverMetrics(

tests/test_metrics.py

@@ -27,7 +27,7 @@ async def test_hasql_context_metrics(pool_manager_factory, pg_dsn):
             metrics = pool_manager.metrics().hasql
             assert metrics == HasqlMetrics(
                 pool=1,
-                acquire=1,
+                acquire={pool_manager.host(pool_manager.pools[0]): 1},
                 pool_time=mock.ANY,
                 acquire_time=mock.ANY,
                 add_connections=mock.ANY,
@@ -39,7 +39,7 @@ async def test_hasql_context_metrics(pool_manager_factory, pg_dsn):
         metrics = pool_manager.metrics().hasql
         assert metrics == HasqlMetrics(
             pool=1,
-            acquire=1,
+            acquire={pool_manager.host(pool_manager.pools[0]): 1},
             pool_time=mock.ANY,
             acquire_time=mock.ANY,
             add_connections=mock.ANY,
@@ -65,7 +65,7 @@ async def test_hasql_metrics(pool_manager_factory, pg_dsn):
         metrics = pool_manager.metrics().hasql
         assert metrics == HasqlMetrics(
             pool=1,
-            acquire=1,
+            acquire={pool_manager.host(pool_manager.pools[0]): 1},
             pool_time=mock.ANY,
             acquire_time=mock.ANY,
             add_connections=mock.ANY,
@@ -79,7 +79,7 @@ async def test_hasql_metrics(pool_manager_factory, pg_dsn):
         metrics = pool_manager.metrics().hasql
         assert metrics == HasqlMetrics(
             pool=1,
-            acquire=1,
+            acquire={pool_manager.host(pool_manager.pools[0]): 1},
             pool_time=mock.ANY,
             acquire_time=mock.ANY,
             add_connections=mock.ANY,
@@ -107,7 +107,7 @@ async def test_hasql_close_metrics(pool_manager_factory, pg_dsn):
         metrics = pool_manager.metrics().hasql
         assert metrics == HasqlMetrics(
             pool=1,
-            acquire=1,
+            acquire={pool_manager.host(pool_manager.pools[0]): 1},
             pool_time=mock.ANY,
             acquire_time=mock.ANY,
             add_connections=mock.ANY,