From 7a254923b9be43383d132f550e654e7534c979b6 Mon Sep 17 00:00:00 2001 From: kevingrismore <146098880+kevingrismore@users.noreply.github.com> Date: Fri, 19 Jan 2024 16:09:19 -0600 Subject: [PATCH] Support MinIO Credentials block as credentials for `push_to_s3` and `pull_from_s3` (#366) Co-authored-by: Chris Guidry --- CHANGELOG.md | 8 +- prefect_aws/deployments/steps.py | 14 +++- .../{deploments => deployments}/test_steps.py | 74 ++++++++++++++++++- 3 files changed, 88 insertions(+), 8 deletions(-) rename tests/{deploments => deployments}/test_steps.py (82%) diff --git a/CHANGELOG.md b/CHANGELOG.md index b40563f9..70227627 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Support MinIO Credentials as `credentials` dict for `push_to_s3` and `pull_from_s3` - [#366](https://github.com/PrefectHQ/prefect-aws/pull/366) + ### Changed - Handle `boto3` clients more efficiently with `lru_cache` - [#361](https://github.com/PrefectHQ/prefect-aws/pull/361) @@ -141,7 +143,7 @@ Released on June 13th, 2023. ### Fixed -- Change prefect.docker import to prefect.utilities.dockerutils to fix a crash when using custom blocks based on S3Bucket - [#273](https://github.com/PrefectHQ/prefect-aws/pull/273) +- Change prefect.docker import to prefect.utilities.dockerutils to fix a crash when using custom blocks based on S3Bucket - [#273](https://github.com/PrefectHQ/prefect-aws/pull/273) ## 0.3.2 @@ -230,7 +232,7 @@ Released on January 4th, 2023. - `list_objects`, `download_object_to_path`, `download_object_to_file_object`, `download_folder_to_path`, `upload_from_path`, `upload_from_file_object`, `upload_from_folder` methods in `S3Bucket` - [#85](https://github.com/PrefectHQ/prefect-aws/pull/175) - `aws_client_parameters` as a field in `AwsCredentials` and `MinioCredentials` blocks - [#175](https://github.com/PrefectHQ/prefect-aws/pull/175) -- `get_client` and `get_s3_client` methods to `AwsCredentials` and `MinioCredentials` blocks - [#175](https://github.com/PrefectHQ/prefect-aws/pull/175) +- `get_client` and `get_s3_client` methods to `AwsCredentials` and `MinioCredentials` blocks - [#175](https://github.com/PrefectHQ/prefect-aws/pull/175) ### Changed @@ -242,7 +244,7 @@ Released on January 4th, 2023. - `endpoint_url` field in S3Bucket; specify `aws_client_parameters` in `AwsCredentials` or `MinIOCredentials` instead - [#175](https://github.com/PrefectHQ/prefect-aws/pull/175) - `basepath` field in S3Bucket; specify `bucket_folder` instead - [#175](https://github.com/PrefectHQ/prefect-aws/pull/175) -- `minio_credentials` and `aws_credentials` field in S3Bucket; use the `credentials` field instead - [#175](https://github.com/PrefectHQ/prefect-aws/pull/175) +- `minio_credentials` and `aws_credentials` field in S3Bucket; use the `credentials` field instead - [#175](https://github.com/PrefectHQ/prefect-aws/pull/175) ## 0.2.1 diff --git a/prefect_aws/deployments/steps.py b/prefect_aws/deployments/steps.py index 7525a5e2..6161dfb8 100644 --- a/prefect_aws/deployments/steps.py +++ b/prefect_aws/deployments/steps.py @@ -62,7 +62,8 @@ def push_to_s3( bucket: The name of the S3 bucket where files will be uploaded. folder: The folder in the S3 bucket where files will be uploaded. credentials: A dictionary of AWS credentials (aws_access_key_id, - aws_secret_access_key, aws_session_token). + aws_secret_access_key, aws_session_token) or MinIO credentials + (minio_root_user, minio_root_password). client_parameters: A dictionary of additional parameters to pass to the boto3 client. ignore_file: The name of the file containing ignore patterns. @@ -139,7 +140,8 @@ def pull_from_s3( bucket: The name of the S3 bucket where files are stored. folder: The folder in the S3 bucket where files are stored. credentials: A dictionary of AWS credentials (aws_access_key_id, - aws_secret_access_key, aws_session_token). + aws_secret_access_key, aws_session_token) or MinIO credentials + (minio_root_user, minio_root_password). client_parameters: A dictionary of additional parameters to pass to the boto3 client. @@ -204,8 +206,12 @@ def get_s3_client( client_parameters = {} # Get credentials from credentials (regardless if block or not) - aws_access_key_id = credentials.get("aws_access_key_id", None) - aws_secret_access_key = credentials.get("aws_secret_access_key", None) + aws_access_key_id = credentials.get( + "aws_access_key_id", credentials.get("minio_root_user", None) + ) + aws_secret_access_key = credentials.get( + "aws_secret_access_key", credentials.get("minio_root_password", None) + ) aws_session_token = credentials.get("aws_session_token", None) # Get remaining session info from credentials, or client_parameters diff --git a/tests/deploments/test_steps.py b/tests/deployments/test_steps.py similarity index 82% rename from tests/deploments/test_steps.py rename to tests/deployments/test_steps.py index 22608bd7..15c4fc25 100644 --- a/tests/deploments/test_steps.py +++ b/tests/deployments/test_steps.py @@ -11,6 +11,16 @@ from prefect_aws.deployments.steps import get_s3_client, pull_from_s3, push_to_s3 +@pytest.fixture(scope="module", autouse=True) +def set_custom_endpoint(): + original = os.environ.get("MOTO_S3_CUSTOM_ENDPOINTS") + os.environ["MOTO_S3_CUSTOM_ENDPOINTS"] = "http://custom.minio.endpoint:9000" + yield + os.environ.pop("MOTO_S3_CUSTOM_ENDPOINTS") + if original is not None: + os.environ["MOTO_S3_CUSTOM_ENDPOINTS"] = original + + @pytest.fixture def s3_setup(): with mock_s3(): @@ -215,8 +225,15 @@ def test_s3_session_with_params(): }, ) get_s3_client(credentials=creds_block.dict()) + get_s3_client( + credentials={ + "minio_root_user": "MY_USER", + "minio_root_password": "MY_PASSWORD", + }, + client_parameters={"endpoint_url": "http://custom.minio.endpoint:9000"}, + ) all_calls = mock_session.mock_calls - assert len(all_calls) == 6 + assert len(all_calls) == 8 assert all_calls[0].kwargs == { "aws_access_key_id": "THE_KEY", "aws_secret_access_key": "SHHH!", @@ -265,6 +282,20 @@ def test_s3_session_with_params(): }.items() <= all_calls[5].kwargs.items() assert all_calls[5].kwargs.get("config").connect_timeout == 123 assert all_calls[5].kwargs.get("config").signature_version is None + assert all_calls[6].kwargs == { + "aws_access_key_id": "MY_USER", + "aws_secret_access_key": "MY_PASSWORD", + "aws_session_token": None, + "profile_name": None, + "region_name": None, + } + assert all_calls[7].args[0] == "s3" + assert { + "api_version": None, + "use_ssl": True, + "verify": None, + "endpoint_url": "http://custom.minio.endpoint:9000", + }.items() <= all_calls[7].kwargs.items() def test_custom_credentials_and_client_parameters(s3_setup, tmp_files): @@ -309,6 +340,47 @@ def test_custom_credentials_and_client_parameters(s3_setup, tmp_files): assert (tmp_path / file.name).exists() +def test_custom_credentials_and_client_parameters_minio(s3_setup, tmp_files): + s3, bucket_name = s3_setup + folder = "my-project" + + # Custom credentials and client parameters + custom_credentials = { + "minio_root_user": "fake_user", + "minio_root_password": "fake_password", + } + + custom_client_parameters = { + "endpoint_url": "http://custom.minio.endpoint:9000", + } + + os.chdir(tmp_files) + + # Test push_to_s3 with custom credentials and client parameters + push_to_s3( + bucket_name, + folder, + credentials=custom_credentials, + client_parameters=custom_client_parameters, + ) + + # Test pull_from_s3 with custom credentials and client parameters + tmp_path = tmp_files / "test_pull" + tmp_path.mkdir(parents=True, exist_ok=True) + os.chdir(tmp_path) + + pull_from_s3( + bucket_name, + folder, + credentials=custom_credentials, + client_parameters=custom_client_parameters, + ) + + for file in tmp_files.iterdir(): + if file.is_file() and file.name != ".prefectignore": + assert (tmp_path / file.name).exists() + + def test_without_prefectignore_file(s3_setup, tmp_files: Path, mock_aws_credentials): s3, bucket_name = s3_setup folder = "my-project"