Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closes #2108. New boto session for each thread #2136

Merged
merged 8 commits into from
Mar 11, 2020
Prev Previous commit
Next Next commit
Use new boto clients if none in context for S3ResultHandler
In retrospect, I probably don't need to get in the cache key but can just go on context directly.
  • Loading branch information
lauralorenz committed Mar 9, 2020
commit a150c107db202fbfb76f0e7f2a66d25d8555b382
4 changes: 2 additions & 2 deletions src/prefect/engine/result_handlers/s3_result_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,9 @@ def initialize_client(self) -> None:

@property
def client(self) -> "boto3.client":
if not prefect.context.caches.get("boto3client"):
if not prefect.context.get("boto3client"):
self.initialize_client()
prefect.context.caches["boto3client"] = self._client
prefect.context["boto3client"] = self._client
return self._client

@client.setter
Expand Down
11 changes: 4 additions & 7 deletions tests/engine/result_handlers/test_result_handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -181,8 +181,7 @@ def test_s3_client_init_uses_secrets(self, session):
assert session.Session().client.called is False

with prefect.context(
secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42)),
caches={},
secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42))
):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_client()
Expand All @@ -195,7 +194,7 @@ def test_s3_client_init_uses_custom_secrets(self, session):
handler = S3ResultHandler(bucket="bob", aws_credentials_secret="MY_FOO")

with prefect.context(
secrets=dict(MY_FOO=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=999)), caches={}
secrets=dict(MY_FOO=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=999))
):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_client()
Expand All @@ -210,8 +209,7 @@ def test_s3_writes_to_blob_prefixed_by_date_suffixed_by_prefect(self, session):
handler = S3ResultHandler(bucket="foo")

with prefect.context(
secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42)),
caches={},
secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42))
):
with set_temporary_config({"cloud.use_local_secrets": True}):
uri = handler.write("so-much-data")
Expand All @@ -238,8 +236,7 @@ def __getstate__(self):
boto3.session.Session().client = client

with prefect.context(
secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42)),
caches={},
secrets=dict(AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42))
):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler = S3ResultHandler(bucket="foo")
Expand Down