refactor: auto-transaction implementation, mypy fixes

Author: Morty Dev

app/features/content/api.py

@@ -5,7 +5,7 @@
 from ...resources.providers import ConnectionProvider
 from .models import Content, ContentIn
 from .tables import CONTENT_TABLE
-from .tasks import fetch_content
+# from .tasks import fetch_content_async
 
 CONNECTION = AppContainer.resources.db.connection  # type: ignore
 API = APIRouter()
@@ -31,12 +31,12 @@ async def create_content(
     async with connection.acquire() as cn, c1.inject(connection), c2.inject(
         connection
     ):
-        assert c1.injected == connection.current == cn
-        assert c2.injected == connection.current == cn
-    content = await connection.one(
-        CONTENT_TABLE.insert(content_in.dict()).returning(CONTENT_TABLE),
-        commit=True,
-    )
-    if content:
-        fetch_content.apply_async((content["id"],))
+        assert c1.current == connection.current == cn
+        assert c2.current == connection.current == cn
+        content = await connection.one(
+            CONTENT_TABLE.insert(content_in.dict()).returning(CONTENT_TABLE),
+            # commit=True,
+        )
+        # if content:
+        # await fetch_content_async(content["id"], connection=connection)
     return content

app/features/content/tasks.py

@@ -15,19 +15,21 @@
 @inject
 async def fetch_content_async(
     content_id,
-    cursor: ConnectionProvider = Provide[AppContainer.resources.db],
+    connection: ConnectionProvider = Provide[
+        AppContainer.resources.db.connection
+    ],
 ):
-    content = await cursor.one(
-        CONTENT_TABLE.select(CONTENT_TABLE.c.id == content_id)
-    )
-    async with httpx.AsyncClient() as cl:
-        body = (await cl.get(content["url"])).text
-    await cursor.execute(
-        CONTENT_TABLE.update(CONTENT_TABLE.c.id == content_id).values(
-            {"body": body}
-        ),
-        commit=True,
-    )
+    async with connection.acquire():
+        content = await connection.one(
+            CONTENT_TABLE.select(CONTENT_TABLE.c.id == content_id)
+        )
+        async with httpx.AsyncClient() as cl:
+            body = (await cl.get(content["url"])).text
+        await connection.execute(
+            CONTENT_TABLE.update(CONTENT_TABLE.c.id == content_id).values(
+                {"body": body}
+            ),
+        )
 
 
 @shared_task(autoretry_for=(Exception,), name="fetch_content")

app/migrations/script.py.mako

@@ -11,9 +11,9 @@ ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.
 revision = ${repr(up_revision)}
-down_revision = ${repr(down_revision)}
-branch_labels = ${repr(branch_labels)}
-depends_on = ${repr(depends_on)}
+down_revision: str | None = ${repr(down_revision)}
+branch_labels: str | None = ${repr(branch_labels)}
+depends_on: str | None = ${repr(depends_on)}
 
 
 def upgrade():

app/migrations/versions/0001_initial_migration.py

@@ -2,17 +2,17 @@
 
 Revision ID: 0001
 Revises:
-Created Date: 2022-05-08 16:55:48.639025
+Created Date: 2022-10-10 00:50:17.533906
 
 """
 import sqlalchemy as sa
 from alembic import op
 
 # revision identifiers, used by Alembic.
 revision = "0001"
-down_revision = None
-branch_labels = None
-depends_on = None
+down_revision: str | None = None
+branch_labels: str | None = None
+depends_on: str | None = None
 
 
 def upgrade():

app/resources/providers.py

@@ -66,8 +66,7 @@ async def somelogic(cn, dal1, dal2):
 
     def __init__(self, engine):
         self.engine = engine
-        self.current = False
-        self.injected = None
+        self.current: AsyncConnection | None = None
 
     @asynccontextmanager
     async def inject(
@@ -77,52 +76,36 @@ async def inject(
             raise ValueError(
                 "Connection already used, you cannot override current one"
             )
-        if self.injected:
-            raise ValueError("Already injected connection exists")
 
-        self.injected = connection.current
+        self.current = connection.current
         yield self
-        self.injected = None
+        self.current = None
 
     @asynccontextmanager
-    async def acquire(
-        self, tx=True
-    ) -> TP.AsyncGenerator[AsyncConnection, None]:
-        if self.current and not self.injected:
-            raise ValueError(
-                "Connection already used, you cannot override current one"
-            )
-
-        if self.injected:
-            yield self.injected
+    async def acquire(self) -> TP.AsyncGenerator[AsyncConnection, None]:
+        if self.current:
+            yield self.current
         else:
-            manager = self.engine.begin if tx else self.engine.connect
-            async with manager() as cn:
+            async with self.engine.begin() as cn:
                 self.current = cn
                 yield cn
                 self.current = None
 
-    async def scalar(self, query, commit: bool = False) -> TP.Any | None:
+    async def execute(self, query) -> None:
+        async with self.acquire() as cn:
+            await cn.execute(query)
+
+    async def scalar(self, query) -> TP.Any | None:
         async with self.acquire() as cn:
             result = (await cn.execute(query)).scalar()
-            if commit:
-                await cn.commit()
-            return result
+        return result
 
-    async def one(
-        self, query, commit: bool = False
-    ) -> dict[str, TP.Any] | None:
+    async def one(self, query) -> dict[str, TP.Any]:
         async with self.acquire() as cn:
             row = (await cn.execute(query)).fetchone()
-            if commit:
-                await cn.commit()
-        return dict(row) if row else None
+        return dict(row) if row else dict()
 
-    async def many(
-        self, query, commit: bool = False
-    ) -> list[dict[str, TP.Any]]:
+    async def many(self, query) -> list[dict[str, TP.Any]]:
         async with self.acquire() as cn:
             rows = (await cn.execute(query)).fetchall()
-            if commit:
-                await cn.commit()
-            return [dict(row) for row in rows]
+        return [dict(row) for row in rows]

mypy.ini

@@ -1,2 +1,5 @@
 [mypy]
 python_version = 3.10
+
+[mypy-asyncpg.*]
+ignore_missing_imports = True

poetry.lock

@@ -7,7 +7,7 @@
 
 @pytest.mark.asyncio
 async def test_content_created(client, cursor):
-    COUNT = 1000
+    COUNT = 100
 
     URLS = [f"https://www.google.com/{i}" for i in range(COUNT)]
     responses = await asyncio.gather(
@@ -28,3 +28,6 @@ async def test_content_created(client, cursor):
         assert resp.status_code == 409, resp.json()
 
     assert len(URLS) == len(list(await cursor.many(CONTENT_TABLE.select())))
+
+    contents = await client.get("/contents")
+    assert len(URLS) == len(contents.json())