def test_s3_with_kwargs(self, session):
from botocore.client import Config
kw = dict(
region_name="us-east-1",
api_version="0.0.1",
endpoint_url="https://minio.local/",
config=Config(
signature_version="s3v4",
s3=dict(
region_name="us-east-1",
addressing_style="path",
inject_host_prefix=False,
),
),
)
handler = S3ResultHandler(bucket="bob",
aws_credentials_secret="AWS_CREDENTIALS",
boto3_kwargs=kw)
with prefect.context(secrets=dict(
AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=999))):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_client()
assert handler.bucket == "bob"
kw.update({
"aws_access_key_id": 1,
"aws_secret_access_key": 999,
})
assert session.Session().client.call_args[1] == kw
def test_serialize(self):
handler = S3ResultHandler(bucket="my-bucket",
aws_credentials_secret="FOO")
serialized = ResultHandlerSchema().dump(handler)
assert serialized["type"] == "S3ResultHandler"
assert serialized["bucket"] == "my-bucket"
assert serialized["aws_credentials_secret"] == "FOO"
def test_basic_conversion_s3_result():
result_handler = S3ResultHandler(bucket="foo",
boto3_kwargs=dict(x=42, y=[1, 2, 3]))
result = ResultHandlerResult.from_result_handler(result_handler)
assert isinstance(result, S3Result)
assert result.bucket == "foo"
assert result.boto3_kwargs == dict(x=42, y=[1, 2, 3])
def test_s3_with_kwargs_invalid_service_name(self, session):
with pytest.raises(AssertionError) as ex:
handler = S3ResultHandler(
bucket="bob",
aws_credentials_secret="AWS_CREDENTIALS",
boto3_kwargs=dict(service_name="s3"),
)
assert str(ex.value) == 'Changing the boto3 "service_name" is not permitted!'
def test_s3_client_init_uses_secrets(self, s3_client):
with prefect.context(secrets=dict(
AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42))):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler = S3ResultHandler(bucket="bob")
assert handler.bucket == "bob"
assert s3_client.call_args[1] == {
"aws_access_key_id": 1,
"aws_secret_access_key": 42,
}
def test_s3_uninitialized_client(self, session):
handler = S3ResultHandler(bucket="bob",
aws_credentials_secret="AWS_CREDENTIALS")
assert handler.bucket == "bob"
with prefect.context(
secrets=dict(
AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=42)),
boto3client="test",
):
with set_temporary_config({"cloud.use_local_secrets": True}):
assert handler.client is not None
def test_roundtrip_never_loads_client(self, monkeypatch):
schema = ResultHandlerSchema()
def raise_me(*args, **kwargs):
raise SyntaxError("oops")
monkeypatch.setattr(S3ResultHandler, "initialize_client", raise_me)
handler = schema.load(
schema.dump(S3ResultHandler(bucket="bucket3", aws_credentials_secret="FOO"))
)
assert isinstance(handler, S3ResultHandler)
assert handler.bucket == "bucket3"
assert handler.aws_credentials_secret == "FOO"
def test_s3_writes_to_blob_prefixed_by_date_suffixed_by_prefect(self, s3_client):
with prefect.context(
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")
uri = handler.write("so-much-data")
used_uri = s3_client.return_value.upload_fileobj.call_args[1]["Key"]
assert used_uri == uri
assert used_uri.startswith(pendulum.now("utc").format("Y/M/D"))
assert used_uri.endswith("prefect_result")
def test_s3_client_init_uses_custom_secrets(self, s3_client):
handler = S3ResultHandler(bucket="bob", aws_credentials_secret="MY_FOO")
with prefect.context(
secrets=dict(MY_FOO=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=999))
):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_client()
assert handler.bucket == "bob"
assert s3_client.call_args[1] == {
"aws_access_key_id": 1,
"aws_secret_access_key": 999,
}
def test_s3_handler_is_pickleable(self, monkeypatch):
class client:
def __init__(self, *args, **kwargs):
pass
def __getstate__(self):
raise ValueError("I cannot be pickled.")
with patch.dict("sys.modules", {"boto3": MagicMock(client=client)}):
with prefect.context(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")
res = cloudpickle.loads(cloudpickle.dumps(handler))
assert isinstance(res, S3ResultHandler)
def __init__(
self,
bucket: str,
aws_access_key_id: str = None,
aws_secret_access_key: str = None,
aws_session_token: str = None,
key: str = None,
) -> None:
self.flows = dict() # type: Dict[str, str]
self._flows = dict() # type: Dict[str, "prefect.core.flow.Flow"]
self.bucket = bucket
self.key = key
self.aws_access_key_id = aws_access_key_id
self.aws_secret_access_key = aws_secret_access_key
self.aws_session_token = aws_session_token
result_handler = S3ResultHandler(bucket=bucket)
super().__init__(result_handler=result_handler)
def __init__(
self,
bucket: str,
aws_access_key_id: str = None,
aws_secret_access_key: str = None,
aws_session_token: str = None,
client_options: dict = None,
key: str = None,
) -> None:
self.flows = dict() # type: Dict[str, str]
self._flows = dict() # type: Dict[str, "Flow"]
self.bucket = bucket
self.key = key
self.aws_access_key_id = aws_access_key_id
self.aws_secret_access_key = aws_secret_access_key
self.aws_session_token = aws_session_token
self.client_options = client_options
result_handler = S3ResultHandler(bucket=bucket)
super().__init__(result_handler=result_handler)
def test_s3_with_kwargs_aws_keys(self, session):
handler = S3ResultHandler(
bucket="bob",
boto3_kwargs=dict(
aws_access_key_id=123,
aws_secret_access_key=456,
),
)
assert ("aws_access_key_id" in handler.boto3_kwargs.keys()
), 'Missing "aws_access_key_id" in boto_kwargs'
assert ("aws_secret_access_key" in handler.boto3_kwargs.keys()
), 'Missing "aws_secret_access_key" in boto_kwargs'
with prefect.context(secrets=dict(
AWS_CREDENTIALS=dict(ACCESS_KEY=1, SECRET_ACCESS_KEY=999))):
with set_temporary_config({"cloud.use_local_secrets": True}):
handler.initialize_client()
assert handler.bucket == "bob"
assert session.Session().client.call_args[1] == {
"aws_access_key_id": 123,
"aws_secret_access_key": 456,
}
def test_roundtrip(self):
schema = ResultHandlerSchema()
handler = schema.load(schema.dump(S3ResultHandler(bucket="bucket3")))
assert isinstance(handler, S3ResultHandler)
assert handler.bucket == "bucket3"
def test_serialize(self, s3_client):
handler = S3ResultHandler(bucket="my-bucket")
serialized = ResultHandlerSchema().dump(handler)
assert serialized["type"] == "S3ResultHandler"
assert serialized["bucket"] == "my-bucket"
@task(
cache_for=datetime.timedelta(days=10),
cache_validator=partial_parameters_only(["x"]),
result_handler=JSONResultHandler(),
)
def cached_task(x, y):
pass
x = Parameter("x")
y = Parameter("y", default=42)
@task(name="Big Name",
checkpoint=True,
result_handler=S3ResultHandler(bucket="blob"))
def terminal_task():
pass
env = RemoteEnvironment(
executor="prefect.engine.executors.DaskExecutor",
executor_kwargs={"scheduler_address": "tcp://"},
)
storage = Docker(
registry_url="prefecthq",
image_name="flows",
image_tag="welcome-flow",
python_dependencies=["boto3"],
)