acryl-datahub 1.2.0.6__py3-none-any.whl → 1.2.0.7__py3-none-any.whl

datahub/sdk/chart.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from datetime import datetime
 from typing import Dict, List, Optional, Type, Union
 
+from deprecated.sphinx import deprecated
 from typing_extensions import Self
 
 import datahub.metadata.schema_classes as models
@@ -23,11 +24,13 @@ from datahub.sdk._shared import (
     HasTerms,
     LinksInputType,
     OwnersInputType,
+    ParentContainerInputType,
     TagsInputType,
     TermsInputType,
 )
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity, ExtraAspectsType
+from datahub.utilities.sentinels import Unset, unset
 
 
 class Chart(
@@ -70,14 +73,15 @@ class Chart(
         last_refreshed: Optional[datetime] = None,
         chart_type: Optional[Union[str, models.ChartTypeClass]] = None,
         access: Optional[str] = None,
+        input_datasets: Optional[List[Union[DatasetUrnOrStr, Dataset]]] = None,
         # Standard aspects.
+        parent_container: ParentContainerInputType | Unset = unset,
         subtype: Optional[str] = None,
         owners: Optional[OwnersInputType] = None,
         links: Optional[LinksInputType] = None,
         tags: Optional[TagsInputType] = None,
         terms: Optional[TermsInputType] = None,
         domain: Optional[DomainInputType] = None,
-        input_datasets: Optional[List[Union[DatasetUrnOrStr, Dataset]]] = None,
         extra_aspects: ExtraAspectsType = None,
     ):
         """Initialize a new Chart instance."""
@@ -91,19 +95,31 @@ class Chart(
 
         self._set_platform_instance(platform, platform_instance)
 
-        # Set additional properties
+        self._ensure_chart_props(display_name=display_name)
+
+        if display_name is not None:
+            self.set_display_name(display_name)
+        if description is not None:
+            self.set_description(description)
         if external_url is not None:
             self.set_external_url(external_url)
         if chart_url is not None:
             self.set_chart_url(chart_url)
         if custom_properties is not None:
             self.set_custom_properties(custom_properties)
+        if last_modified is not None:
+            self.set_last_modified(last_modified)
         if last_refreshed is not None:
             self.set_last_refreshed(last_refreshed)
         if chart_type is not None:
             self.set_chart_type(chart_type)
         if access is not None:
             self.set_access(access)
+        if input_datasets is not None:
+            self.set_input_datasets(input_datasets)
+
+        if parent_container is not unset:
+            self._set_container(parent_container)
         if subtype is not None:
             self.set_subtype(subtype)
         if owners is not None:
@@ -116,14 +132,6 @@ class Chart(
             self.set_terms(terms)
         if domain is not None:
             self.set_domain(domain)
-        if last_modified is not None:
-            self.set_last_modified(last_modified)
-        if input_datasets is not None:
-            self.set_input_datasets(input_datasets)
-        if description is not None:
-            self.set_description(description)
-        if display_name is not None:
-            self.set_display_name(display_name)
 
     @classmethod
     def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
@@ -139,11 +147,13 @@ class Chart(
         assert isinstance(self._urn, ChartUrn)
         return self._urn
 
-    def _ensure_chart_props(self) -> models.ChartInfoClass:
+    def _ensure_chart_props(
+        self, display_name: Optional[str] = None
+    ) -> models.ChartInfoClass:
         """Ensure chart properties exist, using a safer approach."""
         return self._setdefault_aspect(
             models.ChartInfoClass(
-                title=self.urn.chart_id,
+                title=display_name or self.urn.chart_id,
                 description="",
                 lastModified=models.ChangeAuditStampsClass(),
             )
@@ -155,13 +165,15 @@ class Chart(
         return self.urn.chart_id
 
     @property
+    @deprecated("Use display_name instead", version="1.2.0.7")
     def title(self) -> str:
-        """Get the title of the chart."""
-        return self._ensure_chart_props().title
+        """Get the display name of the chart."""
+        return self.display_name
 
+    @deprecated("Use set_display_name instead", version="1.2.0.7")
     def set_title(self, title: str) -> None:
-        """Set the title of the chart."""
-        self._ensure_chart_props().title = title
+        """Set the display name of the chart."""
+        self.set_display_name(title)
 
     @property
     def description(self) -> Optional[str]:
@@ -173,13 +185,13 @@ class Chart(
         self._ensure_chart_props().description = description
 
     @property
-    def display_name(self) -> Optional[str]:
+    def display_name(self) -> str:
         """Get the display name of the chart."""
-        return self.title
+        return self._ensure_chart_props().title
 
     def set_display_name(self, display_name: str) -> None:
         """Set the display name of the chart."""
-        self.set_title(display_name)
+        self._ensure_chart_props().title = display_name
 
     @property
     def external_url(self) -> Optional[str]:
@@ -250,9 +262,11 @@ class Chart(
     def set_chart_type(self, chart_type: Union[str, models.ChartTypeClass]) -> None:
         """Set the type of the chart."""
         if isinstance(chart_type, str):
-            assert chart_type in get_enum_options(models.ChartTypeClass), (
-                f"Invalid chart type: {chart_type}"
-            )
+            chart_type_options = get_enum_options(models.ChartTypeClass)
+            if chart_type not in chart_type_options:
+                raise ValueError(
+                    f"Invalid chart type: {chart_type}; valid types are {chart_type_options}"
+                )
         self._ensure_chart_props().type = chart_type
 
     @property

datahub/sdk/dashboard.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from datetime import datetime
 from typing import Dict, List, Optional, Type, Union
 
+from deprecated.sphinx import deprecated
 from typing_extensions import Self
 
 import datahub.metadata.schema_classes as models
@@ -24,12 +25,14 @@ from datahub.sdk._shared import (
     HasTerms,
     LinksInputType,
     OwnersInputType,
+    ParentContainerInputType,
     TagsInputType,
     TermsInputType,
 )
 from datahub.sdk.chart import Chart
 from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity, ExtraAspectsType
+from datahub.utilities.sentinels import Unset, unset
 
 
 class Dashboard(
@@ -64,7 +67,7 @@ class Dashboard(
         display_name: Optional[str] = None,
         platform_instance: Optional[DataPlatformInstanceUrnOrStr] = None,
         # Dashboard properties.
-        description: str = "",
+        description: Optional[str] = None,
         external_url: Optional[str] = None,
         dashboard_url: Optional[str] = None,
         custom_properties: Optional[Dict[str, str]] = None,
@@ -74,6 +77,7 @@ class Dashboard(
         charts: Optional[List[Union[ChartUrnOrStr, Chart]]] = None,
         dashboards: Optional[List[Union[DashboardUrnOrStr, Dashboard]]] = None,
         # Standard aspects.
+        parent_container: ParentContainerInputType | Unset = unset,
         subtype: Optional[str] = None,
         owners: Optional[OwnersInputType] = None,
         links: Optional[LinksInputType] = None,
@@ -94,18 +98,7 @@ class Dashboard(
         self._set_platform_instance(platform, platform_instance)
 
         # Initialize DashboardInfoClass with default values
-        dashboard_info = models.DashboardInfoClass(
-            title=display_name or name,
-            description=description or "",
-            lastModified=models.ChangeAuditStampsClass(
-                lastModified=None,
-            ),
-            customProperties={},
-            chartEdges=[],
-            datasetEdges=[],
-            dashboards=[],
-        )
-
+        dashboard_info = self._ensure_dashboard_props(display_name=display_name)
         if last_modified:
             dashboard_info.lastModified = models.ChangeAuditStampsClass(
                 lastModified=models.AuditStampClass(
@@ -114,7 +107,6 @@ class Dashboard(
                 ),
             )
 
-        # Set additional properties
         if description is not None:
             self.set_description(description)
         if display_name is not None:
@@ -129,6 +121,15 @@ class Dashboard(
             self.set_last_modified(last_modified)
         if last_refreshed is not None:
             self.set_last_refreshed(last_refreshed)
+        if input_datasets is not None:
+            self.set_input_datasets(input_datasets)
+        if charts is not None:
+            self.set_charts(charts)
+        if dashboards is not None:
+            self.set_dashboards(dashboards)
+
+        if parent_container is not unset:
+            self._set_container(parent_container)
         if subtype is not None:
             self.set_subtype(subtype)
         if owners is not None:
@@ -141,12 +142,6 @@ class Dashboard(
             self.set_terms(terms)
         if domain is not None:
             self.set_domain(domain)
-        if input_datasets is not None:
-            self.set_input_datasets(input_datasets)
-        if charts is not None:
-            self.set_charts(charts)
-        if dashboards is not None:
-            self.set_dashboards(dashboards)
 
     @classmethod
     def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
@@ -162,11 +157,13 @@ class Dashboard(
         assert isinstance(self._urn, DashboardUrn)
         return self._urn
 
-    def _ensure_dashboard_props(self) -> models.DashboardInfoClass:
+    def _ensure_dashboard_props(
+        self, display_name: Optional[str] = None
+    ) -> models.DashboardInfoClass:
         """Get the dashboard properties safely."""
         return self._setdefault_aspect(
             models.DashboardInfoClass(
-                title=self.urn.dashboard_id,
+                title=display_name or self.urn.dashboard_id,
                 description="",
                 lastModified=models.ChangeAuditStampsClass(
                     lastModified=models.AuditStampClass(
@@ -186,60 +183,52 @@ class Dashboard(
         return self.urn.dashboard_id
 
     @property
+    @deprecated("Use display_name instead", version="1.2.0.7")
     def title(self) -> str:
-        """Get the title of the dashboard."""
-        return self._ensure_dashboard_props().title
+        """Get the display name of the dashboard."""
+        return self.display_name
 
+    @deprecated("Use set_display_name instead", version="1.2.0.7")
     def set_title(self, title: str) -> None:
-        """Set the title of the dashboard."""
-        props = self._ensure_dashboard_props()
-        props.title = title
-        self._set_aspect(props)
+        """Set the display name of the dashboard."""
+        self.set_display_name(title)
 
     @property
     def description(self) -> Optional[str]:
         """Get the description of the dashboard."""
-        props = self._ensure_dashboard_props()
-        return props.description
+        # Because description is a required field, we treat "" as None.
+        return self._ensure_dashboard_props().description or None
 
     def set_description(self, description: str) -> None:
         """Set the description of the dashboard."""
-        props = self._ensure_dashboard_props()
-        props.description = description
-        self._set_aspect(props)
+        self._ensure_dashboard_props().description = description
 
     @property
-    def display_name(self) -> Optional[str]:
+    def display_name(self) -> str:
         """Get the display name of the dashboard."""
-        return self.title
+        return self._ensure_dashboard_props().title
 
     def set_display_name(self, display_name: str) -> None:
         """Set the display name of the dashboard."""
-        self.set_title(display_name)
+        self._ensure_dashboard_props().title = display_name
 
     @property
     def external_url(self) -> Optional[str]:
         """Get the external URL of the dashboard."""
-        props = self._ensure_dashboard_props()
-        return props.externalUrl
+        return self._ensure_dashboard_props().externalUrl
 
     def set_external_url(self, external_url: str) -> None:
         """Set the external URL of the dashboard."""
-        props = self._ensure_dashboard_props()
-        props.externalUrl = external_url
-        self._set_aspect(props)
+        self._ensure_dashboard_props().externalUrl = external_url
 
     @property
     def dashboard_url(self) -> Optional[str]:
         """Get the dashboard URL."""
-        props = self._ensure_dashboard_props()
-        return props.dashboardUrl
+        return self._ensure_dashboard_props().dashboardUrl
 
     def set_dashboard_url(self, dashboard_url: str) -> None:
         """Set the dashboard URL."""
-        props = self._ensure_dashboard_props()
-        props.dashboardUrl = dashboard_url
-        self._set_aspect(props)
+        self._ensure_dashboard_props().dashboardUrl = dashboard_url
 
     @property
     def custom_properties(self) -> Dict[str, str]:
@@ -249,9 +238,7 @@ class Dashboard(
 
     def set_custom_properties(self, custom_properties: Dict[str, str]) -> None:
         """Set the custom properties of the dashboard."""
-        props = self._ensure_dashboard_props()
-        props.customProperties = custom_properties
-        self._set_aspect(props)
+        self._ensure_dashboard_props().customProperties = custom_properties
 
     @property
     def last_modified(self) -> Optional[datetime]:
@@ -263,14 +250,12 @@ class Dashboard(
 
     def set_last_modified(self, last_modified: datetime) -> None:
         """Set the last modification timestamp of the dashboard."""
-        props = self._ensure_dashboard_props()
-        props.lastModified = models.ChangeAuditStampsClass(
+        self._ensure_dashboard_props().lastModified = models.ChangeAuditStampsClass(
             lastModified=models.AuditStampClass(
                 time=int(last_modified.timestamp()),
                 actor="urn:li:corpuser:datahub",
             ),
         )
-        self._set_aspect(props)
 
     @property
     def last_refreshed(self) -> Optional[datetime]:
@@ -284,9 +269,7 @@ class Dashboard(
 
     def set_last_refreshed(self, last_refreshed: datetime) -> None:
         """Set the last refresh timestamp of the dashboard."""
-        props = self._ensure_dashboard_props()
-        props.lastRefreshed = int(last_refreshed.timestamp())
-        self._set_aspect(props)
+        self._ensure_dashboard_props().lastRefreshed = int(last_refreshed.timestamp())
 
     @property
     def input_datasets(self) -> List[DatasetUrn]:
@@ -310,7 +293,6 @@ class Dashboard(
                 dataset_urn = DatasetUrn.from_string(dataset)
             dataset_edges.append(models.EdgeClass(destinationUrn=str(dataset_urn)))
         props.datasetEdges = dataset_edges
-        self._set_aspect(props)
 
     def add_input_dataset(self, input_dataset: Union[DatasetUrnOrStr, Dataset]) -> None:
         """Add an input dataset to the dashboard."""
@@ -326,7 +308,6 @@ class Dashboard(
                 models.EdgeClass(destinationUrn=str(input_dataset_urn))
             )
         props.datasetEdges = dataset_edges
-        self._set_aspect(props)
 
     def remove_input_dataset(
         self, input_dataset: Union[DatasetUrnOrStr, Dataset]
@@ -342,7 +323,6 @@ class Dashboard(
             for edge in (props.datasetEdges or [])
             if edge.destinationUrn != str(input_dataset_urn)
         ]
-        self._set_aspect(props)
 
     @property
     def charts(self) -> List[ChartUrn]:
@@ -363,7 +343,6 @@ class Dashboard(
                 chart_urn = ChartUrn.from_string(chart)
             chart_edges.append(models.EdgeClass(destinationUrn=str(chart_urn)))
         props.chartEdges = chart_edges
-        self._set_aspect(props)
 
     def add_chart(self, chart: Union[ChartUrnOrStr, Chart]) -> None:
         """Add a chart to the dashboard."""
@@ -381,7 +360,6 @@ class Dashboard(
         if str(chart_urn) not in existing_urns:
             chart_edges.append(models.EdgeClass(destinationUrn=str(chart_urn)))
         props.chartEdges = chart_edges
-        self._set_aspect(props)
 
     def remove_chart(self, chart: Union[ChartUrnOrStr, Chart]) -> None:
         """Remove a chart from the dashboard."""
@@ -395,7 +373,6 @@ class Dashboard(
             for edge in (props.chartEdges or [])
             if edge.destinationUrn != str(chart_urn)
         ]
-        self._set_aspect(props)
 
     @property
     def dashboards(self) -> List[DashboardUrn]:
@@ -417,7 +394,6 @@ class Dashboard(
             else:
                 dashboard_urn = DashboardUrn.from_string(dashboard)
             props.dashboards.append(models.EdgeClass(destinationUrn=str(dashboard_urn)))
-        self._set_aspect(props)
 
     def add_dashboard(self, dashboard: Union[DashboardUrnOrStr, Dashboard]) -> None:
         """Add a dashboard to the dashboard."""

datahub/sdk/lineage_client.py

@@ -165,11 +165,7 @@ class LineageClient:
         ] = False,
         transformation_text: Optional[str] = None,
     ) -> None:
-        ...
-
-        """
-        Add dataset-to-dataset lineage with column-level mapping.
-        """
+        """Add dataset-to-dataset lineage with column-level mapping."""
 
     @overload
     def add_lineage(
@@ -178,11 +174,7 @@ class LineageClient:
         upstream: Union[DatajobUrnOrStr],
         downstream: DatasetUrnOrStr,
     ) -> None:
-        ...
-
-        """
-        Add dataset-to-datajob or dataset-to-mlmodel lineage.
-        """
+        """Add dataset-to-datajob or dataset-to-mlmodel lineage."""
 
     @overload
     def add_lineage(
@@ -191,11 +183,7 @@ class LineageClient:
         upstream: Union[DatasetUrnOrStr, DatajobUrnOrStr],
         downstream: DatajobUrnOrStr,
     ) -> None:
-        ...
-
-        """
-        Add datajob-to-dataset or datajob-to-datajob lineage.
-        """
+        """Add datajob-to-dataset or datajob-to-datajob lineage."""
 
     @overload
     def add_lineage(
@@ -204,11 +192,7 @@ class LineageClient:
         upstream: Union[DashboardUrnOrStr, DatasetUrnOrStr, ChartUrnOrStr],
         downstream: DashboardUrnOrStr,
     ) -> None:
-        ...
-
-        """
-        Add dashboard-to-dashboard or dashboard-to-dataset lineage.
-        """
+        """Add dashboard-to-dashboard or dashboard-to-dataset lineage."""
 
     @overload
     def add_lineage(
@@ -217,10 +201,7 @@ class LineageClient:
         upstream: DatasetUrnOrStr,
         downstream: ChartUrnOrStr,
     ) -> None:
-        ...
-        """
-        Add dataset-to-chart lineage.
-        """
+        """Add dataset-to-chart lineage."""
 
     # The actual implementation that handles all overloaded cases
     def add_lineage(
@@ -237,8 +218,7 @@ class LineageClient:
         ] = False,
         transformation_text: Optional[str] = None,
     ) -> None:
-        """
-        Add lineage between two entities.
+        """Add lineage between two entities.
 
         This flexible method handles different combinations of entity types:
         - dataset to dataset

datahub/sdk/main_client.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional, overload
+from typing import TYPE_CHECKING, Optional, overload
 
 from datahub.errors import SdkUsageError
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
@@ -9,6 +9,9 @@ from datahub.sdk.entity_client import EntityClient
 from datahub.sdk.lineage_client import LineageClient
 from datahub.sdk.search_client import SearchClient
 
+if TYPE_CHECKING:
+    from datahub.sdk.resolver_client import ResolverClient
+
 
 class DataHubClient:
     """Main client for interacting with DataHub.
@@ -104,13 +107,14 @@ class DataHubClient:
         return EntityClient(self)
 
     @property
-    def resolve(self):  # type: ignore[report-untyped-call]  # Not available due to circular import issues
+    def resolve(self) -> "ResolverClient":
         try:
             from acryl_datahub_cloud.sdk import (  # type: ignore[import-not-found]
                 ResolverClient,
             )
         except ImportError:
-            from datahub.sdk.resolver_client import (  # type: ignore[assignment]  # If the client is not installed, use the one from the SDK
+            # If the client is not installed, use the one from the SDK.
+            from datahub.sdk.resolver_client import (  # type: ignore[assignment]
                 ResolverClient,
             )
         return ResolverClient(self)

datahub/sdk/search_filters.py

@@ -384,6 +384,21 @@ def _filter_discriminator(v: Any) -> Optional[str]:
     return None
 
 
+def _parse_and_like_filter(value: Any) -> Any:
+    # Do not parse if filter is already of type and/or/not or a custom condition
+    # also do not parse container filter if direct_descendants_only is specified
+    if (
+        isinstance(value, dict)
+        and not set(value.keys()).intersection(
+            {"and", "or", "not", "field", "condition", "direct_descendants_only"}
+        )
+        and len(value) > 1
+    ):
+        return {"and": [{k: v} for k, v in value.items()]}
+
+    return value
+
+
 if TYPE_CHECKING or not PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR:
     # The `not TYPE_CHECKING` bit is required to make the linter happy,
     # since we currently only run mypy with pydantic v1.
@@ -445,6 +460,7 @@ else:
             ],
             Discriminator(_filter_discriminator),
         ],
+        pydantic.BeforeValidator(_parse_and_like_filter),
         pydantic.BeforeValidator(_parse_json_from_string),
     ]
 

datahub/specific/aspect_helpers/siblings.py

@@ -0,0 +1,73 @@
+from typing import List
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
+
+
+class HasSiblingsPatch(MetadataPatchProposal):
+    def add_sibling(self, sibling_urn: str, primary: bool = False) -> Self:
+        """Add a sibling relationship to the entity.
+
+        Args:
+            sibling_urn: The URN of the sibling entity to add.
+            primary: Whether this entity should be marked as primary in the relationship.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME,
+            "add",
+            path=("siblings", sibling_urn),
+            value=sibling_urn,
+        )
+
+        # Set primary flag if specified
+        if primary:
+            self._add_patch(
+                Siblings.ASPECT_NAME,
+                "add",
+                path=("primary",),
+                value=primary,
+            )
+
+        return self
+
+    def remove_sibling(self, sibling_urn: str) -> Self:
+        """Remove a sibling relationship from the entity.
+
+        Args:
+            sibling_urn: The URN of the sibling entity to remove.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME,
+            "remove",
+            path=("siblings", sibling_urn),
+            value={},
+        )
+        return self
+
+    def set_siblings(self, sibling_urns: List[str], primary: bool = False) -> Self:
+        """Set the complete list of siblings for the entity.
+
+        This will replace all existing siblings with the new list.
+
+        Args:
+            sibling_urns: The list of sibling URNs to set.
+            primary: Whether this entity should be marked as primary.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME, "add", path=("siblings",), value=sibling_urns
+        )
+
+        self._add_patch(Siblings.ASPECT_NAME, "add", path=("primary",), value=primary)
+
+        return self

datahub/specific/dataset.py

@@ -22,6 +22,7 @@ from datahub.specific.aspect_helpers.fine_grained_lineage import (
     HasFineGrainedLineagePatch,
 )
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
+from datahub.specific.aspect_helpers.siblings import HasSiblingsPatch
 from datahub.specific.aspect_helpers.structured_properties import (
     HasStructuredPropertiesPatch,
 )
@@ -104,6 +105,7 @@ class DatasetPatchBuilder(
     HasTagsPatch,
     HasTermsPatch,
     HasFineGrainedLineagePatch,
+    HasSiblingsPatch,
     MetadataPatchProposal,
 ):
     def __init__(

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -634,6 +634,9 @@ class SqlParsingAggregator(Closeable):
             TableSwap,
         ],
     ) -> None:
+        """
+        This assumes that queries come in order of increasing timestamps.
+        """
         if isinstance(item, KnownQueryLineageInfo):
             self.add_known_query_lineage(item)
         elif isinstance(item, KnownLineageMapping):

datahub/sql_parsing/tool_meta_extractor.py

@@ -208,9 +208,7 @@ class ToolMetaExtractor:
         Returns:
             bool: whether QueryLog entry is that of hex.
         """
-        last_line = _get_last_line(entry.query_text)
-
-        if not last_line.startswith("-- Hex query metadata:"):
+        if "-- Hex query metadata:" not in entry.query_text:
             return False
 
         entry.origin = HEX_PLATFORM_URN