openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0rc2__py3-none-any.whl

metadata/readers/file/config_source_factory.py

@@ -18,6 +18,15 @@ ConfigSource Reader Factory: Helps us choose the reader from
 """
 from typing import Any
 
+from metadata.generated.schema.entity.services.connections.dashboard.powerbi.azureConfig import (
+    AzureConfig as PowerBiAzureConfig,
+)
+from metadata.generated.schema.entity.services.connections.dashboard.powerbi.gcsConfig import (
+    GCSConfig as PowerBiGCSConfig,
+)
+from metadata.generated.schema.entity.services.connections.dashboard.powerbi.s3Config import (
+    S3Config as PowerBiS3Config,
+)
 from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
     AzureConfig,
 )
@@ -58,6 +67,9 @@ CONFIG_SOURCE_READER = {
     DbtAzureConfig.__name__: ADLSReader,
     DbtGcsConfig.__name__: GCSReader,
     DbtS3Config.__name__: S3Reader,
+    PowerBiAzureConfig.__name__: ADLSReader,
+    PowerBiS3Config.__name__: S3Reader,
+    PowerBiGCSConfig.__name__: GCSReader,
 }
 
 

metadata/readers/file/gcs.py

@@ -48,3 +48,21 @@ class GCSReader(Reader):
         to traverse any directories.
         """
         raise NotImplementedError("Not implemented")
+
+    def download(
+        self,
+        path: str,
+        local_file_path: str,
+        *,
+        bucket_name: str = None,
+        verbose: bool = True,
+        **__,
+    ) -> bytes:
+        try:
+            self.client.get_bucket(bucket_name).get_blob(path).download_to_filename(
+                local_file_path
+            )
+        except Exception as err:
+            if verbose:
+                logger.debug(traceback.format_exc())
+            raise ReadException(f"Error downloading file [{path}] from GCS: {err}")

metadata/readers/file/local.py

@@ -78,3 +78,15 @@ class LocalReader(Reader):
                     file_paths.append(f"{root}/{fle}")
 
         return file_paths
+
+    def download(
+        self,
+        path: str,
+        local_file_path: str,
+        *,
+        bucket_name: str = None,
+        verbose: bool = True,
+        **__,
+    ):
+        # Nothing to download for local files
+        pass

metadata/readers/file/s3.py

@@ -46,3 +46,19 @@ class S3Reader(Reader):
         to traverse any directories.
         """
         raise NotImplementedError("Not implemented")
+
+    def download(
+        self,
+        path: str,
+        local_file_path: str,
+        *,
+        bucket_name: str = None,
+        verbose: bool = True,
+        **__,
+    ):
+        try:
+            self.client.download_file(bucket_name, path, local_file_path)
+        except Exception as err:
+            if verbose:
+                logger.debug(traceback.format_exc())
+            raise ReadException(f"Error downloading file [{path}] from S3: {err}")

metadata/utils/constants.py

@@ -24,6 +24,8 @@ from metadata.generated.schema.entity.data.searchIndex import SearchIndex
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.data.topic import Topic
+from metadata.generated.schema.entity.domains.dataProduct import DataProduct
+from metadata.generated.schema.entity.domains.domain import Domain
 from metadata.generated.schema.entity.services.dashboardService import DashboardService
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.entity.services.messagingService import MessagingService
@@ -70,31 +72,37 @@ NO_ACCESS_TOKEN = "no_token"
 
 SAMPLE_DATA_DEFAULT_COUNT = 50
 
-# Mainly used for lineage
-ENTITY_REFERENCE_TYPE_MAP = {
+ENTITY_REFERENCE_CLASS_MAP = {
     # Service Entities
-    DatabaseService.__name__: "databaseService",
-    MessagingService.__name__: "messagingService",
-    DashboardService.__name__: "dashboardService",
-    PipelineService.__name__: "pipelineService",
-    StorageService.__name__: "storageService",
-    MlModelService.__name__: "mlmodelService",
-    MetadataService.__name__: "metadataService",
-    SearchService.__name__: "searchService",
+    "databaseService": DatabaseService,
+    "messagingService": MessagingService,
+    "dashboardService": DashboardService,
+    "pipelineService": PipelineService,
+    "storageService": StorageService,
+    "mlmodelService": MlModelService,
+    "metadataService": MetadataService,
+    "searchService": SearchService,
     # Data Asset Entities
-    Table.__name__: "table",
-    StoredProcedure.__name__: "storedProcedure",
-    Database.__name__: "database",
-    DatabaseSchema.__name__: "databaseSchema",
-    Dashboard.__name__: "dashboard",
-    DashboardDataModel.__name__: "dashboardDataModel",
-    Pipeline.__name__: "pipeline",
-    Chart.__name__: "chart",
-    Topic.__name__: "topic",
-    SearchIndex.__name__: "searchIndex",
-    MlModel.__name__: "mlmodel",
-    Container.__name__: "container",
+    "table": Table,
+    "storedProcedure": StoredProcedure,
+    "database": Database,
+    "databaseSchema": DatabaseSchema,
+    "dashboard": Dashboard,
+    "dashboardDataModel": DashboardDataModel,
+    "pipeline": Pipeline,
+    "chart": Chart,
+    "topic": Topic,
+    "searchIndex": SearchIndex,
+    "mlmodel": MlModel,
+    "container": Container,
     # User Entities
-    User.__name__: "user",
-    Team.__name__: "team",
+    "user": User,
+    "team": Team,
+    # Domain
+    "domain": Domain,
+    "dataProduct": DataProduct,
+}
+
+ENTITY_REFERENCE_TYPE_MAP = {
+    value.__name__: key for key, value in ENTITY_REFERENCE_CLASS_MAP.items()
 }

metadata/utils/datalake/datalake_utils.py

@@ -179,8 +179,10 @@ class GenericDataFrameColumnParser:
         **dict.fromkeys(["float64", "float32", "float"], DataType.FLOAT),
         "bool": DataType.BOOLEAN,
         **dict.fromkeys(
-            ["datetime64", "timedelta[ns]", "datetime64[ns]"], DataType.DATETIME
+            ["datetime64[ns]", "datetime"],
+            DataType.DATETIME,
         ),
+        "timedelta[ns]": DataType.TIME,
         "str": DataType.STRING,
         "bytes": DataType.BYTES,
     }
@@ -253,10 +255,44 @@ class GenericDataFrameColumnParser:
             ):
                 try:
                     # Safely evaluate the input string
-                    df_row_val = data_frame[column_name].dropna().values[0]
-                    parsed_object = ast.literal_eval(str(df_row_val))
+                    df_row_val_list = data_frame[column_name].dropna().values[:1000]
+                    parsed_object_datatype_list = []
+                    for df_row_val in df_row_val_list:
+                        try:
+                            parsed_object_datatype_list.append(
+                                type(ast.literal_eval(str(df_row_val))).__name__.lower()
+                            )
+                        except (ValueError, SyntaxError):
+                            # we try to parse the value as a datetime, if it fails, we fallback to string
+                            # as literal_eval will fail for string
+                            from datetime import datetime
+
+                            from dateutil.parser import ParserError, parse
+
+                            try:
+                                dtype_ = "int64"
+                                if not str(df_row_val).isnumeric():
+                                    # check if the row value is time
+                                    try:
+                                        datetime.strptime(df_row_val, "%H:%M:%S").time()
+                                        dtype_ = "timedelta[ns]"
+                                    except (ValueError, TypeError):
+                                        # check if the row value is date / time / datetime
+                                        type(parse(df_row_val)).__name__.lower()
+                                        dtype_ = "datetime64[ns]"
+                                parsed_object_datatype_list.append(dtype_)
+                            except (ParserError, TypeError):
+                                parsed_object_datatype_list.append("str")
+                        except Exception as err:
+                            logger.debug(
+                                f"Failed to parse datatype for column {column_name}, exc: {err},"
+                                "Falling back to string."
+                            )
+                            parsed_object_datatype_list.append("str")
+
+                    data_type = max(parsed_object_datatype_list)
                     # Determine the data type of the parsed object
-                    data_type = type(parsed_object).__name__.lower()
+
                 except (ValueError, SyntaxError):
                     # Handle any exceptions that may occur
                     data_type = "string"

metadata/utils/elasticsearch.py

@@ -18,8 +18,10 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.analytics.reportData import ReportData
 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
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.glossary import Glossary
@@ -27,6 +29,8 @@ from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
 from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.query import Query
+from metadata.generated.schema.entity.data.searchIndex import SearchIndex
+from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.data.topic import Topic
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
@@ -42,9 +46,13 @@ ES_INDEX_MAP = {
     Database.__name__: "database_search_index",
     DatabaseSchema.__name__: "database_schema_search_index",
     DatabaseService.__name__: "database_service_search_index",
+    StoredProcedure.__name__: "stored_procedure_search_index",
+    DashboardDataModel.__name__: "dashboard_data_model_search_index",
+    SearchIndex.__name__: "search_entity_search_index",
     Team.__name__: "team_search_index",
     User.__name__: "user_search_index",
     Dashboard.__name__: "dashboard_search_index",
+    Chart.__name__: "chart_search_index",
     Topic.__name__: "topic_search_index",
     Pipeline.__name__: "pipeline_search_index",
     Glossary.__name__: "glossary_search_index",

metadata/utils/execution_time_tracker.py

@@ -13,9 +13,11 @@
 ExecutionTimeTracker implementation to help track the execution time of different parts
 of the code.
 """
+import threading
+from copy import deepcopy
 from functools import wraps
 from time import perf_counter
-from typing import Dict, List, Optional
+from typing import List, Optional
 
 from pydantic import BaseModel
 
@@ -34,6 +36,60 @@ class ExecutionTimeTrackerContext(BaseModel):
     stored: bool
 
 
+class ExecutionTimeTrackerContextMap(metaclass=Singleton):
+    """Responsible for managing the ExecutionTimeTracker on different threads."""
+
+    def __init__(self):
+        """Initializes the map."""
+        self.map: dict[int, List[ExecutionTimeTrackerContext]] = {}
+
+    def copy_from_parent(self, parent_thread_id: int, thread_id: Optional[int] = None):
+        """Copy the ExecutionTimeTrackerContext from Parent."""
+        thread_id = thread_id or threading.get_ident()
+
+        self.map[thread_id] = deepcopy(self.map.get(parent_thread_id, []))
+
+    def get_last_stored_context_level(
+        self, thread_id: Optional[int] = None
+    ) -> Optional[str]:
+        """Gets the last stored context level for a given thread."""
+        thread_id = thread_id or threading.get_ident()
+
+        stored_context = [
+            context for context in self.map.get(thread_id, {}) if context.stored
+        ]
+
+        if stored_context:
+            return stored_context[-1].name
+        return None
+
+    def append(
+        self, context: ExecutionTimeTrackerContext, thread_id: Optional[int] = None
+    ):
+        """Appends a new context level for a given thread."""
+        thread_id = thread_id or threading.get_ident()
+        self.map.setdefault(thread_id, []).append(context)
+
+    def pop(self, thread_id: Optional[int] = None) -> ExecutionTimeTrackerContext:
+        """Removes the information of a given thread."""
+        thread_id = thread_id or threading.get_ident()
+        return self.map.get(thread_id, []).pop()
+
+
+class ExecutionTimeTrackerState(metaclass=Singleton):
+    """Tracks the ExecutionTime State across multiple threads."""
+
+    def __init__(self):
+        """Initializes the state and the lock."""
+        self.state = {}
+        self.lock = threading.Lock()
+
+    def add(self, context: ExecutionTimeTrackerContext, elapsed: float):
+        """Updates the State."""
+        with self.lock:
+            self.state[context.name] = self.state.get(context.name, 0) + elapsed
+
+
 class ExecutionTimeTracker(metaclass=Singleton):
     """ExecutionTimeTracker is implemented as a Singleton in order to hold state globally.
 
@@ -60,27 +116,12 @@ class ExecutionTimeTracker(metaclass=Singleton):
             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]
+        self.context_map = ExecutionTimeTrackerContextMap()
+        self.state = ExecutionTimeTrackerState()
 
-        if stored_context:
-            return stored_context[-1].name
-
-        return None
+        self.new_context = ""
+        self.store = True
 
     def __call__(self, context: str, store: bool = True):
         """At every point we open a new Context Manager we can pass the current 'context' and
@@ -92,7 +133,11 @@ class ExecutionTimeTracker(metaclass=Singleton):
             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]
+            [
+                part
+                for part in [self.context_map.get_last_stored_context_level(), context]
+                if part
+            ]
         )
         self.store = store
 
@@ -103,7 +148,7 @@ class ExecutionTimeTracker(metaclass=Singleton):
         ExecutionTimeTrackerContext to the list.
         """
         if self.enabled:
-            self.context.append(
+            self.context_map.append(
                 ExecutionTimeTrackerContext(
                     name=self.new_context, start=perf_counter(), stored=self.store
                 )
@@ -114,7 +159,7 @@ class ExecutionTimeTracker(metaclass=Singleton):
         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)
+            context = self.context_map.pop()
 
             if not context:
                 return
@@ -126,11 +171,7 @@ class ExecutionTimeTracker(metaclass=Singleton):
             )
 
             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
+                self.state.add(context, elapsed)
 
 
 def calculate_execution_time(context: Optional[str] = None, store: bool = True):

metadata/utils/fqn.py

@@ -192,17 +192,34 @@ def _(
 
 @fqn_build_registry.add(DatabaseSchema)
 def _(
-    _: Optional[OpenMetadata],  # ES Search not enabled for Schemas
+    metadata: Optional[OpenMetadata],  # ES Search not enabled for Schemas
     *,
     service_name: str,
-    database_name: str,
+    database_name: Optional[str],
     schema_name: str,
-) -> str:
-    if not service_name or not database_name or not schema_name:
-        raise FQNBuildingException(
-            f"Args should be informed, but got service=`{service_name}`, db=`{database_name}`, schema=`{schema_name}`"
+    skip_es_search: bool = True,
+    fetch_multiple_entities: bool = False,
+) -> Union[Optional[str], Optional[List[str]]]:
+    entity: Optional[Union[DatabaseSchema, List[DatabaseSchema]]] = None
+
+    if not skip_es_search:
+        entity = search_database_schema_from_es(
+            metadata=metadata,
+            database_name=database_name,
+            schema_name=schema_name,
+            fetch_multiple_entities=fetch_multiple_entities,
+            service_name=service_name,
         )
-    return _build(service_name, database_name, schema_name)
+
+    if not entity and database_name:
+        fqn = _build(service_name, database_name, schema_name)
+        return [fqn] if fetch_multiple_entities else fqn
+    if entity and fetch_multiple_entities:
+        return [str(table.fullyQualifiedName.__root__) for table in entity]
+    if entity:
+        return str(entity.fullyQualifiedName.__root__)
+
+    return None
 
 
 @fqn_build_registry.add(Database)
@@ -574,6 +591,43 @@ def build_es_fqn_search_string(
     return fqn_search_string
 
 
+def search_database_schema_from_es(
+    metadata: OpenMetadata,
+    database_name: str,
+    schema_name: str,
+    service_name: str,
+    fetch_multiple_entities: bool = False,
+    fields: Optional[str] = None,
+):
+    """
+    Find database schema entity in elasticsearch index.
+
+    :param metadata: OM Client
+    :param database_name: name of database in which we are searching for database schema
+    :param schema_name: name of schema we are searching for
+    :param service_name: name of service in which we are searching for database schema
+    :param fetch_multiple_entities: should single match be returned or all matches
+    :param fields: additional fields to return
+    :return: entity / entities matching search criteria
+    """
+    if not schema_name:
+        raise FQNBuildingException(
+            f"Schema Name should be informed, but got schema_name=`{schema_name}`"
+        )
+
+    fqn_search_string = _build(service_name or "*", database_name or "*", schema_name)
+
+    es_result = metadata.es_search_from_fqn(
+        entity_type=DatabaseSchema,
+        fqn_search_string=fqn_search_string,
+        fields=fields,
+    )
+
+    return get_entity_from_es_result(
+        entity_list=es_result, fetch_multiple_entities=fetch_multiple_entities
+    )
+
+
 def search_table_from_es(
     metadata: OpenMetadata,
     database_name: str,

metadata/utils/helpers.py

@@ -346,6 +346,8 @@ 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)))
+    if magnitude >= len(units):
+        return f"{int(number / constant_k**magnitude)}e{magnitude*3}"
     return f"{number / constant_k**magnitude:.3f}{units[magnitude]}"
 
 

metadata/utils/importer.py

@@ -196,13 +196,22 @@ def import_connection_fn(connection: BaseModel, function_name: str) -> Callable:
 
     # module building strings read better with .format instead of f-strings
     # pylint: disable=consider-using-f-string
-    _connection_fn = import_from_module(
-        "metadata.ingestion.source.{}.{}.connection.{}".format(
-            service_type.name.lower(),
-            connection_type.value.lower(),
-            function_name,
+
+    if connection.type.value.lower().startswith("custom"):
+        python_class_parts = connection.sourcePythonClass.rsplit(".", 1)
+        python_module_path = ".".join(python_class_parts[:-1])
+
+        _connection_fn = import_from_module(
+            "{}.{}".format(python_module_path, function_name)
+        )
+    else:
+        _connection_fn = import_from_module(
+            "metadata.ingestion.source.{}.{}.connection.{}".format(
+                service_type.name.lower(),
+                connection_type.value.lower(),
+                function_name,
+            )
         )
-    )
 
     return _connection_fn
 

metadata/utils/logger.py

@@ -242,7 +242,7 @@ def _(record: OMetaLifeCycleData) -> str:
     """
     Capture the lifecycle changes of an Entity
     """
-    return f"{type(record.entity).__name__} Lifecycle [{record.entity.name.__root__}]"
+    return f"{record.entity.__name__} Lifecycle [{record.entity_fqn}]"
 
 
 @get_log_name.register

metadata/utils/lru_cache.py

@@ -15,6 +15,8 @@ LRU cache
 
 from collections import OrderedDict
 
+LRU_CACHE_SIZE = 4096
+
 
 class LRUCache:
     """Least Recently Used cache"""

metadata/utils/messaging_utils.py

@@ -0,0 +1,43 @@
+#  Copyright 2024 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.
+
+"""
+Helpers module for messaging sources
+"""
+
+import re
+import traceback
+from typing import Optional
+
+from metadata.utils.logger import utils_logger
+
+logger = utils_logger()
+
+
+def merge_and_clean_protobuf_schema(schema_text: Optional[str]) -> Optional[str]:
+    """
+    Remove the import and extra syntax lines for a schema with references
+    """
+    try:
+        lines = schema_text.splitlines() if schema_text else []
+        new_lines = []
+        for i, line in enumerate(lines):
+            if not re.search(r'import ".*";', line) and not re.search(
+                r"option .*;", line
+            ):
+                if re.search(r'\s*syntax\s*=\s*"proto\d+";\s*', line) and i != 0:
+                    continue
+                new_lines.append(line)
+        return "\n".join(new_lines)
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(f"Failed to merge and clean protobuf schema: {exc}")
+    return None