acryl-datahub-cloud 0.3.7.7rc4__py3-none-any.whl → 0.3.8rc0__py3-none-any.whl

acryl_datahub_cloud/_codegen_config.json

@@ -1,6 +1,6 @@
 {
   "name": "acryl-datahub-cloud",
-  "version": "0.3.7.7rc4",
+  "version": "0.3.8rc0",
   "install_requires": [
     "avro-gen3==0.7.16",
     "acryl-datahub"

acryl_datahub_cloud/api/__init__.py

@@ -0,0 +1 @@
+from acryl_datahub_cloud.api.client import AcrylGraph

acryl_datahub_cloud/api/client.py

@@ -0,0 +1,6 @@
+from acryl_datahub_cloud.api.entity_versioning import EntityVersioningAPI
+from datahub.ingestion.graph.client import DataHubGraph
+
+
+class AcrylGraph(EntityVersioningAPI, DataHubGraph):
+    pass

acryl_datahub_cloud/api/entity_versioning.py

@@ -0,0 +1,167 @@
+import uuid
+from typing import Optional
+
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.metadata.schema_classes import (
+    VersionPropertiesClass,
+    VersionSetPropertiesClass,
+)
+from datahub.metadata.urns import VersionSetUrn
+from datahub.utilities.urns.urn import guess_entity_type
+
+
+class EntityVersioningAPI(DataHubGraph):
+    LINK_VERSION_MUTATION = """
+        mutation($input: LinkVersionInput!) {
+            linkAssetVersion(input: $input)
+        }
+    """
+
+    UNLINK_VERSION_MUTATION = """
+        mutation($input: UnlinkVersionInput!) {
+            unlinkAssetVersion(input: $input)
+        }
+    """
+
+    def link_asset_to_version_set(
+        self,
+        asset_urn: str,
+        version_set_urn: Optional[str],
+        label: str,
+        *,
+        comment: Optional[str] = None,
+    ) -> str:
+        """Sets an entity as the latest version of a version set.
+
+        Can also be used to create a new version set, with `asset_urn` as the first version.
+
+        Args:
+            asset_urn: URN of the entity.
+            version_set_urn: URN of the version set, or None to generate a new version set urn
+            label: Label of the version.
+            comment: Comment about the version.
+
+        Returns:
+            URN of the version set to which `asset_urn` was linked.
+        """
+
+        entity_type = guess_entity_type(asset_urn)
+        if version_set_urn is None:
+            version_set_urn = VersionSetUrn(str(uuid.uuid4()), entity_type).urn()
+        elif guess_entity_type(version_set_urn) != "versionSet":
+            raise ValueError(f"Expected version set URN, got {version_set_urn}")
+
+        entity_version = self.get_aspect(asset_urn, VersionPropertiesClass)
+        if entity_version:
+            raise ValueError(
+                f"Asset {asset_urn} is already a version of {entity_version.versionSet}"
+            )
+
+        variables = {
+            "input": {
+                "versionSet": version_set_urn,
+                "linkedEntity": asset_urn,
+                "version": label,
+                "comment": comment,
+            }
+        }
+        self.execute_graphql(self.LINK_VERSION_MUTATION, variables)
+        return version_set_urn
+
+    def link_asset_to_versioned_asset(
+        self,
+        new_asset_urn: str,
+        old_asset_urn: str,
+        label: str,
+        *,
+        comment: Optional[str] = None,
+    ) -> str:
+        """Sets an entity as the latest version of an existing versioned entity.
+
+        Args:
+            new_asset_urn: URN of the new latest entity.
+            old_asset_urn: URN of an existing versioned entity to link onto.
+            label: Label of the version.
+            comment: Comment about the version.
+
+        Returns:
+            URN of the version set to which `new_asset_urn` was linked.
+        """
+
+        new_entity_type = guess_entity_type(new_asset_urn)
+        old_entity_type = guess_entity_type(old_asset_urn)
+        if new_entity_type != old_entity_type:
+            raise ValueError(
+                f"Expected URNs of the same type, got {new_entity_type} and {old_entity_type}"
+            )
+
+        new_entity_version = self.get_aspect(new_asset_urn, VersionPropertiesClass)
+        if new_entity_version:
+            raise ValueError(
+                f"Asset {new_asset_urn} is already a version of {new_entity_version.versionSet}"
+            )
+        old_entity_version = self.get_aspect(old_asset_urn, VersionPropertiesClass)
+        if not old_entity_version:
+            raise ValueError(f"Asset {old_asset_urn} is not versioned")
+
+        version_set_urn = old_entity_version.versionSet
+        self.link_asset_to_version_set(
+            new_asset_urn, version_set_urn, label, comment=comment
+        )
+        return version_set_urn
+
+    def unlink_asset_from_version_set(self, asset_urn: str) -> Optional[str]:
+        """Unlinks an entity from its version set.
+
+        Args:
+            asset_urn: URN of the entity to unlink from its version set.
+
+        Returns:
+            If successful, the URN of the version set from which `asset_urn` was unlinked.
+        """
+
+        entity_version = self.get_aspect(asset_urn, VersionPropertiesClass)
+        if not entity_version:
+            raise ValueError(f"Asset {asset_urn} is not versioned")
+
+        variables = {
+            "input": {
+                "versionSet": entity_version.versionSet,
+                "unlinkedEntity": asset_urn,
+            }
+        }
+        if self.execute_graphql(self.UNLINK_VERSION_MUTATION, variables):
+            return entity_version.versionSet
+        else:
+            return None
+
+    def unlink_latest_asset_from_version_set(
+        self, version_set_urn: str
+    ) -> Optional[str]:
+        """Unlinks the latest version of a version set.
+
+        Args:
+            version_set_urn: URN of the version set.
+
+        Returns:
+            If successful, the URN of the entity that was unlinked from `version_set_urn`.
+        """
+
+        version_set_properties = self.get_aspect(
+            version_set_urn, VersionSetPropertiesClass
+        )
+        if not version_set_properties:
+            raise ValueError(
+                f"Version set {version_set_urn} does not exist or has no versions"
+            )
+
+        variables = {
+            "input": {
+                "versionSet": version_set_urn,
+                "unlinkedEntity": version_set_properties.latest,
+            }
+        }
+        if self.execute_graphql(self.UNLINK_VERSION_MUTATION, variables):
+            return version_set_properties.latest
+        else:
+            return None

acryl_datahub_cloud/lineage_features/source.py

@@ -83,16 +83,31 @@ class DataHubLineageFeaturesSource(Source):
         query = {
             "query": {
                 "bool": {
-                    "filter": [
-                        {"term": {"source.entityType": "schemaField"}},
-                        {"term": {"destination.entityType": "schemaField"}},
+                    "should": [
+                        {"term": {"relationshipType": "Consumes"}},
                         {"term": {"relationshipType": "DownstreamOf"}},
+                        {"term": {"relationshipType": "TrainedBy"}},
+                        {"term": {"relationshipType": "UsedBy"}},
+                        {"term": {"relationshipType": "MemberOf"}},
+                        {"term": {"relationshipType": "DerivedFrom"}},
+                        {"term": {"relationshipType": "Produces"}},
+                        {"term": {"relationshipType": "DashboardContainsDashboard"}},
+                        {
+                            "bool": {
+                                "must": [
+                                    {"term": {"relationshipType": "Contains"}},
+                                    {"term": {"source.entityType": "dashboard"}},
+                                    {"term": {"destination.entityType": "chart"}},
+                                ]
+                            }
+                        },
                     ],
-                }
+                },
             },
             "sort": [
                 {"source.urn": {"order": "desc"}},
                 {"destination.urn": {"order": "desc"}},
+                {"relationshipType": {"order": "desc"}},
                 {"lifecycleOwner": {"order": "desc"}},
             ],
         }
@@ -124,7 +139,9 @@ class DataHubLineageFeaturesSource(Source):
         for urn in set(self.upstream_counts.keys()).union(
             self.downstream_counts.keys()
         ):
-            print(urn, self.upstream_counts[urn], self.downstream_counts[urn])
+            logger.debug(
+                f"{urn}: {self.upstream_counts[urn]}, {self.downstream_counts[urn]}"
+            )
             yield MetadataChangeProposalWrapper(
                 entityUrn=urn,
                 aspect=LineageFeaturesClass(

acryl_datahub_cloud/metadata/_urns/urn_defs.py

@@ -21,6 +21,60 @@ from datahub.utilities.urns.error import InvalidUrnError
 
 deprecated = functools.partial(_sphinx_deprecated, version="0.12.0.2")
 
+if TYPE_CHECKING:
+    from datahub.metadata.schema_classes import VersionSetKeyClass
+
+class VersionSetUrn(_SpecificUrn):
+    ENTITY_TYPE: ClassVar[str] = "versionSet"
+    URN_PARTS: ClassVar[int] = 2
+
+    def __init__(self, id: str, entity_type: str, *, _allow_coercion: bool = True) -> None:
+        if _allow_coercion:
+            # Field coercion logic (if any is required).
+            id = UrnEncoder.encode_string(id)
+            entity_type = UrnEncoder.encode_string(entity_type)
+
+        # Validation logic.
+        if not id:
+            raise InvalidUrnError("VersionSetUrn id cannot be empty")
+        if UrnEncoder.contains_reserved_char(id):
+            raise InvalidUrnError(f'VersionSetUrn id contains reserved characters')
+        if not entity_type:
+            raise InvalidUrnError("VersionSetUrn entity_type cannot be empty")
+        if UrnEncoder.contains_reserved_char(entity_type):
+            raise InvalidUrnError(f'VersionSetUrn entity_type contains reserved characters')
+
+        super().__init__(self.ENTITY_TYPE, [id, entity_type])
+
+    @classmethod
+    def _parse_ids(cls, entity_ids: List[str]) -> "VersionSetUrn":
+        if len(entity_ids) != cls.URN_PARTS:
+            raise InvalidUrnError(f"VersionSetUrn should have {cls.URN_PARTS} parts, got {len(entity_ids)}: {entity_ids}")
+        return cls(id=entity_ids[0], entity_type=entity_ids[1], _allow_coercion=False)
+
+    @classmethod
+    def underlying_key_aspect_type(cls) -> Type["VersionSetKeyClass"]:
+        from datahub.metadata.schema_classes import VersionSetKeyClass
+
+        return VersionSetKeyClass
+
+    def to_key_aspect(self) -> "VersionSetKeyClass":
+        from datahub.metadata.schema_classes import VersionSetKeyClass
+
+        return VersionSetKeyClass(id=self.id, entityType=self.entity_type)
+
+    @classmethod
+    def from_key_aspect(cls, key_aspect: "VersionSetKeyClass") -> "VersionSetUrn":
+        return cls(id=key_aspect.id, entity_type=key_aspect.entityType)
+
+    @property
+    def id(self) -> str:
+        return self.entity_ids[0]
+
+    @property
+    def entity_type(self) -> str:
+        return self.entity_ids[1]
+
 if TYPE_CHECKING:
     from datahub.metadata.schema_classes import DataHubConnectionKeyClass
 

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/common/__init__.py

@@ -87,6 +87,7 @@ from .....schema_classes import SubTypesClass
 from .....schema_classes import SyncMechanismClass
 from .....schema_classes import TagAssociationClass
 from .....schema_classes import TimeStampClass
+from .....schema_classes import VersionPropertiesClass
 from .....schema_classes import VersionTagClass
 from .....schema_classes import WindowDurationClass
 
@@ -171,6 +172,7 @@ SubTypes = SubTypesClass
 SyncMechanism = SyncMechanismClass
 TagAssociation = TagAssociationClass
 TimeStamp = TimeStampClass
+VersionProperties = VersionPropertiesClass
 VersionTag = VersionTagClass
 WindowDuration = WindowDurationClass
 

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/metadata/key/__init__.py

@@ -65,6 +65,7 @@ from ......schema_classes import SubscriptionKeyClass
 from ......schema_classes import TagKeyClass
 from ......schema_classes import TelemetryKeyClass
 from ......schema_classes import TestKeyClass
+from ......schema_classes import VersionSetKeyClass
 
 
 ActionRequestKey = ActionRequestKeyClass
@@ -125,5 +126,6 @@ SubscriptionKey = SubscriptionKeyClass
 TagKey = TagKeyClass
 TelemetryKey = TelemetryKeyClass
 TestKey = TestKeyClass
+VersionSetKey = VersionSetKeyClass
 
 # fmt: on

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/structured/__init__.py

@@ -13,6 +13,7 @@ from .....schema_classes import StructuredPropertiesClass
 from .....schema_classes import StructuredPropertyDefinitionClass
 from .....schema_classes import StructuredPropertyFilterStatusClass
 from .....schema_classes import StructuredPropertyKeyClass
+from .....schema_classes import StructuredPropertySettingsClass
 from .....schema_classes import StructuredPropertyValueAssignmentClass
 
 
@@ -22,6 +23,7 @@ StructuredProperties = StructuredPropertiesClass
 StructuredPropertyDefinition = StructuredPropertyDefinitionClass
 StructuredPropertyFilterStatus = StructuredPropertyFilterStatusClass
 StructuredPropertyKey = StructuredPropertyKeyClass
+StructuredPropertySettings = StructuredPropertySettingsClass
 StructuredPropertyValueAssignment = StructuredPropertyValueAssignmentClass
 
 # fmt: on

acryl_datahub_cloud/metadata/com/linkedin/pegasus2avro/versionset/__init__.py

@@ -0,0 +1,17 @@
+# mypy: ignore-errors
+# flake8: noqa
+
+# This file is autogenerated by /metadata-ingestion/scripts/avro_codegen.py
+# Do not modify manually!
+
+# pylint: skip-file
+# fmt: off
+# isort: skip_file
+from .....schema_classes import VersionSetPropertiesClass
+from .....schema_classes import VersioningSchemeClass
+
+
+VersionSetProperties = VersionSetPropertiesClass
+VersioningScheme = VersioningSchemeClass
+
+# fmt: on