Added psycopg support. (#70)

Signed-off-by: Pavel Kirilin <win10@list.ru>

Author: s3rius

fastapi_template/cli.py

@@ -1,13 +1,16 @@
 import re
 from argparse import ArgumentParser
 from operator import attrgetter
+from termcolor import cprint
 
 from prompt_toolkit import prompt
 from prompt_toolkit.document import Document
 from prompt_toolkit.shortcuts import checkboxlist_dialog, radiolist_dialog
 from prompt_toolkit.validation import ValidationError, Validator
 
 from fastapi_template.input_model import (
+    SUPPORTED_ORMS,
+    ORMS_WITHOUT_MIGRATIONS,
     ORM,
     BuilderContext,
     DB_INFO,
@@ -199,10 +202,19 @@ def read_user_input(current_context: BuilderContext) -> BuilderContext:
         current_context.orm = radiolist_dialog(
             "ORM",
             text="Which ORM do you want?",
-            values=[(orm, orm.value) for orm in list(ORM) if orm != ORM.none],
+            values=[(orm, orm.value) for orm in SUPPORTED_ORMS[current_context.db]],
         ).run()
         if current_context.orm is None:
             raise KeyboardInterrupt()
+    if (
+        current_context.orm is not None 
+        and current_context.orm != ORM.none 
+        and current_context.orm not in SUPPORTED_ORMS.get(current_context.db, [])
+    ):
+        cprint("This ORM is not supported by chosen database.", "red")
+        raise KeyboardInterrupt()
+    if current_context.orm in ORMS_WITHOUT_MIGRATIONS:
+        current_context.enable_migrations = False
     if current_context.ci_type is None:
         current_context.ci_type = radiolist_dialog(
             "CI",

fastapi_template/input_model.py

@@ -24,6 +24,7 @@ class ORM(enum.Enum):
     ormar = "ormar"
     sqlalchemy = "sqlalchemy"
     tortoise = "tortoise"
+    psycopg = "psycopg"
 
 
 class Database(BaseModel):
@@ -69,6 +70,28 @@ class Database(BaseModel):
     ),
 }
 
+SUPPORTED_ORMS = {
+    DatabaseType.postgresql: [
+        ORM.ormar,
+        ORM.psycopg,
+        ORM.tortoise,
+        ORM.sqlalchemy,
+    ],
+    DatabaseType.sqlite: [
+        ORM.ormar,
+        ORM.tortoise,
+        ORM.sqlalchemy,
+    ],
+    DatabaseType.mysql: [
+        ORM.ormar,
+        ORM.tortoise,
+        ORM.sqlalchemy,
+    ]
+}
+
+ORMS_WITHOUT_MIGRATIONS = [
+    ORM.psycopg,
+]
 
 class BuilderContext(BaseModel):
     """Options for project generation."""

fastapi_template/template/{{cookiecutter.project_name}}/conditional_files.json

@@ -80,6 +80,8 @@
             "{{cookiecutter.project_name}}/db_ormar/models/dummy_model.py",
             "{{cookiecutter.project_name}}/db_tortoise/dao",
             "{{cookiecutter.project_name}}/db_tortoise/models/dummy_model.py",
+            "{{cookiecutter.project_name}}/db_psycopg/dao",
+            "{{cookiecutter.project_name}}/db_psycopg/models/dummy_model.py",
             "{{cookiecutter.project_name}}/tests/test_dummy.py",
             "{{cookiecutter.project_name}}/db_sa/migrations/versions/2021-08-16-16-55_2b7380507a71.py",
             "{{cookiecutter.project_name}}/db_ormar/migrations/versions/2021-08-16-16-55_2b7380507a71.py",
@@ -114,6 +116,12 @@
             "{{cookiecutter.project_name}}/db_ormar"
         ]
     },
+    "PsycoPG": {
+        "enabled": "{{cookiecutter.orm == 'psycopg'}}",
+        "resources": [
+            "{{cookiecutter.project_name}}/db_psycopg"
+        ]
+    },
     "Postgresql DB": {
         "enabled": "{{cookiecutter.db_info.name == 'postgresql'}}",
         "resources": [

fastapi_template/template/{{cookiecutter.project_name}}/pyproject.toml

@@ -67,6 +67,9 @@ aioredis = {version = "^2.0.1", extras = ["hiredis"]}
 {%- if cookiecutter.self_hosted_swagger == 'True' %}
 aiofiles = "^0.8.0"
 {%- endif %}
+{%- if cookiecutter.orm == "psycopg" %}
+psycopg = { version = "^3.0.11", extras = ["binary", "pool"] }
+{%- endif %}
 httptools = "^0.3.0"
 
 [tool.poetry.dev-dependencies]

fastapi_template/template/{{cookiecutter.project_name}}/replaceable_files.json

@@ -2,6 +2,7 @@
     "{{cookiecutter.project_name}}/db": [
         "{{cookiecutter.project_name}}/db_sa",
         "{{cookiecutter.project_name}}/db_ormar",
-        "{{cookiecutter.project_name}}/db_tortoise"
+        "{{cookiecutter.project_name}}/db_tortoise",
+        "{{cookiecutter.project_name}}/db_psycopg"
     ]
 }

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

@@ -31,6 +31,12 @@
 from sqlalchemy.engine import create_engine
 from {{cookiecutter.project_name}}.db.config import database
 from {{cookiecutter.project_name}}.db.utils import create_database, drop_database
+{%- elif cookiecutter.orm == "psycopg" %}
+from psycopg import AsyncConnection
+from psycopg_pool import AsyncConnectionPool
+
+from {{cookiecutter.project_name}}.db.dependencies import get_db_session
+
 {%- endif %}
 
 
@@ -146,6 +152,99 @@ async def initialize_db() -> AsyncGenerator[None, None]:
     await database.disconnect()
     drop_database()
 
+{%- elif cookiecutter.orm == "psycopg" %}
+
+async def drop_db() -> None:
+    """Drops database after tests."""
+    pool = AsyncConnectionPool(conninfo=str(settings.db_url.with_path("/postgres")))
+    await pool.wait()
+    async with pool.connection() as conn:
+        await conn.set_autocommit(True)
+        await conn.execute(
+            "SELECT pg_terminate_backend(pg_stat_activity.pid) "  # noqa: S608
+            "FROM pg_stat_activity "
+            "WHERE pg_stat_activity.datname = %(dbname)s "
+            "AND pid <> pg_backend_pid();",
+            params={
+                "dbname": settings.db_base,
+            }
+        )
+        await conn.execute(
+            f"DROP DATABASE {settings.db_base}",
+        )
+    await pool.close()
+
+
+async def create_db() -> None:  # noqa: WPS217
+    """Creates database for tests."""
+    pool = AsyncConnectionPool(conninfo=str(settings.db_url.with_path("/postgres")))
+    await pool.wait()
+    async with pool.connection() as conn_check:
+        res = await conn_check.execute(
+            "SELECT 1 FROM pg_database WHERE datname=%(dbname)s",
+            params={
+                "dbname": settings.db_base,
+            }
+        )
+        db_exists = False
+        row = await res.fetchone()
+        if row is not None:
+            db_exists = row[0]
+
+    if db_exists:
+        await drop_db()
+
+    async with pool.connection() as conn_create:
+        await conn_create.set_autocommit(True)
+        await conn_create.execute(
+            f"CREATE DATABASE {settings.db_base};",
+        )
+    await pool.close()
+
+
+async def create_tables(connection: AsyncConnection[Any]) -> None:
+    """
+    Create tables for your database.
+
+    Since psycopg doesn't have migration tool,
+    you must create your tables for tests.
+
+    :param connection: connection to database.
+    """
+    {%- if cookiecutter.add_dummy == 'True' %}
+    await connection.execute(
+        "CREATE TABLE dummy ("
+        "id SERIAL primary key,"
+        "name VARCHAR(200)"
+        ");"
+    )
+    {%- endif %}
+    pass  # noqa: WPS420
+
+
+@pytest.fixture
+async def dbsession() -> AsyncGenerator[AsyncConnection[Any], None]:
+    """
+    Creates connection to some test database.
+
+    This connection must be used in tests and for application.
+
+    :yield: connection to database.
+    """
+    await create_db()
+    pool = AsyncConnectionPool(conninfo=str(settings.db_url))
+    await pool.wait()
+
+    async with pool.connection() as create_conn:
+        await create_tables(create_conn)
+
+    try:
+        async with pool.connection() as conn:
+            yield conn
+    finally:
+        await pool.close()
+        await drop_db()
+
 {%- endif %}
 
 
@@ -167,6 +266,8 @@ async def fake_redis() -> AsyncGenerator[FakeRedis, None]:
 def fastapi_app(
     {%- if cookiecutter.orm == "sqlalchemy" %}
     dbsession: AsyncSession,
+    {%- elif cookiecutter.orm == "psycopg" %}
+    dbsession: AsyncConnection[Any],
     {%- endif %}
     {% if cookiecutter.enable_redis == "True" -%}
     fake_redis: FakeRedis,
@@ -178,7 +279,7 @@ def fastapi_app(
     :return: fastapi app with mocked dependencies.
     """
     application = get_app()
-    {% if cookiecutter.orm == "sqlalchemy" -%}
+    {% if cookiecutter.orm in ["sqlalchemy", "psycopg"] -%}
     application.dependency_overrides[get_db_session] = lambda: dbsession
     {%- endif %}
     {%- if cookiecutter.enable_redis == "True" %}

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

@@ -0,0 +1,81 @@
+from termios import OFDEL
+from {{cookiecutter.project_name}}.db.models.dummy_model import DummyModel
+from typing import Any
+
+from fastapi import Depends
+from psycopg import AsyncConnection
+from psycopg.rows import class_row
+from {{cookiecutter.project_name}}.db.dependencies import get_db_session
+from typing import List, Optional
+
+class DummyDAO:
+    """Class for accessing dummy table."""
+
+    def __init__(
+        self,
+        connection: AsyncConnection[Any] = Depends(get_db_session),
+    ):
+        self.connection = connection
+
+
+    async def create_dummy_model(self, name: str) -> None:
+        """
+        Creates new dummy in a database.
+
+        :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 def get_all_dummies(self, limit: int, offset: int) -> List[DummyModel]:
+        """
+        Get all dummy models with limit/offset pagination.
+
+        :param limit: limit of dummies.
+        :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 def filter(
+        self,
+        name: Optional[str] = None,
+    ) -> List[DummyModel]:
+        """
+        Get specific dummy model.
+
+        :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()

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

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

fastapi_template/template/{{cookiecutter.project_name}}/{{cookiecutter.project_name}}/db_psycopg/models/dummy_model.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+class DummyModel(BaseModel):
+    """Dummy model for database."""
+
+    id: int
+    name: str