feat: data contracts models + CLI

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/TestUtils.java

@@ -8,6 +8,7 @@
 import com.datahub.plugins.auth.authorization.Authorizer;
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.entity.ebean.transactions.AspectsBatchImpl;
 import com.linkedin.metadata.models.registry.ConfigEntityRegistry;
@@ -21,6 +22,8 @@
 public class TestUtils {
 
   public static EntityService getMockEntityService() {
+    PathSpecBasedSchemaAnnotationVisitor.class.getClassLoader()
+        .setClassAssertionStatus(PathSpecBasedSchemaAnnotationVisitor.class.getName(), false);
     EntityRegistry registry = new ConfigEntityRegistry(TestUtils.class.getResourceAsStream("/test-entity-registry.yaml"));
     EntityService mockEntityService = Mockito.mock(EntityService.class);
     Mockito.when(mockEntityService.getEntityRegistry()).thenReturn(registry);

datahub-graphql-core/src/test/resources/test-entity-registry.yaml

@@ -181,6 +181,7 @@ entities:
     - assertionInfo
     - dataPlatformInstance
     - assertionRunEvent
+    - assertionActions
     - status
 - name: dataHubRetention
   category: internal
@@ -292,4 +293,11 @@ entities:
   aspects:
     - ownershipTypeInfo
     - status
+- name: dataContract
+  category: core
+  keyAspect: dataContractKey
+  aspects:
+    - dataContractProperties
+    - dataContractStatus
+    - status
 events:

metadata-ingestion/examples/data_contract/pet_of_the_week.dhub.dc.yaml

@@ -0,0 +1,21 @@
+# id: pet_details_dc # Optional: This is the unique identifier for the data contract
+display_name: Data Contract for SampleHiveDataset
+entity: urn:li:dataset:(urn:li:dataPlatform:hive,SampleHiveDataset,PROD)
+freshness:
+  time: 0700
+  granularity: DAILY
+schema:
+  properties:
+    field_foo:
+      type: string
+      native_type: VARCHAR(100)
+    field_bar:
+      type: boolean
+  required:
+    - field_bar
+data_quality:
+  - type: column_range
+    config:
+      column: field_foo
+      min: 0
+      max: 100

metadata-ingestion/src/datahub/api/entities/datacontract/data_quality_assertion.py

@@ -0,0 +1,136 @@
+from typing import List, Optional, Union
+
+import pydantic
+from typing_extensions import Literal
+
+import datahub.emitter.mce_builder as builder
+from datahub.configuration.common import ConfigModel
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.metadata.schema_classes import (
+    AssertionInfoClass,
+    AssertionStdAggregationClass,
+    AssertionStdOperatorClass,
+    AssertionStdParameterClass,
+    AssertionStdParametersClass,
+    AssertionStdParameterTypeClass,
+    AssertionTypeClass,
+    DatasetAssertionInfoClass,
+    DatasetAssertionScopeClass,
+)
+
+
+class IdConfigMixin(ConfigModel):
+    id_raw: Optional[str] = pydantic.Field(
+        default=None,
+        alias="id",
+        description="The id of the assertion. If not provided, one will be generated using the type.",
+    )
+
+    def generate_default_id(self) -> str:
+        raise NotImplementedError
+
+
+class CustomSQLAssertion(IdConfigMixin, ConfigModel):
+    type: Literal["custom_sql"]
+
+    sql: str
+
+    def generate_dataset_assertion_info(
+        self, entity_urn: str
+    ) -> DatasetAssertionInfoClass:
+        return DatasetAssertionInfoClass(
+            dataset=entity_urn,
+            scope=DatasetAssertionScopeClass.UNKNOWN,
+            fields=[],
+            operator=AssertionStdOperatorClass._NATIVE_,
+            aggregation=AssertionStdAggregationClass._NATIVE_,
+            logic=self.sql,
+        )
+
+
+"""
+    def foo():
+        def _gen_column_range_assertion(config: dict) -> DatasetAssertionInfoClass:
+            return DatasetAssertionInfoClass(
+                dataset=self.entity,
+                fields=[
+                    builder.make_schema_field_urn(
+                        self.entity, str(config.get("column"))
+                    )
+                ],
+                operator=AssertionStdOperatorClass.BETWEEN,
+                aggregation=AssertionStdAggregationClass.IDENTITY,
+                scope=DatasetAssertionScopeClass.DATASET_COLUMN,
+                parameters=AssertionStdParametersClass(
+                    minValue=AssertionStdParameterClass(
+                        value=str(config.get("min")),
+                        type=AssertionStdParameterTypeClass.NUMBER,
+                    ),
+                    maxValue=AssertionStdParameterClass(
+                        value=str(config.get("max")),
+                        type=AssertionStdParameterTypeClass.NUMBER,
+                    ),
+                ),
+            )
+
+        dq_type_map = {"column_range": _gen_column_range_assertion}
+"""
+
+
+class ColumnUniqueAssertion(IdConfigMixin, ConfigModel):
+    type: Literal["unique"]
+
+    # TODO: support multiple columns?
+    column: str
+
+    def generate_default_id(self) -> str:
+        return f"{self.type}-{self.column}"
+
+    def generate_dataset_assertion_info(
+        self, entity_urn: str
+    ) -> DatasetAssertionInfoClass:
+        return DatasetAssertionInfoClass(
+            dataset=entity_urn,
+            scope=DatasetAssertionScopeClass.DATASET_COLUMN,
+            fields=[builder.make_schema_field_urn(entity_urn, self.column)],
+            operator=AssertionStdOperatorClass.EQUAL_TO,
+            aggregation=AssertionStdAggregationClass.UNIQUE_PROPOTION,  # purposely using the misspelled version to work with gql
+            parameters=AssertionStdParametersClass(
+                value=AssertionStdParameterClass(
+                    value="1", type=AssertionStdParameterTypeClass.NUMBER
+                )
+            ),
+        )
+
+
+class DataQualityAssertion(ConfigModel):
+    __root__: Union[
+        CustomSQLAssertion,
+        ColumnUniqueAssertion,
+    ] = pydantic.Field(discriminator="type")
+
+    @property
+    def id(self) -> str:
+        if self.__root__.id_raw:
+            return self.__root__.id_raw
+        try:
+            return self.__root__.generate_default_id()
+        except NotImplementedError:
+            return self.__root__.type
+
+    def generate_mcp(
+        self, assertion_urn: str, entity_urn: str
+    ) -> List[MetadataChangeProposalWrapper]:
+        dataset_assertion_info = self.__root__.generate_dataset_assertion_info(
+            entity_urn
+        )
+
+        return [
+            MetadataChangeProposalWrapper(
+                entityUrn=assertion_urn,
+                aspect=AssertionInfoClass(
+                    type=AssertionTypeClass.DATASET,
+                    datasetAssertion=dataset_assertion_info,
+                ),
+            )
+        ]