Bundle analysis: auto open and close DB session

shared/bundle_analysis/report.py

@@ -6,12 +6,27 @@
 
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.orm import Session as SQLAlchemySession
+from sqlalchemy.orm import sessionmaker
 from sqlalchemy.sql import func
 
 from shared.bundle_analysis import models
 from shared.bundle_analysis.parser import Parser
 
 
+def db_query(f):
+    def wrapper(*args, **kw):
+        if not args[0].db_session:
+            raise Exception(f"No DB session found for {f}")
+
+        engine = args[0].db_session.get_bind()
+        Session = sessionmaker(engine)
+        with Session() as session:
+            args[0].db_session = session
+            return f(*args, **kw)
+
+    return wrapper
+
+
 class AssetReport:
     """
     Report wrapper around a single asset (many of which can exist in a single bundle).
@@ -33,6 +48,7 @@ def hashed_name(self):
     def size(self):
         return self.asset.size
 
+    @db_query
     def modules(self):
         return (
             self.db_session.query(models.Module)
@@ -56,6 +72,7 @@ def __init__(self, bundle: models.Bundle):
     def name(self):
         return self.bundle.name
 
+    @db_query
     def asset_reports(self) -> Iterator[AssetReport]:
         assets = (
             self.db_session.query(models.Asset)
@@ -66,6 +83,7 @@ def asset_reports(self) -> Iterator[AssetReport]:
         )
         return (AssetReport(asset) for asset in assets)
 
+    @db_query
     def total_size(self) -> int:
         return (
             self.db_session.query(func.sum(models.Asset.size).label("asset_size"))
@@ -75,6 +93,7 @@ def total_size(self) -> int:
             .scalar()
         ) or 0
 
+    @db_query
     def info(self) -> dict:
         session = (
             self.db_session.query(models.Session)