Skip to content

Commit

Permalink
fix(tableau): fixes wrong argument when reauthenticating
Browse files Browse the repository at this point in the history
  • Loading branch information
sgomezvillamor authored Dec 23, 2024
1 parent 21ddb55 commit 89b0309
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 17 deletions.
52 changes: 38 additions & 14 deletions metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py
Original file line number Diff line number Diff line change
Expand Up @@ -618,6 +618,12 @@ def update_table(
self.parsed_columns = parsed_columns


@dataclass
class SiteIdContentUrl:
site_id: str
site_content_url: str


class TableauSourceReport(StaleEntityRemovalSourceReport):
get_all_datasources_query_failed: bool = False
num_get_datasource_query_failures: int = 0
Expand Down Expand Up @@ -770,7 +776,6 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
config=self.config,
ctx=self.ctx,
site=site,
site_id=site.id,
report=self.report,
server=self.server,
platform=self.platform,
Expand All @@ -789,7 +794,11 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
site_source = TableauSiteSource(
config=self.config,
ctx=self.ctx,
site=site,
site=site
if site
else SiteIdContentUrl(
site_id=self.server.site_id, site_content_url=self.config.site
),
site_id=self.server.site_id,
report=self.report,
server=self.server,
Expand Down Expand Up @@ -823,8 +832,7 @@ def __init__(
self,
config: TableauConfig,
ctx: PipelineContext,
site: Optional[SiteItem],
site_id: Optional[str],
site: Union[SiteItem, SiteIdContentUrl],
report: TableauSourceReport,
server: Server,
platform: str,
Expand All @@ -835,13 +843,18 @@ def __init__(
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

Check warning on line 855 in metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py#L853-L855

Added lines #L853 - L855 were not covered by tests
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] = {}
Expand Down Expand Up @@ -895,10 +908,21 @@ def dataset_browse_prefix(self) -> str:
# datasets also have the env in the browse path
return f"/{self.config.env.lower()}{self.no_env_browse_prefix}"

def _re_authenticate(self):
# 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_id)
def _re_authenticate(self) -> None:
if self.site_content_url:
assert self.site_content_url is not None, "Site Content URL is required"
self.report.info(

Check warning on line 914 in metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py#L912-L914

Added lines #L912 - L914 were not covered by tests
message="Re-authenticating to Tableau",
context=f"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)

Check warning on line 920 in metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py#L920

Added line #L920 was not covered by tests
else:
self.report.warning(

Check warning on line 922 in metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py#L922

Added line #L922 was not covered by tests
message="Site Content URL is not set. Unable to re-authenticate.",
context=f"site_id={self.site_id}, site={self.site}",
)

@property
def site_content_url(self) -> Optional[str]:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from datahub.ingestion.run.pipeline import Pipeline, PipelineContext
from datahub.ingestion.source.tableau import tableau_constant as c
from datahub.ingestion.source.tableau.tableau import (
SiteIdContentUrl,
TableauConfig,
TableauProject,
TableauSiteSource,
Expand Down Expand Up @@ -1008,7 +1009,7 @@ def check_lineage_metadata(
config=config,
ctx=context,
platform="tableau",
site=SiteItem(name="Site 1", content_url="site1"),
site=SiteIdContentUrl(site_id="id1", site_content_url="site1"),
site_id="site1",
report=TableauSourceReport(),
server=Server("https://test-tableau-server.com"),
Expand Down Expand Up @@ -1314,7 +1315,6 @@ def test_permission_warning(pytestconfig, tmp_path, mock_datahub_graph):
config=mock.MagicMock(),
ctx=mock.MagicMock(),
site=mock.MagicMock(),
site_id=None,
server=mock_sdk.return_value,
report=reporter,
)
Expand Down Expand Up @@ -1372,7 +1372,6 @@ def test_extract_project_hierarchy(extract_project_hierarchy, allowed_projects):
ctx=context,
platform="tableau",
site=SiteItem(name="Site 1", content_url="site1"),
site_id="site1",
report=TableauSourceReport(),
server=Server("https://test-tableau-server.com"),
)
Expand Down

0 comments on commit 89b0309

Please sign in to comment.