Fixed closing of client sessions

Author: sveinugu

src/omnipy/data/dataset.py

@@ -43,6 +43,7 @@
 _DatasetT = TypeVar('_DatasetT')
 
 DATA_KEY = 'data'
+ASYNC_LOAD_SLEEP_TIME = 0.05
 
 # def orjson_dumps(v, *, default):
 #     # orjson.dumps returns bytes, to match standard json.dumps we need to decode
@@ -588,19 +589,22 @@ def _load_http_urls(self, http_url_dataset: 'HttpUrlDataset') -> list[asyncio.Ta
 
         async def load_all():
             tasks = []
-            client_sessions = {}
+
             for host in hosts:
-                client_sessions[host] = RateLimitingClientSession(
-                    self.config.http_config_for_host[host].requests_per_time_period,
-                    self.config.http_config_for_host[host].time_period_in_secs)
-
-            for host, indices in hosts.items():
-                task = (
-                    get_json_from_api_endpoint.refine(output_dataset_param='output_dataset').run(
-                        http_url_dataset[indices],
-                        client_session=client_sessions[host],
-                        output_dataset=self))
-                tasks.append(task)
+                async with RateLimitingClientSession(
+                        self.config.http_config_for_host[host].requests_per_time_period,
+                        self.config.http_config_for_host[host].time_period_in_secs
+                ) as client_session:
+                    indices = hosts[host]
+                    task = (
+                        get_json_from_api_endpoint.refine(
+                            output_dataset_param='output_dataset').run(
+                                http_url_dataset[indices],
+                                client_session=client_session,
+                                output_dataset=self))
+                    tasks.append(task)
+                    while not task.done():
+                        await asyncio.sleep(ASYNC_LOAD_SLEEP_TIME)
 
             await asyncio.gather(*tasks)
             return self

src/omnipy/modules/remote/tasks.py

@@ -92,7 +92,8 @@ async def get_json_from_api_endpoint(
             retry_backoff_strategy,
     ):
         async for response in _call_get(url, cast(ClientSession, retry_session)):
-            return JsonModel(await response.json(content_type=None))
+            output_data = JsonModel(await response.json(content_type=None))
+    return output_data
 
 
 @TaskTemplate(iterate_over_data_files=True, output_dataset_cls=StrDataset)
@@ -110,7 +111,8 @@ async def get_str_from_api_endpoint(
             retry_backoff_strategy,
     ):
         async for response in _call_get(url, cast(ClientSession, retry_session)):
-            return StrModel(await response.text())
+            output_data = StrModel(await response.text())
+    return output_data
 
 
 @TaskTemplate(iterate_over_data_files=True, output_dataset_cls=BytesDataset)
@@ -128,7 +130,8 @@ async def get_bytes_from_api_endpoint(
             retry_backoff_strategy,
     ):
         async for response in _call_get(url, cast(ClientSession, retry_session)):
-            return BytesModel(await response.read())
+            output_data = BytesModel(await response.read())
+    return output_data
 
 
 JsonDatasetT = TypeVar('JsonDatasetT', bound=Dataset)