Copy existing code from seed project

Author: andrew-chang-dewitt

db_wrapper/__init__.py

@@ -0,0 +1,22 @@
+"""
+Database module, exposes classes for simplifying working with Postgres.
+
+class.ConnectionParameters:
+    Dataclass to assist with correctly setting up a Postgres connection.
+
+class.Client:
+    The bread & butter of this module. Wraps aiopg to simplify connecting to,
+    executing queries on, & disconnecting from a Postgres database.
+
+class.Model:
+    Not quite an ORM, but close. Used to assist in defining database queries &
+    associating a Model type. Using this class will help manage separation of
+    concerns.
+
+class.ModelData:
+    A simple TypedDict to define the shape of the data for a specific
+    Model instance.
+"""
+from .connection import ConnectionParameters
+from .client import Client
+from .model import Model, ModelData

db_wrapper/client.py

@@ -0,0 +1,113 @@
+"""Wrapper on aiopg to simplify connecting to & interacting with db."""
+
+from __future__ import annotations
+from typing import (
+    Any,
+    TypeVar,
+    Union,
+    Optional,
+    Hashable,
+    List,
+    Dict)
+
+import aiopg  # type: ignore
+from psycopg2.extras import register_uuid
+# importing for the sole purpose of re-exporting
+# pylint: disable=unused-import
+from psycopg2 import sql
+
+from .connection import ConnectionParameters, connect
+
+# add uuid support to psycopg2 & Postgres
+register_uuid()
+
+
+# Generic doesn't need a more descriptive name
+# pylint: disable=invalid-name
+T = TypeVar('T')
+
+# pylint: disable=unsubscriptable-object
+Query = Union[str, sql.Composed]
+
+
+class Client:
+    """Class to manage database connection & expose necessary methods to user.
+
+    Stores connection parameters on init, then exposes methods to
+    asynchronously connect & disconnect the database, as well as execute SQL
+    queries.
+    """
+
+    _connection_params: ConnectionParameters
+    _connection: aiopg.Connection
+
+    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)
+
+    async def disconnect(self) -> None:
+        """Disconnect from the database."""
+        await self._connection.close()
+
+    # PENDS python 3.9 support in pylint
+    # pylint: disable=unsubscriptable-object
+    @staticmethod
+    async def _execute_query(
+        cursor: aiopg.Cursor,
+        query: Query,
+        params: Optional[Dict[Hashable, Any]] = None,
+    ) -> None:
+        if params:
+            await cursor.execute(query, params)
+        else:
+            await cursor.execute(query)
+
+    # PENDS python 3.9 support in pylint
+    # pylint: disable=unsubscriptable-object
+    async def execute(
+        self,
+        query: Query,
+        params: Optional[Dict[Hashable, Any]] = None,
+    ) -> None:
+        """Execute the given SQL query.
+
+        Arguments:
+            query (Query)   -- the SQL query to execute
+            params (dict)  -- a dictionary of parameters to interpolate when
+                              executing the query
+
+        Returns:
+            None
+        """
+        async with self._connection.cursor() as cursor:
+            await self._execute_query(cursor, query, params)
+
+    # PENDS python 3.9 support in pylint
+    # pylint: disable=unsubscriptable-object
+    async def execute_and_return(
+        self,
+        query: Query,
+        params: Optional[Dict[Hashable, Any]] = None,
+    ) -> List[T]:
+        """Execute the given SQL query & return the result.
+
+        Arguments:
+            query (Query)   -- the SQL query to execute
+            params (dict)  -- a dictionary of parameters to interpolate when
+                              executing the query
+
+        Returns:
+            List containing all the rows that matched the query.
+        """
+        async with self._connection.cursor() as cursor:
+            await self._execute_query(cursor, query, params)
+
+            result = []
+
+            async for row in cursor:
+                result.append(row)
+
+            return result

db_wrapper/connection.py

@@ -0,0 +1,75 @@
+"""Interface & function for defining & establishing connection to Postgres."""
+
+import asyncio
+from dataclasses import dataclass
+import logging
+from typing import Optional
+
+from psycopg2 import OperationalError  # type: ignore
+
+# no stubs available, starting out by just determining correct return
+# types & annotating in my wrappers here
+import aiopg  # type: ignore
+
+#
+# Postgres Connection helper
+#
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class ConnectionParameters:
+    """Defines connection parameters for database."""
+
+    host: str
+    user: str
+    password: str
+    database: str
+
+
+async def _try_connect(
+    connection_params: ConnectionParameters,
+    retries: int = 1
+) -> aiopg.Connection:
+    database = connection_params.database
+    user = connection_params.user
+    password = connection_params.password
+    host = connection_params.host
+
+    dsn = f'dbname={database} user={user} password={password} host={host}'
+
+    # PENDS python 3.9 support in pylint
+    # pylint: disable=unsubscriptable-object
+    connection: Optional[aiopg.Connection] = None
+
+    LOGGER.info(
+        f'Attempting to connect to database {database} as {user}@{host}...')
+
+    while connection is None:
+        try:
+            connection = await aiopg.connect(dsn)
+        except OperationalError as err:
+            print(type(err))
+            if retries > 12:
+                raise ConnectionError(
+                    'Max number of connection attempts has been reached (12)'
+                ) from err
+
+            LOGGER.info(
+                f'Connection failed ({retries} time(s))'
+                'retrying again in 5 seconds...')
+
+            await asyncio.sleep(5)
+            return await _try_connect(connection_params, retries + 1)
+
+    return connection
+
+
+# PENDS python 3.9 support in pylint
+# pylint: disable=unsubscriptable-object
+async def connect(
+    connection_params: ConnectionParameters
+) -> aiopg.Connection:
+    """Establish database connection."""
+    return await _try_connect(connection_params)