1515#
1616
1717import logging
18- import os
1918from typing import Annotated
2019
2120from deepdiff import DeepDiff
3332from feeds_operations_gen .models .update_request_gtfs_rt_feed import (
3433 UpdateRequestGtfsRtFeed ,
3534)
36- from helpers .database import start_db_session , refresh_materialized_view
35+ from helpers .database import with_db_session , refresh_materialized_view
3736from helpers .query_helper import query_feed_by_stable_id
3837from .models .update_request_gtfs_rt_feed_impl import UpdateRequestGtfsRtFeedImpl
3938from .request_validator import validate_request
@@ -107,21 +106,21 @@ async def update_gtfs_rt_feed(
107106 - 400: Feed ID not found.
108107 - 500: Internal server error.
109108 """
110- return await self ._update_feed (update_request_gtfs_rt_feed , DataType .GTFS_RT )
109+ return self ._update_feed (update_request_gtfs_rt_feed , DataType .GTFS_RT )
111110
111+ @with_db_session
112112 async def _update_feed (
113113 self ,
114114 update_request_feed : UpdateRequestGtfsFeed | UpdateRequestGtfsRtFeed ,
115115 data_type : DataType ,
116+ db_session ,
116117 ) -> Response :
117118 """
118119 Update the specified feed in the Mobility Database
119120 """
120- session = None
121121 try :
122- session = start_db_session (os .getenv ("FEEDS_DATABASE_URL" ))
123122 feed = await OperationsApiImpl .fetch_feed (
124- data_type , session , update_request_feed
123+ data_type , db_session , update_request_feed
125124 )
126125
127126 logging .info (
@@ -139,14 +138,14 @@ async def _update_feed(
139138 and update_request_feed .operational_status_action != "no_change"
140139 ):
141140 await OperationsApiImpl ._populate_feed_values (
142- feed , impl_class , session , update_request_feed
141+ feed , impl_class , db_session , update_request_feed
143142 )
144- session .flush ()
145- refreshed = refresh_materialized_view (session , t_feedsearch .name )
143+ db_session .flush ()
144+ refreshed = refresh_materialized_view (db_session , t_feedsearch .name )
146145 logging .info (
147146 f"Materialized view { t_feedsearch .name } refreshed: { refreshed } "
148147 )
149- session .commit ()
148+ db_session .commit ()
150149 logging .info (
151150 f"Feed ID: { update_request_feed .id } updated successfully with the following changes: "
152151 f"{ diff .values ()} "
@@ -161,13 +160,9 @@ async def _update_feed(
161160 logging .error (
162161 f"Failed to update feed ID: { update_request_feed .id } . Error: { e } "
163162 )
164- session .rollback ()
165163 if isinstance (e , HTTPException ):
166164 raise e
167165 raise HTTPException (status_code = 500 , detail = f"Internal server error: { e } " )
168- finally :
169- if session :
170- session .close ()
171166
172167 @staticmethod
173168 async def _populate_feed_values (feed , impl_class , session , update_request_feed ):
0 commit comments