Merge pull request #7 from permitio/rk/pg-conn-pool

Postgres Backend: Acquire connections from pool

.github/workflows/test-suite.yml

@@ -15,7 +15,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11-dev"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
 
     services:
       zookeeper:

broadcaster/__init__.py

@@ -1,4 +1,4 @@
 from ._base import Broadcast, Event
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 __all__ = ["Broadcast", "Event"]

broadcaster/_backends/postgres.py

@@ -2,22 +2,38 @@
 from typing import Any
 
 import asyncpg
+import os
 
 from .._base import Event
 from .base import BroadcastBackend
 
+try:
+    POOL_MAX_SIZE = int(os.getenv("BROADCASTER_PG_MAX_POOL_SIZE"))
+except TypeError:
+    POOL_MAX_SIZE = 10
 
 class PostgresBackend(BroadcastBackend):
+    _pools = {}
+    _pools_lock = asyncio.Lock()
+
     def __init__(self, url: str):
         self._url = url
 
+    async def _get_pool(self):
+        async with self.__class__._pools_lock:
+            if self._url not in self.__class__._pools:
+                self.__class__._pools[self._url] = await asyncpg.create_pool(self._url, max_size=POOL_MAX_SIZE)
+            return self.__class__._pools[self._url]
+
     async def connect(self) -> None:
-        self._conn = await asyncpg.connect(self._url)
+        self._conn = await (await self._get_pool()).acquire()
         self._listen_queue: asyncio.Queue = asyncio.Queue()
         self._conn.add_termination_listener(self._termination_listener)
 
     async def disconnect(self) -> None:
-        await self._conn.close()
+        self._conn.remove_termination_listener(self._termination_listener)
+        await (await self._get_pool()).release(self._conn)
+        self._conn = None
 
     async def subscribe(self, channel: str) -> None:
         await self._conn.add_listener(channel, self._listener)