Switch from single connection to pool.

Author: andrew-chang-dewitt

db_wrapper/client/async_client.py

@@ -12,10 +12,10 @@
     Dict)
 
 import aiopg
-from psycopg2.extras import register_uuid, RealDictRow
+from psycopg2.extras import register_uuid, RealDictCursor, RealDictRow  # type: ignore
 from psycopg2 import sql
 
-from db_wrapper.connection import ConnectionParameters, connect
+from db_wrapper.connection import ConnectionParameters, get_pool
 
 # add uuid support to psycopg2 & Postgres
 register_uuid()
@@ -33,18 +33,19 @@ class AsyncClient:
     """
 
     _connection_params: ConnectionParameters
-    _connection: aiopg.Connection
+    _pool: aiopg.Pool
 
     def __init__(self, connection_params: ConnectionParameters) -> None:
         self._connection_params = connection_params
 
     async def connect(self) -> None:
-        """Connect to the database."""
-        self._connection = await connect(self._connection_params)
+        """Create a database connection pool."""
+        self._pool = await get_pool(self._connection_params)
 
     async def disconnect(self) -> None:
-        """Disconnect from the database."""
-        await self._connection.close()
+        """Close database connection pool."""
+        self._pool.close()
+        await self._pool.wait_closed()
 
     # PENDS python 3.9 support in pylint
     # pylint: disable=unsubscriptable-object
@@ -81,7 +82,7 @@ async def execute(
         Returns:
             None
         """
-        async with self._connection.cursor() as cursor:
+        with (await self._pool.cursor(cursor_factory=RealDictCursor) ) as cursor:
             await self._execute_query(cursor, query, params)
 
     # PENDS python 3.9 support in pylint
@@ -101,7 +102,7 @@ async def execute_and_return(
         Returns:
             List containing all the rows that matched the query.
         """
-        async with self._connection.cursor() as cursor:
+        with (await self._pool.cursor(cursor_factory=RealDictCursor) ) as cursor:
             await self._execute_query(cursor, query, params)
 
             result: List[RealDictRow] = await cursor.fetchall()

db_wrapper/connection.py

@@ -41,18 +41,18 @@ async def _try_connect(
     dsn = f"dbname={database} user={user} password={password} " \
           f"host={host} port={port}"
 
+    # return await aiopg.create_pool(dsn)
+
     # PENDS python 3.9 support in pylint
     # pylint: disable=unsubscriptable-object
-    connection: Optional[aiopg.Connection] = None
+    pool: Optional[aiopg.Connection] = None
 
     LOGGER.info(f"Attempting to connect to database {database} as "
                 f"{user}@{host}:{port}...")
 
-    while connection is None:
+    while pool is None:
         try:
-            connection = await aiopg.connect(
-                dsn,
-                cursor_factory=RealDictCursor)
+            pool = await aiopg.create_pool(dsn)
         except psycopg2OpError as err:
             print(type(err))
             if retries > 12:
@@ -67,7 +67,7 @@ async def _try_connect(
             await asyncio.sleep(5)
             return await _try_connect(connection_params, retries + 1)
 
-    return connection
+    return pool
 
 
 def _sync_try_connect(
@@ -112,10 +112,10 @@ def _sync_try_connect(
 
 # PENDS python 3.9 support in pylint
 # pylint: disable=unsubscriptable-object
-async def connect(
+async def get_pool(
     connection_params: ConnectionParameters
-) -> aiopg.Connection:
-    """Establish database connection."""
+) -> aiopg.Pool:
+    """Establish database connection pool."""
     return await _try_connect(connection_params)