acryl-datahub 0.15.0rc25__py3-none-any.whl → 0.15.0.1__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,
@@ -35,7 +35,10 @@ from tableauserverclient import (
     SiteItem,
     TableauAuth,
 )
-from tableauserverclient.server.endpoint.exceptions import NonXMLResponseError
+from tableauserverclient.server.endpoint.exceptions import (
+    InternalServerError,
+    NonXMLResponseError,
+)
 from urllib3 import Retry
 
 import datahub.emitter.mce_builder as builder
@@ -49,6 +52,7 @@ from datahub.configuration.source_common import (
     DatasetSourceConfigMixin,
 )
 from datahub.configuration.validate_field_deprecation import pydantic_field_deprecated
+from datahub.configuration.validate_field_removal import pydantic_removed_field
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import (
     ContainerKey,
@@ -105,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,
@@ -182,6 +187,20 @@ try:
 except ImportError:
     REAUTHENTICATE_ERRORS = (NonXMLResponseError,)
 
+RETRIABLE_ERROR_CODES = [
+    408,  # Request Timeout
+    429,  # Too Many Requests
+    500,  # Internal Server Error
+    502,  # Bad Gateway
+    503,  # Service Unavailable
+    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 |
@@ -283,7 +302,7 @@ class TableauConnectionConfig(ConfigModel):
                 max_retries=Retry(
                     total=self.max_retries,
                     backoff_factor=1,
-                    status_forcelist=[429, 500, 502, 503, 504],
+                    status_forcelist=RETRIABLE_ERROR_CODES,
                 )
             )
             server._session.mount("http://", adapter)
@@ -380,11 +399,6 @@ class TableauConfig(
         description="[advanced] Number of metadata objects (e.g. CustomSQLTable, PublishedDatasource, etc) to query at a time using the Tableau API.",
     )
 
-    fetch_size: int = Field(
-        default=250,
-        description="Specifies the number of records to retrieve in each batch during a query execution.",
-    )
-
     # We've found that even with a small workbook page size (e.g. 10), the Tableau API often
     # returns warnings like this:
     # {
@@ -499,6 +513,10 @@ class TableauConfig(
         "This can only be used with ingest_tags enabled as it will overwrite tags entered from the UI.",
     )
 
+    _fetch_size = pydantic_removed_field(
+        "fetch_size",
+    )
+
     # pre = True because we want to take some decision before pydantic initialize the configuration to default values
     @root_validator(pre=True)
     def projects_backward_compatibility(cls, values: Dict) -> Dict:
@@ -618,6 +636,13 @@ class DatabaseTable:
                 self.parsed_columns = parsed_columns
 
 
+@dataclass
+class SiteIdContentUrl:
+    site_id: str
+    site_content_url: str
+
+
+@dataclass
 class TableauSourceReport(StaleEntityRemovalSourceReport):
     get_all_datasources_query_failed: bool = False
     num_get_datasource_query_failures: int = 0
@@ -634,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:
@@ -645,7 +677,7 @@ def report_user_role(report: TableauSourceReport, server: Server) -> None:
         # the site-role might be different on another site
         logged_in_user: UserInfo = UserInfo.from_server(server=server)
 
-        if not logged_in_user.is_site_administrator_explorer():
+        if not logged_in_user.has_site_administrator_explorer_privileges():
             report.warning(
                 title=title,
                 message=message,
@@ -705,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:
@@ -770,7 +803,6 @@ class TableauSource(StatefulIngestionSourceBase, TestableSource):
                         config=self.config,
                         ctx=self.ctx,
                         site=site,
-                        site_id=site.id,
                         report=self.report,
                         server=self.server,
                         platform=self.platform,
@@ -789,8 +821,14 @@ class TableauSource(StatefulIngestionSourceBase, TestableSource):
                 site_source = TableauSiteSource(
                     config=self.config,
                     ctx=self.ctx,
-                    site=site,
-                    site_id=self.server.site_id,
+                    site=(
+                        site
+                        if site
+                        else SiteIdContentUrl(
+                            site_id=self.server.site_id,
+                            site_content_url=self.config.site,
+                        )
+                    ),
                     report=self.report,
                     server=self.server,
                     platform=self.platform,
@@ -823,8 +861,7 @@ class TableauSiteSource:
         self,
         config: TableauConfig,
         ctx: PipelineContext,
-        site: Optional[SiteItem],
-        site_id: Optional[str],
+        site: Union[SiteItem, SiteIdContentUrl],
         report: TableauSourceReport,
         server: Server,
         platform: str,
@@ -835,13 +872,18 @@ class TableauSiteSource:
         self.ctx: PipelineContext = ctx
         self.platform = platform
 
-        self.site: Optional[SiteItem] = site
-        if site_id is not None:
-            self.site_id: str = site_id
+        self.site: Optional[SiteItem] = None
+        if isinstance(site, SiteItem):
+            self.site = site
+            assert site.id is not None, "Site ID is required"
+            self.site_id = site.id
+            self.site_content_url = site.content_url
+        elif isinstance(site, SiteIdContentUrl):
+            self.site = None
+            self.site_id = site.site_id
+            self.site_content_url = site.site_content_url
         else:
-            assert self.site is not None, "site or site_id is required"
-            assert self.site.id is not None, "site_id is required when site is provided"
-            self.site_id = self.site.id
+            raise AssertionError("site or site id+content_url pair is required")
 
         self.database_tables: Dict[str, DatabaseTable] = {}
         self.tableau_stat_registry: Dict[str, UsageStat] = {}
@@ -895,17 +937,12 @@ class TableauSiteSource:
         # datasets also have the env in the browse path
         return f"/{self.config.env.lower()}{self.no_env_browse_prefix}"
 
-    def _re_authenticate(self):
-        tableau_auth: Union[
-            TableauAuth, PersonalAccessTokenAuth
-        ] = self.config.get_tableau_auth(self.site_id)
-        self.server.auth.sign_in(tableau_auth)
-
-    @property
-    def site_content_url(self) -> Optional[str]:
-        if self.site and self.site.content_url:
-            return self.site.content_url
-        return None
+    def _re_authenticate(self) -> None:
+        logger.info(f"Re-authenticating to Tableau site '{self.site_content_url}'")
+        # 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:
@@ -1148,7 +1185,7 @@ class TableauSiteSource:
         connection_type: str,
         query_filter: str,
         current_cursor: Optional[str],
-        fetch_size: int = 250,
+        fetch_size: int,
         retry_on_auth_error: bool = True,
         retries_remaining: Optional[int] = None,
     ) -> Tuple[dict, Optional[str], int]:
@@ -1171,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,
@@ -1180,24 +1218,56 @@ 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:
+                if retries_remaining <= 0:
+                    raise ise
+                logger.info(f"Retrying query due to error {ise.code}")
+                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=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
@@ -1212,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,
             )
 
@@ -1300,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}")
@@ -1325,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
@@ -1336,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,
@@ -1345,7 +1418,11 @@ class TableauSiteSource:
                     connection_type=connection_type,
                     query_filter=filter_,
                     current_cursor=current_cursor,
-                    fetch_size=self.config.fetch_size,
+                    # `filter_page` contains metadata object IDs (e.g., Project IDs, Field IDs, Sheet IDs, etc.).
+                    # The number of IDs is always less than or equal to page_size.
+                    # If the IDs are primary keys, the number of metadata objects to load matches the number of records to return.
+                    # In our case, mostly, the IDs are primary key, therefore, fetch_size is set equal to page_size.
+                    fetch_size=page_size,
                 )
 
                 yield from connection_objects.get(c.NODES) or []

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/tableau/tableau_constant.py

@@ -82,4 +82,6 @@ PROJECT = "Project"
 SITE = "Site"
 IS_UNSUPPORTED_CUSTOM_SQL = "isUnsupportedCustomSql"
 SITE_PERMISSION = "sitePermission"
-SITE_ROLE = "SiteAdministratorExplorer"
+ROLE_SITE_ADMIN_EXPLORER = "SiteAdministratorExplorer"
+ROLE_SITE_ADMIN_CREATOR = "SiteAdministratorCreator"
+ROLE_SERVER_ADMIN = "ServerAdministrator"

datahub/ingestion/source/tableau/tableau_server_wrapper.py

@@ -11,8 +11,12 @@ class UserInfo:
     site_role: str
     site_id: str
 
-    def is_site_administrator_explorer(self):
-        return self.site_role == c.SITE_ROLE
+    def has_site_administrator_explorer_privileges(self):
+        return self.site_role in [
+            c.ROLE_SITE_ADMIN_EXPLORER,
+            c.ROLE_SITE_ADMIN_CREATOR,
+            c.ROLE_SERVER_ADMIN,
+        ]
 
     @staticmethod
     def from_server(server: Server) -> "UserInfo":

datahub/ingestion/source/tableau/tableau_validation.py

@@ -28,7 +28,7 @@ def check_user_role(
 
     try:
         # TODO: Add check for `Enable Derived Permissions`
-        if not logged_in_user.is_site_administrator_explorer():
+        if not logged_in_user.has_site_administrator_explorer_privileges():
             capability_dict[c.SITE_PERMISSION] = CapabilityReport(
                 capable=False,
                 failure_reason=f"{failure_reason} Their current role is {logged_in_user.site_role}.",

datahub/ingestion/source/unity/proxy.py

@@ -4,7 +4,7 @@ Manage the communication with DataBricks Server and provide equivalent dataclass
 
 import dataclasses
 import logging
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Any, Dict, Iterable, List, Optional, Union, cast
 from unittest.mock import patch
 
@@ -27,6 +27,7 @@ from databricks.sdk.service.sql import (
 from databricks.sdk.service.workspace import ObjectType
 
 import datahub
+from datahub.emitter.mce_builder import parse_ts_millis
 from datahub.ingestion.source.unity.hive_metastore_proxy import HiveMetastoreProxy
 from datahub.ingestion.source.unity.proxy_profiling import (
     UnityCatalogProxyProfilingMixin,
@@ -211,16 +212,8 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                     id=obj.object_id,
                     path=obj.path,
                     language=obj.language,
-                    created_at=(
-                        datetime.fromtimestamp(obj.created_at / 1000, tz=timezone.utc)
-                        if obj.created_at
-                        else None
-                    ),
-                    modified_at=(
-                        datetime.fromtimestamp(obj.modified_at / 1000, tz=timezone.utc)
-                        if obj.modified_at
-                        else None
-                    ),
+                    created_at=parse_ts_millis(obj.created_at),
+                    modified_at=parse_ts_millis(obj.modified_at),
                 )
 
     def query_history(
@@ -452,17 +445,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             properties=obj.properties or {},
             owner=obj.owner,
             generation=obj.generation,
-            created_at=(
-                datetime.fromtimestamp(obj.created_at / 1000, tz=timezone.utc)
-                if obj.created_at
-                else None
-            ),
+            created_at=(parse_ts_millis(obj.created_at) if obj.created_at else None),
             created_by=obj.created_by,
-            updated_at=(
-                datetime.fromtimestamp(obj.updated_at / 1000, tz=timezone.utc)
-                if obj.updated_at
-                else None
-            ),
+            updated_at=(parse_ts_millis(obj.updated_at) if obj.updated_at else None),
             updated_by=obj.updated_by,
             table_id=obj.table_id,
             comment=obj.comment,
@@ -500,12 +485,8 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             query_id=info.query_id,
             query_text=info.query_text,
             statement_type=info.statement_type,
-            start_time=datetime.fromtimestamp(
-                info.query_start_time_ms / 1000, tz=timezone.utc
-            ),
-            end_time=datetime.fromtimestamp(
-                info.query_end_time_ms / 1000, tz=timezone.utc
-            ),
+            start_time=parse_ts_millis(info.query_start_time_ms),
+            end_time=parse_ts_millis(info.query_end_time_ms),
             user_id=info.user_id,
             user_name=info.user_name,
             executed_as_user_id=info.executed_as_user_id,

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/ingestion/source_report/ingestion_stage.py

@@ -14,6 +14,8 @@ LINEAGE_EXTRACTION = "Lineage Extraction"
 USAGE_EXTRACTION_INGESTION = "Usage Extraction Ingestion"
 USAGE_EXTRACTION_OPERATIONAL_STATS = "Usage Extraction Operational Stats"
 USAGE_EXTRACTION_USAGE_AGGREGATION = "Usage Extraction Usage Aggregation"
+EXTERNAL_TABLE_DDL_LINEAGE = "External table DDL Lineage"
+VIEW_PARSING = "View Parsing"
 QUERIES_EXTRACTION = "Queries Extraction"
 PROFILING = "Profiling"
 
@@ -40,4 +42,5 @@ class IngestionStageReport:
             self._timer = PerfTimer()
 
         self.ingestion_stage = f"{stage} at {datetime.now(timezone.utc)}"
+        logger.info(f"Stage started: {self.ingestion_stage}")
         self._timer.start()