refactor db functions from main into db.py

Signed-off-by: Grant Ramsay <seapagan@gmail.com>

Author: seapagan

db.py

@@ -7,3 +7,23 @@
 engine = create_async_engine(DATABASE_URL, echo=False)
 Base = declarative_base()
 async_session = async_sessionmaker(engine, expire_on_commit=False)
+
+
+async def get_db():
+    """Get a database session.
+
+    To be used for dependency injection.
+    """
+    async with async_session() as session:
+        async with session.begin():
+            yield session
+
+
+async def init_models():
+    """Create tables if they don't already exist.
+
+    In a real-life example we would use Alembic to manage migrations.
+    """
+    async with engine.begin() as conn:
+        # await conn.run_sync(Base.metadata.drop_all)
+        await conn.run_sync(Base.metadata.create_all)

main.py

@@ -5,20 +5,10 @@
 from fastapi import Depends, FastAPI
 from sqlalchemy import select
 
-from db import Base, async_session, engine
+from db import get_db, init_models
 from models import User
 
 
-async def init_models():
-    """Create tables if they don't already exist.
-
-    In a real-life example we would use Alembic to manage migrations.
-    """
-    async with engine.begin() as conn:
-        # await conn.run_sync(Base.metadata.drop_all)
-        await conn.run_sync(Base.metadata.create_all)
-
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Run tasks before and after the server starts."""
@@ -29,16 +19,6 @@ async def lifespan(app: FastAPI):
 app = FastAPI(lifespan=lifespan)
 
 
-async def get_db():
-    """Get a database session.
-
-    To be used for dependency injection.
-    """
-    async with async_session() as session:
-        async with session.begin():
-            yield session
-
-
 @app.get("/")
 async def root():
     """Root endpoint."""