openmetadata-ingestion 1.4.6.0rc1__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

metadata/ingestion/ometa/client.py

@@ -11,9 +11,9 @@
 """
 Python API REST wrapper and helpers
 """
-import datetime
 import time
 import traceback
+from datetime import datetime, timezone
 from typing import Callable, Dict, List, Optional, Union
 
 import requests
@@ -21,6 +21,7 @@ from requests.exceptions import HTTPError
 
 from metadata.config.common import ConfigModel
 from metadata.ingestion.ometa.credentials import URL, get_api_version
+from metadata.ingestion.ometa.ttl_cache import TTLCache
 from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import ometa_logger
 
@@ -33,6 +34,12 @@ class RetryException(Exception):
     """
 
 
+class LimitsException(Exception):
+    """
+    API Client Feature Limit exception
+    """
+
+
 class APIError(Exception):
     """
     Represent API related error.
@@ -97,7 +104,8 @@ class ClientConfig(ConfigModel):
     api_version: Optional[str] = "v1"
     retry: Optional[int] = 3
     retry_wait: Optional[int] = 30
-    retry_codes: List[int] = [429, 504]
+    limit_codes: List[int] = [429]
+    retry_codes: List[int] = [504]
     auth_token: Optional[Callable] = None
     access_token: Optional[str] = None
     expires_in: Optional[int] = None
@@ -107,8 +115,10 @@ class ClientConfig(ConfigModel):
     allow_redirects: Optional[bool] = False
     auth_token_mode: Optional[str] = "Bearer"
     verify: Optional[Union[bool, str]] = None
+    ttl_cache: int = 60
 
 
+# pylint: disable=too-many-instance-attributes
 class REST:
     """
     REST client wrapper to manage requests with
@@ -124,20 +134,27 @@ class REST:
         self._retry = self.config.retry
         self._retry_wait = self.config.retry_wait
         self._retry_codes = self.config.retry_codes
+        self._limit_codes = self.config.limit_codes
         self._auth_token = self.config.auth_token
         self._auth_token_mode = self.config.auth_token_mode
         self._verify = self.config.verify
 
-    def _request(
+        self._limits_reached = TTLCache(config.ttl_cache)
+
+    def _request(  # pylint: disable=too-many-arguments
         self,
         method,
         path,
         data=None,
+        json=None,
         base_url: URL = None,
         api_version: str = None,
         headers: dict = None,
     ):
         # pylint: disable=too-many-locals
+        if path in self._limits_reached:
+            raise LimitsException(f"Skipping request - limits reached for {path}")
+
         if not headers:
             headers = {"Content-type": "application/json"}
         base_url = base_url or self._base_url
@@ -145,16 +162,16 @@ class REST:
         url: URL = URL(base_url + "/" + version + path)
         if (
             self.config.expires_in
-            and datetime.datetime.utcnow().timestamp() >= self.config.expires_in
+            and datetime.now(timezone.utc).timestamp() >= self.config.expires_in
             or not self.config.access_token
         ):
             self.config.access_token, expiry = self._auth_token()
             if not self.config.access_token == "no_token":
-                if isinstance(expiry, datetime.datetime):
+                if isinstance(expiry, datetime):
                     self.config.expires_in = expiry.timestamp() - 120
                 else:
                     self.config.expires_in = (
-                        datetime.datetime.utcnow().timestamp() + expiry - 120
+                        datetime.now(timezone.utc).timestamp() + expiry - 120
                     )
 
         headers[self.config.auth_header] = (
@@ -185,12 +202,18 @@ class REST:
 
         method_key = "params" if method.upper() == "GET" else "data"
         opts[method_key] = data
+        if json:
+            opts["json"] = json
 
         total_retries = self._retry if self._retry > 0 else 0
         retry = total_retries
         while retry >= 0:
             try:
                 return self._one_request(method, url, opts, retry)
+            except LimitsException as exc:
+                logger.error(f"Feature limit exceeded for {url}")
+                self._limits_reached.add(path)
+                raise exc
             except RetryException:
                 retry_wait = self._retry_wait * (total_retries - retry + 1)
                 logger.warning(
@@ -214,6 +237,7 @@ class REST:
         Returns the body json in the 200 status.
         """
         retry_codes = self._retry_codes
+        limit_codes = self._limit_codes
         try:
             resp = self._session.request(method, url, **opts)
             resp.raise_for_status()
@@ -231,6 +255,8 @@ class REST:
             # retry if we hit Rate Limit
             if resp.status_code in retry_codes and retry > 0:
                 raise RetryException() from http_error
+            if resp.status_code in limit_codes:
+                raise LimitsException() from http_error
             if "code" in resp.text:
                 error = resp.json()
                 if "code" in error:
@@ -270,7 +296,7 @@ class REST:
         return self._request("GET", path, data)
 
     @calculate_execution_time(context="POST")
-    def post(self, path, data=None):
+    def post(self, path, data=None, json=None):
         """
         POST method
 
@@ -281,7 +307,7 @@ class REST:
         Returns:
             Response
         """
-        return self._request("POST", path, data)
+        return self._request("POST", path, data, json)
 
     @calculate_execution_time(context="PUT")
     def put(self, path, data=None):

metadata/ingestion/ometa/client_utils.py

@@ -18,8 +18,8 @@ from metadata.generated.schema.entity.data.chart import Chart
 from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
     OpenMetadataConnection,
 )
-from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.generated.schema.type.basic import FullyQualifiedEntityName
+from metadata.ingestion.ometa.ometa_api import C, OpenMetadata, T
 from metadata.utils import fqn
 from metadata.utils.logger import ometa_logger
 
@@ -28,7 +28,7 @@ logger = ometa_logger()
 
 def create_ometa_client(
     metadata_config: OpenMetadataConnection,
-) -> OpenMetadata:
+) -> OpenMetadata[T, C]:  # pyright: ignore[reportInvalidTypeVarUse]
     """Create an OpenMetadata client
 
     Args:
@@ -38,7 +38,7 @@ def create_ometa_client(
         OpenMetadata: an OM client
     """
     try:
-        metadata = OpenMetadata(metadata_config)
+        metadata = OpenMetadata[T, C](metadata_config)
         metadata.health_check()
         return metadata
     except Exception as exc:
@@ -49,7 +49,7 @@ def create_ometa_client(
 
 def get_chart_entities_from_id(
     chart_ids: List[str], metadata: OpenMetadata, service_name: str
-) -> List[EntityReference]:
+) -> List[FullyQualifiedEntityName]:
     """
     Method to get the chart entity using get_by_name api
     """
@@ -63,6 +63,5 @@ def get_chart_entities_from_id(
             ),
         )
         if chart:
-            entity = EntityReference(id=chart.id, type="chart")
-            entities.append(entity)
+            entities.append(chart.fullyQualifiedName)
     return entities

metadata/ingestion/ometa/mixins/custom_property_mixin.py

@@ -61,7 +61,7 @@ class OMetaCustomPropertyMixin:
 
         resp = self.client.put(
             f"/metadata/types/{entity_schema.get('id')}",
-            data=ometa_custom_property.createCustomPropertyRequest.json(),
+            data=ometa_custom_property.createCustomPropertyRequest.model_dump_json(),
         )
         return resp
 
@@ -79,9 +79,7 @@ class OMetaCustomPropertyMixin:
         Get the PropertyType for custom properties
         """
         custom_property_type = self.get_custom_property_type(data_type=data_type)
-        return PropertyType(
-            __root__=EntityReference(id=custom_property_type.id, type="type")
-        )
+        return PropertyType(EntityReference(id=custom_property_type.id, type="type"))
 
     def get_entity_custom_properties(self, entity_type: Type[T]) -> Optional[List]:
         """

metadata/ingestion/ometa/mixins/dashboard_mixin.py

@@ -41,7 +41,7 @@ class OMetaDashboardMixin:
         :param dashboard_usage_request: Usage data to add
         """
         resp = self.client.put(
-            f"/usage/dashboard/{dashboard.id.__root__}",
-            data=dashboard_usage_request.json(),
+            f"/usage/dashboard/{dashboard.id.root}",
+            data=dashboard_usage_request.model_dump_json(),
         )
         logger.debug("Published dashboard usage %s", resp)

metadata/ingestion/ometa/mixins/data_insight_mixin.py

@@ -44,7 +44,9 @@ class DataInsightMixin:
             record (ReportData): report data
         """
 
-        resp = self.client.post("/analytics/dataInsights/data", record.json())
+        resp = self.client.post(
+            "/analytics/dataInsights/data", record.model_dump_json()
+        )
 
         return resp
 
@@ -56,7 +58,7 @@ class DataInsightMixin:
             record (ReportData): report data
         """
 
-        resp = self.client.put(f"/kpi/{fqn}/kpiResult", record.json())
+        resp = self.client.put(f"/kpi/{fqn}/kpiResult", record.model_dump_json())
 
         return resp
 
@@ -66,7 +68,9 @@ class DataInsightMixin:
     ) -> List[WebAnalyticEventData]:
         """Get web analytic event"""
 
-        resp = self.client.put("/analytics/web/events/collect", event_data.json())
+        resp = self.client.put(
+            "/analytics/web/events/collect", event_data.model_dump_json()
+        )
 
         return resp
 
@@ -127,7 +131,7 @@ class DataInsightMixin:
             request_params,
         )
 
-        return DataInsightChartResult.parse_obj(resp)
+        return DataInsightChartResult.model_validate(resp)
 
     def get_kpi_result(self, fqn: str, start_ts, end_ts) -> list[KpiResult]:
         """Given FQN return KPI results
@@ -146,9 +150,9 @@ class DataInsightMixin:
         return [KpiResult(**data) for data in resp["data"]]
 
     def create_kpi(self, create: CreateKpiRequest) -> Kpi:
-        resp = self.client.post("/kpi", create.json())
+        resp = self.client.post("/kpi", create.model_dump_json())
 
-        return Kpi.parse_obj(resp)
+        return Kpi.model_validate(resp)
 
     def get_web_analytic_events(
         self, event_type: WebAnalyticEventType, start_ts: int, end_ts: int

metadata/ingestion/ometa/mixins/ingestion_pipeline_mixin.py

@@ -47,7 +47,7 @@ class OMetaIngestionPipelineMixin:
         """
         resp = self.client.put(
             f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
-            data=pipeline_status.json(),
+            data=pipeline_status.model_dump_json(),
         )
         logger.debug(
             f"Created Pipeline Status for pipeline {ingestion_pipeline_fqn}: {resp}"
@@ -104,7 +104,9 @@ class OMetaIngestionPipelineMixin:
         )
 
         if resp:
-            return [PipelineStatus.parse_obj(status) for status in resp.get("data")]
+            return [
+                PipelineStatus.model_validate(status) for status in resp.get("data")
+            ]
         return None
 
     def get_ingestion_pipeline_by_name(

metadata/ingestion/ometa/mixins/lineage_mixin.py

@@ -13,6 +13,7 @@ Mixin class containing Lineage specific methods
 
 To be used by OpenMetadata class
 """
+import functools
 import traceback
 from copy import deepcopy
 from typing import Any, Dict, Generic, List, Optional, Type, TypeVar, Union
@@ -51,11 +52,12 @@ class OMetaLineageMixin(Generic[T]):
     def _merge_column_lineage(
         self, original: List[Dict[str, Any]], updated: List[Dict[str, Any]]
     ):
-        temp_result = []
+        flat_original_result = set()
+        flat_updated_result = set()
         try:
             for column in original or []:
                 if column.get("toColumn") and column.get("fromColumns"):
-                    temp_result.append(
+                    flat_original_result.add(
                         (*column.get("fromColumns", []), column.get("toColumn"))
                     )
             for column in updated or []:
@@ -64,25 +66,28 @@ class OMetaLineageMixin(Generic[T]):
                 else:
                     data = column
                 if data.get("toColumn") and data.get("fromColumns"):
-                    temp_result.append(
+                    flat_updated_result.add(
                         (*data.get("fromColumns", []), data.get("toColumn"))
                     )
         except Exception as exc:
             logger.debug(f"Error while merging column lineage: {exc}")
             logger.debug(traceback.format_exc())
+        union_result = flat_original_result.union(flat_updated_result)
+        if flat_original_result == union_result:
+            return original
         return [
             {"fromColumns": list(col_data[:-1]), "toColumn": col_data[-1]}
-            for col_data in set(temp_result)
+            for col_data in union_result
         ]
 
     def _update_cache(self, request: AddLineageRequest, response: Dict[str, Any]):
         try:
             for res in response.get("downstreamEdges", []):
-                if str(request.edge.toEntity.id.__root__) == res.get("toEntity"):
+                if str(request.edge.toEntity.id.root) == res.get("toEntity"):
                     search_cache.put(
                         (
-                            request.edge.fromEntity.id.__root__,
-                            request.edge.toEntity.id.__root__,
+                            request.edge.fromEntity.id.root,
+                            request.edge.toEntity.id.root,
                         ),
                         {"edge": res.get("lineageDetails")},
                     )
@@ -93,8 +98,8 @@ class OMetaLineageMixin(Generic[T]):
         # discard the cache if failed to update
         search_cache.put(
             (
-                request.edge.fromEntity.id.__root__,
-                request.edge.toEntity.id.__root__,
+                request.edge.fromEntity.id.root,
+                request.edge.toEntity.id.root,
             ),
             None,
         )
@@ -109,8 +114,8 @@ class OMetaLineageMixin(Generic[T]):
         try:
             patch_op_success = False
             if check_patch and data.edge.lineageDetails:
-                from_id = data.edge.fromEntity.id.__root__
-                to_id = data.edge.toEntity.id.__root__
+                from_id = data.edge.fromEntity.id.root
+                to_id = data.edge.toEntity.id.root
                 edge = self.get_lineage_edge(from_id, to_id)
                 if edge:
                     original: AddLineageRequest = deepcopy(data)
@@ -118,7 +123,10 @@ class OMetaLineageMixin(Generic[T]):
                         "columnsLineage", []
                     )
                     original.edge.lineageDetails.pipeline = (
-                        EntityReference(**edge["edge"].get("pipeline"))
+                        EntityReference(
+                            id=edge["edge"]["pipeline"]["id"],
+                            type=edge["edge"]["pipeline"]["type"],
+                        )
                         if edge["edge"].get("pipeline")
                         else None
                     )
@@ -140,24 +148,26 @@ class OMetaLineageMixin(Generic[T]):
                             original.edge.lineageDetails.pipeline
                         )
                     patch = self.patch_lineage_edge(original=original, updated=data)
-                    if patch is not None:
+                    if patch:
                         patch_op_success = True
 
             if patch_op_success is False:
-                self.client.put(self.get_suffix(AddLineageRequest), data=data.json())
+                self.client.put(
+                    self.get_suffix(AddLineageRequest), data=data.model_dump_json()
+                )
 
         except APIError as err:
             logger.debug(traceback.format_exc())
             logger.error(
                 "Error %s trying to PUT lineage for %s: %s",
                 err.status_code,
-                data.json(),
+                data.model_dump_json(),
                 str(err),
             )
             raise err
 
         from_entity_lineage = self.get_lineage_by_id(
-            data.edge.fromEntity.type, str(data.edge.fromEntity.id.__root__)
+            data.edge.fromEntity.type, str(data.edge.fromEntity.id.root)
         )
 
         self._update_cache(data, from_entity_lineage)
@@ -200,7 +210,7 @@ class OMetaLineageMixin(Generic[T]):
         self,
         original: AddLineageRequest,
         updated: AddLineageRequest,
-    ) -> Optional[str]:
+    ) -> Optional[bool]:
         """
         Patches a lineage edge between two entities.
 
@@ -222,18 +232,18 @@ class OMetaLineageMixin(Generic[T]):
             if patch:
                 self.client.patch(
                     f"{self.get_suffix(AddLineageRequest)}/{original.edge.fromEntity.type}/"
-                    f"{original.edge.fromEntity.id.__root__}/{original.edge.toEntity.type}"
-                    f"/{original.edge.toEntity.id.__root__}",
+                    f"{original.edge.fromEntity.id.root}/{original.edge.toEntity.type}"
+                    f"/{original.edge.toEntity.id.root}",
                     data=str(patch),
                 )
-                return str(patch)
+            return True
         except APIError as err:
             logger.debug(traceback.format_exc())
             logger.warning(
                 f"Error Patching Lineage Edge {err.status_code} "
                 f"for {original.edge.fromEntity.fullyQualifiedName}"
             )
-        return None
+        return False
 
     def get_lineage_by_id(
         self,
@@ -312,13 +322,31 @@ class OMetaLineageMixin(Generic[T]):
         """
         try:
             self.client.delete(
-                f"{self.get_suffix(AddLineageRequest)}/{edge.fromEntity.type}/{edge.fromEntity.id.__root__}/"
-                f"{edge.toEntity.type}/{edge.toEntity.id.__root__}"
+                f"{self.get_suffix(AddLineageRequest)}/{edge.fromEntity.type}/{edge.fromEntity.id.root}/"
+                f"{edge.toEntity.type}/{edge.toEntity.id.root}"
             )
         except APIError as err:
             logger.debug(traceback.format_exc())
             logger.error(f"Error {err.status_code} trying to DELETE linage for {edge}")
 
+    @functools.lru_cache(maxsize=LRU_CACHE_SIZE)
+    def delete_lineage_by_source(
+        self, entity_type: str, entity_id: str, source: str
+    ) -> None:
+        """
+        Remove the given Edge
+        """
+        try:
+            self.client.delete(
+                f"{self.get_suffix(AddLineageRequest)}/{entity_type}/{entity_id}/"
+                f"type/{source}"
+            )
+        except APIError as err:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Error {err.status_code} trying to DELETE linage for {entity_id} of type {source}"
+            )
+
     def add_lineage_by_query(
         self,
         database_service: DatabaseService,
@@ -341,7 +369,7 @@ class OMetaLineageMixin(Generic[T]):
             connection_type = database_service.serviceType.value
             add_lineage_request = get_lineage_by_query(
                 metadata=self,
-                service_name=database_service.name.__root__,
+                service_name=database_service.name.root,
                 dialect=ConnectionTypeDialectMapper.dialect_of(connection_type),
                 query=sql,
                 database_name=database_name,

metadata/ingestion/ometa/mixins/mlmodel_mixin.py

@@ -90,7 +90,7 @@ class OMetaMlModelMixin(OMetaLineageMixin):
                 )
             )
 
-        mlmodel_lineage = self.get_lineage_by_id(MlModel, str(model.id.__root__))
+        mlmodel_lineage = self.get_lineage_by_id(MlModel, str(model.id.root))
 
         return mlmodel_lineage
 
@@ -151,7 +151,7 @@ class OMetaMlModelMixin(OMetaLineageMixin):
             mlHyperParameters=[
                 MlHyperParameter(
                     name=key,
-                    value=value,
+                    value=str(value),
                 )
                 for key, value in model.get_params().items()
             ],