Prefer FLYTE_ prefixed AWS creds env vars (#1523)

Signed-off-by: Felix Ruess <[email protected]>

plugins/flytekit-data-fsspec/flytekitplugins/fsspec/persist.py

@@ -8,9 +8,6 @@
 from flytekit.extend import DataPersistence, DataPersistencePlugins
 from flytekit.loggers import logger
 
-S3_ACCESS_KEY_ID_ENV_NAME = "AWS_ACCESS_KEY_ID"
-S3_SECRET_ACCESS_KEY_ENV_NAME = "AWS_SECRET_ACCESS_KEY"
-
 # Refer to https://github.com/fsspec/s3fs/blob/50bafe4d8766c3b2a4e1fc09669cf02fb2d71454/s3fs/core.py#L198
 # for key and secret
 _FSSPEC_S3_KEY_ID = "key"
@@ -19,13 +16,11 @@
 
 def s3_setup_args(s3_cfg: S3Config):
     kwargs = {}
-    if S3_ACCESS_KEY_ID_ENV_NAME not in os.environ:
-        if s3_cfg.access_key_id:
-            kwargs[_FSSPEC_S3_KEY_ID] = s3_cfg.access_key_id
+    if s3_cfg.access_key_id:
+        kwargs[_FSSPEC_S3_KEY_ID] = s3_cfg.access_key_id
 
-    if S3_SECRET_ACCESS_KEY_ENV_NAME not in os.environ:
-        if s3_cfg.secret_access_key:
-            kwargs[_FSSPEC_S3_SECRET] = s3_cfg.secret_access_key
+    if s3_cfg.secret_access_key:
+        kwargs[_FSSPEC_S3_SECRET] = s3_cfg.secret_access_key
 
     # S3fs takes this as a special arg
     if s3_cfg.endpoint is not None:

plugins/flytekit-data-fsspec/tests/test_persist.py

@@ -2,6 +2,7 @@
 import pathlib
 import tempfile
 
+import mock
 from flytekitplugins.fsspec.persist import FSSpecPersistence, s3_setup_args
 from fsspec.implementations.local import LocalFileSystem
 
@@ -19,6 +20,54 @@ def test_s3_setup_args():
     assert kwargs == {"key": "access"}
 
 
+@mock.patch.dict(os.environ, {}, clear=True)
+def test_s3_setup_args_env_empty():
+    kwargs = s3_setup_args(S3Config.auto())
+    assert kwargs == {}
+
+
+@mock.patch.dict(
+    os.environ,
+    {
+        "AWS_ACCESS_KEY_ID": "ignore-user",
+        "AWS_SECRET_ACCESS_KEY": "ignore-secret",
+        "FLYTE_AWS_ACCESS_KEY_ID": "flyte",
+        "FLYTE_AWS_SECRET_ACCESS_KEY": "flyte-secret",
+    },
+    clear=True,
+)
+def test_s3_setup_args_env_both():
+    kwargs = s3_setup_args(S3Config.auto())
+    assert kwargs == {"key": "flyte", "secret": "flyte-secret"}
+
+
+@mock.patch.dict(
+    os.environ,
+    {
+        "FLYTE_AWS_ACCESS_KEY_ID": "flyte",
+        "FLYTE_AWS_SECRET_ACCESS_KEY": "flyte-secret",
+    },
+    clear=True,
+)
+def test_s3_setup_args_env_flyte():
+    kwargs = s3_setup_args(S3Config.auto())
+    assert kwargs == {"key": "flyte", "secret": "flyte-secret"}
+
+
+@mock.patch.dict(
+    os.environ,
+    {
+        "AWS_ACCESS_KEY_ID": "ignore-user",
+        "AWS_SECRET_ACCESS_KEY": "ignore-secret",
+    },
+    clear=True,
+)
+def test_s3_setup_args_env_aws():
+    kwargs = s3_setup_args(S3Config.auto())
+    # not explicitly in kwargs, since fsspec/boto3 will use these env vars by default
+    assert kwargs == {}
+
+
 def test_get_protocol():
     assert FSSpecPersistence.get_protocol("s3://abc") == "s3"
     assert FSSpecPersistence.get_protocol("/abc") == "file"