Repository: reata/sqllineage

Issue:
Metadata Masked When Table was in a previous UPDATE statement
**Describe the bug**
* Updating a table then using it with a join makes the Column Level Lineage to not deduce the source.

**SQL**
```sql
UPDATE `dataset.table_1` SET A = '';

CREATE OR REPLACE TEMP TABLE `table_x` AS
SELECT DISTINCT
    B
FROM `dataset.table_1`
CROSS JOIN `dataset.table_2`
;
```

**To Reproduce**
*Note here we refer to SQL provided in prior step as stored in a file named `test.sql`*
```python
import json
from sqllineage.core.metadata.dummy import DummyMetaDataProvider
from sqllineage.runner import LineageRunner

with open("test.sql") as f:
    sp = f.read()

with open('metadata.json') as file:
    metadata = json.load(file)
provider = DummyMetaDataProvider(metadata)

lineage = LineageRunner(sp, dialect="bigquery", metadata_provider=provider, verbose=True)
lineage.print_column_lineage()
```
Metadata.json:
```
{
  "dataset.table_1": [
    "A",
    "B"
  ],
  "dataset.table_2": [
    "C"
  ]
}
```

**Result**
```
<default>.table_x.b <- b
```

**Expected behavior**
```
<default>.table_x.b <- dataset.table_1.b
```

**Python version (available via `python --version`)**
 - 3.9.12

**SQLLineage version (available via `sqllineage --version`):**
 - 1.5.1

**Notes**
- Removing the UPDATE statement or the JOIN makes it work properly.
- Any kind of JOIN produces the bug


Discussion:
The issue comes from when we handling session metadata. The intention for session metadata is that for temporary table/view created during the session, the schema info won't be available in metadata, as the table/view is not created yet.

But clearly we should limit the scope of session metadata to CREATE (and maybe INSERT). By now means should we register session metadata for UPDATE statement.

In this case, we register session metadata for `dataset.table_1` through UPDATE, which doesn't include any schema information. And in the end this masks the schema info from MetaDataProvider.



Patch:
diff --git a/sqllineage/runner.py b/sqllineage/runner.py
index 335a0eec..9f9e472b 100644
--- a/sqllineage/runner.py
+++ b/sqllineage/runner.py
@@ -199,10 +199,10 @@ def _eval(self):
                 stmt_holder = analyzer.analyze(stmt, session.metadata_provider)
                 if write := stmt_holder.write:
                     tgt_table = next(iter(write))
-                    if isinstance(tgt_table, Table):
-                        session.register_session_metadata(
-                            tgt_table, stmt_holder.get_table_columns(tgt_table)
-                        )
+                    if isinstance(tgt_table, Table) and (
+                        tgt_columns := stmt_holder.get_table_columns(tgt_table)
+                    ):
+                        session.register_session_metadata(tgt_table, tgt_columns)
                 stmt_holders.append(stmt_holder)
             self._stmt_holders = stmt_holders
             self._sql_holder = SQLLineageHolder.of(


Test patch:
diff --git a/tests/helpers.py b/tests/helpers.py
index 2443f46a..006aa1a4 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -1,6 +1,6 @@
 import os.path
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
 from sqlalchemy import (
     Column as SQLAlchemyColumn,
@@ -92,7 +92,7 @@ def assert_column_lineage_equal(
         _assert_column_lineage(lr_sqlfluff, column_lineages)
 
 
-def generate_metadata_providers(test_schemas):
+def generate_metadata_providers(test_schemas) -> List[MetaDataProvider]:
     dummy_provider = DummyMetaDataProvider(test_schemas)
 
     sqlite3_sqlalchemy_provider = SQLAlchemyMetaDataProvider("sqlite:///:memory:")
diff --git a/tests/sql/column/multiple_statements/__init__.py b/tests/sql/column/multiple_statements/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/sql/column/multiple_statements/test_session_metadata.py b/tests/sql/column/multiple_statements/test_session_metadata.py
new file mode 100644
index 00000000..caf84926
--- /dev/null
+++ b/tests/sql/column/multiple_statements/test_session_metadata.py
@@ -0,0 +1,45 @@
+import pytest
+
+from sqllineage.core.metadata_provider import MetaDataProvider
+from sqllineage.utils.entities import ColumnQualifierTuple
+from ....helpers import assert_column_lineage_equal, generate_metadata_providers
+
+
+providers = generate_metadata_providers(
+    {
+        "db1.table1": ["id", "a", "b", "c", "d"],
+        "db2.table2": ["id", "h", "i", "j", "k"],
+        "db3.table3": ["pk", "p", "q", "r"],
+    }
+)
+
+
+@pytest.mark.parametrize("provider", providers)
+def test_do_not_register_session_metadata_for_update_statement(
+    provider: MetaDataProvider,
+):
+    sql = """UPDATE db1.table1 SET a = 1;
+
+CREATE TABLE db1.foo AS
+SELECT a, b, c
+FROM db1.table1 tab1
+INNER JOIN db1.table2 tab2 ON tab1.id = tab2.id
+"""
+    assert_column_lineage_equal(
+        sql,
+        [
+            (
+                ColumnQualifierTuple("a", "db1.table1"),
+                ColumnQualifierTuple("a", "db1.foo"),
+            ),
+            (
+                ColumnQualifierTuple("b", "db1.table1"),
+                ColumnQualifierTuple("b", "db1.foo"),
+            ),
+            (
+                ColumnQualifierTuple("c", "db1.table1"),
+                ColumnQualifierTuple("c", "db1.foo"),
+            ),
+        ],
+        metadata_provider=provider,
+    )