Change psycopg dependencies (#107)

Fixed connection pool dependency for psycopg.

Author: chandr-andr

fastapi_template/template/{{cookiecutter.project_name}}/{{cookiecutter.project_name}}/conftest.py

@@ -51,7 +51,7 @@
 from psycopg import AsyncConnection
 from psycopg_pool import AsyncConnectionPool
 
-from {{cookiecutter.project_name}}.db.dependencies import get_db_session
+from {{cookiecutter.project_name}}.db.dependencies import get_db_pool
 {%- elif cookiecutter.orm == "piccolo" %}
 {%- if cookiecutter.db_info.name == "postgresql" %}
 from piccolo.engine.postgres import PostgresEngine
@@ -245,13 +245,13 @@ async def create_tables(connection: AsyncConnection[Any]) -> None:
 
 
 @pytest.fixture
-async def dbsession() -> AsyncGenerator[AsyncConnection[Any], None]:
+async def dbpool() -> AsyncGenerator[AsyncConnectionPool, None]:
     """
-    Creates connection to some test database.
+    Creates database connections pool to test database.
 
     This connection must be used in tests and for application.
 
-    :yield: connection to database.
+    :yield: database connections pool.
     """
     await create_db()
     pool = AsyncConnectionPool(conninfo=str(settings.db_url))
@@ -261,8 +261,7 @@ async def dbsession() -> AsyncGenerator[AsyncConnection[Any], None]:
         await create_tables(create_conn)
 
     try:
-        async with pool.connection() as conn:
-            yield conn
+        yield pool
     finally:
         await pool.close()
         await drop_db()
@@ -443,7 +442,7 @@ def fastapi_app(
     {%- if cookiecutter.orm == "sqlalchemy" %}
     dbsession: AsyncSession,
     {%- elif cookiecutter.orm == "psycopg" %}
-    dbsession: AsyncConnection[Any],
+    dbpool: AsyncConnectionPool,
     {%- endif %}
     {% if cookiecutter.enable_redis == "True" -%}
     fake_redis: FakeRedis,
@@ -461,8 +460,10 @@ def fastapi_app(
     :return: fastapi app with mocked dependencies.
     """
     application = get_app()
-    {% if cookiecutter.orm in ["sqlalchemy", "psycopg"] -%}
+    {%- if cookiecutter.orm == "sqlalchemy" %}
     application.dependency_overrides[get_db_session] = lambda: dbsession
+    {%- elif cookiecutter.orm == "psycopg" %}
+    application.dependency_overrides[get_db_pool] = lambda: dbpool
     {%- endif %}
     {%- if cookiecutter.enable_redis == "True" %}
     application.dependency_overrides[get_redis_connection] = lambda: fake_redis

fastapi_template/template/{{cookiecutter.project_name}}/{{cookiecutter.project_name}}/db_psycopg/dao/dummy_dao.py

@@ -3,19 +3,19 @@
 from typing import Any
 
 from fastapi import Depends
-from psycopg import AsyncConnection
+from psycopg_pool import AsyncConnectionPool
 from psycopg.rows import class_row
-from {{cookiecutter.project_name}}.db.dependencies import get_db_session
+from {{cookiecutter.project_name}}.db.dependencies import get_db_pool
 from typing import List, Optional
 
 class DummyDAO:
     """Class for accessing dummy table."""
 
     def __init__(
         self,
-        connection: AsyncConnection[Any] = Depends(get_db_session),
+        db_pool: AsyncConnectionPool = Depends(get_db_pool),
     ):
-        self.connection = connection
+        self.db_pool = db_pool
 
 
     async def create_dummy_model(self, name: str) -> None:
@@ -24,15 +24,14 @@ async def create_dummy_model(self, name: str) -> None:
 
         :param name: name of a dummy.
         """
-        async with self.connection.cursor(
-            binary=True,
-        ) as cur:
-            await cur.execute(
-                "INSERT INTO dummy (name) VALUES (%(name)s);",
-                params={
-                    "name": name,
-                }
-            )
+        async with self.db_pool.connection() as connection:
+            async with connection.cursor(binary=True) as cur:
+                await cur.execute(
+                    "INSERT INTO dummy (name) VALUES (%(name)s);",
+                    params={
+                        "name": name,
+                    }
+                )
 
     async def get_all_dummies(self, limit: int, offset: int) -> List[DummyModel]:
         """
@@ -42,18 +41,19 @@ async def get_all_dummies(self, limit: int, offset: int) -> List[DummyModel]:
         :param offset: offset of dummies.
         :return: stream of dummies.
         """
-        async with self.connection.cursor(
-            binary=True,
-            row_factory=class_row(DummyModel)
-        ) as cur:
-            res = await cur.execute(
-                "SELECT id, name FROM dummy LIMIT %(limit)s OFFSET %(offset)s;",
-                params={
-                    "limit": limit,
-                    "offset": offset,
-                }
-            )
-            return await res.fetchall()
+        async with self.db_pool.connection() as connection:
+            async with connection.cursor(
+                binary=True,
+                row_factory=class_row(DummyModel)
+            ) as cur:
+                res = await cur.execute(
+                    "SELECT id, name FROM dummy LIMIT %(limit)s OFFSET %(offset)s;",
+                    params={
+                        "limit": limit,
+                        "offset": offset,
+                    }
+                )
+                return await res.fetchall()
 
     async def filter(
         self,
@@ -65,17 +65,18 @@ async def filter(
         :param name: name of dummy instance.
         :return: dummy models.
         """
-        async with self.connection.cursor(
-            binary=True,
-            row_factory=class_row(DummyModel)
-        ) as cur:
-            if name is not None:
-                res = await cur.execute(
-                    "SELECT id, name FROM dummy WHERE name=%(name)s;",
-                    params={
-                        "name": name,
-                    }
-                )
-            else:
-                res = await cur.execute("SELECT id, name FROM dummy;")
-            return await res.fetchall()
+        async with self.db_pool.connection() as connection:
+            async with connection.cursor(
+                binary=True,
+                row_factory=class_row(DummyModel)
+            ) as cur:
+                if name is not None:
+                    res = await cur.execute(
+                        "SELECT id, name FROM dummy WHERE name=%(name)s;",
+                        params={
+                            "name": name,
+                        }
+                    )
+                else:
+                    res = await cur.execute("SELECT id, name FROM dummy;")
+                return await res.fetchall()

fastapi_template/template/{{cookiecutter.project_name}}/{{cookiecutter.project_name}}/db_psycopg/dependencies.py

@@ -1,20 +1,12 @@
-from typing import Any, AsyncGenerator
-
-from psycopg import AsyncConnection
+from psycopg_pool import AsyncConnectionPool
 from starlette.requests import Request
 
 
-async def get_db_session(
-    request: Request,
-) -> AsyncGenerator[AsyncConnection[Any], None]:
+async def get_db_pool(request: Request) -> AsyncConnectionPool:
     """
-    Create and get database connection.
+    Return database connections pool.
 
     :param request: current request.
-    :yield: database connection.
+    :returns: database connections pool.
     """
-    async with request.app.state.db_pool.connection() as conn:
-        try:
-            yield conn
-        except Exception:  # noqa: S110
-            pass  # noqa: WPS420
+    return request.app.state.db_pool

fastapi_template/template/{{cookiecutter.project_name}}/{{cookiecutter.project_name}}/tests/test_dummy.py

@@ -7,6 +7,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 {%- elif cookiecutter.orm == 'psycopg' %}
 from psycopg.connection_async import AsyncConnection
+from psycopg_pool import AsyncConnectionPool
 {%- endif %}
 from starlette import status
 from {{cookiecutter.project_name}}.db.models.dummy_model import DummyModel
@@ -19,7 +20,7 @@ async def test_creation(
     {%- if cookiecutter.orm == "sqlalchemy" %}
     dbsession: AsyncSession,
     {%- elif cookiecutter.orm == "psycopg" %}
-    dbsession: AsyncConnection[Any],
+    dbpool: AsyncConnectionPool,
     {%- endif %}
 ) -> None:
     """Tests dummy instance creation."""
@@ -43,8 +44,10 @@ async def test_creation(
     )
     {%- endif %}
     assert response.status_code == status.HTTP_200_OK
-    {%- if cookiecutter.orm in ["sqlalchemy", "psycopg"] %}
+    {%- if cookiecutter.orm == "sqlalchemy" %}
     dao = DummyDAO(dbsession)
+    {%- elif cookiecutter.orm == "psycopg" %}
+    dao = DummyDAO(dbpool)
     {%- elif cookiecutter.orm in ["tortoise", "ormar", "piccolo"] %}
     dao = DummyDAO()
     {%- endif %}
@@ -59,12 +62,14 @@ async def test_getting(
     {%- if cookiecutter.orm == "sqlalchemy" %}
     dbsession: AsyncSession,
     {%- elif cookiecutter.orm == "psycopg" %}
-    dbsession: AsyncConnection[Any],
+    dbpool: AsyncConnectionPool,
     {%- endif %}
 ) -> None:
     """Tests dummy instance retrieval."""
-    {%- if cookiecutter.orm in ["sqlalchemy", "psycopg"] %}
+    {%- if cookiecutter.orm == "sqlalchemy" %}
     dao = DummyDAO(dbsession)
+    {%- elif cookiecutter.orm == "psycopg" %}
+    dao = DummyDAO(dbpool)
     {%- elif cookiecutter.orm in ["tortoise", "ormar", "piccolo"] %}
     dao = DummyDAO()
     {%- endif %}