Skip to content

Commit 27ea413

Browse files
committed
iterate on session_token support
1 parent a058be0 commit 27ea413

File tree

1 file changed

+18
-2
lines changed

1 file changed

+18
-2
lines changed

jupyter_drives/manager.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,10 @@ def __init__(self, config: traitlets.config.Config) -> None:
4545
if self._config.provider == 's3':
4646
self._s3_clients = {}
4747
if self._config.access_key_id and self._config.secret_access_key:
48-
self._s3_session = boto3.Session(aws_access_key_id = self._config.access_key_id, aws_secret_access_key = self._config.secret_access_key)
48+
if self._config.session_token is None:
49+
self._s3_session = boto3.Session(aws_access_key_id = self._config.access_key_id, aws_secret_access_key = self._config.secret_access_key)
50+
else:
51+
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)
4952
else:
5053
raise tornado.web.HTTPError(
5154
status_code= httpx.codes.BAD_REQUEST,
@@ -134,7 +137,20 @@ async def mount_drive(self, drive_name, provider, region):
134137
# check if content manager doesn't already exist
135138
if drive_name not in self._content_managers or self._content_managers[drive_name] is None:
136139
if provider == 's3':
137-
store = obs.store.S3Store.from_url("s3://" + drive_name + "/", config = {"aws_access_key_id": self._config.access_key_id, "aws_secret_access_key": self._config.secret_access_key, "aws_region": region})
140+
if self._config.session_token is None:
141+
configuration = {
142+
"aws_access_key_id": self._config.access_key_id,
143+
"aws_secret_access_key": self._config.secret_access_key,
144+
"aws_region": region
145+
}
146+
else:
147+
configuration = {
148+
"aws_access_key_id": self._config.access_key_id,
149+
"aws_secret_access_key": self._config.secret_access_key,
150+
"aws_session_token": self._config.session_token,
151+
"aws_region": region
152+
}
153+
store = obs.store.S3Store.from_url("s3://" + drive_name + "/", config = configuration)
138154
elif provider == 'gcs':
139155
store = obs.store.GCSStore.from_url("gs://" + drive_name + "/", config = {}) # add gcs config
140156
elif provider == 'http':

0 commit comments

Comments
 (0)