feat(db session): adds support for tracking db sessions

src/dispatch/database/core.py

@@ -24,6 +24,7 @@
 from dispatch import config
 from dispatch.exceptions import NotFoundError
 from dispatch.search.fulltext import make_searchable
+from dispatch.database.logging import SessionTracker
 
 
 def create_db_engine(connection_string: str):
@@ -156,7 +157,12 @@ def __repr__(self):
 
 def get_db(request: Request) -> Session:
     """Get database session from request state."""
-    return request.state.db
+    session = request.state.db
+    if not hasattr(session, "_dispatch_session_id"):
+        session._dispatch_session_id = SessionTracker.track_session(
+            session, context="fastapi_request"
+        )
+    return session
 
 
 DbSession = Annotated[Session, Depends(get_db)]
@@ -231,27 +237,32 @@ def refetch_db_session(organization_slug: str) -> Session:
             None: f"dispatch_organization_{organization_slug}",
         }
     )
-    db_session = sessionmaker(bind=schema_engine)()
-    return db_session
+    session = sessionmaker(bind=schema_engine)()
+    session._dispatch_session_id = SessionTracker.track_session(
+        session, context=f"organization_{organization_slug}"
+    )
+    return session
 
 
 @contextmanager
 def get_session() -> Session:
     """Context manager to ensure the session is closed after use."""
     session = SessionLocal()
+    session_id = SessionTracker.track_session(session, context="context_manager")
     try:
         yield session
         session.commit()
     except:
         session.rollback()
         raise
     finally:
+        SessionTracker.untrack_session(session_id)
         session.close()
 
 
 @contextmanager
 def get_organization_session(organization_slug: str) -> Session:
-    """Context manager to ensure the session is closed after use."""
+    """Context manager to ensure the organization session is closed after use."""
     session = refetch_db_session(organization_slug)
     try:
         yield session
@@ -260,4 +271,6 @@ def get_organization_session(organization_slug: str) -> Session:
         session.rollback()
         raise
     finally:
+        if hasattr(session, "_dispatch_session_id"):
+            SessionTracker.untrack_session(session._dispatch_session_id)
         session.close()

src/dispatch/database/logging.py

@@ -0,0 +1,61 @@
+import logging
+import uuid
+from typing import Any
+
+from sqlalchemy.orm import Session
+
+logger = logging.getLogger(__name__)
+
+
+class SessionTracker:
+    """Tracks database session lifecycle events."""
+
+    _sessions: dict[str, dict[str, Any]] = {}
+
+    @classmethod
+    def track_session(cls, session: Session, context: str | None = None) -> str:
+        """Tracks a new database session."""
+        session_id = str(uuid.uuid4())
+        cls._sessions[session_id] = {
+            "session": session,
+            "context": context,
+            "created_at": uuid.uuid1().timestamp(),
+        }
+        logger.info(
+            "Database session created",
+            extra={
+                "session_id": session_id,
+                "context": context,
+                "total_active_sessions": len(cls._sessions),
+            },
+        )
+        return session_id
+
+    @classmethod
+    def untrack_session(cls, session_id: str) -> None:
+        """Untracks a database session."""
+        if session_id in cls._sessions:
+            session_info = cls._sessions.pop(session_id)
+            duration = uuid.uuid1().timestamp() - session_info["created_at"]
+            logger.info(
+                "Database session closed",
+                extra={
+                    "session_id": session_id,
+                    "context": session_info["context"],
+                    "duration_seconds": duration,
+                    "total_active_sessions": len(cls._sessions),
+                },
+            )
+
+    @classmethod
+    def get_active_sessions(cls) -> list[dict[str, Any]]:
+        """Returns information about all active sessions."""
+        current_time = uuid.uuid1().timestamp()
+        return [
+            {
+                "session_id": session_id,
+                "context": info["context"],
+                "age_seconds": current_time - info["created_at"],
+            }
+            for session_id, info in cls._sessions.items()
+        ]