From cb1bb1157a30e2ea731fffd786fdf7f1a852e96a Mon Sep 17 00:00:00 2001
From: treff7es <treff7es@gmail.com>
Date: Fri, 10 Jan 2025 18:32:02 +0100
Subject: [PATCH] Fix TableaUpstream create

---
 .../ingestion/source/tableau/tableau.py       |  2 +-
 .../source/tableau/tableau_common.py          |  7 +++++--
 .../tests/unit/test_tableau_source.py         | 21 +++++++++++++++++++
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py
index d149402741e82f..e8e7f6080e28f6 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau.py
@@ -2661,7 +2661,7 @@ def emit_upstream_tables(self) -> Iterable[MetadataWorkUnit]:
             c.ID_WITH_IN: list(tableau_database_table_id_to_urn_map.keys())
         }
 
-        # Emmitting tables that came from Tableau metadata
+        # Emitting tables that came from Tableau metadata
         for tableau_table in self.get_connection_objects(
             database_tables_graphql_query,
             c.DATABASE_TABLES_CONNECTION,
diff --git a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py
index 8f9d81eb9a18c1..5d5103330fe302 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/tableau/tableau_common.py
@@ -642,8 +642,11 @@ class TableauUpstreamReference:
 
     @classmethod
     def create(
-        cls, d: dict, default_schema_map: Optional[Dict[str, str]] = None
+        cls, d: Dict, default_schema_map: Optional[Dict[str, str]] = None
     ) -> "TableauUpstreamReference":
+        if d is None:
+            raise ValueError("TableauUpstreamReference.create: d is None")
+
         # Values directly from `table` object from Tableau
         database_dict = (
             d.get(c.DATABASE) or {}
@@ -717,7 +720,7 @@ def parse_full_name(full_name: Optional[str]) -> Optional[List[str]]:
         #  schema
 
         # TODO: Validate the startswith check. Currently required for our integration tests
-        if full_name is None or not full_name.startswith("["):
+        if full_name is None:
             return None
 
         return full_name.replace("[", "").replace("]", "").split(".")
diff --git a/metadata-ingestion/tests/unit/test_tableau_source.py b/metadata-ingestion/tests/unit/test_tableau_source.py
index 227519fdb464a8..9521d0a61615a6 100644
--- a/metadata-ingestion/tests/unit/test_tableau_source.py
+++ b/metadata-ingestion/tests/unit/test_tableau_source.py
@@ -5,6 +5,7 @@
 import datahub.ingestion.source.tableau.tableau_constant as c
 from datahub.ingestion.source.tableau.tableau import TableauSiteSource
 from datahub.ingestion.source.tableau.tableau_common import (
+    TableauUpstreamReference,
     get_filter_pages,
     make_filter,
     optimize_query_filter,
@@ -247,3 +248,23 @@ def test_optimize_query_filter_handles_no_duplicates():
     assert len(result) == 2
     assert result[c.ID_WITH_IN] == ["id1", "id2"]
     assert result[c.PROJECT_NAME_WITH_IN] == ["project1", "project2"]
+
+
+def testTableaUpstreamReference():
+    d = {
+        "id": "7127b695-3df5-4a3a-4837-eb0f4b572337",
+        "name": "TABLE1",
+        "database": None,
+        "schema": "SCHEMA1",
+        "fullName": "DB1.SCHEMA1.TABLE1",
+        "connectionType": "snowflake",
+        "description": "",
+        "columnsConnection": {"totalCount": 0},
+    }
+    ref = TableauUpstreamReference.create(d)
+    assert ref
+
+    assert ref.database == "DB1"
+    assert ref.schema == "SCHEMA1"
+    assert ref.table == "TABLE1"
+    assert ref.connection_type == "snowflake"