Session Improvements Async (#168)

* Adds cache decorator & cache or middleware calls

* Remove prints

* Expunge test

* Adds session decorator and middleware cached options

* Change to kwargs instead of args

* Pool info

* IS_DEV change

* improvements

* cache, token reset

* session tests

* exception fix

* thread session

* debug removal

* optional ctx token

* os session lookup

* clean up

* clean up

* PR comments

---------

Co-authored-by: JWittmeyer <[email protected]>

Author: LennartSchmidtKern

business_objects/general.py

@@ -3,27 +3,32 @@
 from sqlalchemy.orm.session import make_transient as make_transient_original
 from ..session import session, engine
 from ..session import request_id_ctx_var
-from ..session import check_session_and_rollback as check_and_roll
+from ..session import check_session_and_rollback
 from ..enums import Tablenames, try_parse_enum_value
 import traceback
 import datetime
 from .. import daemon
 from threading import Lock
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.sql import Select
-
+import os
 
 __THREAD_LOCK = Lock()
 
+IS_DEV = os.getenv("IS_DEV", "false").lower() in {"true", "1", "yes", "y"}
+
 session_lookup = {}
 
 
 def get_ctx_token() -> Any:
     global session_lookup
     session_uuid = str(uuid.uuid4())
     session_id = request_id_ctx_var.set(session_uuid)
-
-    call_stack = "".join(traceback.format_stack()[-5:])
+    if IS_DEV:
+        # traces are usually long running and only useful for debugging
+        call_stack = "".join(traceback.format_stack()[-5:])
+    else:
+        call_stack = "Activate dev mode to see call stack"
     with __THREAD_LOCK:
         session_lookup[session_uuid] = {
             "session_id": session_uuid,
@@ -46,15 +51,18 @@ def get_session_lookup(exclude_last_x_seconds: int = 5) -> Dict[str, Dict[str, A
     ]
 
 
-def reset_ctx_token(
-    ctx_token: Any,
-    remove_db: Optional[bool] = False,
-) -> None:
+def reset_ctx_token(ctx_token: Any = None, remove_db: Optional[bool] = False) -> None:
     if remove_db:
         session.remove()
-    session_uuid = ctx_token.var.get()
 
-    request_id_ctx_var.reset(ctx_token)
+    session_uuid = request_id_ctx_var.get()
+    if session_uuid is None:
+        print("Session not found in context variable", flush=True)
+    if ctx_token:
+        request_id_ctx_var.reset(ctx_token)
+    else:
+        request_id_ctx_var.set(None)
+
     global session_lookup
     with __THREAD_LOCK:
         if session_uuid in session_lookup:
@@ -106,9 +114,13 @@ def commit() -> None:
 
 
 def remove_and_refresh_session(
-    session_token: Any, request_new: bool = False
+    session_token: Any = None, request_new: bool = False
 ) -> Union[Any, None]:
-    check_and_roll()
+    try:
+        check_session_and_rollback()
+    except Exception:
+        print("Error: check_session_and_rollback() failed", flush=True)
+        print(traceback.format_exc(), flush=True)
     reset_ctx_token(session_token, True)
     if request_new:
         return get_ctx_token()

business_objects/organization.py

@@ -4,16 +4,27 @@
 from submodules.model import enums
 
 
-from ..session import session
+from ..session import session, request_id_ctx_var
 from ..models import Organization, Project, User
 from ..business_objects import project, user, general
 from ..util import prevent_sql_injection
+from ..db_cache import TTLCacheDecorator, CacheEnum
 
 
 def get(id: str) -> Organization:
     return session.query(Organization).get(id)
 
 
+@TTLCacheDecorator(CacheEnum.ORGANIZATION, 5, "id")
+def get_org_cached(id: str) -> Organization:
+    o = get(id)
+    if not o:
+        return None
+    general.expunge(o)
+    general.make_transient(o)
+    return o
+
+
 def get_by_name(name: str) -> Organization:
     return session.query(Organization).filter(Organization.name == name).first()
 

business_objects/user.py

@@ -5,6 +5,7 @@
 from typing import List, Optional
 from sqlalchemy import sql
 
+from ..db_cache import TTLCacheDecorator, CacheEnum
 
 from ..util import prevent_sql_injection
 
@@ -13,6 +14,17 @@ def get(user_id: str) -> User:
     return session.query(User).get(user_id)
 
 
+@TTLCacheDecorator(CacheEnum.USER, 5, "user_id")
+def get_user_cached(user_id: str) -> User:
+    user = get(user_id)
+    if not user:
+        return None
+
+    general.expunge(user)
+    general.make_transient(user)
+    return user
+
+
 def get_by_id_list(user_ids: List[str]) -> List[User]:
     return session.query(User).filter(User.id.in_(user_ids)).all()
 

cognition_objects/project.py

@@ -1,13 +1,14 @@
 from typing import List, Optional, Dict, Any, Iterable
 from ..business_objects import general, team_resource, user
 from ..cognition_objects import consumption_log, consumption_summary
-from ..session import session
+from ..session import session, request_id_ctx_var
 from ..models import CognitionProject, TeamMember, TeamResource
 from .. import enums
 from datetime import datetime
 from ..util import prevent_sql_injection
 from sqlalchemy.orm.attributes import flag_modified
 from copy import deepcopy
+from ..db_cache import TTLCacheDecorator, CacheEnum
 
 
 def get(project_id: str) -> CognitionProject:
@@ -18,6 +19,16 @@ def get(project_id: str) -> CognitionProject:
     )
 
 
+@TTLCacheDecorator(CacheEnum.PROJECT, 5, "project_id")
+def get_cached(project_id: str) -> CognitionProject:
+    p = get(project_id)
+    if not p:
+        return None
+    general.expunge(p)
+    general.make_transient(p)
+    return p
+
+
 def get_org_id(project_id: str) -> str:
     if p := get(project_id):
         return str(p.organization_id)
@@ -42,6 +53,16 @@ def get_by_user(project_id: str, user_id: str) -> CognitionProject:
     )
 
 
+@TTLCacheDecorator(CacheEnum.PROJECT, 5, "project_id", "user_id")
+def get_by_user_cached(project_id: str, user_id: str) -> CognitionProject:
+    p = get_by_user(project_id, user_id)
+    if not p:
+        return None
+    general.expunge(p)
+    general.make_transient(p)
+    return p
+
+
 def get_all(org_id: str, order_by_name: bool = False) -> List[CognitionProject]:
     query = session.query(CognitionProject).filter(
         CognitionProject.organization_id == org_id

daemon.py

@@ -1,10 +1,7 @@
 import threading
 from submodules.model.business_objects import general
-from contextvars import ContextVar
 import traceback
 
-thread_session_token = ContextVar("token", default=None)
-
 
 def run_without_db_token(target, *args, **kwargs):
     """
@@ -27,35 +24,22 @@ def run_with_db_token(target, *args, **kwargs):
 
     # this is a workaround to set the token in the actual thread context
     def wrapper():
-        token = general.get_ctx_token()
-        thread_session_token.set(token)
+        general.get_ctx_token()
         try:
             target(*args, **kwargs)
         except Exception:
             print("=== Exception in thread ===", flush=True)
             print(traceback.format_exc(), flush=True)
             print("===========================", flush=True)
         finally:
-            reset_session_token_in_thread(False)
+            general.remove_and_refresh_session()
 
     threading.Thread(
         target=wrapper,
         daemon=True,
     ).start()
 
 
-def reset_session_token_in_thread(request_new: bool = True):
-    token = thread_session_token.get()
-    if not token:
-        # shouldn't happen if used with the run_with_session_token function
-        # so we print where it was called from
-        print(traceback.format_stack())
-        raise ValueError("No token set in thread context")
-    new_token = general.remove_and_refresh_session(token, request_new)
-    if new_token:
-        thread_session_token.set(new_token)
-
-
 def prepare_thread(target, *args, **kwargs) -> threading.Thread:
     return threading.Thread(
         target=target,