openmetadata-ingestion 1.5.3.1__py3-none-any.whl → 1.5.4.0__py3-none-any.whl

metadata/generated/schema/type/entityUsage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityUsage.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/filterPattern.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/filterPattern.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/function.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/function.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/include.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/include.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/jdbcConnection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/jdbcConnection.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/lifeCycle.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/lifeCycle.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/queryParserData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/queryParserData.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2024-09-11T16:38:06+00:00
+#   timestamp: 2024-09-13T12:48:15+00:00
 
 from __future__ import annotations
 

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -16,12 +16,15 @@ To be used by OpenMetadata class
 import functools
 import json
 import traceback
-from typing import Generic, Iterable, List, Optional, Set, Type, TypeVar
+from typing import Generic, Iterable, Iterator, List, Optional, Set, Type, TypeVar
+from urllib.parse import quote_plus
 
-from pydantic import BaseModel
+from pydantic import Field
+from typing_extensions import Annotated
 
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.query import Query
+from metadata.ingestion.models.custom_pydantic import BaseModel
 from metadata.ingestion.ometa.client import REST, APIError
 from metadata.ingestion.ometa.utils import quote
 from metadata.utils.elasticsearch import ES_INDEX_MAP
@@ -32,6 +35,45 @@ logger = ometa_logger()
 T = TypeVar("T", bound=BaseModel)
 
 
+class TotalModel(BaseModel):
+    """Elasticsearch total model"""
+
+    relation: str
+    value: int
+
+
+class HitsModel(BaseModel):
+    """Elasticsearch hits model"""
+
+    index: Annotated[str, Field(description="Index name", alias="_index")]
+    type: Annotated[
+        Optional[str],
+        Field(default=None, description="Type of the document", alias="_type"),
+    ]
+    id: Annotated[str, Field(description="Document ID", alias="_id")]
+    score: Annotated[
+        Optional[float], Field(description="Score of the document", alias="_score")
+    ]
+    source: Annotated[dict, Field(description="Document source", alias="_source")]
+    sort: Annotated[
+        List[str],
+        Field(description="Sort field. Used internally to get the next page FQN"),
+    ]
+
+
+class ESHits(BaseModel):
+    """Elasticsearch hits model"""
+
+    total: Annotated[TotalModel, Field(description="Total matched elements")]
+    hits: Annotated[List[HitsModel], Field(description="List of matched elements")]
+
+
+class ESResponse(BaseModel):
+    """Elasticsearch response model"""
+
+    hits: ESHits
+
+
 class ESMixin(Generic[T]):
     """
     OpenMetadata API methods related to Elasticsearch.
@@ -46,6 +88,12 @@ class ESMixin(Generic[T]):
         "&size={size}&index={index}"
     )
 
+    # sort_field needs to be unique for the pagination to work, so we can use the FQN
+    paginate_query = (
+        "/search/query?q=&size={size}&deleted=false{filter}&index={index}"
+        "&sort_field=fullyQualifiedName{after}"
+    )
+
     @functools.lru_cache(maxsize=512)
     def _search_es_entity(
         self,
@@ -252,3 +300,65 @@ class ESMixin(Generic[T]):
             logger.debug(traceback.format_exc())
             logger.warning(f"Unknown error extracting results from ES query [{err}]")
             return None
+
+    def paginate_es(
+        self,
+        entity: Type[T],
+        query_filter: Optional[str] = None,
+        size: int = 100,
+        fields: Optional[List[str]] = None,
+    ) -> Iterator[T]:
+        """Paginate through the ES results, ignoring individual errors"""
+        after: Optional[str] = None
+        error_pages = 0
+        query = functools.partial(
+            self.paginate_query.format,
+            index=ES_INDEX_MAP[entity.__name__],
+            filter="&query_filter=" + quote_plus(query_filter) if query_filter else "",
+            size=size,
+        )
+        while True:
+            query_string = query(
+                after="&search_after=" + quote_plus(after) if after else ""
+            )
+            response = self._get_es_response(query_string)
+
+            # Allow 3 errors getting pages before getting out of the loop
+            if not response:
+                error_pages += 1
+                if error_pages < 3:
+                    continue
+                else:
+                    break
+
+            # Get the data
+            for hit in response.hits.hits:
+                try:
+                    yield self.get_by_name(
+                        entity=entity,
+                        fqn=hit.source["fullyQualifiedName"],
+                        fields=fields,
+                        nullable=False,  # Raise an error if we don't find the Entity
+                    )
+                except Exception as exc:
+                    logger.warning(
+                        f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
+                    )
+
+            # Get next page
+            last_hit = response.hits.hits[-1] if response.hits.hits else None
+            if not last_hit or not last_hit.sort:
+                logger.info("No more pages to fetch")
+                break
+
+            after = ",".join(last_hit.sort)
+
+    def _get_es_response(self, query_string: str) -> Optional[ESResponse]:
+        """Get the Elasticsearch response"""
+        try:
+            response = self.client.get(query_string)
+            return ESResponse.model_validate(response)
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error while getting ES response: {exc}")
+        return None

metadata/utils/constants.py

@@ -106,3 +106,5 @@ ENTITY_REFERENCE_CLASS_MAP = {
 ENTITY_REFERENCE_TYPE_MAP = {
     value.__name__: key for key, value in ENTITY_REFERENCE_CLASS_MAP.items()
 }
+
+CUSTOM_CONNECTOR_PREFIX = "custom"

metadata/utils/importer.py

@@ -27,6 +27,8 @@ from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.metadataIngestion.workflow import Sink as WorkflowSink
 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.client_version import get_client_version
+from metadata.utils.constants import CUSTOM_CONNECTOR_PREFIX
 from metadata.utils.logger import utils_logger
 from metadata.utils.singleton import Singleton
 
@@ -43,6 +45,38 @@ class DynamicImportException(Exception):
     Raise it when having issues dynamically importing objects
     """
 
+    def __init__(self, module: str, key: str = None, cause: Exception = None):
+        self.module = module
+        self.key = key
+        self.cause = cause
+
+    def __str__(self):
+        import_path = self.module
+        if self.key:
+            import_path += f".{self.key}"
+        return f"Cannot import {import_path} due to {self.cause}"
+
+
+class MissingPluginException(Exception):
+    """
+    An excpetion that captures a missing openmetadata-ingestion plugin for a specific connector.
+    """
+
+    def __init__(self, plugin: str):
+        self.plugin = plugin
+
+    def __str__(self):
+        try:
+            version = "==" + get_client_version()
+        except Exception:
+            logger.warning("unable to get client version")
+            logger.debug(traceback.format_exc())
+            version = ""
+        return (
+            f"You might be missing the plugin [{self.plugin}]. Try:\n"
+            f'pip install "openmetadata-ingestion[{self.plugin}]{version}"'
+        )
+
 
 def get_module_dir(type_: str) -> str:
     """
@@ -93,13 +127,13 @@ def import_from_module(key: str) -> Type[Any]:
     Dynamically import an object from a module path
     """
 
+    module_name, obj_name = key.rsplit(MODULE_SEPARATOR, 1)
     try:
-        module_name, obj_name = key.rsplit(MODULE_SEPARATOR, 1)
         obj = getattr(importlib.import_module(module_name), obj_name)
         return obj
     except Exception as err:
         logger.debug(traceback.format_exc())
-        raise DynamicImportException(f"Cannot load object from {key} due to {err}")
+        raise DynamicImportException(module=module_name, key=obj_name, cause=err)
 
 
 # module building strings read better with .format instead of f-strings
@@ -200,7 +234,7 @@ 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
 
-    if connection.type.value.lower().startswith("custom"):
+    if connection.type.value.lower().startswith(CUSTOM_CONNECTOR_PREFIX):
         python_class_parts = connection.sourcePythonClass.rsplit(".", 1)
         python_module_path = ".".join(python_class_parts[:-1])
 
@@ -261,9 +295,7 @@ class SideEffectsLoader(metaclass=Singleton):
                     SideEffectsLoader.modules.add(module.__name__)
                 except Exception as err:
                     logger.debug(traceback.format_exc())
-                    raise DynamicImportException(
-                        f"Cannot load object from {module} due to {err}"
-                    )
+                    raise DynamicImportException(module=module, cause=err)
             else:
                 logger.debug(f"Module {module} already imported")
 

metadata/utils/logger.py

@@ -170,10 +170,14 @@ def set_loggers_level(level: Union[int, str] = logging.INFO):
 
 
 def log_ansi_encoded_string(
-    color: Optional[ANSI] = None, bold: bool = False, message: str = ""
+    color: Optional[ANSI] = None,
+    bold: bool = False,
+    message: str = "",
+    level=logging.INFO,
 ):
-    utils_logger().info(
-        f"{ANSI.BOLD.value if bold else ''}{color.value if color else ''}{message}{ANSI.ENDC.value}"
+    utils_logger().log(
+        level=level,
+        msg=f"{ANSI.BOLD.value if bold else ''}{color.value if color else ''}{message}{ANSI.ENDC.value}",
     )
 
 

metadata/workflow/ingestion.py

@@ -21,7 +21,7 @@ To be extended by any other workflow:
 """
 import traceback
 from abc import ABC, abstractmethod
-from typing import List, Tuple, cast
+from typing import List, Tuple, Type, cast
 
 from metadata.config.common import WorkflowExecutionError
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
@@ -46,6 +46,13 @@ from metadata.utils.class_helper import (
     get_service_class_from_service_type,
     get_service_type_from_source_type,
 )
+from metadata.utils.constants import CUSTOM_CONNECTOR_PREFIX
+from metadata.utils.importer import (
+    DynamicImportException,
+    MissingPluginException,
+    import_from_module,
+    import_source_class,
+)
 from metadata.utils.logger import ingestion_logger
 from metadata.workflow.base import BaseWorkflow, InvalidWorkflowJSONException
 from metadata.workflow.workflow_status_mixin import SUCCESS_THRESHOLD_VALUE
@@ -218,3 +225,22 @@ class IngestionWorkflow(BaseWorkflow, ABC):
             raise WorkflowExecutionError(
                 f"Profiler is not supported for the service connection: {self.config.source.serviceConnection}"
             )
+
+    def import_source_class(self) -> Type[Source]:
+        source_type = self.config.source.type.lower()
+        try:
+            return (
+                import_from_module(
+                    self.config.source.serviceConnection.root.config.sourcePythonClass
+                )
+                if source_type.startswith(CUSTOM_CONNECTOR_PREFIX)
+                else import_source_class(
+                    service_type=self.service_type, source_type=source_type
+                )
+            )
+        except DynamicImportException as e:
+            if source_type.startswith(CUSTOM_CONNECTOR_PREFIX):
+                raise e
+            logger.debug(traceback.format_exc())
+            logger.error(f"Failed to import source of type '{source_type}'")
+            raise MissingPluginException(source_type)

metadata/workflow/metadata.py

@@ -14,11 +14,7 @@ Workflow definition for metadata related ingestions: metadata and lineage.
 
 from metadata.config.common import WorkflowExecutionError
 from metadata.ingestion.api.steps import Sink, Source
-from metadata.utils.importer import (
-    import_from_module,
-    import_sink_class,
-    import_source_class,
-)
+from metadata.utils.importer import import_sink_class
 from metadata.utils.logger import ingestion_logger
 from metadata.workflow.ingestion import IngestionWorkflow
 
@@ -47,15 +43,7 @@ class MetadataWorkflow(IngestionWorkflow):
                 "configuration here: https://docs.open-metadata.org/connectors"
             )
 
-        source_class = (
-            import_from_module(
-                self.config.source.serviceConnection.root.config.sourcePythonClass
-            )
-            if source_type.startswith("custom")
-            else import_source_class(
-                service_type=self.service_type, source_type=source_type
-            )
-        )
+        source_class = self.import_source_class()
 
         pipeline_name = (
             self.ingestion_pipeline.fullyQualifiedName.root

metadata/workflow/profiler.py

@@ -44,6 +44,7 @@ class ProfilerWorkflow(IngestionWorkflow):
 
     def _get_source_class(self):
         if self.config.source.serviceName:
+            self.import_source_class()
             return OpenMetadataSource
         logger.info(
             "Database Service name not provided, we will scan all the tables "

metadata/workflow/usage.py

@@ -16,9 +16,7 @@ from metadata.config.common import WorkflowExecutionError
 from metadata.ingestion.api.steps import BulkSink, Processor, Source, Stage
 from metadata.utils.importer import (
     import_bulk_sink_type,
-    import_from_module,
     import_processor_class,
-    import_source_class,
     import_stage_class,
 )
 from metadata.utils.logger import ingestion_logger
@@ -51,16 +49,7 @@ class UsageWorkflow(IngestionWorkflow):
                 "configuration here: https://docs.open-metadata.org/connectors"
             )
 
-        source_class = (
-            import_from_module(
-                self.config.source.serviceConnection.root.config.sourcePythonClass
-            )
-            if source_type.startswith("custom")
-            else import_source_class(
-                service_type=self.service_type, source_type=source_type
-            )
-        )
-
+        source_class = self.import_source_class()
         source: Source = source_class.create(
             self.config.source.model_dump(), self.metadata
         )

metadata/workflow/workflow_init_error_handler.py

@@ -12,6 +12,7 @@
 """
 Module handles the init error messages from different workflows
 """
+import logging
 import traceback
 from pathlib import Path
 from typing import Any, Dict, Optional, Type, Union
@@ -25,7 +26,7 @@ from metadata.ingestion.api.parser import (
     ParsingConfigurationError,
 )
 from metadata.utils.constants import UTF_8
-from metadata.utils.logger import ANSI, log_ansi_encoded_string
+from metadata.utils.logger import ANSI, log_ansi_encoded_string, utils_logger
 
 EXAMPLES_WORKFLOW_PATH: Path = Path(__file__).parent / "../examples" / "workflows"
 
@@ -74,10 +75,10 @@ class WorkflowInitErrorHandler:
                 source_type_name, pipeline_type
             )
         else:
+            utils_logger().debug(traceback.format_exc())
             WorkflowInitErrorHandler._print_error_msg(
                 f"\nError initializing {pipeline_type.name}: {exc}"
             )
-            WorkflowInitErrorHandler._print_error_msg(traceback.format_exc())
 
         WorkflowInitErrorHandler._print_more_info(pipeline_type)
 
@@ -151,4 +152,15 @@ class WorkflowInitErrorHandler:
         """
         Print message with error style
         """
-        log_ansi_encoded_string(color=ANSI.BRIGHT_RED, bold=False, message=f"{msg}")
+        log_ansi_encoded_string(
+            color=ANSI.BRIGHT_RED, bold=False, message=f"{msg}", level=logging.ERROR
+        )
+
+    @staticmethod
+    def _print_debug_msg(msg: str) -> None:
+        """
+        Print message with error style
+        """
+        log_ansi_encoded_string(
+            color=ANSI.YELLOW, bold=False, message=f"{msg}", level=logging.DEBUG
+        )