Abstract Secret Management Support

.github/workflows/pull_request_push_test.yml

@@ -123,7 +123,8 @@ jobs:
           COSMOS1_KEY: ${{secrets.COSMOS1_KEY}}
           SQL1_USER: ${{secrets.SQL1_USER}}
           SQL1_PASSWORD: ${{secrets.SQL1_PASSWORD}}
-
+          AWS_ACCESS_KEY_ID: ${{secrets.AWS_ACCESS_KEY_ID}}
+          AWS_SECRET_ACCESS_KEY: ${{secrets.AWS_SECRET_ACCESS_KEY}}
         run: |
           # run only test with databricks. run in 4 parallel jobs
           pytest -n 6 feathr_project/test/
@@ -193,6 +194,8 @@ jobs:
           COSMOS1_KEY: ${{secrets.COSMOS1_KEY}}
           SQL1_USER: ${{secrets.SQL1_USER}}
           SQL1_PASSWORD: ${{secrets.SQL1_PASSWORD}}
+          AWS_ACCESS_KEY_ID: ${{secrets.AWS_ACCESS_KEY_ID}}
+          AWS_SECRET_ACCESS_KEY: ${{secrets.AWS_SECRET_ACCESS_KEY}}
         run: |
           # skip databricks related test as we just ran the test; also seperate databricks and synapse test to make sure there's no write conflict
           # run in 4 parallel jobs to make the time shorter

docs/how-to-guides/feathr-configuration-and-env.md

@@ -27,8 +27,60 @@ This allows end users to store the configuration in a secure way, say in Kuberne
 Feathr will get the configurations in the following order:
 
 1. If the key is set in the environment variable, Feathr will use the value of that environment variable
-2. If it's not set in the environment, then a value is retrieved from the feathr_config.yaml file with the same config key.
-3. If it's not available in the feathr_config.yaml file, Feathr will try to retrieve the value from a key vault service. Currently only Azure Key Vault is supported.
+2. If it's not set in the environment, then a value is retrieved from the `feathr_config.yaml` file with the same config key.
+3. If it's not available in the `feathr_config.yaml` file, Feathr will try to retrieve the value from a key vault service. Currently both Azure Key Vault and AWS Secrets Manager are supported.
+
+# Using Secret Management Service in Feathr
+
+Feathr supports using a Secret Management service for all the credentials and environment variables. Currently the supported secret management services are Azure Key Vault and AWS Secrets Manager. 
+
+In order to use those secret management service, there are two steps:
+
+Step 1: Tell Feathr which secret management service to use, and what is the corresponding namespace.
+
+If using Azure Key Vault:
+```yaml
+secrets:
+  azure_key_vault:
+    name: feathrazuretest3-kv
+```
+
+If using AWS Secret Manager, users should put the corresponding secret_id in the `feathr_config.yaml` section, like below, so that Feathr knows which secret_id to use to retrieve the required credentials.
+```yaml
+secrets:
+  aws_secrets_manager:
+    secret_id: feathrsecret_namespace
+```
+
+Step 2: Initialize a secret management client and pass it to Feathr.
+
+For Azure Key Vault:
+```python
+from azure.keyvault.secrets import SecretClient
+secret_client = SecretClient(
+                vault_url = f"https://<replace_with_key_vault_name>.vault.azure.net",
+                credential=DefaultAzureCredential()
+            )
+feathr_client = FeathrClient(..., secret_manager_client = secret_client)
+```
+
+For AWS Secrets Manager, users need to create a SecretCache object and pass it to Feathr client, like below:
+```python
+import botocore 
+import botocore.session 
+from aws_secretsmanager_caching import SecretCache, SecretCacheConfig 
+
+client = botocore.session.get_session().create_client(
+    service_name='secretsmanager',
+    aws_access_key_id = '<replace_your_aws_access_key_id>',
+    aws_secret_access_key= '<replace_your_aws_secret_access_key>',
+    region_name=region_name
+)
+cache_config = SecretCacheConfig()
+cache = SecretCache( config = cache_config, client = client)
+feathr_client = FeathrClient(..., secret_manager_client = cache)
+
+```
 
 # A list of environment variables that Feathr uses
 

feathr_project/feathr/client.py

@@ -50,17 +50,14 @@ class FeathrClient(object):
         config_path (str, optional): config path. See [Feathr Config Template](https://github.com/feathr-ai/feathr/blob/main/feathr_project/feathrcli/data/feathr_user_workspace/feathr_config.yaml) for more details.  Defaults to "./feathr_config.yaml".
         local_workspace_dir (str, optional): set where is the local work space dir. If not set, Feathr will create a temporary folder to store local workspace related files.
         credential (optional): credential to access cloud resources,  most likely to be the returned result of DefaultAzureCredential(). If not set, Feathr will initialize DefaultAzureCredential() inside the __init__ function to get credentials.
-        project_registry_tag (Dict[str, str]): adding tags for project in Feathr registry. This might be useful if you want to tag your project as deprecated, or allow certain customizations on project leve. Default is empty
-
-    Raises:
-        RuntimeError: Fail to create the client since necessary environment variables are not set for Redis
-        client creation.
+        project_registry_tag (Dict[str, str]): adding tags for project in Feathr registry. This might be useful if you want to tag your project as deprecated, or allow certain customizations on project level. Default is empty.
+        secret_manager_client: the secret manager client initialized outside of Feathr.  End users need to initialize the secret manager outside of Feathr and pass it to Feathr so Feathr can use it to get required secrets.
     """
-    def __init__(self, config_path:str = "./feathr_config.yaml", local_workspace_dir: str = None, credential=None, project_registry_tag: Dict[str, str]=None):
+    def __init__(self, config_path:str = "./feathr_config.yaml", local_workspace_dir: str = None, credential=None, project_registry_tag: Dict[str, str]=None, secret_manager_client = None):
         self.logger = logging.getLogger(__name__)
         # Redis key separator
         self._KEY_SEPARATOR = ':'
-        self.envutils = _EnvVaraibleUtil(config_path)
+        self.envutils = _EnvVaraibleUtil(config_path, secret_manager_client)
         if local_workspace_dir:
             self.local_workspace_dir = local_workspace_dir
         else:
@@ -171,17 +168,6 @@ def __init__(self, config_path:str = "./feathr_config.yaml", local_workspace_dir
         # initialize registry
         self.registry = default_registry_client(self.project_name, config_path=config_path, credential=self.credential)
 
-    def _check_required_environment_variables_exist(self):
-        """Checks if the required environment variables(form feathr_config.yaml) is set.
-
-        Some required information has to be set via environment variables so the client can work.
-        """
-        props = self.secret_names
-        for required_field in (self.required_fields + props):
-            if required_field not in os.environ:
-                raise RuntimeError(f'{required_field} is not set in environment variable. All required environment '
-                                   f'variables are: {self.required_fields}.')
-
     def register_features(self, from_context: bool = True):
         """Registers features based on the current workspace
 
@@ -400,13 +386,12 @@ def _construct_redis_client(self):
         port = self.redis_port
         ssl_enabled = self.redis_ssl_enabled
 
-        redis_client = redis.Redis(
+        self.redis_client = redis.Redis(
             host=host,
             port=port,
             password=password,
             ssl=ssl_enabled)
         self.logger.info('Redis connection is successful and completed.')
-        self.redis_client = redis_client
 
 
     def get_offline_features(self,
@@ -571,7 +556,7 @@ def monitor_features(self, settings: MonitoringSettings, execution_configuration
     # Should search in both 'derived_feature_list' and 'anchor_list'
     # Return related keys(key_column list) or None if cannot find the feature
     def _get_feature_key(self, feature_name: str):
-        features = []
+        features: List[FeatureBase] = []
         if 'derived_feature_list' in dir(self):
             features += self.derived_feature_list
         if 'anchor_list' in dir(self):
@@ -747,7 +732,7 @@ def _get_s3_config_str(self):
         # keys can't be only accessed through environment
         access_key = self.envutils.get_environment_variable('S3_ACCESS_KEY')
         secret_key = self.envutils.get_environment_variable('S3_SECRET_KEY')
-        # HOCCON format will be parsed by the Feathr job
+        # HOCON format will be parsed by the Feathr job
         config_str = """
             S3_ENDPOINT: {S3_ENDPOINT}
             S3_ACCESS_KEY: "{S3_ACCESS_KEY}"
@@ -762,7 +747,7 @@ def _get_adls_config_str(self):
         # if ADLS Account is set in the feathr_config, then we need other environment variables
         # keys can't be only accessed through environment
         key = self.envutils.get_environment_variable('ADLS_KEY')
-        # HOCCON format will be parsed by the Feathr job
+        # HOCON format will be parsed by the Feathr job
         config_str = """
             ADLS_ACCOUNT: {ADLS_ACCOUNT}
             ADLS_KEY: "{ADLS_KEY}"
@@ -776,7 +761,7 @@ def _get_blob_config_str(self):
         # if BLOB Account is set in the feathr_config, then we need other environment variables
         # keys can't be only accessed through environment
         key = self.envutils.get_environment_variable('BLOB_KEY')
-        # HOCCON format will be parsed by the Feathr job
+        # HOCON format will be parsed by the Feathr job
         config_str = """
             BLOB_ACCOUNT: {BLOB_ACCOUNT}
             BLOB_KEY: "{BLOB_KEY}"
@@ -792,7 +777,7 @@ def _get_sql_config_str(self):
         driver = self.envutils.get_environment_variable('JDBC_DRIVER')
         auth_flag = self.envutils.get_environment_variable('JDBC_AUTH_FLAG')
         token = self.envutils.get_environment_variable('JDBC_TOKEN')
-        # HOCCON format will be parsed by the Feathr job
+        # HOCON format will be parsed by the Feathr job
         config_str = """
             JDBC_TABLE: {JDBC_TABLE}
             JDBC_USER: {JDBC_USER}
@@ -809,7 +794,7 @@ def _get_monitoring_config_str(self):
         user = self.envutils.get_environment_variable_with_default('monitoring', 'database', 'sql', 'user')
         password = self.envutils.get_environment_variable('MONITORING_DATABASE_SQL_PASSWORD')
         if url:
-            # HOCCON format will be parsed by the Feathr job
+            # HOCON format will be parsed by the Feathr job
             config_str = """
                 MONITORING_DATABASE_SQL_URL: "{url}"
                 MONITORING_DATABASE_SQL_USER: {user}
@@ -826,7 +811,7 @@ def _get_snowflake_config_str(self):
         sf_user = self.envutils.get_environment_variable_with_default('offline_store', 'snowflake', 'user')
         sf_role = self.envutils.get_environment_variable_with_default('offline_store', 'snowflake', 'role')
         sf_password = self.envutils.get_environment_variable('JDBC_SF_PASSWORD')
-        # HOCCON format will be parsed by the Feathr job
+        # HOCON format will be parsed by the Feathr job
         config_str = """
             JDBC_SF_URL: {JDBC_SF_URL}
             JDBC_SF_USER: {JDBC_SF_USER}
@@ -839,7 +824,7 @@ def _get_kafka_config_str(self):
         """Construct the Kafka config string. The endpoint, access key, secret key, and other parameters can be set via
         environment variables."""
         sasl = self.envutils.get_environment_variable('KAFKA_SASL_JAAS_CONFIG')
-        # HOCCON format will be parsed by the Feathr job
+        # HOCON format will be parsed by the Feathr job
         config_str = """
             KAFKA_SASL_JAAS_CONFIG: "{sasl}"
             """.format(sasl=sasl)

feathr_project/feathr/secrets/abc.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+
+from typing import Any, Dict, List, Optional, Tuple
+
+
+class FeathrSecretsManagementClient(ABC):
+    """This is the abstract class for all the secrets management service, which are used to store the credentials that Feathr might use.
+    """
+
+    @abstractmethod
+    def __init__(self, secret_namespace: str, secret_client) -> None:
+        """Initialize the FeathrSecretsManagementClient class.
+
+        Args:
+            secret_namespace (str): a namespace that Feathr needs to get secrets from. 
+            For Azure Key Vault, it is something like the key vault name. 
+            For AWS secrets manager, it is something like a secret name.
+
+            secret_client: A client that will be used to retrieve Feathr secrets.
+        """
+        pass
+
+    @abstractmethod
+    def get_feathr_secret(self, secret_name: str) -> str:
+        """Get Feathr Secrets from a certain secret management service, such as Azure Key Vault or AWS Secrets Manager. 
+
+        Returns:
+            str: returned secret from secret management service
+        """
+        pass

feathr_project/feathr/secrets/akv_client.py

@@ -1,31 +1,37 @@
 from azure.keyvault.secrets import SecretClient
-from azure.identity import DefaultAzureCredential
 from loguru import logger
 from azure.core.exceptions import ResourceNotFoundError
+from feathr.secrets.abc import FeathrSecretsManagementClient
 
-class AzureKeyVaultClient:
-    def __init__(self, akv_name: str):
-        self.akv_name = akv_name
-        self.secret_client = None
 
-    def get_feathr_akv_secret(self, secret_name: str):
+class AzureKeyVaultClient(FeathrSecretsManagementClient):
+    def __init__(self, secret_namespace: str, secret_client: SecretClient = None):
+        """Initializes the AzureKeyVaultClient. Note that `secret_namespace` is not used, since the namespace information will be included in secret_client.
+        """
+        self.secret_client = secret_client
+        if self.secret_client is not None and not isinstance(secret_client, SecretClient):
+            raise RuntimeError(
+                "You need to pass an azure.keyvault.secrets.SecretClient instance.")
+
+    def get_feathr_secret(self, secret_name: str) -> str:
         """Get Feathr Secrets from Azure Key Vault. Note that this function will replace '_' in `secret_name` with '-' since Azure Key Vault doesn't support it
 
         Returns:
-            _type_: _description_
+            str: returned secret from secret management service
         """
         if self.secret_client is None:
-            self.secret_client = SecretClient(
-                vault_url = f"https://{self.akv_name}.vault.azure.net",
-                credential=DefaultAzureCredential()
-            )
+            raise RuntimeError("You need to pass an azure.keyvault.secrets.SecretClient instance when initializing FeathrClient.")
+
         try:
             # replace '_' with '-' since Azure Key Vault doesn't support it
-            variable_replaced = secret_name.replace('_','-') #.upper()
-            logger.info('Fetching the secret {} from Key Vault {}.', variable_replaced, self.akv_name)
+            variable_replaced = secret_name.replace('_', '-')  # .upper()
+            logger.info('Fetching the secret {} from Key Vault {}.',
+                        variable_replaced, self.secret_client.vault_url)
             secret = self.secret_client.get_secret(variable_replaced)
-            logger.info('Secret {} fetched from Key Vault {}.', variable_replaced, self.akv_name)
+            logger.info('Secret {} fetched from Key Vault {}.',
+                        variable_replaced, self.secret_client.vault_url)
             return secret.value
-        except ResourceNotFoundError as e:
-            logger.error(f"Secret {secret_name} cannot be found in Key Vault {self.akv_name}.")
-            raise
+        except ResourceNotFoundError:
+            logger.error(
+                f"Secret {secret_name} cannot be found in Key Vault {self.secret_client.vault_url}.")
+            raise

feathr_project/feathr/secrets/aws_secretmanager.py

@@ -0,0 +1,33 @@
+from loguru import logger
+import json
+from feathr.secrets.abc import FeathrSecretsManagementClient
+from aws_secretsmanager_caching.secret_cache import SecretCache
+
+
+class AWSSecretManagerClient(FeathrSecretsManagementClient):
+    def __init__(self, secret_namespace: str = None, secret_client: SecretCache = None):
+        self.secret_id = secret_namespace
+        self.secret_client = secret_client
+        # make sure secret_client is a SecretCache type
+        if secret_client is not None and not isinstance(secret_client, SecretCache):
+            raise RuntimeError(
+                "You need to pass a aws_secretsmanager_caching.secret_cache.SecretCache instance. Please refer to https://docs.aws.amazon.com/secretsmanager/latest/userguide/retrieving-secrets_cache-python.html for more details.")
+
+    def get_feathr_secret(self, secret_name: str):
+        """Get Feathr Secrets from AWS Secrets manager. It's also recommended that the client passes a cache objects to reduce cost.
+        See more details here: https://docs.aws.amazon.com/secretsmanager/latest/userguide/retrieving-secrets_cache-python.html
+        """
+        if self.secret_client is None:
+            raise RuntimeError(
+                "You need to pass a aws_secretsmanager_caching.secret_cache.SecretCache instance when initializing FeathrClient.")
+
+        try:
+            get_secret_value_response = self.secret_client.get_secret_string(
+                self.secret_id)
+            # result is in str format, so we need to load it as a dict
+            secret = json.loads(get_secret_value_response)
+            return secret[secret_name]
+        except KeyError as e:
+            logger.error(
+                f"Secret {secret_name} cannot be found in secretsmanager {self.secret_id}.")
+            raise e

feathr_project/feathr/spark_provider/_abc.py

@@ -2,6 +2,7 @@
 from typing import Dict, List, Optional, Tuple
 
 
+
 class SparkJobLauncher(ABC):
     """This is the abstract class for all the spark launchers. All the Spark launcher should implement those interfaces
     """
@@ -52,8 +53,5 @@ def get_status(self) -> str:
 
         Returns:
             str: Status of the current job
-
-        Returns:
-            str: _description_
         """
         pass