From 9590e9484dbf0b4a5b20752f9276f230a407ebf2 Mon Sep 17 00:00:00 2001
From: JerrySentry <jerry.feng@sentry.io>
Date: Wed, 21 Feb 2024 14:16:18 -0500
Subject: [PATCH] Bundle analysis: auto open and close DB session

---
 shared/bundle_analysis/report.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/shared/bundle_analysis/report.py b/shared/bundle_analysis/report.py
index 6a6cc1064..6aacccca7 100644
--- a/shared/bundle_analysis/report.py
+++ b/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)