From a8eec38beb91c4f1e4a4ac7f92f1a6e5b7ce9032 Mon Sep 17 00:00:00 2001
From: Drew Bednar <drew@androiddrew.com>
Date: Sun, 13 Feb 2022 13:44:04 -0500
Subject: [PATCH] Now with async sqlalchemy

---
 README.md                                     |  4 ++
 fastapi-sqlmodel-alembic/docker-compose.yml   |  5 +-
 fastapi-sqlmodel-alembic/project/app/db.py    | 50 ++++++++++++++++---
 fastapi-sqlmodel-alembic/project/app/main.py  | 16 +++---
 .../project/requirements.txt                  |  1 +
 5 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index e69de29..042de0a 100644
--- a/README.md
+++ b/README.md
@@ -0,0 +1,4 @@
+# Asyncio Learning
+
+This repo contains a number for examples and experiments related to python3's asyncio functionality.
+
diff --git a/fastapi-sqlmodel-alembic/docker-compose.yml b/fastapi-sqlmodel-alembic/docker-compose.yml
index c83f69f..3526fd5 100644
--- a/fastapi-sqlmodel-alembic/docker-compose.yml
+++ b/fastapi-sqlmodel-alembic/docker-compose.yml
@@ -10,7 +10,10 @@ services:
     ports:
       - 8004:8000
     environment:
-      - DATABASE_URL=postgresql://postgres:postgres@db:5432/foo
+#      This is the synchronous uri we started with
+#      - DATABASE_URL=postgresql://postgres:postgres@db:5432/foo
+#      This is the async uri after installing asyncpg
+      - DATABASE_URL=postgresql+asyncpg://postgres:postgres@db:5432/foo
     depends_on:
       - db
 
diff --git a/fastapi-sqlmodel-alembic/project/app/db.py b/fastapi-sqlmodel-alembic/project/app/db.py
index 949e991..954216c 100644
--- a/fastapi-sqlmodel-alembic/project/app/db.py
+++ b/fastapi-sqlmodel-alembic/project/app/db.py
@@ -1,22 +1,56 @@
 import os
 
-from sqlmodel import create_engine, SQLModel, Session
+from sqlmodel import SQLModel
 
+# We used the SQLAlchemy constructs -- e.g., create_async_engine and AsyncSession
+# -- since SQLModel 0.0.4 does not have wrappers for them as of writing.
+from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
 
 DATABASE_URL = os.environ.get("DATABASE_URL")
 
-
 # The major differences between SQLModel's create_engine and SQLAlchemy's
 # version is that the SQLModel version adds type annotations (for editor support)
 # and enables the SQLAlchemy "2.0" style of engines and connections.
-engine = create_engine(DATABASE_URL, echo=True)  # Will want a "debug mode where it doesn't echo for prod
 
+# We need the 2.0 style for async support.
+# engine = create_engine(DATABASE_URL, echo=True)
+engine = create_async_engine(DATABASE_URL, echo=True, future=True)  # Will want a "debug mode where it doesn't echo for prod
+
+
+async def init_db():
+    """
+    Our new Async init_db was
+
+    def init_db():
+        SQLModel.metadata.create_all(engine)
+
+    metadata.create_all doesn't execute asynchronously,
+    so we used run_sync to execute it synchronously within
+    the async function.
+    """
+    async with engine.begin() as conn:
+        # await conn.run_sync(SQLModel.metadata.drop_all)
+        await conn.run_sync(SQLModel.metadata.create_all)
 
-def init_db():
-    SQLModel.metadata.create_all(engine)
 
+async def get_session() -> AsyncSession:
+    """
+    Our new async get_session
 
-def get_session():
-    """A coroutine function for getting a database session."""
-    with Session(engine) as session:
+    Was
+
+    def get_session():
+        with Session(engine) as session:
+            yield session
+
+    We disabled expire on commit behavior by passing in expire_on_commit=False.
+    I believe this is how flask-sqlalchemy also works
+    """
+    async_session = sessionmaker(
+        engine, class_=AsyncSession, expire_on_commit=False
+    )
+    async with async_session() as session:
         yield session
+
+
diff --git a/fastapi-sqlmodel-alembic/project/app/main.py b/fastapi-sqlmodel-alembic/project/app/main.py
index 74fd66c..fccffe9 100644
--- a/fastapi-sqlmodel-alembic/project/app/main.py
+++ b/fastapi-sqlmodel-alembic/project/app/main.py
@@ -1,6 +1,6 @@
 from fastapi import Depends, FastAPI
 from sqlalchemy import select
-from sqlmodel import Session
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from app.db import get_session, init_db
 from app.models import Song, SongCreate
@@ -9,13 +9,13 @@ app = FastAPI()
 
 
 @app.on_event("startup")
-def on_startup():
+async def on_startup():
     """
     It's worth noting that from app.models import Song is required.
     Without it, the song table will not be created. Kind of like
     how I have to do my flask sql alchemy models
     """
-    init_db()
+    await init_db()
 
 
 @app.get("/ping")
@@ -24,16 +24,16 @@ async def pong():
 
 
 @app.get("/songs", response_model=list[Song])
-def get_songs(session: Session = Depends(get_session)):
-    result = session.execute(select(Song))
+async def get_songs(session: AsyncSession = Depends(get_session)):
+    result = await session.execute(select(Song))
     songs = result.scalars().all()
     return [Song(name=song.name, artist=song.artist, id=song.id) for song in songs]
 
 
 @app.post("/songs")
-def add_song(song: SongCreate, session: Session = Depends(get_session)):
+async def add_song(song: SongCreate, session: AsyncSession = Depends(get_session)):
     song = Song(name=song.name, artist=song.artist)
     session.add(song)
-    session.commit()
-    session.refresh(song)
+    await session.commit()
+    await session.refresh(song)  # refreshed becuase of the async nature
     return song
diff --git a/fastapi-sqlmodel-alembic/project/requirements.txt b/fastapi-sqlmodel-alembic/project/requirements.txt
index 0074634..33e6815 100644
--- a/fastapi-sqlmodel-alembic/project/requirements.txt
+++ b/fastapi-sqlmodel-alembic/project/requirements.txt
@@ -1,3 +1,4 @@
+asyncpg==0.24.0
 fastapi==0.68.1
 psycopg2-binary==2.9.1
 sqlmodel==0.0.4