From 8a046875174871e0bde8a82a0f2d0c1d4398f72e Mon Sep 17 00:00:00 2001
From: James Robinson <james.robinson@noteable.io>
Date: Sun, 22 Jan 2023 16:11:54 -0500
Subject: [PATCH 1/2] Manage more than a single asyncpg pool within app.state.

---
 fastapi_asyncpg/__init__.py | 41 ++++++++++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/fastapi_asyncpg/__init__.py b/fastapi_asyncpg/__init__.py
index f569094..958dd6b 100644
--- a/fastapi_asyncpg/__init__.py
+++ b/fastapi_asyncpg/__init__.py
@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 from fastapi import FastAPI
 
 import asyncpg
@@ -46,19 +45,20 @@ async def on_connect(self):
         the db"""
         # if the pool is comming from outside (tests), don't connect it
         if self._pool:
-            self.app.state.pool = self._pool
+            self._get_pool_manager_from_app().put(self.dsn, self._pool)
             return
         pool = await asyncpg.create_pool(dsn=self.dsn, **self.con_opts)
         async with pool.acquire() as db:
             await self.init_db(db)
-        self.app.state.pool = pool
+
+        self._get_pool_manager_from_app().put(self.dsn, pool)
 
     async def on_disconnect(self):
         # if the pool is comming from outside, don't desconnect it
         # someone else will do (usualy a pytest fixture)
         if self._pool:
             return
-        await self.app.state.pool.close()
+        await self.pool.close()
 
     def on_init(self, func):
         self.init_db = func
@@ -66,7 +66,10 @@ def on_init(self, func):
 
     @property
     def pool(self):
-        return self.app.state.pool
+        """Fetch the connection pool associated with our DSN from the
+        pool manager stashed within app.state.
+        """
+        return self._get_pool_manager_from_app().get(self.dsn)
 
     async def connection(self):
         """
@@ -103,9 +106,37 @@ async def get_content(db = Depens(db.transaction)):
             else:
                 await txn.commit()
 
+    def _get_pool_manager_from_app(self):
+        """Find or create singleton AppPoolManager instance within self.app.state"""
+        if not (
+            pool_manager := getattr(self.app.state, "fastapi_asyncpg_pool_manager")
+        ):
+            self.app.state.fastapi_asyncpg_pool_manager = AppPoolManager()
+
+        return pool_manager
+
     atomic = transaction
 
 
+class AppPoolManager:
+    """Object placed into fastapi app.state to manage one or more
+    asyncpg.pool.Pool instances within the fastapi app.
+
+    If the app uses more than one asyncpg database, then there
+    will be more than one pool. We separate them by the
+    connection DSN.
+    """
+
+    def __init__(self):
+        self._pool_by_dsn = {}
+
+    def get(self, dsn):
+        return self._pool_by_dsn[dsn]
+
+    def put(self, dsn, pool):
+        self._pool_by_dsn[dsn] = pool
+
+
 class SingleConnectionTestingPool:
     """A fake pool that simulates pooling, but runs on
     a single transaction that it's rolled back after

From cca9aaf6f2b16293995156d54c81e63a2c6fe9f5 Mon Sep 17 00:00:00 2001
From: James Robinson <james.robinson@noteable.io>
Date: Sun, 22 Jan 2023 16:18:56 -0500
Subject: [PATCH 2/2] Avoid getattr on self.app.state.

---
 fastapi_asyncpg/__init__.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/fastapi_asyncpg/__init__.py b/fastapi_asyncpg/__init__.py
index 958dd6b..c19060e 100644
--- a/fastapi_asyncpg/__init__.py
+++ b/fastapi_asyncpg/__init__.py
@@ -108,12 +108,10 @@ async def get_content(db = Depens(db.transaction)):
 
     def _get_pool_manager_from_app(self):
         """Find or create singleton AppPoolManager instance within self.app.state"""
-        if not (
-            pool_manager := getattr(self.app.state, "fastapi_asyncpg_pool_manager")
-        ):
+        if not hasattr(self.app.state, "fastapi_asyncpg_pool_manager"):
             self.app.state.fastapi_asyncpg_pool_manager = AppPoolManager()
 
-        return pool_manager
+        return self.app.state.fastapi_asyncpg_pool_manager
 
     atomic = transaction