From 189d029f93549fb39044c42f35b65a5b571c2b11 Mon Sep 17 00:00:00 2001
From: albertkun <albertk@gmx.com>
Date: Wed, 31 Jan 2024 12:43:23 -0800
Subject: [PATCH] refactor: Swiftly API connection to use async/await

---
 .../app/utils/gtfs_rt_helper.py               | 66 +++++++++----------
 data-loading-service/requirements.txt         |  4 +-
 fastapi/app/requirements.txt                  |  3 +-
 fastapi/requirements.txt                      |  3 +-
 requirements.txt                              |  4 +-
 5 files changed, 43 insertions(+), 37 deletions(-)

diff --git a/data-loading-service/app/utils/gtfs_rt_helper.py b/data-loading-service/app/utils/gtfs_rt_helper.py
index 8e9a8999..07b7c94f 100644
--- a/data-loading-service/app/utils/gtfs_rt_helper.py
+++ b/data-loading-service/app/utils/gtfs_rt_helper.py
@@ -43,21 +43,16 @@
 SWIFTLY_GTFS_RT_TRIP_UPDATES = 'gtfs-rt-trip-updates'
 SWIFTLY_GTFS_RT_VEHICLE_POSITIONS = 'gtfs-rt-vehicle-positions'
 
-# engine = create_engine(Config.API_DB_URI, echo=False,executemany_mode="values")
-
-# Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
-
-# insp = inspect(engine)
-# session = Session()
-
 SERVICE_DICT = {
     'LACMTA': 'lametro',
     'LACMTA_Rail': 'lametro-rail'
 }
 
 SWIFTLY_AGENCY_IDS = ['LACMTA', 'LACMTA_Rail']
+import asyncio
 
 # Connect to the database
+
 def connect_to_db():
     try:
         print('Connecting to the database')
@@ -69,29 +64,6 @@ def connect_to_db():
     finally:
         session.close()
 
-def connect_to_swiftly(service, endpoint):
-    swiftly_endpoint = ''
-    swiftly_endpoint = SWIFTLY_API_REALTIME + service + '/' + endpoint
-
-    if (service == 'lametro'):
-        key = Config.SWIFTLY_AUTH_KEY_BUS
-    elif (service == 'lametro-rail'):
-        key = Config.SWIFTLY_AUTH_KEY_RAIL
-    header = { 
-        "Authorization": key
-    }
-    try:
-        print('Connecting to Swiftly API: ' + swiftly_endpoint)
-        response = requests.get(swiftly_endpoint, headers=header)
-        print('Response status code: ' + str(response.status_code))
-        if (response.status_code == 200):
-            return response.content
-        else:
-            return False
-    except Exception as e:
-        print.exception('Error connecting to Swiftly API: ' + str(e))
-        return False
-
 def get_agency_id(service):
     if (service == 'bus'):
         return 'LACMTA'
@@ -124,13 +96,41 @@ def get_route_code_from_trip_route_id(trip_id,agency_id):
         val = str(trip_id).split('-')[0]
     return val
 
-def update_gtfs_realtime_data():
+import aiohttp
+import asyncio
+
+async def connect_to_swiftly(service, endpoint):
+    swiftly_endpoint = SWIFTLY_API_REALTIME + service + '/' + endpoint
+
+    if (service == 'lametro'):
+        key = Config.SWIFTLY_AUTH_KEY_BUS
+    elif (service == 'lametro-rail'):
+        key = Config.SWIFTLY_AUTH_KEY_RAIL
+    header = { 
+        "Authorization": key
+    }
+    try:
+        print('Connecting to Swiftly API: ' + swiftly_endpoint)
+        async with aiohttp.ClientSession() as session:
+            async with session.get(swiftly_endpoint, headers=header) as response:
+                print('Response status code: ' + str(response.status))
+                if (response.status == 200):
+                    return await response.read()
+                else:
+                    return False
+    except Exception as e:
+        print.exception('Error connecting to Swiftly API: ' + str(e))
+        return False
+
+async def update_gtfs_realtime_data():
     process_start = timeit.default_timer()
-    connect_to_db()
+    await connect_to_db()
     combined_trip_update_dataframes = []
     combined_stop_time_dataframes = []
     combined_vehicle_position_dataframes = []
-    
+
+    tasks = [process_agency(agency) for agency in SWIFTLY_AGENCY_IDS]
+    await asyncio.gather(*tasks)
     for agency in SWIFTLY_AGENCY_IDS:
         feed = FeedMessage()
         response_data = connect_to_swiftly(SERVICE_DICT[agency], SWIFTLY_GTFS_RT_TRIP_UPDATES)
diff --git a/data-loading-service/requirements.txt b/data-loading-service/requirements.txt
index d84d5e20..4c70eae6 100644
--- a/data-loading-service/requirements.txt
+++ b/data-loading-service/requirements.txt
@@ -24,4 +24,6 @@ protobuf==3.19.6
 polyline
 versiontag
 gitpython
-crython
\ No newline at end of file
+crython
+asyncio
+aiohttp
\ No newline at end of file
diff --git a/fastapi/app/requirements.txt b/fastapi/app/requirements.txt
index aa9554a6..8bc3bce3 100644
--- a/fastapi/app/requirements.txt
+++ b/fastapi/app/requirements.txt
@@ -33,4 +33,5 @@ gitpython
 fastapi-pagination
 geojson
 pytest
-pickle
\ No newline at end of file
+pickle
+httpx
\ No newline at end of file
diff --git a/fastapi/requirements.txt b/fastapi/requirements.txt
index 15ddf754..67330efb 100644
--- a/fastapi/requirements.txt
+++ b/fastapi/requirements.txt
@@ -32,4 +32,5 @@ gitpython
 fastapi-pagination
 geojson
 pytest
-asyncpg
\ No newline at end of file
+asyncpg
+httpx
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 9ab0e194..141619c0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -32,4 +32,6 @@ fastapi-cache2[redis]
 aioredis
 versiontag
 gitpython
-crython
\ No newline at end of file
+crython
+aiohttp
+asyncio
\ No newline at end of file