openmetadata-ingestion 1.3.0.1__py3-none-any.whl → 1.3.1.0__py3-none-any.whl

metadata/utils/secrets/aws_secrets_manager.py

@@ -17,6 +17,9 @@ from typing import Optional
 
 from botocore.exceptions import ClientError
 
+from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
+    SecretsManagerClientLoader,
+)
 from metadata.generated.schema.security.secrets.secretsManagerProvider import (
     SecretsManagerProvider,
 )
@@ -32,8 +35,10 @@ class AWSSecretsManager(AWSBasedSecretsManager):
     Secrets Manager Implementation Class
     """
 
-    def __init__(self, credentials: Optional["AWSCredentials"]):
-        super().__init__(credentials, "secretsmanager", SecretsManagerProvider.aws)
+    def __init__(self, loader: SecretsManagerClientLoader):
+        super().__init__(
+            client="secretsmanager", provider=SecretsManagerProvider.aws, loader=loader
+        )
 
     def get_string_value(self, secret_id: str) -> Optional[str]:
         """

metadata/utils/secrets/aws_ssm_secrets_manager.py

@@ -17,6 +17,9 @@ from typing import Optional
 
 from botocore.exceptions import ClientError
 
+from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
+    SecretsManagerClientLoader,
+)
 from metadata.generated.schema.security.secrets.secretsManagerProvider import (
     SecretsManagerProvider,
 )
@@ -32,8 +35,10 @@ class AWSSSMSecretsManager(AWSBasedSecretsManager):
     AWS SSM Parameter Store Secret Manager Class
     """
 
-    def __init__(self, credentials: Optional["AWSCredentials"]):
-        super().__init__(credentials, "ssm", SecretsManagerProvider.aws)
+    def __init__(self, loader: SecretsManagerClientLoader):
+        super().__init__(
+            client="ssm", provider=SecretsManagerProvider.aws, loader=loader
+        )
 
     def get_string_value(self, secret_id: str) -> Optional[str]:
         """

metadata/utils/secrets/azure_kv_secrets_manager.py

@@ -0,0 +1,148 @@
+#  Copyright 2022 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+Abstract class for AWS based secrets manager implementations
+"""
+import os
+import traceback
+from abc import ABC
+from typing import Optional
+
+from azure.identity import ClientSecretCredential, DefaultAzureCredential
+from azure.keyvault.secrets import KeyVaultSecret, SecretClient
+
+from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
+    SecretsManagerClientLoader,
+)
+from metadata.generated.schema.security.secrets.secretsManagerProvider import (
+    SecretsManagerProvider,
+)
+from metadata.utils.dispatch import enum_register
+from metadata.utils.logger import utils_logger
+from metadata.utils.secrets.external_secrets_manager import (
+    SECRET_MANAGER_AIRFLOW_CONF,
+    ExternalSecretsManager,
+    SecretsManagerConfigException,
+)
+
+logger = utils_logger()
+
+secrets_manager_client_loader = enum_register()
+
+
+# pylint: disable=import-outside-toplevel
+@secrets_manager_client_loader.add(SecretsManagerClientLoader.noop.value)
+def _() -> None:
+    return None
+
+
+@secrets_manager_client_loader.add(SecretsManagerClientLoader.airflow.value)
+def _() -> Optional["AzureCredentials"]:
+    from airflow.configuration import conf
+
+    from metadata.generated.schema.security.credentials.azureCredentials import (
+        AzureCredentials,
+    )
+
+    key_vault_name = conf.get(
+        SECRET_MANAGER_AIRFLOW_CONF, "azure_key_vault_name", fallback=None
+    )
+    if not key_vault_name:
+        raise ValueError(
+            "Missing `azure_key_vault_name` config for Azure Key Vault Secrets Manager Provider."
+        )
+
+    tenant_id = conf.get(SECRET_MANAGER_AIRFLOW_CONF, "azure_tenant_id", fallback=None)
+    client_id = conf.get(SECRET_MANAGER_AIRFLOW_CONF, "azure_client_id", fallback=None)
+    client_secret = conf.get(
+        SECRET_MANAGER_AIRFLOW_CONF, "azure_client_secret", fallback=None
+    )
+
+    return AzureCredentials(
+        clientId=client_id,
+        clientSecret=client_secret,
+        tenantId=tenant_id,
+        vaultName=key_vault_name,
+    )
+
+
+@secrets_manager_client_loader.add(SecretsManagerClientLoader.env.value)
+def _() -> Optional["AzureCredentials"]:
+    from metadata.generated.schema.security.credentials.azureCredentials import (
+        AzureCredentials,
+    )
+
+    # Load only the AZURE_KEY_VAULT_NAME (required) variable and use the
+    # Default Auth chain
+    # https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#defaultazurecredential
+    key_vault_name = os.getenv("AZURE_KEY_VAULT_NAME")
+
+    if not key_vault_name:
+        raise ValueError(
+            "Missing `azure_key_vault_name` config for Azure Key Vault Secrets Manager Provider."
+        )
+
+    return AzureCredentials(vaultName=key_vault_name)
+
+
+class AzureKVSecretsManager(ExternalSecretsManager, ABC):
+    """
+    Azure Key Vault Secrets Manager class
+    """
+
+    def __init__(
+        self,
+        loader: SecretsManagerClientLoader,
+    ):
+        super().__init__(provider=SecretsManagerProvider.azure_kv, loader=loader)
+
+        if (
+            self.credentials.tenantId
+            and self.credentials.clientId
+            and self.credentials.clientSecret
+        ):
+            azure_identity = ClientSecretCredential(
+                tenant_id=self.credentials.tenantId,
+                client_id=self.credentials.clientId,
+                client_secret=self.credentials.clientSecret.get_secret_value(),
+            )
+        else:
+            azure_identity = DefaultAzureCredential()
+
+        self.client = SecretClient(
+            vault_url=f"https://{self.credentials.vaultName}.vault.azure.net/",
+            credential=azure_identity,
+        )
+
+    def get_string_value(self, secret_id: str) -> str:
+        """
+        :param secret_id: The secret id to retrieve
+        :return: The value of the secret
+        """
+        try:
+            secret: KeyVaultSecret = self.client.get_secret(secret_id)
+            logger.debug(f"Got value for secret {secret_id}")
+            return secret.value
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Could not get the secret value of {secret_id} due to [{exc}]"
+            )
+            raise exc
+
+    def load_credentials(self) -> Optional["AzureCredentials"]:
+        """Load the provider credentials based on the loader type"""
+        try:
+            loader_fn = secrets_manager_client_loader.registry.get(self.loader.value)
+            return loader_fn()
+        except Exception as err:
+            raise SecretsManagerConfigException(f"Error loading credentials - [{err}]")

metadata/utils/secrets/external_secrets_manager.py

@@ -12,18 +12,40 @@
 """
 Abstract class for third party secrets' manager implementations
 """
-from abc import ABC
+from abc import ABC, abstractmethod
+from typing import Any
 
+from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
+    SecretsManagerClientLoader,
+)
 from metadata.generated.schema.security.secrets.secretsManagerProvider import (
     SecretsManagerProvider,
 )
 from metadata.utils.secrets.secrets_manager import SecretsManager
 
+SECRET_MANAGER_AIRFLOW_CONF = "openmetadata_secrets_manager"
+
+
+class SecretsManagerConfigException(Exception):
+    """
+    Invalid config that does not allow us to create
+    the SecretsManagerFactory
+    """
+
 
 class ExternalSecretsManager(SecretsManager, ABC):
     """
     Abstract class for third party secrets' manager implementations
     """
 
-    def __init__(self, provider: SecretsManagerProvider):
-        self.provider = provider.name
+    def __init__(
+        self, provider: SecretsManagerProvider, loader: SecretsManagerClientLoader
+    ):
+        self.provider = provider
+        self.loader = loader
+
+        self.credentials = self.load_credentials()
+
+    @abstractmethod
+    def load_credentials(self) -> Any:
+        """Load the provider credentials based on the loader type"""

metadata/utils/secrets/secrets_manager_factory.py

@@ -12,7 +12,7 @@
 """
 Secrets manager factory module
 """
-from typing import Any, Optional
+from typing import Optional
 
 from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
     SecretsManagerClientLoader,
@@ -22,19 +22,12 @@ from metadata.generated.schema.security.secrets.secretsManagerProvider import (
 )
 from metadata.utils.secrets.aws_secrets_manager import AWSSecretsManager
 from metadata.utils.secrets.aws_ssm_secrets_manager import AWSSSMSecretsManager
-from metadata.utils.secrets.client.loader import secrets_manager_client_loader
-from metadata.utils.secrets.noop_secrets_manager import DBSecretsManager
+from metadata.utils.secrets.azure_kv_secrets_manager import AzureKVSecretsManager
+from metadata.utils.secrets.db_secrets_manager import DBSecretsManager
 from metadata.utils.secrets.secrets_manager import SecretsManager
 from metadata.utils.singleton import Singleton
 
 
-class SecretsManagerConfigException(Exception):
-    """
-    Invalid config that does not allow us to create
-    the SecretsManagerFactory
-    """
-
-
 class SecretsManagerFactory(metaclass=Singleton):
     """
     Singleton factory to initialize a secret manager. It will return always the same secret manager instance.
@@ -56,11 +49,9 @@ class SecretsManagerFactory(metaclass=Singleton):
         self._secrets_manager_provider = secrets_manager_provider
         self._secrets_manager_loader = secrets_manager_loader
 
-        credentials = self._load_secrets_manager_credentials()
-
         self.secrets_manager = self._get_secrets_manager(
             secrets_manager_provider,
-            credentials,
+            secrets_manager_loader,
         )
 
     @property
@@ -74,13 +65,12 @@ class SecretsManagerFactory(metaclass=Singleton):
     def _get_secrets_manager(
         self,
         secrets_manager_provider: SecretsManagerProvider,
-        credentials: Any = None,
+        secrets_manager_loader: SecretsManagerClientLoader,
     ) -> SecretsManager:
         """
         Method to get the secrets manager based on the arguments passed
         :param secrets_manager_provider: the secrets' manager provider
-        :param credentials: optional credentials that could be required by the clients of the secrets' manager
-                            implementations
+        :param secrets_manager_loader: how to retrieve the secrets manager keys from the environment
         :return: a secrets manager
         """
         if (
@@ -92,25 +82,18 @@ class SecretsManagerFactory(metaclass=Singleton):
             SecretsManagerProvider.aws,
             SecretsManagerProvider.managed_aws,
         ):
-            return AWSSecretsManager(credentials)
+            return AWSSecretsManager(secrets_manager_loader)
         if secrets_manager_provider in (
             SecretsManagerProvider.aws_ssm,
             SecretsManagerProvider.managed_aws_ssm,
         ):
-            return AWSSSMSecretsManager(credentials)
+            return AWSSSMSecretsManager(secrets_manager_loader)
+        if secrets_manager_provider in (
+            SecretsManagerProvider.azure_kv,
+            SecretsManagerProvider.managed_azure_kv,
+        ):
+            return AzureKVSecretsManager(secrets_manager_loader)
         raise NotImplementedError(f"[{secrets_manager_provider}] is not implemented.")
 
     def get_secrets_manager(self):
         return self.secrets_manager
-
-    def _load_secrets_manager_credentials(self) -> Optional["AWSCredentials"]:
-        if not self.secrets_manager_loader:
-            return None
-
-        try:
-            loader_fn = secrets_manager_client_loader.registry.get(
-                self.secrets_manager_loader.value
-            )
-            return loader_fn(self.secrets_manager_provider)
-        except Exception as err:
-            raise SecretsManagerConfigException(f"Error loading credentials - [{err}]")

metadata/workflow/base.py

@@ -44,6 +44,7 @@ from metadata.utils.class_helper import (
     get_reference_type_from_service_type,
     get_service_class_from_service_type,
 )
+from metadata.utils.execution_time_tracker import ExecutionTimeTracker
 from metadata.utils.helpers import datetime_to_ts
 from metadata.utils.logger import ingestion_logger, set_loggers_level
 from metadata.workflow.output_handler import report_ingestion_status
@@ -92,6 +93,9 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
         self._timer: Optional[RepeatedTimer] = None
         self._ingestion_pipeline: Optional[IngestionPipeline] = None
         self._start_ts = datetime_to_ts(datetime.now())
+        self._execution_time_tracker = ExecutionTimeTracker(
+            log_level == LogLevels.DEBUG
+        )
 
         set_loggers_level(log_level.value)
 

metadata/workflow/output_handler.py

@@ -26,6 +26,8 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 )
 from metadata.generated.schema.metadataIngestion.workflow import LogLevels
 from metadata.ingestion.api.step import Summary
+from metadata.utils.execution_time_tracker import ExecutionTimeTracker
+from metadata.utils.helpers import pretty_print_time_duration
 from metadata.utils.logger import ANSI, log_ansi_encoded_string
 
 WORKFLOW_FAILURE_MESSAGE = "Workflow finished with failures"
@@ -137,6 +139,25 @@ def is_debug_enabled(workflow) -> bool:
     )
 
 
+def print_execution_time_summary():
+    """Log the ExecutionTimeTracker Summary."""
+    tracker = ExecutionTimeTracker()
+
+    summary_table = {
+        "Context": [],
+        "Execution Time Aggregate": [],
+    }
+
+    for key in sorted(tracker.state.keys()):
+        summary_table["Context"].append(key)
+        summary_table["Execution Time Aggregate"].append(
+            pretty_print_time_duration(tracker.state[key])
+        )
+
+    log_ansi_encoded_string(bold=True, message="Execution Time Summary")
+    log_ansi_encoded_string(message=f"\n{tabulate(summary_table, tablefmt='grid')}")
+
+
 def print_workflow_summary(workflow: "BaseWorkflow") -> None:
     """
     Args:
@@ -148,6 +169,7 @@ def print_workflow_summary(workflow: "BaseWorkflow") -> None:
 
     if is_debug_enabled(workflow):
         print_workflow_status_debug(workflow)
+        print_execution_time_summary()
 
     failures = []
     total_records = 0