⚡️Cached Sessions (#8)

Author: priyanshu-panwar

fastapi_utilities/__init__.py

@@ -2,3 +2,5 @@
 from .repeat.repeat_at import repeat_at
 
 from .timer import add_timer_middleware
+
+from .session import FastAPISessionMaker

fastapi_utilities/session/__init__.py

@@ -0,0 +1 @@
+from .session import FastAPISessionMaker

fastapi_utilities/session/session.py

@@ -0,0 +1,57 @@
+from contextlib import contextmanager
+from sqlalchemy import create_engine
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session, sessionmaker
+
+
+class FastAPISessionMaker:
+    """
+    This will allow us to create a cached session maker that can be used as context manager.
+    """
+
+    def __init__(self, db_url: str):
+        """
+        `db_url` should be any sqlalchemy-compatible database URI.
+        """
+        self.db_url = db_url
+        self._cached_engine: Engine = None
+        self._cached_session_maker: Session = None
+
+    def _get_db(self):
+        """
+        This will return a cached db connection Session.
+        """
+        if self._cached_engine is None:
+            self._cached_engine = create_engine(self.db_url)
+        if self._cached_session_maker is None:
+            self._cached_session_maker = sessionmaker(
+                autocommit=False, autoflush=False, bind=self._cached_engine
+            )
+        db = self._cached_session_maker()
+        try:
+            yield db
+            db.commit()
+        except Exception as e:
+            db.rollback()
+            raise e
+        finally:
+            db.close()
+
+    @contextmanager
+    def context_session(self):
+        """
+        A context manager that works for `get_db` dependency injection.
+
+        Usage:
+            session_maker = FastAPISessionMaker(database_uri)
+            with session_maker.context_session() as session:
+                session.query(...)
+        """
+        yield from self._get_db()
+
+    def reset_session(self):
+        """
+        This will reset the sessionmaker and engine.
+        """
+        self._cached_session_maker = None
+        self._cached_engine = None

tests/test_session.py

@@ -0,0 +1,59 @@
+import pytest
+from pathlib import Path
+import random
+from _pytest.capture import CaptureFixture
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from fastapi_utilities import FastAPISessionMaker
+
+# setup
+db_path = Path("./db.sqlite3")
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+DB_URL = f"sqlite:///{db_path}"
+engine = create_engine(DB_URL, connect_args={"check_same_thread": False})
+Base = declarative_base()
+
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
+from sqlalchemy import Column, Integer
+
+
+class User(Base):
+    __tablename__ = "users"
+
+    id = Column(Integer, primary_key=True)
+
+
+app = FastAPI()
+Base.metadata.create_all(bind=engine)
+
+# done
+
+
+def test_session():
+    sm = FastAPISessionMaker(DB_URL)
+    with sm.context_session() as session:
+        x = User(id=random.randint(0, 10000))
+        session.add(x)
+
+
+def test_reset_session():
+    sm = FastAPISessionMaker(DB_URL)
+    sm.reset_session()
+
+
+def test_session_raise_error(capsys: CaptureFixture[str]) -> None:
+    sm = FastAPISessionMaker(DB_URL)
+    try:
+        with sm.context_session() as session:
+            x = User(id=1)
+            session.add(x)
+            x = User(id=1)
+            session.add(x)
+    except Exception:
+        out, err = capsys.readouterr()
+        assert out == ""