From d1f7a4ba6766f7c1a186780b69b193ea0483f2fe Mon Sep 17 00:00:00 2001
From: Jingsi Lu <jingsi@mobilitydata.org>
Date: Tue, 17 Dec 2024 14:15:57 -0500
Subject: [PATCH] lint error fixes

---
 .../feed_sync_process_transitland/src/main.py |  2 +-
 functions-python/helpers/database.py          |  2 +-
 .../impl/feeds_operations_impl.py             | 23 ++++++++-----------
 3 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/functions-python/feed_sync_process_transitland/src/main.py b/functions-python/feed_sync_process_transitland/src/main.py
index 1e938c48e..d895d516a 100644
--- a/functions-python/feed_sync_process_transitland/src/main.py
+++ b/functions-python/feed_sync_process_transitland/src/main.py
@@ -25,7 +25,7 @@
 from sqlalchemy.orm import Session
 
 from helpers.database import Database, configure_polymorphic_mappers
-from helpers.logger import Logger, StableIdFilter
+from helpers.logger import Logger
 from database_gen.sqlacodegen_models import Feed
 from helpers.feed_sync.models import TransitFeedSyncPayload as FeedPayload
 from .feed_processor_utils import check_url_status, create_new_feed
diff --git a/functions-python/helpers/database.py b/functions-python/helpers/database.py
index c9107bc75..3cd09e6f7 100644
--- a/functions-python/helpers/database.py
+++ b/functions-python/helpers/database.py
@@ -18,7 +18,7 @@
 import logging
 import os
 import threading
-from typing import Final, Optional
+from typing import Optional
 
 from sqlalchemy import create_engine, text, event, Engine
 from sqlalchemy.orm import sessionmaker, Session, mapper, class_mapper
diff --git a/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py b/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py
index 0cf88530b..5923acc10 100644
--- a/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py
+++ b/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py
@@ -15,7 +15,6 @@
 #
 
 import logging
-import os
 from typing import Annotated
 
 from deepdiff import DeepDiff
@@ -33,7 +32,7 @@
 from feeds_operations_gen.models.update_request_gtfs_rt_feed import (
     UpdateRequestGtfsRtFeed,
 )
-from helpers.database import start_db_session, refresh_materialized_view
+from helpers.database import with_db_session, refresh_materialized_view
 from helpers.query_helper import query_feed_by_stable_id
 from .models.update_request_gtfs_rt_feed_impl import UpdateRequestGtfsRtFeedImpl
 from .request_validator import validate_request
@@ -107,21 +106,21 @@ async def update_gtfs_rt_feed(
             - 400: Feed ID not found.
             - 500: Internal server error.
         """
-        return await self._update_feed(update_request_gtfs_rt_feed, DataType.GTFS_RT)
+        return self._update_feed(update_request_gtfs_rt_feed, DataType.GTFS_RT)
 
+    @with_db_session
     async def _update_feed(
         self,
         update_request_feed: UpdateRequestGtfsFeed | UpdateRequestGtfsRtFeed,
         data_type: DataType,
+        db_session,
     ) -> Response:
         """
         Update the specified feed in the Mobility Database
         """
-        session = None
         try:
-            session = start_db_session(os.getenv("FEEDS_DATABASE_URL"))
             feed = await OperationsApiImpl.fetch_feed(
-                data_type, session, update_request_feed
+                data_type, db_session, update_request_feed
             )
 
             logging.info(
@@ -139,14 +138,14 @@ async def _update_feed(
                 and update_request_feed.operational_status_action != "no_change"
             ):
                 await OperationsApiImpl._populate_feed_values(
-                    feed, impl_class, session, update_request_feed
+                    feed, impl_class, db_session, update_request_feed
                 )
-                session.flush()
-                refreshed = refresh_materialized_view(session, t_feedsearch.name)
+                db_session.flush()
+                refreshed = refresh_materialized_view(db_session, t_feedsearch.name)
                 logging.info(
                     f"Materialized view {t_feedsearch.name} refreshed: {refreshed}"
                 )
-                session.commit()
+                db_session.commit()
                 logging.info(
                     f"Feed ID: {update_request_feed.id} updated successfully with the following changes: "
                     f"{diff.values()}"
@@ -161,13 +160,9 @@ async def _update_feed(
             logging.error(
                 f"Failed to update feed ID: {update_request_feed.id}. Error: {e}"
             )
-            session.rollback()
             if isinstance(e, HTTPException):
                 raise e
             raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
-        finally:
-            if session:
-                session.close()
 
     @staticmethod
     async def _populate_feed_values(feed, impl_class, session, update_request_feed):