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

metadata/ingestion/source/database/dbt/metadata.py

@@ -31,8 +31,6 @@ from metadata.generated.schema.entity.services.databaseService import DatabaseSe
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
-from metadata.generated.schema.entity.teams.team import Team
-from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
@@ -140,39 +138,29 @@ class DbtSource(DbtServiceSource):
 
     def get_dbt_owner(
         self, manifest_node: dict, catalog_node: Optional[dict]
-    ) -> Optional[str]:
+    ) -> Optional[EntityReference]:
         """
         Returns dbt owner
         """
-        owner = None
-        dbt_owner = None
-        if catalog_node:
-            dbt_owner = catalog_node.metadata.owner
-        if manifest_node:
-            dbt_owner = manifest_node.meta.get(DbtCommonEnum.OWNER.value)
-        if dbt_owner:
-            owner_name = dbt_owner
-            user_owner_fqn = fqn.build(
-                self.metadata, entity_type=User, user_name=owner_name
-            )
-            if user_owner_fqn:
-                owner = self.metadata.get_entity_reference(
-                    entity=User, fqn=user_owner_fqn
-                )
-            else:
-                team_owner_fqn = fqn.build(
-                    self.metadata, entity_type=Team, team_name=owner_name
-                )
-                if team_owner_fqn:
-                    owner = self.metadata.get_entity_reference(
-                        entity=Team, fqn=team_owner_fqn
-                    )
-                else:
+        try:
+            owner = None
+            dbt_owner = None
+            if catalog_node:
+                dbt_owner = catalog_node.metadata.owner
+            if manifest_node:
+                dbt_owner = manifest_node.meta.get(DbtCommonEnum.OWNER.value)
+            if dbt_owner:
+                owner = self.metadata.get_reference_by_name(name=dbt_owner)
+                if not owner:
                     logger.warning(
                         "Unable to ingest owner from DBT since no user or"
                         f" team was found with name {dbt_owner}"
                     )
-        return owner
+            return owner
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Unable to ingest owner from DBT due to: {exc}")
+        return None
 
     def check_columns(self, catalog_node):
         for catalog_key, catalog_column in catalog_node.get("columns").items():

metadata/ingestion/source/database/oracle/queries.py

@@ -172,14 +172,14 @@ SELECT
   SQL_FULLTEXT AS query_text,
   TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') AS start_time,
   ELAPSED_TIME / 1000 AS duration,
-  TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000, 'SECOND') AS end_time
+  TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000000, 'SECOND') AS end_time
 FROM gv$sql
 WHERE OBJECT_STATUS = 'VALID' 
   {filters}
   AND SQL_FULLTEXT NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
   AND SQL_FULLTEXT NOT LIKE '/* {{"app": "dbt", %%}} */%%'
   AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') >= TO_TIMESTAMP('{start_time}', 'yy-MM-dd HH24:MI:SS')
-  AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000, 'SECOND') 
+  AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000000, 'SECOND') 
   < TO_TIMESTAMP('{end_time}', 'yy-MM-dd HH24:MI:SS')
 ORDER BY FIRST_LOAD_TIME DESC 
 OFFSET 0 ROWS FETCH NEXT {result_limit} ROWS ONLY

metadata/ingestion/source/messaging/common_broker_source.py

@@ -226,14 +226,16 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
         if topic_entity and self.generate_sample_data:
             topic_name = topic_details.topic_name
             sample_data = []
+            messages = None
             try:
-                self.consumer_client.subscribe(
-                    [topic_name], on_assign=on_partitions_assignment_to_consumer
-                )
-                logger.info(
-                    f"Broker consumer polling for sample messages in topic {topic_name}"
-                )
-                messages = self.consumer_client.consume(num_messages=10, timeout=10)
+                if self.consumer_client:
+                    self.consumer_client.subscribe(
+                        [topic_name], on_assign=on_partitions_assignment_to_consumer
+                    )
+                    logger.info(
+                        f"Broker consumer polling for sample messages in topic {topic_name}"
+                    )
+                    messages = self.consumer_client.consume(num_messages=10, timeout=10)
             except Exception as exc:
                 yield Either(
                     left=StackTraceError(

metadata/ingestion/source/messaging/kafka/connection.py

@@ -15,7 +15,7 @@ Source connection handler
 from dataclasses import dataclass
 from typing import Optional, Union
 
-from confluent_kafka.admin import AdminClient
+from confluent_kafka.admin import AdminClient, KafkaException
 from confluent_kafka.avro import AvroConsumer
 from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistryClient
 
@@ -35,6 +35,21 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
+class InvalidKafkaCreds(Exception):
+    """
+    Class to indicate invalid kafka credentials exception
+    """
+
+
+class SchemaRegistryException(Exception):
+    """
+    Class to indicate invalid schema registry not initialized
+    """
+
+
+TIMEOUT_SECONDS = 10
+
+
 @dataclass
 class KafkaClient:
     def __init__(self, admin_client, schema_registry_client, consumer_client) -> None:
@@ -60,6 +75,14 @@ def get_connection(
         if connection.saslMechanism:
             connection.consumerConfig["sasl.mechanism"] = connection.saslMechanism.value
 
+        if (
+            connection.consumerConfig.get("security.protocol") is None
+            and connection.securityProtocol
+        ):
+            connection.consumerConfig[
+                "security.protocol"
+            ] = connection.securityProtocol.value
+
     if connection.basicAuthUserInfo:
         connection.schemaRegistryConfig = connection.schemaRegistryConfig or {}
         connection.schemaRegistryConfig[
@@ -109,9 +132,27 @@ def test_connection(
     """
 
     def custom_executor():
-        _ = client.admin_client.list_topics().topics
-
-    test_fn = {"GetTopics": custom_executor}
+        try:
+            client.admin_client.list_topics(timeout=TIMEOUT_SECONDS).topics
+        except KafkaException as err:
+            raise InvalidKafkaCreds(
+                f"Failed to fetch topics due to: {err}. "
+                "Please validate credentials and check if you are using correct security protocol"
+            )
+
+    def schema_registry_test():
+        if client.schema_registry_client:
+            client.schema_registry_client.get_subjects()
+        else:
+            raise SchemaRegistryException(
+                "Schema Registry not initialized, please provide schema registry "
+                "credentials in case you want topic schema and sample data to be ingested"
+            )
+
+    test_fn = {
+        "GetTopics": custom_executor,
+        "CheckSchemaRegistry": schema_registry_test,
+    }
 
     test_connection_steps(
         metadata=metadata,

metadata/ingestion/source/mlmodel/sagemaker/metadata.py

@@ -24,6 +24,9 @@ from metadata.generated.schema.entity.data.mlmodel import (
 from metadata.generated.schema.entity.services.connections.mlmodel.sageMakerConnection import (
     SageMakerConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
@@ -128,14 +131,23 @@ class SagemakerSource(MlModelServiceSource):
         """
         Prepare the Request model
         """
-        mlmodel_request = CreateMlModelRequest(
-            name=model.name,
-            algorithm=self._get_algorithm(),  # Setting this to a constant
-            mlStore=self._get_ml_store(model.name),
-            service=self.context.mlmodel_service,
-        )
-        yield mlmodel_request
-        self.register_record(mlmodel_request=mlmodel_request)
+        try:
+            mlmodel_request = CreateMlModelRequest(
+                name=model.name,
+                algorithm=self._get_algorithm(),  # Setting this to a constant
+                mlStore=self._get_ml_store(model.name),
+                service=self.context.mlmodel_service,
+            )
+            yield Either(right=mlmodel_request)
+            self.register_record(mlmodel_request=mlmodel_request)
+        except Exception as exc:  # pylint: disable=broad-except
+            yield Either(
+                left=StackTraceError(
+                    name=model.name,
+                    error=f"Error creating mlmodel: {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
 
     def _get_ml_store(  # pylint: disable=arguments-differ
         self,

metadata/profiler/orm/functions/conn_test.py

@@ -48,6 +48,7 @@ def _(*_, **__):
 
 @compiles(ConnTestFn, Dialects.Db2)
 @compiles(ConnTestFn, Dialects.IbmDbSa)
+@compiles(ConnTestFn, Dialects.Ibmi)
 def _(*_, **__):
     return "SELECT 42 FROM SYSIBM.SYSDUMMY1;"
 

metadata/profiler/orm/functions/sum.py

@@ -40,6 +40,7 @@ def _(element, compiler, **kw):
 
 
 @compiles(SumFn, Dialects.BigQuery)
+@compiles(SumFn, Dialects.Postgres)
 def _(element, compiler, **kw):
     """Handle case where column type is INTEGER but SUM returns a NUMBER"""
     proc = compiler.process(element.clauses, **kw)

metadata/profiler/orm/registry.py

@@ -64,6 +64,7 @@ class Dialects(Enum):
     Hive = "hive"
     Impala = "impala"
     IbmDbSa = "ibm_db_sa"
+    Ibmi = "ibmi"
     MariaDB = "mariadb"
     MSSQL = "mssql"
     MySQL = "mysql"

metadata/profiler/processor/core.py

@@ -53,7 +53,7 @@ from metadata.profiler.metrics.static.row_count import RowCount
 from metadata.profiler.orm.registry import NOT_COMPUTE
 from metadata.profiler.processor.sample_data_handler import upload_sample_data
 from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
-from metadata.utils.helpers import calculate_execution_time
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -550,7 +550,7 @@ class Profiler(Generic[TMetric]):
 
         return table_profile
 
-    @calculate_execution_time
+    @calculate_execution_time(store=False)
     def generate_sample_data(self) -> Optional[TableData]:
         """Fetch and ingest sample data
 

metadata/utils/datalake/datalake_utils.py

@@ -169,6 +169,7 @@ class GenericDataFrameColumnParser:
             ["datetime64", "timedelta[ns]", "datetime64[ns]"], DataType.DATETIME
         ),
         "str": DataType.STRING,
+        "bytes": DataType.BYTES,
     }
 
     def __init__(self, data_frame: "DataFrame"):
@@ -247,8 +248,13 @@ class GenericDataFrameColumnParser:
                     data_type = "string"
 
             data_type = cls._data_formats.get(
-                data_type or data_frame[column_name].dtypes.name, DataType.STRING
+                data_type or data_frame[column_name].dtypes.name,
             )
+            if not data_type:
+                logger.debug(
+                    f"unknown data type {data_frame[column_name].dtypes.name}. resolving to string."
+                )
+            data_type = data_type or DataType.STRING
         except Exception as err:
             logger.warning(
                 f"Failed to distinguish data type for column {column_name}, Falling back to {data_type}, exc: {err}"

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
@@ -101,55 +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):
-        # 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).
-        generator = func(*args, **kwargs)
-
-        while True:
-            start = perf_counter()
-
-            try:
-                element = next(generator)
-            except StopIteration:
-                return
-
-            end = perf_counter()
-
-            logger.debug(
-                f"{func.__name__} executed in { pretty_print_time_duration(end - start)}"
-            )
-
-            yield element
-
-    return calculate_debug_time
-
-
 def pretty_print_time_duration(duration: Union[int, float]) -> str:
     """
     Method to format and display the time

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}]")