From 35d0c85bf9f0b432b45004078d1a58dbb8b60763 Mon Sep 17 00:00:00 2001 From: anitarua Date: Wed, 19 Nov 2025 16:50:05 -0800 Subject: [PATCH 1/7] feat: new credential provider methods for accepting global api keys --- src/momento/auth/credential_provider.py | 41 ++++++++++ .../momento/auth/test_credential_provider.py | 81 +++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index eb16000f..49d3d7be 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -102,3 +102,44 @@ def _obscure(self, value: str) -> str: def get_auth_token(self) -> str: return self.auth_token + + @staticmethod + def global_key_from_string(api_key: str, endpoint: str) -> CredentialProvider: + """Creates a CredentialProvider from a global API key and endpoint. + + Args: + api_key (str): The global API key. + endpoint (str): The Momento service endpoint. + + Returns: + CredentialProvider + """ + if len(api_key) == 0: + raise RuntimeError("API key cannot be empty.") + if len(endpoint) == 0: + raise RuntimeError("Endpoint cannot be empty.") + return CredentialProvider( + auth_token=api_key, + control_endpoint=momento_endpoint_resolver._MOMENTO_CONTROL_ENDPOINT_PREFIX + endpoint, + cache_endpoint=momento_endpoint_resolver._MOMENTO_CACHE_ENDPOINT_PREFIX + endpoint, + token_endpoint=momento_endpoint_resolver._MOMENTO_TOKEN_ENDPOINT_PREFIX + endpoint, + port=443, + ) + + @staticmethod + def global_key_from_environment_variable(env_var_name: str, endpoint: str) -> CredentialProvider: + """Creates a CredentialProvider from an endpoint and a global API key stored in an environment variable. + + Args: + env_var_name (str): Name of the environment variable from which the global API key will be read. + endpoint (str): The Momento service endpoint. + + Returns: + CredentialProvider + """ + if len(env_var_name) == 0: + raise RuntimeError("Environment variable name cannot be empty.") + api_key = os.getenv(env_var_name) + if not api_key: + raise RuntimeError(f"Missing required environment variable {env_var_name}") + return CredentialProvider.global_key_from_string(api_key, endpoint) diff --git a/tests/momento/auth/test_credential_provider.py b/tests/momento/auth/test_credential_provider.py index 2258e584..855c52b9 100644 --- a/tests/momento/auth/test_credential_provider.py +++ b/tests/momento/auth/test_credential_provider.py @@ -97,3 +97,84 @@ def test_endpoints(provider: CredentialProvider, auth_token: str, control_endpoi def test_env_token_raises_if_not_exists() -> None: with pytest.raises(RuntimeError, match=r"Missing required environment variable"): CredentialProvider.from_environment_variable(env_var_name=uuid_str()) + + +# Global API Key Tests +test_global_api_key = "testToken" +test_global_endpoint = "testEndpoint" +test_global_env_var_name = "MOMENTO_TEST_GLOBAL_API_KEY" +os.environ[test_global_env_var_name] = test_global_api_key + + +@pytest.mark.parametrize( + "provider, expected_api_key, expected_control_endpoint, expected_cache_endpoint, expected_token_endpoint", + [ + # global_key_from_string - basic usage + ( + CredentialProvider.global_key_from_string( + api_key=test_global_api_key, + endpoint=test_global_endpoint, + ), + test_global_api_key, + f"control.{test_global_endpoint}", + f"cache.{test_global_endpoint}", + f"token.{test_global_endpoint}", + ), + # global_key_from_environment_variable - basic usage + ( + CredentialProvider.global_key_from_environment_variable( + env_var_name=test_global_env_var_name, + endpoint=test_global_endpoint, + ), + test_global_api_key, + f"control.{test_global_endpoint}", + f"cache.{test_global_endpoint}", + f"token.{test_global_endpoint}", + ), + ], +) +def test_global_api_key_endpoints( + provider: CredentialProvider, + expected_api_key: str, + expected_control_endpoint: str, + expected_cache_endpoint: str, + expected_token_endpoint: str, +) -> None: + assert provider.auth_token == expected_api_key + assert provider.control_endpoint == expected_control_endpoint + assert provider.cache_endpoint == expected_cache_endpoint + assert provider.token_endpoint == expected_token_endpoint + + +def test_global_key_from_string_raises_if_api_key_empty() -> None: + with pytest.raises(RuntimeError, match=r"API key cannot be empty"): + CredentialProvider.global_key_from_string(api_key="", endpoint=test_global_endpoint) + + +def test_global_key_from_string_raises_if_endpoint_empty() -> None: + with pytest.raises(RuntimeError, match=r"Endpoint cannot be empty"): + CredentialProvider.global_key_from_string(api_key=test_global_api_key, endpoint="") + + +def test_global_key_from_env_raises_if_env_var_name_empty() -> None: + with pytest.raises(RuntimeError, match=r"Environment variable name cannot be empty"): + CredentialProvider.global_key_from_environment_variable(env_var_name="", endpoint=test_global_endpoint) + + +def test_global_key_from_env_raises_if_env_var_missing() -> None: + with pytest.raises(RuntimeError, match=r"Missing required environment variable"): + CredentialProvider.global_key_from_environment_variable(env_var_name=uuid_str(), endpoint=test_global_endpoint) + + +def test_global_key_from_env_raises_if_endpoint_empty() -> None: + with pytest.raises(RuntimeError, match=r"Endpoint cannot be empty"): + CredentialProvider.global_key_from_environment_variable(env_var_name=test_global_env_var_name, endpoint="") + + +def test_global_key_from_env_raises_if_api_key_empty_string() -> None: + empty_api_key_env_var = uuid_str() + os.environ[empty_api_key_env_var] = "" + with pytest.raises(RuntimeError, match=r"Missing required environment variable"): + CredentialProvider.global_key_from_environment_variable( + env_var_name=empty_api_key_env_var, endpoint=test_global_endpoint + ) From 6fcc96e52aeed0a41f621dec434c23ebffdd5aa9 Mon Sep 17 00:00:00 2001 From: anitarua Date: Mon, 1 Dec 2025 12:13:54 -0800 Subject: [PATCH 2/7] check key type to provide useful errors --- src/momento/auth/credential_provider.py | 16 ++++ src/momento/auth/momento_endpoint_resolver.py | 10 +++ .../momento/auth/test_credential_provider.py | 73 +++++++++++++++---- 3 files changed, 86 insertions(+), 13 deletions(-) diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index 49d3d7be..ea85df72 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -118,6 +118,14 @@ def global_key_from_string(api_key: str, endpoint: str) -> CredentialProvider: raise RuntimeError("API key cannot be empty.") if len(endpoint) == 0: raise RuntimeError("Endpoint cannot be empty.") + if momento_endpoint_resolver._is_base64(api_key): + raise RuntimeError( + "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_string()` instead?" + ) + if not momento_endpoint_resolver._is_global_api_key(api_key): + raise RuntimeError( + "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_string()` instead?" + ) return CredentialProvider( auth_token=api_key, control_endpoint=momento_endpoint_resolver._MOMENTO_CONTROL_ENDPOINT_PREFIX + endpoint, @@ -142,4 +150,12 @@ def global_key_from_environment_variable(env_var_name: str, endpoint: str) -> Cr api_key = os.getenv(env_var_name) if not api_key: raise RuntimeError(f"Missing required environment variable {env_var_name}") + if momento_endpoint_resolver._is_base64(api_key): + raise RuntimeError( + "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + ) + if not momento_endpoint_resolver._is_global_api_key(api_key): + raise RuntimeError( + "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + ) return CredentialProvider.global_key_from_string(api_key, endpoint) diff --git a/src/momento/auth/momento_endpoint_resolver.py b/src/momento/auth/momento_endpoint_resolver.py index 8285c55c..1e9fab03 100644 --- a/src/momento/auth/momento_endpoint_resolver.py +++ b/src/momento/auth/momento_endpoint_resolver.py @@ -14,6 +14,8 @@ _MOMENTO_TOKEN_ENDPOINT_PREFIX = "token." _CONTROL_ENDPOINT_CLAIM_ID = "cp" _CACHE_ENDPOINT_CLAIM_ID = "c" +_API_KEY_TYPE_CLAIM_ID = "t" +_GLOBAL_API_KEY_TYPE = "g" @dataclass @@ -67,3 +69,11 @@ def _is_base64(value: Union[bytes, str]) -> bool: return base64.b64encode(base64.b64decode(value)) == value except Exception: return False + + +def _is_global_api_key(value: str) -> bool: + try: + claims = jwt.decode(value, options={"verify_signature": False}) # type: ignore[misc] + return _API_KEY_TYPE_CLAIM_ID in claims and claims[_API_KEY_TYPE_CLAIM_ID] == _GLOBAL_API_KEY_TYPE # type: ignore[misc] + except DecodeError: + return False diff --git a/tests/momento/auth/test_credential_provider.py b/tests/momento/auth/test_credential_provider.py index 855c52b9..d912782a 100644 --- a/tests/momento/auth/test_credential_provider.py +++ b/tests/momento/auth/test_credential_provider.py @@ -1,6 +1,7 @@ import base64 import json import os +import re import jwt import pytest @@ -23,6 +24,13 @@ os.environ[test_env_var_name] = test_token os.environ[test_v1_env_var_name] = test_encoded_v1_token.decode("utf-8") +# For global API key tests +test_global_key_message = {"t": "g"} +test_global_api_key = jwt.encode(test_global_key_message, "secret", algorithm="HS512") # TODO: use ES256? +test_global_endpoint = "testEndpoint" +test_global_env_var_name = uuid_str() +os.environ[test_global_env_var_name] = test_global_api_key + @pytest.mark.parametrize( "provider, auth_token, control_endpoint, cache_endpoint", @@ -99,13 +107,6 @@ def test_env_token_raises_if_not_exists() -> None: CredentialProvider.from_environment_variable(env_var_name=uuid_str()) -# Global API Key Tests -test_global_api_key = "testToken" -test_global_endpoint = "testEndpoint" -test_global_env_var_name = "MOMENTO_TEST_GLOBAL_API_KEY" -os.environ[test_global_env_var_name] = test_global_api_key - - @pytest.mark.parametrize( "provider, expected_api_key, expected_control_endpoint, expected_cache_endpoint, expected_token_endpoint", [ @@ -147,34 +148,80 @@ def test_global_api_key_endpoints( def test_global_key_from_string_raises_if_api_key_empty() -> None: - with pytest.raises(RuntimeError, match=r"API key cannot be empty"): + with pytest.raises(RuntimeError, match="API key cannot be empty"): CredentialProvider.global_key_from_string(api_key="", endpoint=test_global_endpoint) def test_global_key_from_string_raises_if_endpoint_empty() -> None: - with pytest.raises(RuntimeError, match=r"Endpoint cannot be empty"): + with pytest.raises(RuntimeError, match="Endpoint cannot be empty"): CredentialProvider.global_key_from_string(api_key=test_global_api_key, endpoint="") def test_global_key_from_env_raises_if_env_var_name_empty() -> None: - with pytest.raises(RuntimeError, match=r"Environment variable name cannot be empty"): + with pytest.raises(RuntimeError, match="Environment variable name cannot be empty"): CredentialProvider.global_key_from_environment_variable(env_var_name="", endpoint=test_global_endpoint) def test_global_key_from_env_raises_if_env_var_missing() -> None: - with pytest.raises(RuntimeError, match=r"Missing required environment variable"): + with pytest.raises(RuntimeError, match="Missing required environment variable"): CredentialProvider.global_key_from_environment_variable(env_var_name=uuid_str(), endpoint=test_global_endpoint) def test_global_key_from_env_raises_if_endpoint_empty() -> None: - with pytest.raises(RuntimeError, match=r"Endpoint cannot be empty"): + with pytest.raises(RuntimeError, match="Endpoint cannot be empty"): CredentialProvider.global_key_from_environment_variable(env_var_name=test_global_env_var_name, endpoint="") def test_global_key_from_env_raises_if_api_key_empty_string() -> None: empty_api_key_env_var = uuid_str() os.environ[empty_api_key_env_var] = "" - with pytest.raises(RuntimeError, match=r"Missing required environment variable"): + with pytest.raises(RuntimeError, match="Missing required environment variable"): CredentialProvider.global_key_from_environment_variable( env_var_name=empty_api_key_env_var, endpoint=test_global_endpoint ) + + +def test_global_key_from_string_raises_if_base64_api_key() -> None: + with pytest.raises( + RuntimeError, + match=re.escape( + "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_string()` instead?" + ), + ): + CredentialProvider.global_key_from_string( + api_key=test_encoded_v1_token.decode("utf-8"), endpoint=test_global_endpoint + ) + + +def test_global_key_from_env_raises_if_base64_api_key() -> None: + with pytest.raises( + RuntimeError, + match=re.escape( + "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + ), + ): + CredentialProvider.global_key_from_environment_variable( + env_var_name=test_v1_env_var_name, endpoint=test_global_endpoint + ) + + +def test_global_key_from_string_raises_if_pre_v1_token() -> None: + with pytest.raises( + RuntimeError, + match=re.escape( + "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_string()` instead?" + ), + ): + CredentialProvider.global_key_from_string(api_key=test_token, endpoint=test_global_endpoint) + + +def test_global_key_from_env_raises_if_pre_v1_token() -> None: + with pytest.raises( + RuntimeError, + match=re.escape( + "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + ), + ): + CredentialProvider.global_key_from_environment_variable( + env_var_name=test_env_var_name, endpoint=test_global_endpoint + ) From 4b3cf4a3b4ca5b8e7be5c2eb8347db81db04bc99 Mon Sep 17 00:00:00 2001 From: anitarua Date: Mon, 1 Dec 2025 14:03:53 -0800 Subject: [PATCH 3/7] use correct errors and test non-global api key path --- src/momento/auth/credential_provider.py | 29 ++++++++++++------- src/momento/auth/momento_endpoint_resolver.py | 5 ++++ .../momento/auth/test_credential_provider.py | 27 ++++++++++++----- 3 files changed, 42 insertions(+), 19 deletions(-) diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index ea85df72..38445233 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -5,6 +5,9 @@ from dataclasses import dataclass from typing import Dict, Optional +from momento.errors.exceptions import InvalidArgumentException +from momento.internal.services import Service + from . import momento_endpoint_resolver @@ -115,16 +118,18 @@ def global_key_from_string(api_key: str, endpoint: str) -> CredentialProvider: CredentialProvider """ if len(api_key) == 0: - raise RuntimeError("API key cannot be empty.") + raise InvalidArgumentException("API key cannot be empty.", Service.AUTH) if len(endpoint) == 0: - raise RuntimeError("Endpoint cannot be empty.") + raise InvalidArgumentException("Endpoint cannot be empty.", Service.AUTH) if momento_endpoint_resolver._is_base64(api_key): - raise RuntimeError( - "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_string()` instead?" + raise InvalidArgumentException( + "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_string()` instead?", + Service.AUTH, ) if not momento_endpoint_resolver._is_global_api_key(api_key): - raise RuntimeError( - "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_string()` instead?" + raise InvalidArgumentException( + "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_string()` instead?", + Service.AUTH, ) return CredentialProvider( auth_token=api_key, @@ -146,16 +151,18 @@ def global_key_from_environment_variable(env_var_name: str, endpoint: str) -> Cr CredentialProvider """ if len(env_var_name) == 0: - raise RuntimeError("Environment variable name cannot be empty.") + raise InvalidArgumentException("Environment variable name cannot be empty.", Service.AUTH) api_key = os.getenv(env_var_name) if not api_key: raise RuntimeError(f"Missing required environment variable {env_var_name}") if momento_endpoint_resolver._is_base64(api_key): - raise RuntimeError( - "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + raise InvalidArgumentException( + "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?", + Service.AUTH, ) if not momento_endpoint_resolver._is_global_api_key(api_key): - raise RuntimeError( - "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + raise InvalidArgumentException( + "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?", + Service.AUTH, ) return CredentialProvider.global_key_from_string(api_key, endpoint) diff --git a/src/momento/auth/momento_endpoint_resolver.py b/src/momento/auth/momento_endpoint_resolver.py index 1e9fab03..98f1a896 100644 --- a/src/momento/auth/momento_endpoint_resolver.py +++ b/src/momento/auth/momento_endpoint_resolver.py @@ -46,6 +46,11 @@ def resolve(auth_token: str) -> _TokenAndEndpoints: auth_token=info["api_key"], # type: ignore[misc] ) else: + if _is_global_api_key(auth_token): + raise InvalidArgumentException( + "Received a global API key. Are you using the correct key? Or did you mean to use `global_key_from_string()` or `global_key_from_environment_variable()` instead?", + Service.AUTH, + ) return _get_endpoint_from_token(auth_token) diff --git a/tests/momento/auth/test_credential_provider.py b/tests/momento/auth/test_credential_provider.py index d912782a..6f85a8c2 100644 --- a/tests/momento/auth/test_credential_provider.py +++ b/tests/momento/auth/test_credential_provider.py @@ -7,6 +7,7 @@ import pytest from momento.auth.credential_provider import CredentialProvider from momento.auth.momento_endpoint_resolver import _Base64DecodedV1Token +from momento.errors.exceptions import InvalidArgumentException from tests.utils import uuid_str @@ -148,17 +149,17 @@ def test_global_api_key_endpoints( def test_global_key_from_string_raises_if_api_key_empty() -> None: - with pytest.raises(RuntimeError, match="API key cannot be empty"): + with pytest.raises(InvalidArgumentException, match="API key cannot be empty"): CredentialProvider.global_key_from_string(api_key="", endpoint=test_global_endpoint) def test_global_key_from_string_raises_if_endpoint_empty() -> None: - with pytest.raises(RuntimeError, match="Endpoint cannot be empty"): + with pytest.raises(InvalidArgumentException, match="Endpoint cannot be empty"): CredentialProvider.global_key_from_string(api_key=test_global_api_key, endpoint="") def test_global_key_from_env_raises_if_env_var_name_empty() -> None: - with pytest.raises(RuntimeError, match="Environment variable name cannot be empty"): + with pytest.raises(InvalidArgumentException, match="Environment variable name cannot be empty"): CredentialProvider.global_key_from_environment_variable(env_var_name="", endpoint=test_global_endpoint) @@ -168,7 +169,7 @@ def test_global_key_from_env_raises_if_env_var_missing() -> None: def test_global_key_from_env_raises_if_endpoint_empty() -> None: - with pytest.raises(RuntimeError, match="Endpoint cannot be empty"): + with pytest.raises(InvalidArgumentException, match="Endpoint cannot be empty"): CredentialProvider.global_key_from_environment_variable(env_var_name=test_global_env_var_name, endpoint="") @@ -183,7 +184,7 @@ def test_global_key_from_env_raises_if_api_key_empty_string() -> None: def test_global_key_from_string_raises_if_base64_api_key() -> None: with pytest.raises( - RuntimeError, + InvalidArgumentException, match=re.escape( "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_string()` instead?" ), @@ -195,7 +196,7 @@ def test_global_key_from_string_raises_if_base64_api_key() -> None: def test_global_key_from_env_raises_if_base64_api_key() -> None: with pytest.raises( - RuntimeError, + InvalidArgumentException, match=re.escape( "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" ), @@ -207,7 +208,7 @@ def test_global_key_from_env_raises_if_base64_api_key() -> None: def test_global_key_from_string_raises_if_pre_v1_token() -> None: with pytest.raises( - RuntimeError, + InvalidArgumentException, match=re.escape( "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_string()` instead?" ), @@ -217,7 +218,7 @@ def test_global_key_from_string_raises_if_pre_v1_token() -> None: def test_global_key_from_env_raises_if_pre_v1_token() -> None: with pytest.raises( - RuntimeError, + InvalidArgumentException, match=re.escape( "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" ), @@ -225,3 +226,13 @@ def test_global_key_from_env_raises_if_pre_v1_token() -> None: CredentialProvider.global_key_from_environment_variable( env_var_name=test_env_var_name, endpoint=test_global_endpoint ) + + +def test_global_key_provided_to_from_string() -> None: + with pytest.raises( + InvalidArgumentException, + match=re.escape( + "Received a global API key. Are you using the correct key? Or did you mean to use `global_key_from_string()` or `global_key_from_environment_variable()` instead?" + ), + ): + CredentialProvider.from_string(auth_token=test_global_api_key) From 595ef73ad705bdae662d35975726da80426f97ac Mon Sep 17 00:00:00 2001 From: anitarua Date: Tue, 9 Dec 2025 11:28:17 -0800 Subject: [PATCH 4/7] rename global to v2, add disposable token method --- src/momento/auth/credential_provider.py | 73 ++++++--- src/momento/auth/momento_endpoint_resolver.py | 10 +- .../momento/auth/test_credential_provider.py | 145 +++++++++++------- 3 files changed, 145 insertions(+), 83 deletions(-) diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index 38445233..770d9967 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -4,6 +4,7 @@ import os from dataclasses import dataclass from typing import Dict, Optional +from warnings import warn from momento.errors.exceptions import InvalidArgumentException from momento.internal.services import Service @@ -45,6 +46,7 @@ def from_environment_variable( Returns: CredentialProvider """ + warn("from_environment_variable is deprecated, use from_env_var_v2 instead", DeprecationWarning, stacklevel=2) api_key = os.getenv(env_var_name) if not api_key: raise RuntimeError(f"Missing required environment variable {env_var_name}") @@ -71,6 +73,11 @@ def from_string( Returns: CredentialProvider """ + warn( + "from_string is deprecated, use from_api_key_v2 or from_disposable_token instead", + DeprecationWarning, + stacklevel=2, + ) token_and_endpoints = momento_endpoint_resolver.resolve(auth_token) control_endpoint = control_endpoint or token_and_endpoints.control_endpoint cache_endpoint = cache_endpoint or token_and_endpoints.cache_endpoint @@ -107,7 +114,7 @@ def get_auth_token(self) -> str: return self.auth_token @staticmethod - def global_key_from_string(api_key: str, endpoint: str) -> CredentialProvider: + def from_api_key_v2(api_key: str, endpoint: str) -> CredentialProvider: """Creates a CredentialProvider from a global API key and endpoint. Args: @@ -121,14 +128,10 @@ def global_key_from_string(api_key: str, endpoint: str) -> CredentialProvider: raise InvalidArgumentException("API key cannot be empty.", Service.AUTH) if len(endpoint) == 0: raise InvalidArgumentException("Endpoint cannot be empty.", Service.AUTH) - if momento_endpoint_resolver._is_base64(api_key): - raise InvalidArgumentException( - "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_string()` instead?", - Service.AUTH, - ) - if not momento_endpoint_resolver._is_global_api_key(api_key): + + if not momento_endpoint_resolver._is_v2_api_key(api_key): raise InvalidArgumentException( - "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_string()` instead?", + "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_string()` with a legacy key instead?", Service.AUTH, ) return CredentialProvider( @@ -140,29 +143,53 @@ def global_key_from_string(api_key: str, endpoint: str) -> CredentialProvider: ) @staticmethod - def global_key_from_environment_variable(env_var_name: str, endpoint: str) -> CredentialProvider: + def from_env_var_v2(api_key_env_var: str, endpoint_env_var: str) -> CredentialProvider: """Creates a CredentialProvider from an endpoint and a global API key stored in an environment variable. Args: - env_var_name (str): Name of the environment variable from which the global API key will be read. - endpoint (str): The Momento service endpoint. + api_key_env_var (str): Name of the environment variable from which the global API key will be read. + endpoint_env_var (str): Name of the environment variable from which the Momento service endpoint will be read. Returns: CredentialProvider """ - if len(env_var_name) == 0: - raise InvalidArgumentException("Environment variable name cannot be empty.", Service.AUTH) - api_key = os.getenv(env_var_name) + if len(api_key_env_var) == 0: + raise InvalidArgumentException("API key environment variable name cannot be empty.", Service.AUTH) + if len(endpoint_env_var) == 0: + raise InvalidArgumentException("Endpoint environment variable name cannot be empty.", Service.AUTH) + + api_key = os.getenv(api_key_env_var) if not api_key: - raise RuntimeError(f"Missing required environment variable {env_var_name}") - if momento_endpoint_resolver._is_base64(api_key): - raise InvalidArgumentException( - "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?", - Service.AUTH, - ) - if not momento_endpoint_resolver._is_global_api_key(api_key): + raise RuntimeError(f"Missing required environment variable {api_key_env_var}") + endpoint = os.getenv(endpoint_env_var) + if not endpoint: + raise RuntimeError(f"Missing required environment variable {endpoint_env_var}") + + if not momento_endpoint_resolver._is_v2_api_key(api_key): raise InvalidArgumentException( - "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?", + "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` with a legacy key instead?", Service.AUTH, ) - return CredentialProvider.global_key_from_string(api_key, endpoint) + return CredentialProvider.from_api_key_v2(api_key, endpoint) + + @staticmethod + def from_disposable_token(auth_token: str) -> CredentialProvider: + """Reads and parses a Momento disposable auth token. + + Args: + auth_token (str): the Momento disposable auth token + + Returns: + CredentialProvider + """ + if len(auth_token) == 0: + raise InvalidArgumentException("Disposable token cannot be empty.", Service.AUTH) + token_and_endpoints = momento_endpoint_resolver.resolve(auth_token) + auth_token = token_and_endpoints.auth_token + return CredentialProvider( + auth_token, + token_and_endpoints.control_endpoint, + token_and_endpoints.cache_endpoint, + token_and_endpoints.token_endpoint, + 443, + ) diff --git a/src/momento/auth/momento_endpoint_resolver.py b/src/momento/auth/momento_endpoint_resolver.py index 98f1a896..000a2161 100644 --- a/src/momento/auth/momento_endpoint_resolver.py +++ b/src/momento/auth/momento_endpoint_resolver.py @@ -46,9 +46,9 @@ def resolve(auth_token: str) -> _TokenAndEndpoints: auth_token=info["api_key"], # type: ignore[misc] ) else: - if _is_global_api_key(auth_token): + if _is_v2_api_key(auth_token): raise InvalidArgumentException( - "Received a global API key. Are you using the correct key? Or did you mean to use `global_key_from_string()` or `global_key_from_environment_variable()` instead?", + "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?", Service.AUTH, ) return _get_endpoint_from_token(auth_token) @@ -76,9 +76,11 @@ def _is_base64(value: Union[bytes, str]) -> bool: return False -def _is_global_api_key(value: str) -> bool: +def _is_v2_api_key(key: str) -> bool: + if _is_base64(key): + return False try: - claims = jwt.decode(value, options={"verify_signature": False}) # type: ignore[misc] + claims = jwt.decode(key, options={"verify_signature": False}) # type: ignore[misc] return _API_KEY_TYPE_CLAIM_ID in claims and claims[_API_KEY_TYPE_CLAIM_ID] == _GLOBAL_API_KEY_TYPE # type: ignore[misc] except DecodeError: return False diff --git a/tests/momento/auth/test_credential_provider.py b/tests/momento/auth/test_credential_provider.py index 6f85a8c2..c4844684 100644 --- a/tests/momento/auth/test_credential_provider.py +++ b/tests/momento/auth/test_credential_provider.py @@ -25,12 +25,14 @@ os.environ[test_env_var_name] = test_token os.environ[test_v1_env_var_name] = test_encoded_v1_token.decode("utf-8") -# For global API key tests -test_global_key_message = {"t": "g"} -test_global_api_key = jwt.encode(test_global_key_message, "secret", algorithm="HS512") # TODO: use ES256? -test_global_endpoint = "testEndpoint" -test_global_env_var_name = uuid_str() -os.environ[test_global_env_var_name] = test_global_api_key +# For v2 API key tests +test_v2_key_message = {"t": "g", "id": "some-id"} +test_v2_api_key = jwt.encode(test_v2_key_message, "secret", algorithm="HS512") +test_v2_key_env_var_name = uuid_str() +test_v2_endpoint = "testEndpoint" +test_v2_endpoint_env_var_name = uuid_str() +os.environ[test_v2_key_env_var_name] = test_v2_api_key +os.environ[test_v2_endpoint_env_var_name] = test_v2_endpoint @pytest.mark.parametrize( @@ -111,31 +113,31 @@ def test_env_token_raises_if_not_exists() -> None: @pytest.mark.parametrize( "provider, expected_api_key, expected_control_endpoint, expected_cache_endpoint, expected_token_endpoint", [ - # global_key_from_string - basic usage + # v2_key_from_string - basic usage ( - CredentialProvider.global_key_from_string( - api_key=test_global_api_key, - endpoint=test_global_endpoint, + CredentialProvider.from_api_key_v2( + api_key=test_v2_api_key, + endpoint=test_v2_endpoint, ), - test_global_api_key, - f"control.{test_global_endpoint}", - f"cache.{test_global_endpoint}", - f"token.{test_global_endpoint}", + test_v2_api_key, + f"control.{test_v2_endpoint}", + f"cache.{test_v2_endpoint}", + f"token.{test_v2_endpoint}", ), - # global_key_from_environment_variable - basic usage + # v2_key_from_environment_variable - basic usage ( - CredentialProvider.global_key_from_environment_variable( - env_var_name=test_global_env_var_name, - endpoint=test_global_endpoint, + CredentialProvider.from_env_var_v2( + api_key_env_var=test_v2_key_env_var_name, + endpoint_env_var=test_v2_endpoint_env_var_name, ), - test_global_api_key, - f"control.{test_global_endpoint}", - f"cache.{test_global_endpoint}", - f"token.{test_global_endpoint}", + test_v2_api_key, + f"control.{test_v2_endpoint}", + f"cache.{test_v2_endpoint}", + f"token.{test_v2_endpoint}", ), ], ) -def test_global_api_key_endpoints( +def test_v2_api_key_endpoints( provider: CredentialProvider, expected_api_key: str, expected_control_endpoint: str, @@ -148,91 +150,122 @@ def test_global_api_key_endpoints( assert provider.token_endpoint == expected_token_endpoint -def test_global_key_from_string_raises_if_api_key_empty() -> None: +def test_v2_key_from_string_raises_if_api_key_empty() -> None: with pytest.raises(InvalidArgumentException, match="API key cannot be empty"): - CredentialProvider.global_key_from_string(api_key="", endpoint=test_global_endpoint) + CredentialProvider.from_api_key_v2(api_key="", endpoint=test_v2_endpoint) -def test_global_key_from_string_raises_if_endpoint_empty() -> None: +def test_v2_key_from_string_raises_if_endpoint_empty() -> None: with pytest.raises(InvalidArgumentException, match="Endpoint cannot be empty"): - CredentialProvider.global_key_from_string(api_key=test_global_api_key, endpoint="") + CredentialProvider.from_api_key_v2(api_key=test_v2_api_key, endpoint="") -def test_global_key_from_env_raises_if_env_var_name_empty() -> None: - with pytest.raises(InvalidArgumentException, match="Environment variable name cannot be empty"): - CredentialProvider.global_key_from_environment_variable(env_var_name="", endpoint=test_global_endpoint) +def test_v2_key_from_env_raises_if_env_var_name_empty() -> None: + with pytest.raises(InvalidArgumentException, match="API key environment variable name cannot be empty"): + CredentialProvider.from_env_var_v2(api_key_env_var="", endpoint_env_var=test_v2_endpoint_env_var_name) -def test_global_key_from_env_raises_if_env_var_missing() -> None: +def test_v2_key_from_env_raises_if_env_var_missing() -> None: with pytest.raises(RuntimeError, match="Missing required environment variable"): - CredentialProvider.global_key_from_environment_variable(env_var_name=uuid_str(), endpoint=test_global_endpoint) + CredentialProvider.from_env_var_v2(api_key_env_var=uuid_str(), endpoint_env_var=test_v2_endpoint_env_var_name) -def test_global_key_from_env_raises_if_endpoint_empty() -> None: - with pytest.raises(InvalidArgumentException, match="Endpoint cannot be empty"): - CredentialProvider.global_key_from_environment_variable(env_var_name=test_global_env_var_name, endpoint="") +def test_v2_key_from_env_raises_if_endpoint_empty() -> None: + with pytest.raises(InvalidArgumentException, match="Endpoint environment variable name cannot be empty"): + CredentialProvider.from_env_var_v2(api_key_env_var=test_v2_key_env_var_name, endpoint_env_var="") -def test_global_key_from_env_raises_if_api_key_empty_string() -> None: +def test_v2_key_from_env_raises_if_api_key_empty_string() -> None: empty_api_key_env_var = uuid_str() os.environ[empty_api_key_env_var] = "" with pytest.raises(RuntimeError, match="Missing required environment variable"): - CredentialProvider.global_key_from_environment_variable( - env_var_name=empty_api_key_env_var, endpoint=test_global_endpoint + CredentialProvider.from_env_var_v2( + api_key_env_var=empty_api_key_env_var, endpoint_env_var=test_v2_endpoint_env_var_name ) -def test_global_key_from_string_raises_if_base64_api_key() -> None: +def test_v2_key_from_string_raises_if_base64_api_key() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_string()` instead?" + "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_string()` with a legacy key instead?" ), ): - CredentialProvider.global_key_from_string( - api_key=test_encoded_v1_token.decode("utf-8"), endpoint=test_global_endpoint + CredentialProvider.from_api_key_v2( + api_key=test_encoded_v1_token.decode("utf-8"), endpoint=test_v2_endpoint_env_var_name ) -def test_global_key_from_env_raises_if_base64_api_key() -> None: +def test_v2_key_from_env_raises_if_base64_api_key() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Did not expect global API key to be base64 encoded. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` with a legacy key instead?" ), ): - CredentialProvider.global_key_from_environment_variable( - env_var_name=test_v1_env_var_name, endpoint=test_global_endpoint + CredentialProvider.from_env_var_v2( + api_key_env_var=test_v1_env_var_name, endpoint_env_var=test_v2_endpoint_env_var_name ) -def test_global_key_from_string_raises_if_pre_v1_token() -> None: +def test_v2_key_from_string_raises_if_pre_v1_token() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_string()` instead?" + "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_string()` with a legacy key instead?" ), ): - CredentialProvider.global_key_from_string(api_key=test_token, endpoint=test_global_endpoint) + CredentialProvider.from_api_key_v2(api_key=test_token, endpoint=test_v2_endpoint_env_var_name) -def test_global_key_from_env_raises_if_pre_v1_token() -> None: +def test_v2_key_from_env_raises_if_pre_v1_token() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Provided API key is not a valid global API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` instead?" + "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` with a legacy key instead?" ), ): - CredentialProvider.global_key_from_environment_variable( - env_var_name=test_env_var_name, endpoint=test_global_endpoint + CredentialProvider.from_env_var_v2( + api_key_env_var=test_env_var_name, endpoint_env_var=test_v2_endpoint_env_var_name ) -def test_global_key_provided_to_from_string() -> None: +def test_v2_key_provided_to_from_string() -> None: + with pytest.raises( + InvalidArgumentException, + match=re.escape( + "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?" + ), + ): + CredentialProvider.from_string(auth_token=test_v2_api_key) + + +def test_v2_key_provided_to_from_disposable_token() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Received a global API key. Are you using the correct key? Or did you mean to use `global_key_from_string()` or `global_key_from_environment_variable()` instead?" + "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?" ), ): - CredentialProvider.from_string(auth_token=test_global_api_key) + CredentialProvider.from_disposable_token(auth_token=test_v2_api_key) + + +def test_from_disposable_token_raises_if_token_empty() -> None: + with pytest.raises(InvalidArgumentException, match="Disposable token cannot be empty."): + CredentialProvider.from_disposable_token(auth_token="") + + +def test_from_disposable_token_accepts_v1_api_key() -> None: + provider = CredentialProvider.from_disposable_token(auth_token=test_encoded_v1_token.decode("utf-8")) + assert provider.auth_token == test_v1_api_key + assert provider.control_endpoint == "control.test.momentohq.com" + assert provider.cache_endpoint == "cache.test.momentohq.com" + assert provider.token_endpoint == "token.test.momentohq.com" + + +def test_from_disposable_token_accepts_pre_v1_token() -> None: + provider = CredentialProvider.from_disposable_token(auth_token=test_token) + assert provider.auth_token == test_token + assert provider.control_endpoint == test_control_endpoint + assert provider.cache_endpoint == test_cache_endpoint + assert provider.token_endpoint == f"token.{test_cache_endpoint}" From 4acc1d45907902b77907a703a3cd16597607676b Mon Sep 17 00:00:00 2001 From: anitarua Date: Fri, 12 Dec 2025 17:03:26 -0800 Subject: [PATCH 5/7] from_env_var_v2 optional args --- src/momento/auth/credential_provider.py | 20 ++++++---- src/momento/auth/momento_endpoint_resolver.py | 2 +- .../momento/auth/test_credential_provider.py | 37 ++++++++++++------- 3 files changed, 37 insertions(+), 22 deletions(-) diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index 770d9967..07180de2 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -46,7 +46,11 @@ def from_environment_variable( Returns: CredentialProvider """ - warn("from_environment_variable is deprecated, use from_env_var_v2 instead", DeprecationWarning, stacklevel=2) + warn( + "from_environment_variable is deprecated, use from_environment_variables_v2 instead", + DeprecationWarning, + stacklevel=2, + ) api_key = os.getenv(env_var_name) if not api_key: raise RuntimeError(f"Missing required environment variable {env_var_name}") @@ -115,10 +119,10 @@ def get_auth_token(self) -> str: @staticmethod def from_api_key_v2(api_key: str, endpoint: str) -> CredentialProvider: - """Creates a CredentialProvider from a global API key and endpoint. + """Creates a CredentialProvider from a v2 API key and endpoint. Args: - api_key (str): The global API key. + api_key (str): The v2 API key. endpoint (str): The Momento service endpoint. Returns: @@ -143,12 +147,14 @@ def from_api_key_v2(api_key: str, endpoint: str) -> CredentialProvider: ) @staticmethod - def from_env_var_v2(api_key_env_var: str, endpoint_env_var: str) -> CredentialProvider: - """Creates a CredentialProvider from an endpoint and a global API key stored in an environment variable. + def from_environment_variables_v2( + api_key_env_var: str = "MOMENTO_API_KEY", endpoint_env_var: str = "MOMENTO_ENDPOINT" + ) -> CredentialProvider: + """Creates a CredentialProvider from an endpoint and v2 API key stored in the environment variables MOMENTO_API_KEY and MOMENTO_ENDPOINT. Args: - api_key_env_var (str): Name of the environment variable from which the global API key will be read. - endpoint_env_var (str): Name of the environment variable from which the Momento service endpoint will be read. + api_key_env_var (str): Optionally provide an alternate environment variable name from which the v2 API key will be read. + endpoint_env_var (str): Optionally provide an alternate environment variable name from which the Momento service endpoint will be read. Returns: CredentialProvider diff --git a/src/momento/auth/momento_endpoint_resolver.py b/src/momento/auth/momento_endpoint_resolver.py index 000a2161..20c86ff0 100644 --- a/src/momento/auth/momento_endpoint_resolver.py +++ b/src/momento/auth/momento_endpoint_resolver.py @@ -48,7 +48,7 @@ def resolve(auth_token: str) -> _TokenAndEndpoints: else: if _is_v2_api_key(auth_token): raise InvalidArgumentException( - "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?", + "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?", Service.AUTH, ) return _get_endpoint_from_token(auth_token) diff --git a/tests/momento/auth/test_credential_provider.py b/tests/momento/auth/test_credential_provider.py index c4844684..8d81256c 100644 --- a/tests/momento/auth/test_credential_provider.py +++ b/tests/momento/auth/test_credential_provider.py @@ -26,11 +26,11 @@ os.environ[test_v1_env_var_name] = test_encoded_v1_token.decode("utf-8") # For v2 API key tests -test_v2_key_message = {"t": "g", "id": "some-id"} +test_v2_key_message = {"t": "g", "jti": "some-id"} test_v2_api_key = jwt.encode(test_v2_key_message, "secret", algorithm="HS512") -test_v2_key_env_var_name = uuid_str() +test_v2_key_env_var_name = "MOMENTO_API_KEY" test_v2_endpoint = "testEndpoint" -test_v2_endpoint_env_var_name = uuid_str() +test_v2_endpoint_env_var_name = "MOMENTO_ENDPOINT" os.environ[test_v2_key_env_var_name] = test_v2_api_key os.environ[test_v2_endpoint_env_var_name] = test_v2_endpoint @@ -113,7 +113,6 @@ def test_env_token_raises_if_not_exists() -> None: @pytest.mark.parametrize( "provider, expected_api_key, expected_control_endpoint, expected_cache_endpoint, expected_token_endpoint", [ - # v2_key_from_string - basic usage ( CredentialProvider.from_api_key_v2( api_key=test_v2_api_key, @@ -124,9 +123,8 @@ def test_env_token_raises_if_not_exists() -> None: f"cache.{test_v2_endpoint}", f"token.{test_v2_endpoint}", ), - # v2_key_from_environment_variable - basic usage ( - CredentialProvider.from_env_var_v2( + CredentialProvider.from_environment_variables_v2( api_key_env_var=test_v2_key_env_var_name, endpoint_env_var=test_v2_endpoint_env_var_name, ), @@ -135,6 +133,13 @@ def test_env_token_raises_if_not_exists() -> None: f"cache.{test_v2_endpoint}", f"token.{test_v2_endpoint}", ), + ( + CredentialProvider.from_environment_variables_v2(), + test_v2_api_key, + f"control.{test_v2_endpoint}", + f"cache.{test_v2_endpoint}", + f"token.{test_v2_endpoint}", + ), ], ) def test_v2_api_key_endpoints( @@ -162,24 +167,28 @@ def test_v2_key_from_string_raises_if_endpoint_empty() -> None: def test_v2_key_from_env_raises_if_env_var_name_empty() -> None: with pytest.raises(InvalidArgumentException, match="API key environment variable name cannot be empty"): - CredentialProvider.from_env_var_v2(api_key_env_var="", endpoint_env_var=test_v2_endpoint_env_var_name) + CredentialProvider.from_environment_variables_v2( + api_key_env_var="", endpoint_env_var=test_v2_endpoint_env_var_name + ) def test_v2_key_from_env_raises_if_env_var_missing() -> None: with pytest.raises(RuntimeError, match="Missing required environment variable"): - CredentialProvider.from_env_var_v2(api_key_env_var=uuid_str(), endpoint_env_var=test_v2_endpoint_env_var_name) + CredentialProvider.from_environment_variables_v2( + api_key_env_var=uuid_str(), endpoint_env_var=test_v2_endpoint_env_var_name + ) def test_v2_key_from_env_raises_if_endpoint_empty() -> None: with pytest.raises(InvalidArgumentException, match="Endpoint environment variable name cannot be empty"): - CredentialProvider.from_env_var_v2(api_key_env_var=test_v2_key_env_var_name, endpoint_env_var="") + CredentialProvider.from_environment_variables_v2(api_key_env_var=test_v2_key_env_var_name, endpoint_env_var="") def test_v2_key_from_env_raises_if_api_key_empty_string() -> None: empty_api_key_env_var = uuid_str() os.environ[empty_api_key_env_var] = "" with pytest.raises(RuntimeError, match="Missing required environment variable"): - CredentialProvider.from_env_var_v2( + CredentialProvider.from_environment_variables_v2( api_key_env_var=empty_api_key_env_var, endpoint_env_var=test_v2_endpoint_env_var_name ) @@ -203,7 +212,7 @@ def test_v2_key_from_env_raises_if_base64_api_key() -> None: "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` with a legacy key instead?" ), ): - CredentialProvider.from_env_var_v2( + CredentialProvider.from_environment_variables_v2( api_key_env_var=test_v1_env_var_name, endpoint_env_var=test_v2_endpoint_env_var_name ) @@ -225,7 +234,7 @@ def test_v2_key_from_env_raises_if_pre_v1_token() -> None: "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_environment_variable()` with a legacy key instead?" ), ): - CredentialProvider.from_env_var_v2( + CredentialProvider.from_environment_variables_v2( api_key_env_var=test_env_var_name, endpoint_env_var=test_v2_endpoint_env_var_name ) @@ -234,7 +243,7 @@ def test_v2_key_provided_to_from_string() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?" + "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?" ), ): CredentialProvider.from_string(auth_token=test_v2_api_key) @@ -244,7 +253,7 @@ def test_v2_key_provided_to_from_disposable_token() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_env_var_v2()` instead?" + "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?" ), ): CredentialProvider.from_disposable_token(auth_token=test_v2_api_key) From 68d4d438811b9c24dd1789ab9f4336dc7ab7da26 Mon Sep 17 00:00:00 2001 From: anitarua Date: Tue, 16 Dec 2025 15:06:47 -0800 Subject: [PATCH 6/7] update docs and hints --- src/momento/auth/credential_provider.py | 6 +++++- src/momento/auth/momento_endpoint_resolver.py | 2 +- tests/momento/auth/test_credential_provider.py | 8 ++++---- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/momento/auth/credential_provider.py b/src/momento/auth/credential_provider.py index 07180de2..ef874f45 100644 --- a/src/momento/auth/credential_provider.py +++ b/src/momento/auth/credential_provider.py @@ -31,6 +31,8 @@ def from_environment_variable( ) -> CredentialProvider: """Reads and parses a Momento auth token stored as an environment variable. + Deprecated as of v1.28.0. Use from_environment_variables_v2 instead. + Args: env_var_name (str): Name of the environment variable from which the API key will be read control_endpoint (Optional[str], optional): Optionally overrides the default control endpoint. @@ -65,6 +67,8 @@ def from_string( ) -> CredentialProvider: """Reads and parses a Momento auth token. + Deprecated as of v1.28.0. Use from_api_key_v2 or from_disposable_token instead. + Args: auth_token (str): the Momento API key (previously: auth token) control_endpoint (Optional[str], optional): Optionally overrides the default control endpoint. @@ -135,7 +139,7 @@ def from_api_key_v2(api_key: str, endpoint: str) -> CredentialProvider: if not momento_endpoint_resolver._is_v2_api_key(api_key): raise InvalidArgumentException( - "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_string()` with a legacy key instead?", + "Received an invalid v2 API key. Are you using the correct key and the correct CredentialProvider method?", Service.AUTH, ) return CredentialProvider( diff --git a/src/momento/auth/momento_endpoint_resolver.py b/src/momento/auth/momento_endpoint_resolver.py index 20c86ff0..11718e55 100644 --- a/src/momento/auth/momento_endpoint_resolver.py +++ b/src/momento/auth/momento_endpoint_resolver.py @@ -48,7 +48,7 @@ def resolve(auth_token: str) -> _TokenAndEndpoints: else: if _is_v2_api_key(auth_token): raise InvalidArgumentException( - "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?", + "Unexpectedly received a v2 API key. Are you using the correct key and the correct CredentialProvider method?", Service.AUTH, ) return _get_endpoint_from_token(auth_token) diff --git a/tests/momento/auth/test_credential_provider.py b/tests/momento/auth/test_credential_provider.py index 8d81256c..139ed753 100644 --- a/tests/momento/auth/test_credential_provider.py +++ b/tests/momento/auth/test_credential_provider.py @@ -197,7 +197,7 @@ def test_v2_key_from_string_raises_if_base64_api_key() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_string()` with a legacy key instead?" + "Received an invalid v2 API key. Are you using the correct key and the correct CredentialProvider method?" ), ): CredentialProvider.from_api_key_v2( @@ -221,7 +221,7 @@ def test_v2_key_from_string_raises_if_pre_v1_token() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Received an invalid v2 API key. Are you using the correct key? Or did you mean to use `from_string()` with a legacy key instead?" + "Received an invalid v2 API key. Are you using the correct key and the correct CredentialProvider method?" ), ): CredentialProvider.from_api_key_v2(api_key=test_token, endpoint=test_v2_endpoint_env_var_name) @@ -243,7 +243,7 @@ def test_v2_key_provided_to_from_string() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?" + "Unexpectedly received a v2 API key. Are you using the correct key and the correct CredentialProvider method?" ), ): CredentialProvider.from_string(auth_token=test_v2_api_key) @@ -253,7 +253,7 @@ def test_v2_key_provided_to_from_disposable_token() -> None: with pytest.raises( InvalidArgumentException, match=re.escape( - "Received a v2 API key. Are you using the correct key? Or did you mean to use `from_api_key_v2()` or `from_environment_variables_v2()` instead?" + "Unexpectedly received a v2 API key. Are you using the correct key and the correct CredentialProvider method?" ), ): CredentialProvider.from_disposable_token(auth_token=test_v2_api_key) From c6bb5ea516eaf3fb0f4516329dfe8949bb03a913 Mon Sep 17 00:00:00 2001 From: anitarua Date: Tue, 16 Dec 2025 15:09:37 -0800 Subject: [PATCH 7/7] add resolve docstring --- src/momento/auth/momento_endpoint_resolver.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/momento/auth/momento_endpoint_resolver.py b/src/momento/auth/momento_endpoint_resolver.py index 11718e55..f6cc07ce 100644 --- a/src/momento/auth/momento_endpoint_resolver.py +++ b/src/momento/auth/momento_endpoint_resolver.py @@ -33,6 +33,14 @@ class _Base64DecodedV1Token: def resolve(auth_token: str) -> _TokenAndEndpoints: + """Helper function used by from_string and from_disposable_token to parse legacy and v1 auth tokens. + + Args: + auth_token (str): The auth token to be resolved. + + Returns: + _TokenAndEndpoints + """ if not auth_token: raise InvalidArgumentException("malformed auth token", Service.AUTH)