acryl-datahub 0.15.0.1rc13__py3-none-any.whl → 0.15.0.1rc15__py3-none-any.whl

datahub/ingestion/source/tableau/tableau.py

@@ -2,9 +2,9 @@ import json
 import logging
 import re
 import time
-from collections import OrderedDict
-from dataclasses import dataclass
-from datetime import datetime
+from collections import OrderedDict, defaultdict
+from dataclasses import dataclass, field as dataclass_field
+from datetime import datetime, timedelta, timezone
 from functools import lru_cache
 from typing import (
     Any,
@@ -109,6 +109,7 @@ from datahub.ingestion.source.tableau.tableau_common import (
     make_filter,
     make_fine_grained_lineage_class,
     make_upstream_class,
+    optimize_query_filter,
     published_datasource_graphql_query,
     query_metadata_cursor_based_pagination,
     sheet_graphql_query,
@@ -195,6 +196,11 @@ RETRIABLE_ERROR_CODES = [
     504,  # Gateway Timeout
 ]
 
+# From experience, this expiry time typically ranges from 50 minutes
+# to 2 hours but might as well be configurable. We will allow upto
+# 10 minutes of such expiry time
+REGULAR_AUTH_EXPIRY_PERIOD = timedelta(minutes=10)
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 # Replace / with |
@@ -636,6 +642,7 @@ class SiteIdContentUrl:
     site_content_url: str
 
 
+@dataclass
 class TableauSourceReport(StaleEntityRemovalSourceReport):
     get_all_datasources_query_failed: bool = False
     num_get_datasource_query_failures: int = 0
@@ -652,7 +659,14 @@ class TableauSourceReport(StaleEntityRemovalSourceReport):
     num_upstream_table_lineage_failed_parse_sql: int = 0
     num_upstream_fine_grained_lineage_failed_parse_sql: int = 0
     num_hidden_assets_skipped: int = 0
-    logged_in_user: List[UserInfo] = []
+    logged_in_user: List[UserInfo] = dataclass_field(default_factory=list)
+    last_authenticated_at: Optional[datetime] = None
+
+    num_expected_tableau_metadata_queries: int = 0
+    num_actual_tableau_metadata_queries: int = 0
+    tableau_server_error_stats: Dict[str, int] = dataclass_field(
+        default_factory=(lambda: defaultdict(int))
+    )
 
 
 def report_user_role(report: TableauSourceReport, server: Server) -> None:
@@ -723,6 +737,7 @@ class TableauSource(StatefulIngestionSourceBase, TestableSource):
         try:
             logger.info(f"Authenticated to Tableau site: '{site_content_url}'")
             self.server = self.config.make_tableau_client(site_content_url)
+            self.report.last_authenticated_at = datetime.now(timezone.utc)
             report_user_role(report=self.report, server=self.server)
         # Note that we're not catching ConfigurationError, since we want that to throw.
         except ValueError as e:
@@ -806,10 +821,13 @@ class TableauSource(StatefulIngestionSourceBase, TestableSource):
                 site_source = TableauSiteSource(
                     config=self.config,
                     ctx=self.ctx,
-                    site=site
-                    if site
-                    else SiteIdContentUrl(
-                        site_id=self.server.site_id, site_content_url=self.config.site
+                    site=(
+                        site
+                        if site
+                        else SiteIdContentUrl(
+                            site_id=self.server.site_id,
+                            site_content_url=self.config.site,
+                        )
                     ),
                     report=self.report,
                     server=self.server,
@@ -924,6 +942,7 @@ class TableauSiteSource:
         # Sign-in again may not be enough because Tableau sometimes caches invalid sessions
         # so we need to recreate the Tableau Server object
         self.server = self.config.make_tableau_client(self.site_content_url)
+        self.report.last_authenticated_at = datetime.now(timezone.utc)
 
     def _populate_usage_stat_registry(self) -> None:
         if self.server is None:
@@ -1189,6 +1208,7 @@ class TableauSiteSource:
         )
         try:
             assert self.server is not None
+            self.report.num_actual_tableau_metadata_queries += 1
             query_data = query_metadata_cursor_based_pagination(
                 server=self.server,
                 main_query=query,
@@ -1198,25 +1218,36 @@ class TableauSiteSource:
                 qry_filter=query_filter,
             )
 
-        except REAUTHENTICATE_ERRORS:
-            if not retry_on_auth_error:
+        except REAUTHENTICATE_ERRORS as e:
+            self.report.tableau_server_error_stats[e.__class__.__name__] += 1
+            if not retry_on_auth_error or retries_remaining <= 0:
                 raise
 
-            # If ingestion has been running for over 2 hours, the Tableau
-            # temporary credentials will expire. If this happens, this exception
-            # will be thrown, and we need to re-authenticate and retry.
-            self._re_authenticate()
+            # We have been getting some irregular authorization errors like below well before the expected expiry time
+            # - within few seconds of initial authentication . We'll retry without re-auth for such cases.
+            # <class 'tableauserverclient.server.endpoint.exceptions.NonXMLResponseError'>:
+            # b'{"timestamp":"xxx","status":401,"error":"Unauthorized","path":"/relationship-service-war/graphql"}'
+            if self.report.last_authenticated_at and (
+                datetime.now(timezone.utc) - self.report.last_authenticated_at
+                > REGULAR_AUTH_EXPIRY_PERIOD
+            ):
+                # If ingestion has been running for over 2 hours, the Tableau
+                # temporary credentials will expire. If this happens, this exception
+                # will be thrown, and we need to re-authenticate and retry.
+                self._re_authenticate()
+
             return self.get_connection_object_page(
                 query=query,
                 connection_type=connection_type,
                 query_filter=query_filter,
                 fetch_size=fetch_size,
                 current_cursor=current_cursor,
-                retry_on_auth_error=False,
+                retry_on_auth_error=True,
                 retries_remaining=retries_remaining - 1,
             )
 
         except InternalServerError as ise:
+            self.report.tableau_server_error_stats[InternalServerError.__name__] += 1
             # In some cases Tableau Server returns 504 error, which is a timeout error, so it worths to retry.
             # Extended with other retryable errors.
             if ise.code in RETRIABLE_ERROR_CODES:
@@ -1229,13 +1260,14 @@ class TableauSiteSource:
                     query_filter=query_filter,
                     fetch_size=fetch_size,
                     current_cursor=current_cursor,
-                    retry_on_auth_error=False,
+                    retry_on_auth_error=True,
                     retries_remaining=retries_remaining - 1,
                 )
             else:
                 raise ise
 
         except OSError:
+            self.report.tableau_server_error_stats[OSError.__name__] += 1
             # In tableauseverclient 0.26 (which was yanked and released in 0.28 on 2023-10-04),
             # the request logic was changed to use threads.
             # https://github.com/tableau/server-client-python/commit/307d8a20a30f32c1ce615cca7c6a78b9b9bff081
@@ -1250,7 +1282,7 @@ class TableauSiteSource:
                 query_filter=query_filter,
                 fetch_size=fetch_size,
                 current_cursor=current_cursor,
-                retry_on_auth_error=False,
+                retry_on_auth_error=True,
                 retries_remaining=retries_remaining - 1,
             )
 
@@ -1338,7 +1370,7 @@ class TableauSiteSource:
                         query_filter=query_filter,
                         fetch_size=fetch_size,
                         current_cursor=current_cursor,
-                        retry_on_auth_error=False,
+                        retry_on_auth_error=True,
                         retries_remaining=retries_remaining,
                     )
                 raise RuntimeError(f"Query {connection_type} error: {errors}")
@@ -1363,6 +1395,8 @@ class TableauSiteSource:
         query_filter: dict = {},
         page_size_override: Optional[int] = None,
     ) -> Iterable[dict]:
+        query_filter = optimize_query_filter(query_filter)
+
         # Calls the get_connection_object_page function to get the objects,
         # and automatically handles pagination.
         page_size = page_size_override or self.config.page_size
@@ -1374,6 +1408,7 @@ class TableauSiteSource:
             while has_next_page:
                 filter_: str = make_filter(filter_page)
 
+                self.report.num_expected_tableau_metadata_queries += 1
                 (
                     connection_objects,
                     current_cursor,

datahub/ingestion/source/tableau/tableau_common.py

@@ -1,3 +1,4 @@
+import copy
 import html
 import json
 import logging
@@ -35,6 +36,7 @@ from datahub.metadata.schema_classes import (
     UpstreamClass,
 )
 from datahub.sql_parsing.sqlglot_lineage import ColumnLineageInfo, SqlParsingResult
+from datahub.utilities.ordered_set import OrderedSet
 
 logger = logging.getLogger(__name__)
 
@@ -1000,3 +1002,19 @@ def get_filter_pages(query_filter: dict, page_size: int) -> List[dict]:
         ]
 
     return filter_pages
+
+
+def optimize_query_filter(query_filter: dict) -> dict:
+    """
+    Duplicates in the filter cause duplicates in the result,
+    leading to entities/aspects being emitted multiple times unnecessarily
+    """
+    optimized_query = copy.deepcopy(query_filter)
+
+    if query_filter.get(c.ID_WITH_IN):
+        optimized_query[c.ID_WITH_IN] = list(OrderedSet(query_filter[c.ID_WITH_IN]))
+    if query_filter.get(c.PROJECT_NAME_WITH_IN):
+        optimized_query[c.PROJECT_NAME_WITH_IN] = list(
+            OrderedSet(query_filter[c.PROJECT_NAME_WITH_IN])
+        )
+    return optimized_query

datahub/ingestion/source/usage/usage_common.py

@@ -54,6 +54,20 @@ def default_user_urn_builder(email: str) -> str:
     return builder.make_user_urn(email.split("@")[0])
 
 
+def extract_user_email(user: str) -> Optional[str]:
+    """Extracts user email from user input
+
+    >>> extract_user_email('urn:li:corpuser:abc@xyz.com')
+    'abc@xyz.com'
+    >>> extract_user_email('urn:li:corpuser:abc')
+    >>> extract_user_email('abc@xyz.com')
+    'abc@xyz.com'
+    """
+    if user.startswith(("urn:li:corpuser:", "urn:li:corpGroup:")):
+        user = user.split(":")[-1]
+    return user if "@" in user else None
+
+
 def make_usage_workunit(
     bucket_start_time: datetime,
     resource: ResourceType,
@@ -104,7 +118,7 @@ def make_usage_workunit(
             DatasetUserUsageCountsClass(
                 user=user_urn_builder(user),
                 count=count,
-                userEmail=user if "@" in user else None,
+                userEmail=extract_user_email(user),
             )
             for user, count in user_freq
         ],

datahub/specific/aspect_helpers/custom_properties.py

@@ -0,0 +1,79 @@
+from abc import abstractmethod
+from typing import Dict, Optional, Tuple
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
+
+
+class HasCustomPropertiesPatch(MetadataPatchProposal):
+    @classmethod
+    @abstractmethod
+    def _custom_properties_location(self) -> Tuple[str, PatchPath]:
+        ...
+
+    def add_custom_property(self, key: str, value: str) -> Self:
+        """Add a custom property to the entity.
+
+        Args:
+            key: The key of the custom property.
+            value: The value of the custom property.
+
+        Returns:
+            The patch builder instance.
+        """
+        aspect_name, path = self._custom_properties_location()
+        self._add_patch(
+            aspect_name,
+            "add",
+            path=(*path, key),
+            value=value,
+        )
+        return self
+
+    def add_custom_properties(
+        self, custom_properties: Optional[Dict[str, str]] = None
+    ) -> Self:
+        if custom_properties is not None:
+            for key, value in custom_properties.items():
+                self.add_custom_property(key, value)
+        return self
+
+    def remove_custom_property(self, key: str) -> Self:
+        """Remove a custom property from the entity.
+
+        Args:
+            key: The key of the custom property to remove.
+
+        Returns:
+            The patch builder instance.
+        """
+        aspect_name, path = self._custom_properties_location()
+        self._add_patch(
+            aspect_name,
+            "remove",
+            path=(*path, key),
+            value={},
+        )
+        return self
+
+    def set_custom_properties(self, custom_properties: Dict[str, str]) -> Self:
+        """Sets the custom properties of the entity.
+
+        This method replaces all existing custom properties with the given dictionary.
+
+        Args:
+            custom_properties: A dictionary containing the custom properties to be set.
+
+        Returns:
+            The patch builder instance.
+        """
+
+        aspect_name, path = self._custom_properties_location()
+        self._add_patch(
+            aspect_name,
+            "add",
+            path=path,
+            value=custom_properties,
+        )
+        return self

datahub/specific/aspect_helpers/ownership.py

@@ -0,0 +1,67 @@
+from typing import List, Optional
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.schema_classes import (
+    OwnerClass,
+    OwnershipClass,
+    OwnershipTypeClass,
+)
+
+
+class HasOwnershipPatch(MetadataPatchProposal):
+    def add_owner(self, owner: OwnerClass) -> Self:
+        """Add an owner to the entity.
+
+        Args:
+            owner: The Owner object to add.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            OwnershipClass.ASPECT_NAME,
+            "add",
+            path=("owners", owner.owner, str(owner.type)),
+            value=owner,
+        )
+        return self
+
+    def remove_owner(
+        self, owner: str, owner_type: Optional[OwnershipTypeClass] = None
+    ) -> Self:
+        """Remove an owner from the entity.
+
+        If owner_type is not provided, the owner will be removed regardless of ownership type.
+
+        Args:
+            owner: The owner to remove.
+            owner_type: The ownership type of the owner (optional).
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            OwnershipClass.ASPECT_NAME,
+            "remove",
+            path=("owners", owner) + ((str(owner_type),) if owner_type else ()),
+            value=owner,
+        )
+        return self
+
+    def set_owners(self, owners: List[OwnerClass]) -> Self:
+        """Set the owners of the entity.
+
+        This will effectively replace all existing owners with the new list - it doesn't really patch things.
+
+        Args:
+            owners: The list of owners to set.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            OwnershipClass.ASPECT_NAME, "add", path=("owners",), value=owners
+        )
+        return self

datahub/specific/aspect_helpers/structured_properties.py

@@ -0,0 +1,72 @@
+from typing import List, Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.schema_classes import (
+    StructuredPropertiesClass,
+    StructuredPropertyValueAssignmentClass,
+)
+from datahub.utilities.urns.structured_properties_urn import (
+    make_structured_property_urn,
+)
+
+
+class HasStructuredPropertiesPatch(MetadataPatchProposal):
+    def set_structured_property(
+        self, key: str, value: Union[str, float, List[Union[str, float]]]
+    ) -> Self:
+        """Add or update a structured property.
+
+        Args:
+            key: the name of the property (either bare or urn form)
+            value: the value of the property (for multi-valued properties, this can be a list)
+
+        Returns:
+            The patch builder instance.
+        """
+        self.remove_structured_property(key)
+        self.add_structured_property(key, value)
+        return self
+
+    def remove_structured_property(self, key: str) -> Self:
+        """Remove a structured property.
+
+        Args:
+            key: the name of the property (either bare or urn form)
+
+        Returns:
+            The patch builder instance.
+        """
+
+        self._add_patch(
+            StructuredPropertiesClass.ASPECT_NAME,
+            "remove",
+            path=("properties", make_structured_property_urn(key)),
+            value={},
+        )
+        return self
+
+    def add_structured_property(
+        self, key: str, value: Union[str, float, List[Union[str, float]]]
+    ) -> Self:
+        """Add a structured property.
+
+        Args:
+            key: the name of the property (either bare or urn form)
+            value: the value of the property (for multi-valued properties, this value will be appended to the list)
+
+        Returns:
+            The patch builder instance.
+        """
+
+        self._add_patch(
+            StructuredPropertiesClass.ASPECT_NAME,
+            "add",
+            path=("properties", make_structured_property_urn(key)),
+            value=StructuredPropertyValueAssignmentClass(
+                propertyUrn=make_structured_property_urn(key),
+                values=value if isinstance(value, list) else [value],
+            ),
+        )
+        return self

datahub/specific/aspect_helpers/tags.py

@@ -0,0 +1,42 @@
+from typing import Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.schema_classes import (
+    GlobalTagsClass as GlobalTags,
+    TagAssociationClass as Tag,
+)
+from datahub.metadata.urns import TagUrn, Urn
+
+
+class HasTagsPatch(MetadataPatchProposal):
+    def add_tag(self, tag: Tag) -> Self:
+        """Adds a tag to the entity.
+
+        Args:
+            tag: The Tag object representing the tag to be added.
+
+        Returns:
+            The patch builder instance.
+        """
+
+        # TODO: Make this support raw strings, in addition to Tag objects.
+        self._add_patch(
+            GlobalTags.ASPECT_NAME, "add", path=("tags", tag.tag), value=tag
+        )
+        return self
+
+    def remove_tag(self, tag: Union[str, Urn]) -> Self:
+        """Removes a tag from the entity.
+
+        Args:
+            tag: The tag to remove, specified as a string or Urn object.
+
+        Returns:
+            The patch builder instance.
+        """
+        if isinstance(tag, str) and not tag.startswith("urn:li:tag:"):
+            tag = TagUrn.create_from_id(tag)
+        self._add_patch(GlobalTags.ASPECT_NAME, "remove", path=("tags", tag), value={})
+        return self

datahub/specific/aspect_helpers/terms.py

@@ -0,0 +1,43 @@
+from typing import Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.schema_classes import (
+    GlossaryTermAssociationClass as Term,
+    GlossaryTermsClass,
+)
+from datahub.metadata.urns import GlossaryTermUrn, Urn
+
+
+class HasTermsPatch(MetadataPatchProposal):
+    def add_term(self, term: Term) -> Self:
+        """Adds a glossary term to the entity.
+
+        Args:
+            term: The Term object representing the glossary term to be added.
+
+        Returns:
+            The patch builder instance.
+        """
+        # TODO: Make this support raw strings, in addition to Term objects.
+        self._add_patch(
+            GlossaryTermsClass.ASPECT_NAME, "add", path=("terms", term.urn), value=term
+        )
+        return self
+
+    def remove_term(self, term: Union[str, Urn]) -> Self:
+        """Removes a glossary term from the entity.
+
+        Args:
+            term: The term to remove, specified as a string or Urn object.
+
+        Returns:
+            The patch builder instance.
+        """
+        if isinstance(term, str) and not term.startswith("urn:li:glossaryTerm:"):
+            term = GlossaryTermUrn(term)
+        self._add_patch(
+            GlossaryTermsClass.ASPECT_NAME, "remove", path=("terms", term), value={}
+        )
+        return self