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

feat(python): Support passing aws_profile in storage_options #20965

Merged
merged 2 commits into from
Jan 29, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 42 additions & 3 deletions py-polars/polars/io/cloud/credential_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, TypedDict, Union

from polars._utils.various import issue_warning

if TYPE_CHECKING:
if sys.version_info >= (3, 10):
from typing import TypeAlias
Expand Down Expand Up @@ -484,7 +486,7 @@ def _maybe_init_credential_provider(
elif k in {"azure_use_azure_cli", "use_azure_cli"}:
continue
elif k in OBJECT_STORE_CLIENT_OPTIONS:
pass
continue
else:
# We assume some sort of access key was given, so we
# just dispatch to the rust side.
Expand All @@ -503,7 +505,9 @@ def _maybe_init_credential_provider(
)
elif _is_aws_cloud(scheme):
region = None
profile = os.getenv("AWS_PROFILE")
default_region = None
unhandled_key = None

if storage_options is not None:
for k, v in storage_options.items():
Expand All @@ -514,14 +518,49 @@ def _maybe_init_credential_provider(
region = v
elif k in {"aws_default_region", "default_region"}:
default_region = v
elif k in {"aws_profile", "profile"}:
profile = v
elif k in OBJECT_STORE_CLIENT_OPTIONS:
continue
else:
# We assume some sort of access key was given, so we
# just dispatch to the rust side.
return None
unhandled_key = k

to_silence_this_warning = (
"To silence this warning, pass 'aws_profile': None in "
"storage_options, or unset the AWS_PROFILE environment flag."
)

if unhandled_key is not None:
if profile is not None:
msg = (
f"the configured AWS profile '{profile}' may be ignored "
"as it is not compatible with the provided "
f"storage_option key '{unhandled_key}'. "
f"{to_silence_this_warning}"
)
issue_warning(msg, UserWarning)

return None

try:
provider = CredentialProviderAWS(
profile_name=profile, region_name=region or default_region
)
except ImportError:
if profile is not None:
msg = (
f"the configured AWS profile '{profile}' may not "
"be used as boto3 is not installed. "
f"{to_silence_this_warning}"
)
# Conservatively warn instead of hard error. It could just be
# set as a default environment flag.
issue_warning(msg, UserWarning)
# Note: Enclosing scope is also within a try-except.
raise

provider = CredentialProviderAWS(region_name=region or default_region)
elif storage_options is not None and any(
key.lower() not in OBJECT_STORE_CLIENT_OPTIONS for key in storage_options
):
Expand Down
1 change: 1 addition & 0 deletions py-polars/tests/unit/dataframe/test_df.py
Original file line number Diff line number Diff line change
Expand Up @@ -2387,6 +2387,7 @@ def test_selection_regex_and_multicol() -> None:


@pytest.mark.parametrize("subset", ["a", cs.starts_with("x", "a")])
@pytest.mark.may_fail_auto_streaming # Flaky in CI, see https://github.com/pola-rs/polars/issues/20943
def test_unique_on_sorted(subset: Any) -> None:
df = pl.DataFrame(data={"a": [1, 1, 3], "b": [1, 2, 3]})

Expand Down
Loading