openmetadata-ingestion 1.8.9.0__py3-none-any.whl → 1.9.0.0__py3-none-any.whl

metadata/ingestion/ometa/mixins/data_contract_mixin.py

@@ -0,0 +1,205 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  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.
+"""
+Mixin class containing Data Contract specific methods
+
+To be used by OpenMetadata class
+"""
+import traceback
+from typing import Optional
+
+from metadata.generated.schema.entity.data.dataContract import DataContract
+from metadata.generated.schema.entity.datacontract.dataContractResult import (
+    DataContractResult,
+)
+from metadata.generated.schema.type.basic import Uuid
+from metadata.ingestion.ometa.client import REST
+from metadata.ingestion.ometa.utils import model_str
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+
+class OMetaDataContractMixin:
+    """
+    OpenMetadata API methods related to Data Contracts.
+    To be inherited by OpenMetadata
+    """
+
+    client: REST
+
+    def put_data_contract_result(
+        self, data_contract_id: Uuid, result: DataContractResult
+    ) -> Optional[DataContractResult]:
+        """
+        Create or update a data contract execution result
+
+        Args:
+            data_contract_id: UUID of the data contract
+            result: DataContractResult object containing execution details
+
+        Returns:
+            DataContractResult if successful, None otherwise
+        """
+        try:
+            resp = self.client.put(
+                f"{self.get_suffix(DataContract)}/{model_str(data_contract_id)}/results",
+                data=result.model_dump_json(),
+            )
+            if resp:
+                return DataContractResult(**resp)
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error creating data contract result for {model_str(data_contract_id)}: {err}"
+            )
+        return None
+
+    def get_data_contract_results(
+        self,
+        data_contract_id: Uuid,
+        limit: int = 10,
+        start_ts: Optional[int] = None,
+        end_ts: Optional[int] = None,
+    ) -> Optional[list]:
+        """
+        Get data contract execution results
+
+        Args:
+            data_contract_id: UUID of the data contract
+            limit: Maximum number of results to return
+            start_ts: Start timestamp filter
+            end_ts: End timestamp filter
+
+        Returns:
+            List of DataContractResult objects if successful, None otherwise
+        """
+        try:
+            # Build query parameters
+            query_params = [f"limit={limit}"]
+            if start_ts:
+                query_params.append(f"startTs={start_ts}")
+            if end_ts:
+                query_params.append(f"endTs={end_ts}")
+
+            query_string = "&".join(query_params)
+            url = f"{self.get_suffix(DataContract)}/{model_str(data_contract_id)}/results?{query_string}"
+
+            resp = self.client.get(url)
+            if resp:
+                return [DataContractResult(**result) for result in resp.get("data", [])]
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error getting data contract results for {model_str(data_contract_id)}: {err}"
+            )
+        return None
+
+    def get_latest_data_contract_result(
+        self, data_contract_id: Uuid
+    ) -> Optional[DataContractResult]:
+        """
+        Get the latest data contract execution result
+
+        Args:
+            data_contract_id: UUID of the data contract
+
+        Returns:
+            DataContractResult if successful, None otherwise
+        """
+        try:
+            resp = self.client.get(
+                f"{self.get_suffix(DataContract)}/{model_str(data_contract_id)}/results/latest"
+            )
+            if resp:
+                return DataContractResult(**resp)
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error getting latest data contract result for {model_str(data_contract_id)}: {err}"
+            )
+        return None
+
+    def get_data_contract_result_by_id(
+        self, data_contract_id: Uuid, result_id: Uuid
+    ) -> Optional[DataContractResult]:
+        """
+        Get a specific data contract execution result by ID
+
+        Args:
+            data_contract_id: UUID of the data contract
+            result_id: UUID of the specific result
+
+        Returns:
+            DataContractResult if successful, None otherwise
+        """
+        try:
+            resp = self.client.get(
+                f"{self.get_suffix(DataContract)}/{model_str(data_contract_id)}/results/{model_str(result_id)}"
+            )
+            if resp:
+                return DataContractResult(**resp)
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error getting data contract result {model_str(result_id)} for {model_str(data_contract_id)}: {err}"
+            )
+        return None
+
+    def delete_data_contract_result(
+        self, data_contract_id: Uuid, timestamp: int
+    ) -> bool:
+        """
+        Delete a data contract result at a specific timestamp
+
+        Args:
+            data_contract_id: UUID of the data contract
+            timestamp: Timestamp of the result to delete
+
+        Returns:
+            True if successful, False otherwise
+        """
+        try:
+            self.client.delete(
+                f"{self.get_suffix(DataContract)}/{model_str(data_contract_id)}/results/{timestamp}"
+            )
+            return True
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error deleting data contract result at {timestamp} for {model_str(data_contract_id)}: {err}"
+            )
+        return False
+
+    def validate_data_contract(
+        self, data_contract_id: Uuid
+    ) -> Optional[DataContractResult]:
+        """
+        Trigger on-demand validation of a data contract
+
+        Args:
+            data_contract_id: UUID of the data contract to validate
+
+        Returns:
+            DataContractResult if successful, None otherwise
+        """
+        try:
+            resp = self.client.post(
+                f"{self.get_suffix(DataContract)}/{model_str(data_contract_id)}/validate"
+            )
+            if resp:
+                return DataContractResult(**resp)
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error validating data contract {model_str(data_contract_id)}: {err}"
+            )
+        return None

metadata/ingestion/ometa/mixins/patch_mixin.py

@@ -25,7 +25,6 @@ from metadata.generated.schema.entity.automations.workflow import (
 )
 from metadata.generated.schema.entity.automations.workflow import WorkflowStatus
 from metadata.generated.schema.entity.data.table import Column, Table, TableConstraint
-from metadata.generated.schema.entity.domains.domain import Domain
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
 )
@@ -34,7 +33,6 @@ from metadata.generated.schema.entity.services.ingestionPipelines.reverseIngesti
 )
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
 from metadata.generated.schema.type.basic import EntityLink, Markdown
-from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.generated.schema.type.entityReferenceList import EntityReferenceList
 from metadata.generated.schema.type.lifeCycle import LifeCycle
 from metadata.generated.schema.type.tagLabel import TagLabel
@@ -124,6 +122,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         restrict_update_fields: Optional[List] = None,
         array_entity_fields: Optional[List] = None,
         override_metadata: Optional[bool] = False,
+        skip_on_failure: Optional[bool] = True,
     ) -> Optional[T]:
         """
         Given an Entity type and Source entity and Destination entity,
@@ -135,6 +134,9 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             destination: payload with changes applied to the source.
             allowed_fields: List of field names to filter from source and destination models
             restrict_update_fields: List of field names which will only support add operation
+            array_entity_fields: List of array fields to sort for consistent patching
+            override_metadata: Whether to override existing metadata fields
+            skip_on_failure: Whether to skip the patch operation on failure (default: True)
 
         Returns
             Updated Entity
@@ -147,6 +149,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
                 restrict_update_fields=restrict_update_fields,
                 array_entity_fields=array_entity_fields,
                 override_metadata=override_metadata,
+                skip_on_failure=skip_on_failure,
             )
 
             if not patch:
@@ -160,9 +163,19 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Error trying to PATCH {get_log_name(source)}: {exc}")
-
-        return None
+            if skip_on_failure:
+                entity_name = get_log_name(source)
+                logger.warning(
+                    f"Failed to update {entity_name}. The patch operation was skipped. "
+                    f"Reason: {exc}"
+                )
+                return None
+            else:
+                entity_name = get_log_name(source)
+                raise RuntimeError(
+                    f"Failed to update {entity_name}. The patch operation failed. "
+                    f"Set 'skip_on_failure=True' to skip failed patches. Error: {exc}"
+                ) from exc
 
     def patch_description(
         self,
@@ -170,6 +183,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         source: T,
         description: str,
         force: bool = False,
+        skip_on_failure: bool = True,
     ) -> Optional[T]:
         """
         Given an Entity type and ID, JSON PATCH the description.
@@ -180,33 +194,51 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             description: new description to add
             force: if True, we will patch any existing description. Otherwise, we will maintain
                 the existing data.
+            skip_on_failure: if True, return None on failure instead of raising exception
         Returns
             Updated Entity
         """
-        if isinstance(source, TestCase):
-            instance: Optional[T] = self._fetch_entity_if_exists(
-                entity=entity,
-                entity_id=source.id,
-                fields=["testDefinition", "testSuite"],
-            )
-        else:
-            instance: Optional[T] = self._fetch_entity_if_exists(
-                entity=entity, entity_id=source.id
-            )
+        try:
+            if isinstance(source, TestCase):
+                instance: Optional[T] = self._fetch_entity_if_exists(
+                    entity=entity,
+                    entity_id=source.id,
+                    fields=["testDefinition", "testSuite"],
+                )
+            else:
+                instance: Optional[T] = self._fetch_entity_if_exists(
+                    entity=entity, entity_id=source.id
+                )
 
-        if not instance:
-            return None
+            if not instance:
+                return None
 
-        if instance.description and not force:
-            # If the description is already present and force is not passed,
-            # description will not be overridden
-            return None
+            if instance.description and not force:
+                # If the description is already present and force is not passed,
+                # description will not be overridden
+                return None
 
-        # https://docs.pydantic.dev/latest/usage/exporting_models/#modelcopy
-        destination = source.model_copy(deep=True)
-        destination.description = Markdown(description)
+            # https://docs.pydantic.dev/latest/usage/exporting_models/#modelcopy
+            destination = source.model_copy(deep=True)
+            destination.description = Markdown(description)
 
-        return self.patch(entity=entity, source=source, destination=destination)
+            return self.patch(
+                entity=entity,
+                source=source,
+                destination=destination,
+                skip_on_failure=skip_on_failure,
+            )
+        except Exception as exc:
+            if skip_on_failure:
+                logger.debug(traceback.format_exc())
+                entity_name = get_log_name(source)
+                logger.warning(
+                    f"Failed to patch description for {entity_name}. The patch operation was skipped. "
+                    f"Reason: {exc}"
+                )
+                return None
+            else:
+                raise
 
     def patch_table_constraints(
         self,
@@ -275,6 +307,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         operation: Union[
             PatchOperation.ADD, PatchOperation.REMOVE
         ] = PatchOperation.ADD,
+        skip_on_failure: bool = True,
     ) -> Optional[T]:
         """
         Given an Entity type and ID, JSON PATCH the tag.
@@ -284,29 +317,47 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             source: Source entity object
             tag_label: TagLabel to add or remove
             operation: Patch Operation to add or remove the tag.
+            skip_on_failure: if True, return None on failure instead of raising exception
         Returns
             Updated Entity
         """
-        instance: Optional[T] = self._fetch_entity_if_exists(
-            entity=entity, entity_id=source.id, fields=["tags"]
-        )
-        if not instance:
-            return None
+        try:
+            instance: Optional[T] = self._fetch_entity_if_exists(
+                entity=entity, entity_id=source.id, fields=["tags"]
+            )
+            if not instance:
+                return None
 
-        # Initialize empty tag list or the last updated tags
-        source.tags = instance.tags or []
-        destination = source.model_copy(deep=True)
+            # Initialize empty tag list or the last updated tags
+            source.tags = instance.tags or []
+            destination = source.model_copy(deep=True)
 
-        tag_fqns = {label.tagFQN.root for label in tag_labels}
+            tag_fqns = {label.tagFQN.root for label in tag_labels}
 
-        if operation == PatchOperation.REMOVE:
-            for tag in destination.tags:
-                if tag.tagFQN.root in tag_fqns:
-                    destination.tags.remove(tag)
-        else:
-            destination.tags.extend(tag_labels)
+            if operation == PatchOperation.REMOVE:
+                for tag in destination.tags:
+                    if tag.tagFQN.root in tag_fqns:
+                        destination.tags.remove(tag)
+            else:
+                destination.tags.extend(tag_labels)
 
-        return self.patch(entity=entity, source=source, destination=destination)
+            return self.patch(
+                entity=entity,
+                source=source,
+                destination=destination,
+                skip_on_failure=skip_on_failure,
+            )
+        except Exception as exc:
+            if skip_on_failure:
+                logger.debug(traceback.format_exc())
+                entity_name = get_log_name(source)
+                logger.warning(
+                    f"Failed to patch tags for {entity_name}. The patch operation was skipped. "
+                    f"Reason: {exc}"
+                )
+                return None
+            else:
+                raise
 
     def patch_tag(
         self,
@@ -316,11 +367,16 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         operation: Union[
             PatchOperation.ADD, PatchOperation.REMOVE
         ] = PatchOperation.ADD,
+        skip_on_failure: bool = True,
     ) -> Optional[T]:
         """Will be deprecated in 1.3"""
         logger.warning("patch_tag will be deprecated in 1.3. Use `patch_tags` instead.")
         return self.patch_tags(
-            entity=entity, source=source, tag_labels=[tag_label], operation=operation
+            entity=entity,
+            source=source,
+            tag_labels=[tag_label],
+            operation=operation,
+            skip_on_failure=skip_on_failure,
         )
 
     def patch_owner(
@@ -554,17 +610,37 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             )
             return None
 
-    def patch_domain(self, entity: Entity, domain: Domain) -> Optional[Entity]:
-        """Patch domain data for an Entity"""
-        try:
-            destination: Entity = entity.model_copy(deep=True)
-            destination.domain = EntityReference(id=domain.id, type="domain")
-            return self.patch(
-                entity=type(entity), source=entity, destination=destination
-            )
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.warning(
-                f"Error trying to Patch Domain for {entity.fullyQualifiedName.root}: {exc}"
-            )
+    def patch_domain(
+        self,
+        entity: Type[T],
+        source: T,
+        domains: EntityReferenceList = None,
+        force: bool = False,
+    ) -> Optional[T]:
+        """
+        Given an Entity type and ID, JSON PATCH the owner. If not owner Entity type and
+        not owner ID are provided, the owner is removed.
+
+        Args
+            entity (T): Entity Type of the entity to be patched
+            entity_id: ID of the entity to be patched
+            owner: Entity Reference of the owner. If None, the owner will be removed
+            force: if True, we will patch any existing owner. Otherwise, we will maintain
+                the existing data.
+        Returns
+            Updated Entity
+        """
+        instance: Optional[T] = self._fetch_entity_if_exists(
+            entity=entity, entity_id=source.id, fields=["domains"]
+        )
+
+        if not instance:
             return None
+
+        if instance.domains and instance.domains.root and not force:
+            return None
+
+        destination = deepcopy(instance)
+        destination.domains = domains
+
+        return self.patch(entity=entity, source=instance, destination=destination)

metadata/ingestion/ometa/ometa_api.py

@@ -37,6 +37,7 @@ from metadata.ingestion.ometa.mixins.custom_property_mixin import (
     OMetaCustomPropertyMixin,
 )
 from metadata.ingestion.ometa.mixins.dashboard_mixin import OMetaDashboardMixin
+from metadata.ingestion.ometa.mixins.data_contract_mixin import OMetaDataContractMixin
 from metadata.ingestion.ometa.mixins.data_insight_mixin import DataInsightMixin
 from metadata.ingestion.ometa.mixins.domain_mixin import OMetaDomainMixin
 from metadata.ingestion.ometa.mixins.es_mixin import ESMixin
@@ -116,6 +117,7 @@ class OpenMetadata(
     ESMixin,
     OMetaServerMixin,
     OMetaDashboardMixin,
+    OMetaDataContractMixin,
     OMetaPatchMixin,
     OMetaTestsMixin,
     DataInsightMixin,
@@ -289,6 +291,7 @@ class OpenMetadata(
             .replace("storedprocedure", "storedProcedure")
             .replace("ingestionpipeline", "ingestionPipeline")
             .replace("dataproduct", "dataProduct")
+            .replace("datacontract", "dataContract")
         )
         class_path = ".".join(
             filter(
@@ -325,7 +328,7 @@ class OpenMetadata(
         resp = fn(
             # this might be a regular pydantic model so we build the context manually
             self.get_suffix(entity),
-            data=data.model_dump_json(context={"mask_secrets": False}),
+            data=data.model_dump_json(context={"mask_secrets": False}, by_alias=True),
         )
         if not resp:
             raise EmptyPayloadException(

metadata/ingestion/ometa/routes.py

@@ -38,6 +38,11 @@ from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequ
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
+from metadata.generated.schema.api.data.createDataContract import (
+    CreateDataContractRequest,
+)
+from metadata.generated.schema.api.data.createDirectory import CreateDirectoryRequest
+from metadata.generated.schema.api.data.createFile import CreateFileRequest
 from metadata.generated.schema.api.data.createGlossary import CreateGlossaryRequest
 from metadata.generated.schema.api.data.createGlossaryTerm import (
     CreateGlossaryTermRequest,
@@ -52,11 +57,15 @@ from metadata.generated.schema.api.data.createQueryCostRecord import (
 from metadata.generated.schema.api.data.createSearchIndex import (
     CreateSearchIndexRequest,
 )
+from metadata.generated.schema.api.data.createSpreadsheet import (
+    CreateSpreadsheetRequest,
+)
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
 from metadata.generated.schema.api.data.createTopic import CreateTopicRequest
+from metadata.generated.schema.api.data.createWorksheet import CreateWorksheetRequest
 from metadata.generated.schema.api.docStore.createDocument import CreateDocumentRequest
 from metadata.generated.schema.api.domains.createDataProduct import (
     CreateDataProductRequest,
@@ -74,6 +83,9 @@ from metadata.generated.schema.api.services.createDashboardService import (
 from metadata.generated.schema.api.services.createDatabaseService import (
     CreateDatabaseServiceRequest,
 )
+from metadata.generated.schema.api.services.createDriveService import (
+    CreateDriveServiceRequest,
+)
 from metadata.generated.schema.api.services.createMessagingService import (
     CreateMessagingServiceRequest,
 )
@@ -133,6 +145,9 @@ 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.dataContract import DataContract
+from metadata.generated.schema.entity.data.directory import Directory
+from metadata.generated.schema.entity.data.file import File
 from metadata.generated.schema.entity.data.glossary import Glossary
 from metadata.generated.schema.entity.data.glossaryTerm import GlossaryTerm
 from metadata.generated.schema.entity.data.metric import Metric
@@ -142,9 +157,11 @@ from metadata.generated.schema.entity.data.query import Query
 from metadata.generated.schema.entity.data.queryCostRecord import QueryCostRecord
 from metadata.generated.schema.entity.data.report import Report
 from metadata.generated.schema.entity.data.searchIndex import SearchIndex
+from metadata.generated.schema.entity.data.spreadsheet import Spreadsheet
 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.data.worksheet import Worksheet
 from metadata.generated.schema.entity.docStore.document import Document
 from metadata.generated.schema.entity.domains.dataProduct import DataProduct
 from metadata.generated.schema.entity.domains.domain import Domain
@@ -156,6 +173,7 @@ from metadata.generated.schema.entity.services.connections.testConnectionDefinit
 )
 from metadata.generated.schema.entity.services.dashboardService import DashboardService
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.services.driveService import DriveService
 from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipeline import (
     IngestionPipeline,
 )
@@ -213,6 +231,15 @@ ROUTES = {
     CreateAPICollectionRequest.__name__: "/apiCollections",
     Document.__name__: "/docStore",
     CreateDocumentRequest.__name__: "/docStore",
+    # Drive
+    Directory.__name__: "/drives/directories",
+    CreateDirectoryRequest.__name__: "/drives/directories",
+    Spreadsheet.__name__: "/drives/spreadsheets",
+    CreateSpreadsheetRequest.__name__: "/drives/spreadsheets",
+    File.__name__: "/drives/files",
+    CreateFileRequest.__name__: "/drives/files",
+    Worksheet.__name__: "/drives/worksheets",
+    CreateWorksheetRequest.__name__: "/drives/worksheets",
     # Classifications
     Tag.__name__: "/tags",
     CreateTagRequest.__name__: "/tags",
@@ -242,6 +269,8 @@ ROUTES = {
     Workflow.__name__: "/automations/workflows",
     CreateWorkflowRequest.__name__: "/automations/workflows",
     # Services
+    DriveService.__name__: "/services/driveServices",
+    CreateDriveServiceRequest.__name__: "/services/driveServices",
     ApiService.__name__: "/services/apiServices",
     CreateApiServiceRequest.__name__: "/services/apiServices",
     DatabaseService.__name__: "/services/databaseServices",
@@ -295,4 +324,7 @@ ROUTES = {
     # Query Cost
     QueryCostRecord.__name__: "/queryCostRecord",
     CreateQueryCostRecordRequest.__name__: "/queryCostRecord",
+    # Data Contracts
+    DataContract.__name__: "/dataContracts",
+    CreateDataContractRequest.__name__: "/dataContracts",
 }