feat: add base entity class with deferred reflection (#7)

Author: hadrien

README.md

@@ -22,3 +22,10 @@ SqlAlchemy integration for FastAPI®
   app = FastAPI()
   fastapi_sqla.setup(app)
   ```
+* Adding a new entity class:
+  ```python
+  from fastapi_sqla import Base
+
+  class Entity(Base):
+      __tablename__ = "table-name-in-db"
+  ```

fastapi_sqla.py

@@ -3,9 +3,10 @@
 import structlog
 from fastapi import FastAPI
 from sqlalchemy import engine_from_config
+from sqlalchemy.ext.declarative import DeferredReflection, declarative_base
 from sqlalchemy.orm.session import sessionmaker
 
-__all__ = ["setup"]
+__all__ = ["Base", "setup"]
 
 logger = structlog.get_logger(__name__)
 
@@ -18,5 +19,11 @@ def setup(app: FastAPI):
 
 def startup():
     engine = engine_from_config(os.environ, prefix="sqlalchemy_")
+    Base.metadata.bind = engine
+    Base.prepare(engine)
     _Session.configure(bind=engine)
     logger.info("startup", engine=engine)
+
+
+class Base(declarative_base(cls=DeferredReflection)):  # type: ignore
+    __abstract__ = True

tests/conftest.py

@@ -1,3 +1,4 @@
+import importlib
 import os
 from unittest.mock import patch
 
@@ -13,5 +14,13 @@ def db_uri():
 @fixture(autouse=True)
 def environ(db_uri):
     values = {"sqlalchemy_url": db_uri}
-    with patch.dict("os.environ", values=values, clear=True) as environ:
-        yield environ
+    with patch.dict("os.environ", values=values, clear=True):
+        yield values
+
+
+@fixture(autouse=True)
+def tear_down():
+    import fastapi_sqla
+
+    # reload fastapi_sqla to clear sqla deferred reflection mapping stored in Base
+    importlib.reload(fastapi_sqla)

tests/test_base.py

@@ -0,0 +1,40 @@
+from pytest import fixture, raises
+from sqlalchemy import engine_from_config
+from sqlalchemy.exc import NoSuchTableError
+from sqlalchemy.orm.session import close_all_sessions
+
+
+@fixture(autouse=True)
+def setup_tear_down(environ):
+    engine = engine_from_config(environ, prefix="sqlalchemy_")
+
+    engine.execute("CREATE TABLE IF NOT EXISTS test_table (id integer primary key)")
+    yield
+    close_all_sessions()
+    engine.execute("DROP TABLE test_table")
+
+
+def test_startup_reflect_test_table():
+    from fastapi_sqla import Base, _Session, startup
+
+    class TestTable(Base):
+        __tablename__ = "test_table"
+
+    startup()
+
+    session = _Session()
+    session.add(TestTable(id=1))
+    session.add(TestTable(id=2))
+    session.commit()
+
+    assert session.query(TestTable).count() == 2
+
+
+def test_startup_fails_when_table_doesnt_exist():
+    from fastapi_sqla import Base, startup
+
+    class TestTable(Base):
+        __tablename__ = "does_not_exist"
+
+    with raises(NoSuchTableError):
+        startup()

tests/test_startup.py

@@ -1,15 +1,13 @@
 import httpx
 from asgi_lifespan import LifespanManager
+from fastapi import FastAPI
 from pytest import fixture, mark
 from sqlalchemy.orm.session import close_all_sessions
 
 
 @fixture(autouse=True)
 def setup_tear_down():
-    from fastapi_sqla import _Session
-
     yield
-    _Session.configure(bind=None)
     close_all_sessions()
 
 
@@ -25,8 +23,6 @@ def test_startup():
 
 @mark.asyncio
 async def test_fastapi_integration():
-    from fastapi import FastAPI
-
     from fastapi_sqla import _Session, setup
 
     app = FastAPI()