fix(ingest/sql): improve handling of views with dots in their names (#…

…9183)
datahub-project · Nov 28, 2023 · 3d7962c · 3d7962c
1 parent 966cb17
commit 3d7962c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
@@ -1054,7 +1054,7 @@ def _run_sql_parser(
         return view_definition_lineage_helper(raw_lineage, view_urn)
 
     def get_db_schema(self, dataset_identifier: str) -> Tuple[Optional[str], str]:
-        database, schema, _view = dataset_identifier.split(".")
+        database, schema, _view = dataset_identifier.split(".", 2)
         return database, schema
 
     def get_profiler_instance(self, inspector: Inspector) -> "DatahubGEProfiler":

diff --git a/metadata-ingestion/tests/unit/test_sql_common.py b/metadata-ingestion/tests/unit/test_sql_common.py
@@ -102,3 +102,17 @@ def test_use_source_schema_for_foreign_key_if_not_specified():
 def test_get_platform_from_sqlalchemy_uri(uri: str, expected_platform: str) -> None:
     platform: str = get_platform_from_sqlalchemy_uri(uri)
     assert platform == expected_platform
+
+
+def test_get_db_schema_with_dots_in_view_name():
+    config: SQLCommonConfig = _TestSQLAlchemyConfig()
+    ctx: PipelineContext = PipelineContext(run_id="test_ctx")
+    platform: str = "TEST"
+    source = _TestSQLAlchemySource(config=config, ctx=ctx, platform=platform)
+
+    database, schema = source.get_db_schema(
+        dataset_identifier="database.schema.long.view.name1"
+    )
+
+    assert database == "database"
+    assert schema == "schema"