openmetadata-ingestion 1.7.0.0rc2__py3-none-any.whl → 1.7.1.0__py3-none-any.whl

metadata/ingestion/source/dashboard/tableau/client.py

@@ -13,23 +13,22 @@ Wrapper module of TableauServerConnection client
 """
 import math
 import traceback
-from typing import Any, Callable, Dict, List, Optional
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
-import pandas as pd
 import validators
 from cached_property import cached_property
-from packaging import version
-from tableau_api_lib import TableauServerConnection
-from tableau_api_lib.utils import extract_pages, flatten_dict_column
-from tableau_api_lib.utils.querying import get_views_dataframe
-
-from metadata.ingestion.source.dashboard.tableau import (
-    TABLEAU_GET_VIEWS_PARAM_DICT,
-    TABLEAU_GET_WORKBOOKS_PARAM_DICT,
+from tableauserverclient import (
+    Pager,
+    PersonalAccessTokenAuth,
+    Server,
+    TableauAuth,
+    ViewItem,
 )
+
 from metadata.ingestion.source.dashboard.tableau.models import (
     CustomSQLTablesResponse,
     DataSource,
+    TableauBaseModel,
     TableauChart,
     TableauDashboard,
     TableauDatasources,
@@ -45,6 +44,18 @@ from metadata.utils.logger import ometa_logger
 logger = ometa_logger()
 
 
+class TableauWorkBookException(Exception):
+    """
+    Raise when Workbooks information is not retrieved from the Tableau APIs
+    """
+
+
+class TableauChartsException(Exception):
+    """
+    Raise when Charts information is not retrieved from the Tableau APIs
+    """
+
+
 class TableauOwnersNotFound(Exception):
     """
     Raise when Owner information is not retrieved from the Tableau APIs
@@ -62,183 +73,153 @@ class TableauClient:
     Wrapper to TableauServerConnection
     """
 
-    _client: TableauServerConnection
-
     def __init__(
         self,
-        tableau_server_config: Dict[str, Dict[str, Any]],
+        tableau_server_auth: Union[PersonalAccessTokenAuth, TableauAuth],
         config,
-        env: str,
-        ssl_verify: bool,
+        verify_ssl,
         pagination_limit: int,
     ):
-        # ssl_verify is typed as a `bool` in TableauServerConnection
-        # However, it is passed as `verify=self.ssl_verify` in each `requests` call.
-        # In requests (https://requests.readthedocs.io/en/latest/user/advanced.html?highlight=ssl#ssl-cert-verification)
-        # the param can be None, False to ignore HTTPS certs or a string with the path to the cert.
-        self._client = TableauServerConnection(
-            config_json=tableau_server_config,
-            env=env,
-            ssl_verify=ssl_verify,
-        )
+        self.tableau_server = Server(str(config.hostPort), use_server_version=True)
+        self.tableau_server.add_http_options({"verify": verify_ssl})
+        self.tableau_server.auth.sign_in(tableau_server_auth)
         self.config = config
-        self._client.sign_in().json()
         self.pagination_limit = pagination_limit
         self.custom_sql_table_queries: Dict[str, List[str]] = {}
         self.usage_metrics: Dict[str, int] = {}
 
     @cached_property
     def server_info(self) -> Callable:
-        return self._client.server_info
+        return self.tableau_server.server_info.get
 
-    @property
     def server_api_version(self) -> str:
-        return self.server_info().json()["serverInfo"]["restApiVersion"]
+        return self.tableau_server.version
 
     @property
     def site_id(self) -> str:
-        return self._client.site_id
-
-    @property
-    def query_workbooks_for_site(self) -> Callable:
-        return self._client.query_workbooks_for_site
-
-    @property
-    def query_views_for_site(self) -> Callable:
-        return self._client.query_views_for_site
-
-    def get_owners(self) -> Optional[List[TableauOwner]]:
-        owners = [workbook.owner for workbook in self.get_workbooks()]
-        if len(owners) > 0:
-            return owners
-        raise TableauOwnersNotFound(
-            "Unable to fetch Dashboard Owners from tableau\n"
-            "Please check if the user has permissions to access the Owner information"
-        )
+        return self.tableau_server.site_id
 
-    def get_all_workbook_usage_metrics(self) -> None:
-        """
-        Get the usage metrics for all workbook and store it in self.usage_metrics
-        """
+    def get_tableau_owner(self, owner_id: str) -> Optional[TableauOwner]:
         try:
-            views_df = get_views_dataframe(self._client, all_fields=False)
-            if views_df.empty:
-                logger.debug("No views data found to process usage metrics.")
-                return
-
-            if "workbook" not in views_df.columns:
-                logger.debug("Expected 'workbook' column not found in views data.")
-                return
-
-            usage_views_df = flatten_dict_column(
-                df=views_df, keys=["id"], col_name="workbook"
-            )
-
-            if (
-                "workbook_id" not in usage_views_df.columns
-                or "usage_totalViewCount" not in usage_views_df.columns
-            ):
-                logger.debug(
-                    "Expected columns 'workbook_id' or 'usage_totalViewCount' not found after flattening."
-                )
-                return
-
-            usage_views_df = usage_views_df[["workbook_id", "usage_totalViewCount"]]
-            usage_views_df = (
-                usage_views_df.groupby("workbook_id", as_index=False)[
-                    "usage_totalViewCount"
-                ]
-                .apply(lambda x: pd.to_numeric(x, errors="coerce").astype(int).sum())
-                .reset_index()
-            )
-            self.usage_metrics.update(
-                usage_views_df.set_index("workbook_id")[
-                    "usage_totalViewCount"
-                ].to_dict()
+            owner = self.tableau_server.users.get_by_id(owner_id) if owner_id else None
+            if owner and owner.email:
+                return TableauOwner(id=owner.id, name=owner.name, email=owner.email)
+        except Exception as err:
+            logger.warning(
+                f"Failed to fetch owner details for ID {owner_id}: {str(err)}"
             )
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.warning(f"Error processing workbook usage metrics: {exc}")
+        return None
 
-    def get_workbook_view_count_by_id(self, workbook_id: str) -> Optional[int]:
+    def get_workbook_charts_and_user_count(
+        self, views: List[ViewItem]
+    ) -> Optional[Tuple[Optional[int], Optional[List[TableauChart]]]]:
         """
-        Get a workbook view count by dashboard_id
+        Fetches workbook charts and dashboard user view count
         """
-        if not self.usage_metrics:
-            self.get_all_workbook_usage_metrics()
-        return self.usage_metrics.get(workbook_id)
+        view_count = 0
+        charts: Optional[List[TableauChart]] = []
+        for view in views or []:
+            try:
+                charts.append(
+                    TableauChart(
+                        id=view.id,
+                        name=view.name,
+                        tags=view.tags,
+                        owner=self.get_tableau_owner(view.owner_id),
+                        contentUrl=view.content_url,
+                        sheetType=view.sheet_type,
+                    )
+                )
+                view_count += view.total_views
+            except AttributeError as e:
+                logger.warning(
+                    f"Failed to process view due to missing attribute: {str(e)}"
+                )
+                continue
+            except Exception as e:
+                logger.warning(f"Failed to process view: {str(e)}")
+                continue
 
-    def get_workbooks(self) -> List[TableauDashboard]:
-        return [
-            TableauDashboard(**workbook)
-            for workbook in extract_pages(
-                self.query_workbooks_for_site,
-                parameter_dict=TABLEAU_GET_WORKBOOKS_PARAM_DICT,
-            )
-        ]
-
-    def get_charts(self) -> List[TableauChart]:
-        # For charts, we can also pick up usage as a field
-        return [
-            TableauChart(**chart)
-            for chart in extract_pages(
-                self.query_views_for_site,
-                content_id=self.site_id,
-                parameter_dict=TABLEAU_GET_VIEWS_PARAM_DICT,
-            )
-        ]
+        return charts, view_count
 
-    def get_workbook_charts(self, dashboard_id: str) -> Optional[List[TableauChart]]:
+    def get_workbooks(self) -> List[TableauDashboard]:
         """
-        Get the charts for a workbook
+        Fetch all tableau workbooks
         """
-        try:
-            return [
-                TableauChart(**chart)
-                for chart in self._client.query_views_for_workbook(
-                    workbook_id=dashboard_id,
-                    parameter_dict=TABLEAU_GET_VIEWS_PARAM_DICT,
-                ).json()["views"]["view"]
-            ]
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.warning(
-                f"Error processing charts for dashboard [{dashboard_id}]: {exc}"
-            )
-        return None
+        workbooks: Optional[List[TableauDashboard]] = []
+        self.cache_custom_sql_tables()
+        for workbook in Pager(self.tableau_server.workbooks):
+            try:
+                self.tableau_server.workbooks.populate_views(workbook, usage=True)
+                charts, user_views = self.get_workbook_charts_and_user_count(
+                    workbook.views
+                )
+                workbooks.append(
+                    TableauDashboard(
+                        id=workbook.id,
+                        name=workbook.name,
+                        project=TableauBaseModel(
+                            id=workbook.project_id, name=workbook.project_name
+                        ),
+                        description=workbook.description,
+                        owner=self.get_tableau_owner(workbook.owner_id),
+                        tags=workbook.tags,
+                        webpageUrl=workbook.webpage_url,
+                        charts=charts,
+                        dataModels=self.get_datasources(dashboard_id=workbook.id),
+                        user_views=user_views,
+                    )
+                )
+            except AttributeError as err:
+                logger.warning(
+                    f"Failed to process workbook due to missing attribute: {str(err)}"
+                )
+                continue
+            except Exception as err:
+                logger.warning(f"Failed to process workbook: {str(err)}")
+                continue
+        return workbooks
+
+    def test_get_workbooks(self):
+        for workbook in Pager(self.tableau_server.workbooks):
+            if workbook.id is not None:
+                self.tableau_server.workbooks.populate_views(workbook, usage=True)
+                return workbook
+            break
+        raise TableauWorkBookException(
+            "Unable to fetch Dashboards from tableau\n"
+            "Please check if the user has permissions to access the Dashboards information"
+        )
 
-    def test_api_version(self):
-        """
-        Method to validate the declared v/s server tableau rest api version
-        """
-        server_api_version = version.parse(self.server_api_version)
-        declared_api_version = version.parse(self.config.apiVersion)
-        if declared_api_version > server_api_version:
-            raise ValueError(
-                f"""
-            Your API version of '{declared_api_version}' is too damn high!
-            The server you are establishing a connection with is using REST API version '{server_api_version}'.
-            """
-            )
-        if declared_api_version < server_api_version:
-            raise ValueError(
-                f"""
-            The Tableau Server REST API version you specified is lower than the version your server uses.
-            Your Tableau Server is on REST API version {server_api_version}.
-            The REST API version you specified is {declared_api_version}.
-            For optimal results, please change the 'api_version' config variable to {server_api_version}.
-            """
-            )
+    def test_get_workbook_views(self):
+        workbook = self.test_get_workbooks()
+        charts, _ = self.get_workbook_charts_and_user_count(workbook.views)
+        if charts:
+            return True
+        raise TableauChartsException(
+            "Unable to fetch Charts from tableau\n"
+            "Please check if the user has permissions to access the Charts information"
+        )
+
+    def test_get_owners(self) -> Optional[List[TableauOwner]]:
+        workbook = self.test_get_workbooks()
+        owners = self.get_tableau_owner(workbook.owner_id)
+        if owners is not None:
+            return owners
+        raise TableauOwnersNotFound(
+            "Unable to fetch Dashboard Owners from tableau\n"
+            "Please check if the user has permissions to access the Owner information"
+        )
 
     def test_site_url(self):
         """
         Method to test the site url and site name fields
         """
-        validation = validators.url(self.config.siteUrl)
+        validation = validators.url(self.config.siteName)
         if validation:
             raise ValueError(
                 f"""
-            The site url "{self.config.siteUrl}" is in incorrect format.
+            The site url "{self.config.siteName}" is in incorrect format.
             If "https://xxx.tableau.com/#/site/MarketingTeam/home" represents the homepage url for your tableau site,
             the "MarketingTeam" from the url should be entered in the Site Name and Site Url fields.
             """
@@ -249,17 +230,16 @@ class TableauClient:
         """
         Method to test the datamodels
         """
-        workbooks = self.get_workbooks()
+        workbook = self.test_get_workbooks()
 
-        if len(workbooks) == 0:
+        if workbook.id is None:
             raise TableauDataModelsException(
                 "Unable to get any workbooks to fetch tableau data sources"
             )
 
         # Take the 1st workbook's id and pass to the graphql query
-        test_workbook = workbooks[0]
         data = self._query_datasources(
-            dashboard_id=test_workbook.id, entities_per_page=1, offset=0
+            dashboard_id=workbook.id, entities_per_page=1, offset=0
         )
         if data:
             return data
@@ -278,16 +258,17 @@ class TableauClient:
         Method to query the graphql endpoint to get data sources
         """
         try:
-            datasources_graphql_result = self._client.metadata_graphql_query(
+            datasources_graphql_result = self.tableau_server.metadata.query(
                 query=TABLEAU_DATASOURCES_QUERY.format(
                     workbook_id=dashboard_id, first=entities_per_page, offset=offset
                 )
             )
             if datasources_graphql_result:
-                resp = datasources_graphql_result.json()
-                if resp and resp.get("data"):
+                if datasources_graphql_result and datasources_graphql_result.get(
+                    "data"
+                ):
                     tableau_datasource_connection = TableauDatasourcesConnection(
-                        **resp["data"]["workbooks"][0]
+                        **datasources_graphql_result["data"]["workbooks"][0]
                     )
                     return tableau_datasource_connection.embeddedDatasourcesConnection
         except Exception:
@@ -330,15 +311,15 @@ class TableauClient:
             logger.warning("Unable to fetch Data Sources")
         return None
 
-    def get_custom_sql_table_queries(self, dashboard_id: str) -> Optional[List[str]]:
+    def get_custom_sql_table_queries(self, datasource_id: str) -> Optional[List[str]]:
         """
         Get custom SQL table queries for a specific dashboard/workbook ID
         """
-        logger.debug(f"Getting custom SQL table queries for dashboard {dashboard_id}")
+        logger.debug(f"Getting custom SQL table queries for datasource {datasource_id}")
 
-        if dashboard_id in self.custom_sql_table_queries:
-            logger.debug(f"Found cached queries for dashboard {dashboard_id}")
-            return self.custom_sql_table_queries[dashboard_id]
+        if datasource_id in self.custom_sql_table_queries:
+            logger.debug(f"Found cached queries for datasource {datasource_id}")
+            return self.custom_sql_table_queries[datasource_id]
 
         return None
 
@@ -347,30 +328,30 @@ class TableauClient:
         Fetch all custom SQL tables and cache their queries by workbook ID
         """
         try:
-            result = self._client.metadata_graphql_query(
+            result = self.tableau_server.metadata.query(
                 query=TALEAU_GET_CUSTOM_SQL_QUERY
             )
-            if not result or not (response_json := result.json()):
+            if not result:
                 logger.debug("No result returned from GraphQL query")
                 return
 
-            response = CustomSQLTablesResponse(**response_json)
+            response = CustomSQLTablesResponse(**result)
             if not response.data:
                 logger.debug("No data found in GraphQL response")
                 return
 
             for tables in response.data.values():
                 for table in tables:
-                    if not (table.query and table.downstreamWorkbooks):
+                    if not (table.query and table.downstreamDatasources):
                         logger.debug(
                             f"Skipping table {table} - missing query or workbooks"
                         )
                         continue
 
                     query = table.query
-                    for workbook in table.downstreamWorkbooks:
+                    for datasource in table.downstreamDatasources:
                         self.custom_sql_table_queries.setdefault(
-                            workbook.luid, []
+                            datasource.id, []
                         ).append(query)
 
         except Exception:
@@ -378,4 +359,4 @@ class TableauClient:
             logger.warning("Unable to fetch Custom SQL Tables")
 
     def sign_out(self) -> None:
-        self._client.sign_out()
+        self.tableau_server.auth.sign_out()

metadata/ingestion/source/dashboard/tableau/connection.py

@@ -13,10 +13,9 @@
 Source connection handler
 """
 import traceback
-from functools import partial
 from typing import Any, Dict, Optional
 
-from tableau_api_lib.utils import extract_pages
+import tableauserverclient as TSC
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -36,10 +35,6 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.dashboard.tableau import (
-    TABLEAU_GET_VIEWS_PARAM_DICT,
-    TABLEAU_GET_WORKBOOKS_PARAM_DICT,
-)
 from metadata.ingestion.source.dashboard.tableau.client import TableauClient
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
@@ -52,14 +47,13 @@ def get_connection(connection: TableauConnection) -> TableauClient:
     """
     Create connection
     """
-    tableau_server_config = build_server_config(connection)
+    tableau_server_auth = build_server_config(connection)
     get_verify_ssl = get_verify_ssl_fn(connection.verifySSL)
     try:
         return TableauClient(
-            tableau_server_config=tableau_server_config,
+            tableau_server_auth=tableau_server_auth,
             config=connection,
-            env=connection.env,
-            ssl_verify=get_verify_ssl(connection.sslConfig),
+            verify_ssl=get_verify_ssl(connection.sslConfig),
             pagination_limit=connection.paginationLimit,
         )
     except Exception as exc:
@@ -83,26 +77,11 @@ def test_connection(
 
     test_fn = {
         "ServerInfo": client.server_info,
-        # The Tableau server_info API doesn't provide direct access to the API version.
-        # This is due to the "api_version" being a mandatory field for the tableau library's connection class.
-        # Without this information, requests to the Tableau server cannot be made,
-        # including fetching the server info containing the "api_version".
-        # Consequently, we'll compare the declared api_version with the server's api_version during the test connection
-        # once the tableau library's connection class is initialized.
-        "ValidateApiVersion": client.test_api_version,
+        "ValidateApiVersion": client.server_api_version,
         "ValidateSiteUrl": client.test_site_url,
-        "GetWorkbooks": partial(
-            extract_pages,
-            query_func=client.query_workbooks_for_site,
-            parameter_dict=TABLEAU_GET_WORKBOOKS_PARAM_DICT,
-        ),
-        "GetViews": partial(
-            extract_pages,
-            query_func=client.query_views_for_site,
-            content_id=client.site_id,
-            parameter_dict=TABLEAU_GET_VIEWS_PARAM_DICT,
-        ),
-        "GetOwners": client.get_owners,
+        "GetWorkbooks": client.test_get_workbooks,
+        "GetViews": client.test_get_workbook_views,
+        "GetOwners": client.test_get_owners,
         "GetDataModels": client.test_get_datamodels,
     }
 
@@ -123,24 +102,20 @@ def build_server_config(connection: TableauConnection) -> Dict[str, Dict[str, An
     Returns:
         Client configuration
     """
-    tableau_server_config = {
-        f"{connection.env}": {
-            "server": str(connection.hostPort),
-            "api_version": str(connection.apiVersion),
-            "site_name": connection.siteName if connection.siteName else "",
-            "site_url": connection.siteUrl if connection.siteUrl else "",
-        }
-    }
+
     if isinstance(connection.authType, BasicAuth):
-        tableau_server_config[connection.env]["username"] = connection.authType.username
-        tableau_server_config[connection.env][
-            "password"
-        ] = connection.authType.password.get_secret_value()
+        tableau_auth = TSC.TableauAuth(
+            username=connection.authType.username,
+            password=connection.authType.password.get_secret_value(),
+            site_id=connection.siteName if connection.siteName else "",
+        )
     elif isinstance(connection.authType, AccessTokenAuth):
-        tableau_server_config[connection.env][
-            "personal_access_token_name"
-        ] = connection.authType.personalAccessTokenName
-        tableau_server_config[connection.env][
-            "personal_access_token_secret"
-        ] = connection.authType.personalAccessTokenSecret.get_secret_value()
-    return tableau_server_config
+        tableau_auth = TSC.PersonalAccessTokenAuth(
+            token_name=connection.authType.personalAccessTokenName,
+            personal_access_token=connection.authType.personalAccessTokenSecret.get_secret_value(),
+            site_id=connection.siteName if connection.siteName else "",
+        )
+    else:
+        raise ValueError("Unsupported authentication type")
+
+    return tableau_auth