Track DB usage metrics

run_server.py

@@ -3,6 +3,9 @@
 import cherrypy
 
 import uber.server
+import logging
+logging.basicConfig()
+logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
 
 if __name__ == '__main__':
     cherrypy.engine.start()

uber/models/__init__.py

@@ -32,7 +32,7 @@
 from sqlalchemy.util import immutabledict
 
 import uber
-from uber.config import c, create_namespace_uuid
+from uber.config import c, create_namespace_uuid, threadlocal
 from uber.errors import HTTPRedirect
 from uber.decorators import cost_property, department_id_adapter, presave_adjustment, suffix_property
 from uber.models.types import Choice, DefaultColumn as Column, MultiChoice, utcnow
@@ -44,7 +44,6 @@
 def _make_getter(model):
     def getter(
             self, params=None, *, bools=(), checkgroups=(), allowed=(), restricted=False, ignore_csrf=False, **query):
-
         if query:
             return self.query(model).filter_by(**query).one()
         elif isinstance(params, str):
@@ -630,6 +629,26 @@ def minutestr(dt):
 from uber.models.tracking import Tracking  # noqa: E402
 
 
+import sqlalchemy.event
+import traceback
+import os
+
+base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+
+def catch_queries_callback(connection, cursor, query, *_):
+    stacktrace = [
+        (file_path, val1, val2, line_contents)
+        for (file_path, val1, val2, line_contents) in traceback.extract_stack()
+        if file_path.startswith(base_path)
+    ]
+
+    print("\n\n")
+    for i in stacktrace:
+        print(f"\t{i}")
+    print()
+
+
+
 class Session(SessionManager):
     # This looks strange, but `sqlalchemy.create_engine` will throw an error
     # if it's passed arguments that aren't supported by the given DB engine.
@@ -642,6 +661,7 @@ class Session(SessionManager):
         ('pool_pre_ping', True),
         ('pool_recycle', c.SQLALCHEMY_POOL_RECYCLE)] if v > -1)
     engine = sqlalchemy.create_engine(c.SQLALCHEMY_URL, **_engine_kwargs)
+    sqlalchemy.event.listen(engine, "before_cursor_execute", catch_queries_callback)
 
     @classmethod
     def initialize_db(cls, modify_tables=False, drop=False, initialize=False):
@@ -747,8 +767,14 @@ def iexact(self, **filters):
 
     class SessionMixin:
         def current_admin_account(self):
-            if getattr(cherrypy, 'session', {}).get('account_id'):
-                return self.admin_account(cherrypy.session.get('account_id'))
+            admin_account = threadlocal.get("current_admin_account")
+            if admin_account is None:
+                if getattr(cherrypy, 'session', {}).get('account_id'):
+                    admin_account = self.admin_account(cherrypy.session.get('account_id'))
+                    threadlocal.set("current_admin_account", admin_account)
+            else:
+                self.add(admin_account)
+            return admin_account
 
         def admin_attendee(self):
             if getattr(cherrypy, 'session', {}).get('account_id'):

uber/models/admin.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta
 
 import cherrypy
+import functools
 from pockets import classproperty, listify
 from pytz import UTC
 from residue import CoerceUTF8 as UnicodeText, UTCDateTime, UUID
@@ -10,7 +11,7 @@
 from sqlalchemy.schema import ForeignKey, Table, UniqueConstraint, Index
 from sqlalchemy.types import Boolean, Date, Integer
 
-from uber.config import c
+from uber.config import c, threadlocal
 from uber.decorators import presave_adjustment
 from uber.models import MagModel
 from uber.models.types import default_relationship as relationship, utcnow, DefaultColumn as Column
@@ -30,7 +31,6 @@
     Index('ix_admin_access_group_access_group_id', 'access_group_id'),
 )
 
-
 class AdminAccount(MagModel):
     attendee_id = Column(UUID, ForeignKey('attendee.id'), unique=True)
     access_groups = relationship(
@@ -75,14 +75,23 @@ def admin_email():
 
     @staticmethod
     def get_access_set(id=None, include_read_only=False):
+        cache_key = f"get_access_list({id}, {include_read_only})"
+        cached = threadlocal.get(cache_key)
+        if cached is not None:
+            return cached
         try:
             from uber.models import Session
             with Session() as session:
-                id = id or cherrypy.session.get('account_id')
-                account = session.admin_account(id)
+                if id:
+                    account = session.admin_account(id)
+                else:
+                    account = session.current_admin_account()
                 if include_read_only:
-                    return account.read_or_write_access_set
-                return account.write_access_set
+                    access_set = account.read_or_write_access_set
+                else:
+                    access_set = account.write_access_set
+            threadlocal.set(cache_key, access_set)
+            return access_set
         except Exception:
             return set()