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

metadata/utils/execution_time_tracker.py

@@ -0,0 +1,199 @@
+#  Copyright 2021 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.
+
+"""
+ExecutionTimeTracker implementation to help track the execution time of different parts
+of the code.
+"""
+from functools import wraps
+from time import perf_counter
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from metadata.utils.helpers import pretty_print_time_duration
+from metadata.utils.logger import utils_logger
+from metadata.utils.singleton import Singleton
+
+logger = utils_logger()
+
+
+class ExecutionTimeTrackerContext(BaseModel):
+    """Small Model to hold the ExecutionTimeTracker context."""
+
+    name: str
+    start: float
+    stored: bool
+
+
+class ExecutionTimeTracker(metaclass=Singleton):
+    """ExecutionTimeTracker is implemented as a Singleton in order to hold state globally.
+
+    It works as a Context Manager in order to track and log execution times.
+
+    Example:
+
+        def my_function():
+            tracker = ExecutionTimeTracker()
+
+            with tracker(context="MyFunction", store=True):
+                other_opeartion()
+
+    """
+
+    def __init__(self, enabled: bool = False):
+        """When instantiated we can pass if we want it enabled or disabled in order to
+        avoid overhead when not needed.
+
+        Attrs
+        ------
+            enabled: Defines if it will be enabled or not.
+            context: Keeps track of the context levels and their state.
+            state: Keeps track of the global state for the Execution Time Tracker.
+        """
+        self.enabled: bool = enabled
+        self.context: List[ExecutionTimeTrackerContext] = []
+        self.state: Dict[str, float] = {}
+        self.new_context = None
+        self.store = True
+
+    @property
+    def last_stored_context_level(self) -> Optional[str]:
+        """Returns the last stored context level.
+
+        In order to provide better logs and keep track where in the code the time is being
+        measured we keep track of nested contexts.
+
+        If a given context is not stored it will only log to debug but won't be part of the
+        global state.
+        """
+        stored_context = [context for context in self.context if context.stored]
+
+        if stored_context:
+            return stored_context[-1].name
+
+        return None
+
+    def __call__(self, context: str, store: bool = True):
+        """At every point we open a new Context Manager we can pass the current 'context' and
+        if we want to 'store' it.
+
+        Sets the temporary attributes used within the context:
+
+            new_context: Full Context name, appending the given context to the last stored context level.
+            store: If True, it will take part of the global state. Otherwise it will only log to debug.
+        """
+        self.new_context = ".".join(
+            [part for part in [self.last_stored_context_level, context] if part]
+        )
+        self.store = store
+
+        return self
+
+    def __enter__(self):
+        """If enabled, when entering the context, we append a new
+        ExecutionTimeTrackerContext to the list.
+        """
+        if self.enabled:
+            self.context.append(
+                ExecutionTimeTrackerContext(
+                    name=self.new_context, start=perf_counter(), stored=self.store
+                )
+            )
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """If enabled, when exiting the context, we calculate the elapsed time and log to debug.
+        If the context.stored is True, we also save it to the global state."""
+        if self.enabled:
+            stop = perf_counter()
+            context = self.context.pop(-1)
+
+            if not context:
+                return
+
+            elapsed = stop - context.start
+
+            logger.debug(
+                "%s executed in %s", context.name, pretty_print_time_duration(elapsed)
+            )
+
+            if context.stored:
+                self._save(context, elapsed)
+
+    def _save(self, context: ExecutionTimeTrackerContext, elapsed: float):
+        """Small utility to save the new measure to the global accumulator."""
+        self.state[context.name] = self.state.get(context.name, 0) + elapsed
+
+
+def calculate_execution_time(context: Optional[str] = None, store: bool = True):
+    """Utility decorator to be able to use the ExecutionTimeTracker on a function.
+
+    It receives the context and if it should store it.
+
+    Example:
+
+        @calculate_execution_time(context="MyContext")
+        def my_function():
+            ...
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def inner(*args, **kwargs):
+            execution_time = ExecutionTimeTracker()
+
+            with execution_time(context or func.__name__, store):
+                result = func(*args, **kwargs)
+
+            return result
+
+        return inner
+
+    return decorator
+
+
+def calculate_execution_time_generator(
+    context: Optional[str] = None, store: bool = True
+):
+    """Utility decorator to be able to use the ExecutionTimeTracker on a generator function.
+
+    It receives the context and if it should store it.
+
+    Example:
+
+        @calculate_execution_time_generator(context="MyContext")
+        def my_generator():
+            ...
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def inner(*args, **kwargs):
+            # NOTE: We are basically implementing by hand a simplified version of 'yield from'
+            # in order to be able to calculate the time difference correctly.
+            # The 'while True' loop allows us to guarantee we are iterating over all thje values
+            # from func(*args, **kwargs).
+            execution_time = ExecutionTimeTracker()
+
+            generator = func(*args, **kwargs)
+
+            while True:
+                with execution_time(context or func.__name__, store):
+                    try:
+                        element = next(generator)
+                    except StopIteration:
+                        return
+
+                yield element
+
+        return inner
+
+    return decorator

metadata/utils/filters.py

@@ -54,6 +54,10 @@ def _filter(filter_pattern: Optional[FilterPattern], name: str) -> bool:
         # No filter pattern, nothing to filter
         return False
 
+    if filter_pattern and not name:
+        # Filter pattern is present but not the name so we'll filter it out
+        return True
+
     if filter_pattern.includes:
         validate_regex(filter_pattern.includes)
         return not any(  # pylint: disable=use-a-generator

metadata/utils/helpers.py

@@ -20,10 +20,8 @@ import re
 import shutil
 import sys
 from datetime import datetime, timedelta
-from functools import wraps
 from math import floor, log
 from pathlib import Path
-from time import perf_counter
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import sqlparse
@@ -31,7 +29,9 @@ from sqlparse.sql import Statement
 
 from metadata.generated.schema.entity.data.chart import ChartType
 from metadata.generated.schema.entity.data.table import Column, Table
+from metadata.generated.schema.entity.feed.suggestion import Suggestion, SuggestionType
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.type.basic import EntityLink
 from metadata.generated.schema.type.tagLabel import TagLabel
 from metadata.utils.constants import DEFAULT_DATABASE
 from metadata.utils.logger import utils_logger
@@ -99,40 +99,6 @@ om_chart_type_dict = {
 }
 
 
-def calculate_execution_time(func):
-    """
-    Method to calculate workflow execution time
-    """
-
-    @wraps(func)
-    def calculate_debug_time(*args, **kwargs):
-        start = perf_counter()
-        result = func(*args, **kwargs)
-        end = perf_counter()
-        logger.debug(
-            f"{func.__name__} executed in { pretty_print_time_duration(end - start)}"
-        )
-        return result
-
-    return calculate_debug_time
-
-
-def calculate_execution_time_generator(func):
-    """
-    Generator method to calculate workflow execution time
-    """
-
-    def calculate_debug_time(*args, **kwargs):
-        start = perf_counter()
-        yield from func(*args, **kwargs)
-        end = perf_counter()
-        logger.debug(
-            f"{func.__name__} executed in { pretty_print_time_duration(end - start)}"
-        )
-
-    return calculate_debug_time
-
-
 def pretty_print_time_duration(duration: Union[int, float]) -> str:
     """
     Method to format and display the time
@@ -227,12 +193,38 @@ def find_in_iter(element: Any, container: Iterable[Any]) -> Optional[Any]:
     return next((elem for elem in container if elem == element), None)
 
 
-def find_column_in_table(column_name: str, table: Table) -> Optional[Column]:
+def find_column_in_table(
+    column_name: str, table: Table, case_sensitive: bool = True
+) -> Optional[Column]:
     """
     If the column exists in the table, return it
     """
+
+    def equals(first: str, second: str) -> bool:
+        if case_sensitive:
+            return first == second
+        return first.lower() == second.lower()
+
     return next(
-        (col for col in table.columns if col.name.__root__ == column_name), None
+        (col for col in table.columns if equals(col.name.__root__, column_name)), None
+    )
+
+
+def find_suggestion(
+    suggestions: List[Suggestion],
+    suggestion_type: SuggestionType,
+    entity_link: EntityLink,
+) -> Optional[Suggestion]:
+    """Given a list of suggestions, a suggestion type and an entity link, find
+    one suggestion in the list that matches the criteria
+    """
+    return next(
+        (
+            sugg
+            for sugg in suggestions
+            if sugg.type == suggestion_type and sugg.entityLink == entity_link
+        ),
+        None,
     )
 
 

metadata/utils/secrets/aws_based_secrets_manager.py

@@ -12,20 +12,75 @@
 """
 Abstract class for AWS based secrets manager implementations
 """
+import os
 from abc import ABC, abstractmethod
 from typing import Optional
 
 from metadata.clients.aws_client import AWSClient
+from metadata.generated.schema.security.secrets.secretsManagerClientLoader import (
+    SecretsManagerClientLoader,
+)
 from metadata.generated.schema.security.secrets.secretsManagerProvider import (
     SecretsManagerProvider,
 )
+from metadata.ingestion.models.custom_pydantic import CustomSecretStr
+from metadata.utils.dispatch import enum_register
 from metadata.utils.logger import utils_logger
-from metadata.utils.secrets.external_secrets_manager import ExternalSecretsManager
+from metadata.utils.secrets.external_secrets_manager import (
+    SECRET_MANAGER_AIRFLOW_CONF,
+    ExternalSecretsManager,
+    SecretsManagerConfigException,
+)
 
 logger = utils_logger()
 
 NULL_VALUE = "null"
 
+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["AWSCredentials"]:
+    from airflow.configuration import conf
+
+    from metadata.generated.schema.security.credentials.awsCredentials import (
+        AWSCredentials,
+    )
+
+    aws_region = conf.get(SECRET_MANAGER_AIRFLOW_CONF, "aws_region", fallback=None)
+    if aws_region:
+        credentials = AWSCredentials(awsRegion=aws_region)
+        credentials.awsAccessKeyId = conf.get(
+            SECRET_MANAGER_AIRFLOW_CONF, "aws_access_key_id", fallback=""
+        )
+        credentials.awsSecretAccessKey = CustomSecretStr(
+            conf.get(SECRET_MANAGER_AIRFLOW_CONF, "aws_secret_access_key", fallback="")
+        )
+        return credentials
+
+    return None
+
+
+@secrets_manager_client_loader.add(SecretsManagerClientLoader.env.value)
+def _() -> Optional["AWSCredentials"]:
+    from metadata.generated.schema.security.credentials.awsCredentials import (
+        AWSCredentials,
+    )
+
+    # Loading the env vars required by boto3
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html
+    aws_region = os.getenv("AWS_DEFAULT_REGION")
+    if aws_region:
+        return AWSCredentials(awsRegion=aws_region)
+
+    return None
+
 
 class AWSBasedSecretsManager(ExternalSecretsManager, ABC):
     """
@@ -34,12 +89,12 @@ class AWSBasedSecretsManager(ExternalSecretsManager, ABC):
 
     def __init__(
         self,
-        credentials: Optional["AWSCredentials"],
         client: str,
         provider: SecretsManagerProvider,
+        loader: SecretsManagerClientLoader,
     ):
-        super().__init__(provider)
-        self.client = AWSClient(credentials).get_client(client)
+        super().__init__(provider=provider, loader=loader)
+        self.client = AWSClient(self.credentials).get_client(client)
 
     @abstractmethod
     def get_string_value(self, secret_id: str) -> str:
@@ -47,3 +102,11 @@ class AWSBasedSecretsManager(ExternalSecretsManager, ABC):
         :param secret_id: The secret id to retrieve
         :return: The value of the secret
         """
+
+    def load_credentials(self) -> Optional["AWSCredentials"]:
+        """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/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"""