Don't use database as fallback when no schema parsed. (#32959)

Signed-off-by: Jakub Dardzinski <[email protected]>
apache · Aug 1, 2023 · 2a39914 · 2a39914
1 parent 9b7ed43
commit 2a39914
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/airflow/providers/openlineage/sqlparser.py b/airflow/providers/openlineage/sqlparser.py
@@ -291,6 +291,6 @@ def _get_tables_hierarchy(
             else:
                 db = None
             schemas = hierarchy.setdefault(normalize_name(db) if db else db, {})
-            tables = schemas.setdefault(normalize_name(table.schema) if table.schema else db, [])
+            tables = schemas.setdefault(normalize_name(table.schema) if table.schema else None, [])
             tables.append(table.name)
         return hierarchy
diff --git a/tests/providers/openlineage/utils/test_sqlparser.py b/tests/providers/openlineage/utils/test_sqlparser.py
@@ -89,6 +89,14 @@ def test_get_tables_hierarchy(self):
             is_cross_db=True,
         ) == {"db": {"schema1": ["Table2"]}, "db2": {"schema1": ["Table1"]}}
 
+        # cross db, no db & schema parsed
+        assert SQLParser._get_tables_hierarchy(
+            [DbTableMeta("Table1"), DbTableMeta("Table2")],
+            normalize_name_lower,
+            database="Db",
+            is_cross_db=True,
+        ) == {"db": {None: ["Table1", "Table2"]}}
+
     def test_normalize_sql(self):
         assert SQLParser.normalize_sql("select * from asdf") == "select * from asdf"