openmetadata-ingestion 1.3.0.0.dev0__py3-none-any.whl → 1.3.1.0__py3-none-any.whl

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,26 +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/application.py

@@ -15,9 +15,6 @@ from abc import ABC, abstractmethod
 from typing import List, Optional
 
 from metadata.config.common import WorkflowExecutionError
-from metadata.generated.schema.entity.applications.configuration.applicationConfig import (
-    AppConfig,
-)
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
 )
@@ -50,9 +47,14 @@ class AppRunner(Step, ABC):
     """Class that knows how to execute the Application logic."""
 
     def __init__(
-        self, config: AppConfig.__fields__["__root__"].type_, metadata: OpenMetadata
+        self,
+        config: OpenMetadataApplicationConfig,
+        metadata: OpenMetadata,
     ):
-        self.config = config
+        self.app_config = config.appConfig.__root__ if config.appConfig else None
+        self.private_config = (
+            config.appPrivateConfig.__root__ if config.appPrivateConfig else None
+        )
         self.metadata = metadata
 
         super().__init__()
@@ -67,7 +69,8 @@ class AppRunner(Step, ABC):
 
     @classmethod
     def create(cls, config_dict: dict, metadata: OpenMetadata) -> "Step":
-        return cls(config=config_dict, metadata=metadata)
+        config = OpenMetadataApplicationConfig.parse_obj(config_dict)
+        return cls(config=config, metadata=metadata)
 
 
 class ApplicationWorkflow(BaseWorkflow, ABC):
@@ -77,7 +80,6 @@ class ApplicationWorkflow(BaseWorkflow, ABC):
     runner: Optional[AppRunner]
 
     def __init__(self, config_dict: dict):
-
         self.runner = None  # Will be passed in post-init
         # TODO: Create a parse_gracefully method
         self.config = OpenMetadataApplicationConfig.parse_obj(config_dict)
@@ -114,9 +116,7 @@ class ApplicationWorkflow(BaseWorkflow, ABC):
 
         try:
             self.runner = runner_class(
-                config=self.config.appConfig.__root__
-                if self.config.appConfig
-                else None,
+                config=self.config,
                 metadata=self.metadata,
             )
         except Exception as exc:

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)
 
@@ -235,7 +239,6 @@ class BaseWorkflow(ABC, WorkflowStatusMixin):
             service = self._get_ingestion_pipeline_service()
 
             if service is not None:
-
                 return self.metadata.create_or_update(
                     CreateIngestionPipelineRequest(
                         name=pipeline_name,

metadata/workflow/metadata.py

@@ -31,7 +31,6 @@ class MetadataWorkflow(IngestionWorkflow):
     """
 
     def set_steps(self):
-
         # We keep the source registered in the workflow
         self.source = self._get_source()
         sink = self._get_sink()

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

metadata/workflow/usage.py

@@ -33,7 +33,6 @@ class UsageWorkflow(IngestionWorkflow):
     """
 
     def set_steps(self):
-
         # We keep the source registered in the workflow
         self.source = self._get_source()
         processor = self._get_processor()