switch to async logic for region extraction

Author: DenisaCG

jupyter_drives/manager.py

@@ -16,7 +16,7 @@
 from libcloud.storage.types import Provider
 from libcloud.storage.providers import get_driver
 import pyarrow
-import boto3
+from aiobotocore.session import get_session
 
 from .log import get_logger
 from .base import DrivesConfig
@@ -42,19 +42,10 @@ def __init__(self, config: traitlets.config.Config) -> None:
         self._content_managers = {}
         self._max_files_listed = 1000
 
-         # initiate boto3 session if we are dealing with S3 drives
+         # initiate aiobotocore session if we are dealing with S3 drives
         if self._config.provider == 's3':
             self._s3_clients = {}
-            if self._config.access_key_id and self._config.secret_access_key:
-                if self._config.session_token is None:
-                    self._s3_session = boto3.Session(aws_access_key_id = self._config.access_key_id, aws_secret_access_key = self._config.secret_access_key)
-                else:
-                    self._s3_session = boto3.Session(aws_access_key_id = self._config.access_key_id, aws_secret_access_key = self._config.secret_access_key, aws_session_token = self._config.session_token)
-            else:
-                raise tornado.web.HTTPError(
-                status_code= httpx.codes.BAD_REQUEST,
-                reason="No credentials specified. Please set them in your user jupyter_server_config file.",
-                )
+            self._s3_session = get_session()
 
     @property
     def base_api_url(self) -> str:
@@ -149,7 +140,7 @@ async def mount_drive(self, drive_name, provider):
             if drive_name not in self._content_managers or self._content_managers[drive_name] is None:
                 if provider == 's3':
                     # get region of drive
-                    region = self._get_drive_location(drive_name)
+                    region = await self._get_drive_location(drive_name)
                     if self._config.session_token is None:
                         configuration = {
                             "aws_access_key_id": self._config.access_key_id,
@@ -555,7 +546,7 @@ async def presigned_link(self, drive_name, path):
             }
         return response
 
-    def _get_drive_location(self, drive_name):
+    async def _get_drive_location(self, drive_name):
         """Helping function for getting drive region.
 
         Args:
@@ -564,10 +555,9 @@ def _get_drive_location(self, drive_name):
         location = 'eu-north-1'
         try:
             # set temporary client for location extraction
-            s3 = self._s3_session.client('s3')
-            result = s3.get_bucket_location(Bucket = drive_name)
-
-            location = result['LocationConstraint']
+            async with self._s3_session.create_client('s3', aws_secret_access_key=self._config.secret_access_key, aws_access_key_id=self._config.access_key_id, aws_session_token=self._config.session_token) as client:
+                result = await client.get_bucket_location(Bucket=drive_name)
+                location = result['LocationConstraint']
         except Exception as e:
              raise tornado.web.HTTPError(
             status_code= httpx.codes.BAD_REQUEST,