refactored cloud functions db session management

Author: qcdyx

functions-python/gbfs_validator/src/main.py

@@ -18,7 +18,7 @@
     PipelineStage,
     MaxExecutionsReachedError,
 )
-from helpers.database import start_db_session
+from helpers.database import Database
 from helpers.logger import Logger, StableIdFilter
 from helpers.parser import jsonify_pubsub
 from .gbfs_utils import (
@@ -34,19 +34,18 @@
 
 
 def fetch_all_gbfs_feeds() -> List[Gbfsfeed]:
-    session = None
+    db = Database(database_url=os.getenv("FEEDS_DATABASE_URL"))
     try:
-        session = start_db_session(os.getenv("FEEDS_DATABASE_URL"))
-        gbfs_feeds = (
-            session.query(Gbfsfeed).options(joinedload(Gbfsfeed.gbfsversions)).all()
-        )
-        return gbfs_feeds
+        with db.start_db_session() as session:
+            gbfs_feeds = (
+                session.query(Gbfsfeed).options(joinedload(Gbfsfeed.gbfsversions)).all()
+            )
+            return gbfs_feeds
     except Exception as e:
         logging.error(f"Error fetching all GBFS feeds: {e}")
         raise e
     finally:
-        if session:
-            session.close()
+        pass
 
 
 @functions_framework.cloud_event
@@ -92,7 +91,6 @@ def gbfs_validator_pubsub(cloud_event: CloudEvent):
             save_trace_with_error(trace, error_message, trace_service)
             return error_message
 
-        session = None
         try:
             storage_client = storage.Client()
             bucket = storage_client.bucket(BUCKET_NAME)
@@ -108,17 +106,16 @@ def gbfs_validator_pubsub(cloud_event: CloudEvent):
         try:
             snapshot = validator.create_snapshot(feed_id)
             validation_results = validator.validate_gbfs_feed(bucket)
-            session = start_db_session(os.getenv("FEEDS_DATABASE_URL"))
-            save_snapshot_and_report(session, snapshot, validation_results)
-
+            db = Database(database_url=os.getenv("FEEDS_DATABASE_URL"))
+            with db.start_db_session() as session:
+                save_snapshot_and_report(session, snapshot, validation_results)
         except Exception as e:
             error_message = f"Error validating GBFS feed: {e}"
             logging.error(f"{error_message}\nTraceback:\n{traceback.format_exc()}")
             save_trace_with_error(trace, error_message, trace_service)
             return error_message
         finally:
-            if session:
-                session.close()
+            pass
 
         trace.status = Status.SUCCESS
         trace_service.save(trace)

functions-python/helpers/feed_sync/feed_sync_dispatcher.py

@@ -18,6 +18,7 @@
 import os
 import logging
 
+from helpers.database import Database
 from helpers.feed_sync.feed_sync_common import FeedSyncProcessor
 from helpers.pub_sub import get_pubsub_client, publish
 

functions-python/validation_report_processor/src/main.py

@@ -18,7 +18,7 @@
 import logging
 from datetime import datetime
 import requests
-from database.database import Database
+from helpers.database import Database
 import functions_framework
 from database_gen.sqlacodegen_models import (
     Validationreport,