openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__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/fqn.py

@@ -28,6 +28,7 @@ from metadata.generated.antlr.FqnLexer import FqnLexer
 from metadata.generated.antlr.FqnParser import FqnParser
 from metadata.generated.schema.entity.classification.tag import Tag
 from metadata.generated.schema.entity.data.chart import Chart
+from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.dashboard import Dashboard
 from metadata.generated.schema.entity.data.dashboardDataModel import DashboardDataModel
 from metadata.generated.schema.entity.data.database import Database
@@ -284,6 +285,25 @@ def _(
     return _build(service_name, topic_name)
 
 
+@fqn_build_registry.add(Container)
+def _(
+    _: Optional[OpenMetadata],  # ES Index not necessary for Container FQN building
+    *,
+    service_name: str,
+    parent_container: str,
+    container_name: str,
+) -> str:
+    if not service_name or not container_name:
+        raise FQNBuildingException(
+            f"Args should be informed, but got service=`{service_name}`, container=`{container_name}``"
+        )
+    return (
+        _build(parent_container, container_name, quote=False)
+        if parent_container
+        else (_build(service_name, container_name))
+    )
+
+
 @fqn_build_registry.add(SearchIndex)
 def _(
     _: Optional[OpenMetadata],  # ES Index not necessary for Search Index FQN building

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,
     )
 
 
@@ -354,7 +346,7 @@ def format_large_string_numbers(number: Union[float, int]) -> str:
     units = ["", "K", "M", "B", "T"]
     constant_k = 1000.0
     magnitude = int(floor(log(abs(number), constant_k)))
-    return f"{number / constant_k**magnitude:.2f}{units[magnitude]}"
+    return f"{number / constant_k**magnitude:.3f}{units[magnitude]}"
 
 
 def clean_uri(uri: str) -> str:
@@ -430,6 +422,9 @@ def is_safe_sql_query(sql_query: str) -> bool:
         "SET TRANSACTION",
     }
 
+    if sql_query is None:
+        return True
+
     parsed_queries: Tuple[Statement] = sqlparse.parse(sql_query)
     for parsed_query in parsed_queries:
         validation = [

metadata/utils/importer.py

@@ -14,7 +14,7 @@ Helpers to import python classes and modules dynamically
 import importlib
 import traceback
 from enum import Enum
-from typing import Callable, Optional, Type, TypeVar
+from typing import Any, Callable, Optional, Type, TypeVar
 
 from pydantic import BaseModel
 
@@ -23,7 +23,6 @@ from metadata.generated.schema.entity.services.connections.metadata.openMetadata
 )
 from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.metadataIngestion.workflow import Sink as WorkflowSink
-from metadata.ingestion.api.step import Step
 from metadata.ingestion.api.steps import BulkSink, Processor, Sink, Source, Stage
 from metadata.utils.class_helper import get_service_type_from_source_type
 from metadata.utils.logger import utils_logger
@@ -86,7 +85,7 @@ def get_class_name_root(type_: str) -> str:
     )
 
 
-def import_from_module(key: str) -> Type[Step]:
+def import_from_module(key: str) -> Type[Any]:
     """
     Dynamically import an object from a module path
     """

metadata/utils/life_cycle_utils.py

@@ -27,12 +27,12 @@ QUERY_TYPES_DICT = {
 }
 
 
-select_pattern = re.compile(r"^\s*(SELECT|SHOW|DESCRIBE)", re.IGNORECASE)
-create_pattern = re.compile(r"^\s*CREATE", re.IGNORECASE)
+select_pattern = re.compile(r".*\s*(SELECT|SHOW|DESCRIBE)", re.IGNORECASE)
+create_pattern = re.compile(r".*\s*CREATE", re.IGNORECASE)
 update_pattern = re.compile(
-    r"^\s*(UPDATE|INSERT|DELETE|MERGE|TRUNCATE_TABLE|ALTER)", re.IGNORECASE
+    r".*\s*(UPDATE|INSERT|DELETE|MERGE|TRUNCATE_TABLE|ALTER)", re.IGNORECASE
 )
-drop_pattern = re.compile(r"^\s*DROP", re.IGNORECASE)
+drop_pattern = re.compile(r".*\s*DROP", re.IGNORECASE)
 
 logger = utils_logger()
 

metadata/utils/logger.py

@@ -54,6 +54,7 @@ class Loggers(Enum):
     TEST_SUITE = "TestSuite"
     DATA_INSIGHT = "DataInsight"
     QUERY_RUNNER = "QueryRunner"
+    APP = "App"
 
     @DynamicClassAttribute
     def value(self):
@@ -153,6 +154,14 @@ def great_expectations_logger():
     return logging.getLogger(Loggers.GREAT_EXPECTATIONS.value)
 
 
+def app_logger():
+    """
+    Method to get the APP logger
+    """
+
+    return logging.getLogger(Loggers.APP.value)
+
+
 def query_runner_logger():
     """
     Method to get the QUERY_RUNNER logger
@@ -245,9 +254,11 @@ def _(record: TableAndTests) -> str:
 
 
 @get_log_name.register
-def _(_: TestCaseResults) -> Optional[str]:
+def _(record: TestCaseResults) -> str:
     """We don't want to log this in the status"""
-    return None
+    return ",".join(
+        set(result.testCase.name.__root__ for result in record.test_results)
+    )
 
 
 @get_log_name.register

metadata/utils/partition.py

@@ -15,6 +15,7 @@ from typing import Optional
 
 from metadata.generated.schema.entity.data.table import (
     IntervalType,
+    PartitionIntervalUnit,
     PartitionProfilerConfig,
     Table,
 )
@@ -47,8 +48,10 @@ def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
                 return PartitionProfilerConfig(
                     enablePartitioning=True,
                     partitionColumnName=entity.tablePartition.columns[0],
-                    partitionIntervalUnit=entity.tablePartition.interval,
-                    partitionInterval=30,
+                    partitionIntervalUnit=PartitionIntervalUnit.DAY
+                    if entity.tablePartition.interval != "HOUR"
+                    else entity.tablePartition.interval,
+                    partitionInterval=1,
                     partitionIntervalType=entity.tablePartition.intervalType.value,
                     partitionValues=None,
                     partitionIntegerRangeStart=None,
@@ -60,8 +63,10 @@ def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
                     partitionColumnName="_PARTITIONDATE"
                     if entity.tablePartition.interval == "DAY"
                     else "_PARTITIONTIME",
-                    partitionIntervalUnit=entity.tablePartition.interval,
-                    partitionInterval=30,
+                    partitionIntervalUnit=PartitionIntervalUnit.DAY
+                    if entity.tablePartition.interval != "HOUR"
+                    else entity.tablePartition.interval,
+                    partitionInterval=1,
                     partitionIntervalType=entity.tablePartition.intervalType.value,
                     partitionValues=None,
                     partitionIntegerRangeStart=None,
@@ -72,7 +77,7 @@ def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
                     enablePartitioning=True,
                     partitionColumnName=entity.tablePartition.columns[0],
                     partitionIntervalUnit=None,
-                    partitionInterval=30,
+                    partitionInterval=None,
                     partitionIntervalType=entity.tablePartition.intervalType.value,
                     partitionValues=None,
                     partitionIntegerRangeStart=1,

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]:
         """