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"], )