openmetadata-ingestion 1.5.13.1__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/utils/dispatch.py

@@ -14,9 +14,7 @@ Helper that implements custom dispatcher logic
 """
 
 from collections import namedtuple
-from functools import update_wrapper
-from types import MappingProxyType
-from typing import Any, Callable, Type, TypeVar
+from typing import Type, TypeVar
 
 from pydantic import BaseModel
 
@@ -56,53 +54,3 @@ def class_register():
 
     Register = namedtuple("Register", ["add", "registry"])
     return Register(add, registry)
-
-
-def valuedispatch(func) -> Callable:
-    """Value dispatch for methods and functions
-
-    Args:
-        func (_type_): function to run
-
-    Returns:
-        Callable: wrapper
-    """
-
-    registry = {}
-
-    def _is_valid_dispatch(value):
-        return isinstance(value, str)
-
-    def dispatch(value: str) -> Callable:
-        try:
-            impl = registry[value]
-        except KeyError:
-            impl = registry[object]
-        return impl
-
-    def register(value, func=None) -> Callable:
-        if _is_valid_dispatch(value):
-            if func is None:
-                return lambda f: register(value, f)
-        else:
-            raise TypeError(
-                "Invalid first argument to reigister()." f"{value} is not a string."
-            )
-
-        registry[value] = func
-        return func
-
-    def wrapper(*args, **kwargs) -> Any:
-        if not args:
-            raise TypeError(f"{func_name} requires at least 1 argument")
-        if isinstance(args[0], (str, bytes)):
-            return dispatch(str(args[0]))(*args, **kwargs)
-        return dispatch(args[1])(*args, **kwargs)
-
-    func_name = getattr(func, "__name__", "method value dispatch")
-    registry[object] = func
-    wrapper.register = register
-    wrapper.dispatch = dispatch
-    wrapper.registry = MappingProxyType(registry)  # making registry read only
-    update_wrapper(wrapper, func)
-    return wrapper

metadata/utils/entity_link.py

@@ -19,19 +19,23 @@ from antlr4.CommonTokenStream import CommonTokenStream
 from antlr4.error.ErrorStrategy import BailErrorStrategy
 from antlr4.InputStream import InputStream
 from antlr4.tree.Tree import ParseTreeWalker
-from pydantic import BaseModel
 from requests.compat import unquote_plus
 
 from metadata.antlr.split_listener import EntityLinkSplitListener
 from metadata.generated.antlr.EntityLinkLexer import EntityLinkLexer
 from metadata.generated.antlr.EntityLinkParser import EntityLinkParser
 from metadata.generated.schema.entity.data.table import Table
+from metadata.ingestion.models.custom_pydantic import BaseModel
 from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
 from metadata.utils.dispatch import class_register
 
 T = TypeVar("T", bound=BaseModel)
 
 
+class CustomColumnName(BaseModel):
+    root: str
+
+
 class EntityLinkBuildingException(Exception):
     """
     Raise for inconsistencies when building the EntityLink
@@ -69,7 +73,9 @@ def get_decoded_column(entity_link: str) -> str:
         entity_link: entity link
     """
 
-    return unquote_plus(entity_link.split("::")[-1].replace(">", ""))
+    return CustomColumnName(
+        root=unquote_plus(entity_link.split("::")[-1].replace(">", ""))
+    ).root
 
 
 def get_table_fqn(entity_link: str) -> str:

metadata/utils/entity_utils.py

@@ -0,0 +1,93 @@
+#  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.
+"""Entity Utilities"""
+from enum import Enum
+from typing import Type
+
+from metadata.generated.schema.entity.services.apiService import (
+    ApiService,
+    ApiServiceType,
+)
+from metadata.generated.schema.entity.services.dashboardService import (
+    DashboardService,
+    DashboardServiceType,
+)
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseService,
+    DatabaseServiceType,
+)
+from metadata.generated.schema.entity.services.messagingService import (
+    MessagingService,
+    MessagingServiceType,
+)
+from metadata.generated.schema.entity.services.metadataService import (
+    MetadataService,
+    MetadataServiceType,
+)
+from metadata.generated.schema.entity.services.mlmodelService import (
+    MlModelService,
+    MlModelServiceType,
+)
+from metadata.generated.schema.entity.services.pipelineService import (
+    PipelineService,
+    PipelineServiceType,
+)
+from metadata.generated.schema.entity.services.searchService import (
+    SearchService,
+    SearchServiceType,
+)
+from metadata.generated.schema.entity.services.storageService import (
+    StorageService,
+    StorageServiceType,
+)
+
+
+class ServiceClass(Enum):
+    """Kinds of supported services"""
+
+    DATABASE = DatabaseService
+    API = ApiService
+    DASHBOARD = DashboardService
+    PIPELINE = PipelineService
+    MESSAGING = MessagingService
+    METADATA = MetadataService
+    ML_MODEL = MlModelService
+    SEARCH = SearchService
+    STORAGE = StorageService
+
+
+SERVICE_TYPE_MAP = {
+    DatabaseService: DatabaseServiceType,
+    ApiService: ApiServiceType,
+    DashboardService: DashboardServiceType,
+    PipelineService: PipelineServiceType,
+    MessagingService: MessagingServiceType,
+    MetadataService: MetadataServiceType,
+    MlModelService: MlModelServiceType,
+    SearchService: SearchServiceType,
+    StorageService: StorageServiceType,
+}
+
+
+def service_class(service_type) -> Type:
+    """Get the service class based on the service type
+
+    Args:
+        service_type (str): Service type
+    Returns:
+        str
+    """
+    for service in ServiceClass:
+        if service_type.casefold() in {
+            key.casefold() for key in SERVICE_TYPE_MAP[service.value].__members__
+        }:
+            return service.value
+    raise ValueError(f"Unsupported service type: {service_type}")

metadata/utils/filters.py

@@ -280,3 +280,18 @@ def filter_by_search_index(
     :return: True for filtering, False otherwise
     """
     return _filter(search_index_filter_pattern, search_index_name)
+
+
+def filter_by_classification(
+    classification_pattern: Optional[FilterPattern], classification_name: str
+) -> bool:
+    """
+    Return True if the models needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param search_index_filter_pattern: Model defining search index filtering logic
+    :param search_index_name: search index name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(classification_pattern, classification_name)

metadata/utils/fqn.py

@@ -27,6 +27,7 @@ from metadata.antlr.split_listener import FqnSplitListener
 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.apiCollection import APICollection
 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
@@ -250,6 +251,20 @@ def _(
     return _build(service_name, dashboard_name)
 
 
+@fqn_build_registry.add(APICollection)
+def _(
+    _: Optional[OpenMetadata],  # ES Index not necessary for dashboard FQN building
+    *,
+    service_name: str,
+    api_collection_name: str,
+) -> str:
+    if not service_name or not api_collection_name:
+        raise FQNBuildingException(
+            f"Args should be informed, but got service=`{service_name}`, collection=`{api_collection_name}``"
+        )
+    return _build(service_name, api_collection_name)
+
+
 @fqn_build_registry.add(Chart)
 def _(
     _: Optional[OpenMetadata],  # ES Index not necessary for dashboard FQN building

metadata/utils/helpers.py

@@ -87,9 +87,11 @@ om_chart_type_dict = {
     "dual_line": ChartType.Line,
     "line_multi": ChartType.Line,
     "table": ChartType.Table,
+    "levelTable": ChartType.Table,
     "dist_bar": ChartType.Bar,
     "bar": ChartType.Bar,
     "box_plot": ChartType.BoxPlot,
+    "box": ChartType.BoxPlot,
     "boxplot": ChartType.BoxPlot,
     "histogram": ChartType.Histogram,
     "treemap": ChartType.Area,
@@ -358,8 +360,7 @@ def clean_uri(uri: Union[str, Url]) -> str:
     make it http://localhost:9000
     """
     # force a string of the given Uri if needed
-    if isinstance(uri, Url):
-        uri = str(uri)
+    uri = str(uri)
     return uri[:-1] if uri.endswith("/") else uri
 
 

metadata/utils/importer.py

@@ -25,7 +25,7 @@ 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.steps import BulkSink, Processor, Sink, Source, Stage
+from metadata.ingestion.api.steps import BulkSink, Processor, Sink, 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
@@ -126,37 +126,21 @@ def import_from_module(key: str) -> Type[Any]:
     """
     Dynamically import an object from a module path
     """
-
+    logger.debug("Importing: %s", key)
     module_name, obj_name = key.rsplit(MODULE_SEPARATOR, 1)
     try:
         obj = getattr(importlib.import_module(module_name), obj_name)
         return obj
-    except Exception as err:
+    except (ModuleNotFoundError, ImportError) as err:
         logger.debug(traceback.format_exc())
         raise DynamicImportException(module=module_name, key=obj_name, cause=err)
 
 
-# module building strings read better with .format instead of f-strings
-# pylint: disable=consider-using-f-string
-def import_source_class(
-    service_type: ServiceType, source_type: str, from_: str = "ingestion"
-) -> Type[Source]:
-    return import_from_module(
-        "metadata.{}.source.{}.{}.{}.{}Source".format(
-            from_,
-            service_type.name.lower(),
-            get_module_dir(source_type),
-            get_source_module_name(source_type),
-            get_class_name_root(source_type),
-        )
-    )
-
-
 def import_processor_class(
     processor_type: str, from_: str = "ingestion"
 ) -> Type[Processor]:
     return import_from_module(
-        "metadata.{}.processor.{}.{}Processor".format(
+        "metadata.{}.processor.{}.{}Processor".format(  # pylint: disable=consider-using-f-string
             from_,
             get_module_name(processor_type),
             get_class_name_root(processor_type),
@@ -166,7 +150,7 @@ def import_processor_class(
 
 def import_stage_class(stage_type: str, from_: str = "ingestion") -> Type[Stage]:
     return import_from_module(
-        "metadata.{}.stage.{}.{}Stage".format(
+        "metadata.{}.stage.{}.{}Stage".format(  # pylint: disable=consider-using-f-string
             from_,
             get_module_name(stage_type),
             get_class_name_root(stage_type),
@@ -176,7 +160,7 @@ def import_stage_class(stage_type: str, from_: str = "ingestion") -> Type[Stage]
 
 def import_sink_class(sink_type: str, from_: str = "ingestion") -> Type[Sink]:
     return import_from_module(
-        "metadata.{}.sink.{}.{}Sink".format(
+        "metadata.{}.sink.{}.{}Sink".format(  # pylint: disable=consider-using-f-string
             from_,
             get_module_name(sink_type),
             get_class_name_root(sink_type),
@@ -188,7 +172,7 @@ def import_bulk_sink_type(
     bulk_sink_type: str, from_: str = "ingestion"
 ) -> Type[BulkSink]:
     return import_from_module(
-        "metadata.{}.bulksink.{}.{}BulkSink".format(
+        "metadata.{}.bulksink.{}.{}BulkSink".format(  # pylint: disable=consider-using-f-string
             from_,
             get_module_name(bulk_sink_type),
             get_class_name_root(bulk_sink_type),
@@ -273,7 +257,7 @@ def import_test_case_class(
         test_definition[0].upper() + test_definition[1:]
     )  # change test names to camel case
     return import_from_module(
-        "metadata.data_quality.validations.{}.{}.{}.{}Validator".format(
+        "metadata.data_quality.validations.{}.{}.{}.{}Validator".format(  # pylint: disable=consider-using-f-string
             test_type.lower(),
             runner_type,
             test_definition,
@@ -302,3 +286,7 @@ class SideEffectsLoader(metaclass=Singleton):
 
 def import_side_effects(*modules):
     SideEffectsLoader().import_side_effects(*modules)
+
+
+def get_class_path(module):
+    return module.__module__ + "." + module.__name__

metadata/utils/logger.py

@@ -13,10 +13,11 @@ Module centralising logger configs
 """
 
 import logging
+from copy import deepcopy
 from enum import Enum
 from functools import singledispatch
 from types import DynamicClassAttribute
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 
 from metadata.data_quality.api.models import (
     TableAndTests,
@@ -40,6 +41,8 @@ BASE_LOGGING_FORMAT = (
 )
 logging.basicConfig(format=BASE_LOGGING_FORMAT, datefmt="%Y-%m-%d %H:%M:%S")
 
+REDACTED_KEYS = {"serviceConnection", "securityConfig"}
+
 
 class Loggers(Enum):
     """
@@ -49,6 +52,7 @@ class Loggers(Enum):
     OMETA = "OMetaAPI"
     CLI = "Metadata"
     PROFILER = "Profiler"
+    SAMPLER = "Sampler"
     PII = "PII"
     INGESTION = "Ingestion"
     UTILS = "Utils"
@@ -100,6 +104,14 @@ def profiler_logger():
     return logging.getLogger(Loggers.PROFILER.value)
 
 
+def sampler_logger():
+    """
+    Method to get the SAMPLER logger
+    """
+
+    return logging.getLogger(Loggers.SAMPLER.value)
+
+
 def pii_logger():
     """
     Method to get the PROFILER logger
@@ -285,3 +297,21 @@ def _(record: QueryParserData) -> str:
         mask_query(query.sql, query.dialect) for query in record.parsedData
     )
     return f"Usage ParsedData [{queries}]"
+
+
+def redacted_config(config: Dict[str, Union[str, dict]]) -> Dict[str, Union[str, dict]]:
+    config_copy = deepcopy(config)
+
+    def traverse_and_modify(obj):
+        if isinstance(obj, dict):
+            for key, value in obj.items():
+                if key in REDACTED_KEYS:
+                    obj[key] = "REDACTED"
+                else:
+                    traverse_and_modify(value)
+        elif isinstance(obj, list):
+            for item in obj:
+                traverse_and_modify(item)
+
+    traverse_and_modify(config_copy)
+    return config_copy

metadata/utils/lru_cache.py

@@ -13,47 +13,108 @@
 LRU cache
 """
 
+import threading
 from collections import OrderedDict
+from typing import Callable, Generic, TypeVar
 
 LRU_CACHE_SIZE = 4096
 
+T = TypeVar("T")
 
-class LRUCache:
+
+class LRUCache(Generic[T]):
     """Least Recently Used cache"""
 
     def __init__(self, capacity: int) -> None:
         self._cache = OrderedDict()
         self.capacity = capacity
+        self.lock = threading.Lock()
 
     def clear(self):
-        self._cache = OrderedDict()
+        with self.lock:
+            self._cache = OrderedDict()
 
-    def get(self, key):
+    def get(self, key) -> T:
         """
         Returns the value associated to `key` if it exists,
         updating the cache usage.
         Raises `KeyError` if `key doesn't exist in the cache.
+
+        Args:
+            key: The key to get the value for
+
+        Returns:
+            The value associated to `key`
         """
-        self._cache.move_to_end(key)
-        return self._cache[key]
+        with self.lock:
+            self._cache.move_to_end(key)
+            return self._cache[key]
 
-    def put(self, key, value) -> None:
+    def put(self, key: str, value: T) -> None:
         """
         Assigns `value` to `key`, overwriting `key` if it already exists
         in the cache and updating the cache usage.
         If the size of the cache grows above capacity, pops the least used
         element.
+
+        Args:
+            key: The key to assign the value to
+            value: The value to assign to the key
         """
-        self._cache[key] = value
-        self._cache.move_to_end(key)
-        if len(self._cache) > self.capacity:
-            self._cache.popitem(last=False)
+        with self.lock:
+            self._cache[key] = value
+            self._cache.move_to_end(key)
+            if len(self._cache) > self.capacity:
+                self._cache.popitem(last=False)
 
     def __contains__(self, key) -> bool:
-        if key not in self._cache:
-            return False
-        self._cache.move_to_end(key)
-        return True
+        with self.lock:
+            if key not in self._cache:
+                return False
+            self._cache.move_to_end(key)
+            return True
 
     def __len__(self) -> int:
-        return len(self._cache)
+        with self.lock:
+            return len(self._cache)
+
+    def wrap(self, key_func: Callable[..., str]):
+        """Decorator to cache the result of a function based on its arguments.
+
+        Example:
+        ```python
+        import time
+        from metadata.utils.lru_cache import LRUCache
+        cache = LRUCache(4096)
+
+        @cache.wrap(lambda x, y: f"{x}-{y}")
+        def add(x, y):
+            time.sleep(1)
+            return x + y
+        start1 = time.time()
+        add(1, 2)  # This will be cached and take 1 second
+        print('took', time.time() - start1, 'seconds')
+        start2 = time.time()
+        add(1, 2)  # This will return the cached value and take no time
+        print('took', time.time() - start2, 'seconds')
+        ```
+        Args:
+            key_func: A function that generates a key based on the arguments
+                of the decorated function.
+
+        Returns:
+            A decorator that caches the result of the decorated function.
+        """
+
+        def wrapper(func: Callable[..., T]):
+            def wrapped(*args, **kwargs) -> T:
+                key = key_func(*args, **kwargs)
+                if key in self:
+                    return self.get(key)
+                value = func(*args, **kwargs)
+                self.put(key, value)
+                return value
+
+            return wrapped
+
+        return wrapper

metadata/utils/profiler_utils.py

@@ -20,6 +20,11 @@ from typing import Optional, Tuple
 import sqlparse
 from pydantic import BaseModel
 
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -34,20 +39,12 @@ class QueryResult(BaseModel):
     schema_name: str
     table_name: str
     query_type: str
-    timestamp: datetime
+    start_time: datetime
     query_id: Optional[str] = None
     query_text: Optional[str] = None
     rows: Optional[int] = None
 
 
-class SnowflakeQueryResult(QueryResult):
-    """Snowflake system metric query result"""
-
-    rows_inserted: Optional[int] = None
-    rows_updated: Optional[int] = None
-    rows_deleted: Optional[int] = None
-
-
 def clean_up_query(query: str) -> str:
     """remove comments and newlines from query"""
     return sqlparse.format(query, strip_comments=True).replace("\\n", "")
@@ -108,3 +105,40 @@ def set_cache(cache: defaultdict, key: str, value):
             cache[key_] = value
             break
         cache = cache[key_]
+
+
+def get_context_entities(
+    entity: Table, metadata: OpenMetadata
+) -> Tuple[DatabaseSchema, Database, DatabaseService]:
+    """Based on the table, get all the parent entities"""
+    schema_entity = None
+    database_entity = None
+    db_service = None
+
+    if entity.databaseSchema:
+        schema_entity_list = metadata.es_search_from_fqn(
+            entity_type=DatabaseSchema,
+            fqn_search_string=entity.databaseSchema.fullyQualifiedName,
+            fields="databaseSchemaProfilerConfig",
+        )
+        if schema_entity_list:
+            schema_entity = schema_entity_list[0]
+
+    if entity.database:
+        database_entity_list = metadata.es_search_from_fqn(
+            entity_type=Database,
+            fqn_search_string=entity.database.fullyQualifiedName,
+            fields="databaseProfilerConfig",
+        )
+        if database_entity_list:
+            database_entity = database_entity_list[0]
+
+    if entity.service:
+        db_service_list = metadata.es_search_from_fqn(
+            entity_type=DatabaseService,
+            fqn_search_string=entity.service.fullyQualifiedName,
+        )
+        if db_service_list:
+            db_service = db_service_list[0]
+
+    return schema_entity, database_entity, db_service

metadata/utils/service_spec/__init__.py

@@ -0,0 +1,5 @@
+"""Module for the OpenMetadat Ingestion Service Specification (ServiceSpec)"""
+
+from metadata.utils.service_spec.service_spec import BaseSpec
+
+__all__ = ["BaseSpec"]

metadata/utils/service_spec/default.py

@@ -0,0 +1,21 @@
+"""
+Default service specs for services.
+"""
+
+from typing import Optional
+
+from metadata.data_quality.interface.sqlalchemy.sqa_test_suite_interface import (
+    SQATestSuiteInterface,
+)
+from metadata.profiler.interface.sqlalchemy.profiler_interface import (
+    SQAProfilerInterface,
+)
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+from metadata.utils.importer import get_class_path
+from metadata.utils.service_spec.service_spec import BaseSpec
+
+
+class DefaultDatabaseSpec(BaseSpec):
+    profiler_class: Optional[str] = get_class_path(SQAProfilerInterface)
+    sampler_class: Optional[str] = get_class_path(SQASampler)
+    test_suite_class: Optional[str] = get_class_path(SQATestSuiteInterface)