Repository: tobymao/sqlglot

Issue:
lineage returns sql comments as part of node name
Having a snowflake query, sqlglot lineage parser returns nearby sql comments as part of node name - is this expected behaviour ?.
Below a test case and screen shot:

###  pip3 install sqlglot
import sys
import os

root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(root_dir)
current_file_path = os.path.abspath(__file__)
current_filename = os.path.basename(current_file_path)
print("** Testfile:", current_filename, '**')

import unittest
import sqlglot
from sqlglot.lineage import lineage
from sqlglot import  exp

class Test_SQLGLOT_COMMENT_BUG(unittest.TestCase):
    sql_content = """        
        CREATE TRANSIENT TABLE
            AAA.temp_table_123 DATA_RETENTION_TIME_IN_DAYS = 0
        AS
        SELECT DISTINCT
            *
        FROM (
            SELECT
                -- NO PROBLEM WITH THIS SQL COMMENT --
                MD5(NVL(UPPER(TRIM(some_column_B)), '-1') || '~' || NVL(UPPER(TRIM(another_column_A)), '-1') || '~' || NVL(UPPER(TRIM(comment_error_here)), '-1') || '~' || NVL(UPPER(TRIM(id)), '-1')) AS no_error_id
                -- CUSTOM SQL END --
                , xx_yy_number
                , some_ID_FIELD
                , some_column_B
                , another_column_A
                , comment_error_here
                -- ERROR SQL COMMENT START ERROR HERE HERE HERE HERE !!!! ---------------------------------------------------------------
                , MD5(NVL(UPPER(TRIM(some_column_B)), '-1') || '~' || NVL(UPPER(TRIM(another_column_A)), '-1') || '~' || NVL(UPPER(TRIM(comment_error_here)), '-1') || '~' || NVL(UPPER(TRIM(id)), '-1')) AS some_hash
                -- ERROR SQL COMMENT END --
                , 1 AS status
                , extra_column
            FROM
                domain.some_table_x x_table
            WHERE
                some_ID_FIELD IN (10,11,12,13,14)
            QUALIFY
                xx_yy_number = MAX(xx_yy_number) OVER (PARTITION BY some_column_B, another_column_A, comment_error_here)
        ) src
        WHERE
            NOT EXISTS (
                SELECT
                    1
                FROM
                    AAA.abc_dfg exists_relation
                WHERE
                    exists_relation.no_error_id = src.no_error_id
                    AND exists_relation.status = 1
            );
    """

    def test_bug_demonstration(self):
        dialect = "snowflake"
        ast = sqlglot.parse_one(self.sql_content, read=dialect)
        sql = ast.sql()
        for select in ast.find_all(exp.Select):
            for projection in select.expressions:
                column_lineage = {}
                column_name = projection.output_name
                try:
                    column_lineage = lineage((column_name), sql, dialect=dialect, trim_selects=False)
                except Exception as exception:
                    print(f'column_name: <' + column_name + f'> not parseable by lineage(): {exception}')
                if hasattr(column_lineage, 'downstream'):
                    for node in column_lineage.downstream:
                        if hasattr(node, 'downstream'):
                            if len(node.downstream) > 0:
                                print("node name:","<",str(node.downstream[0].name),">")



if __name__ == '__main__':
    unittest.main()

![image](https://github.com/tobymao/sqlglot/assets/14150443/47564e0c-7895-4908-b474-745b9b8ade8a)



Discussion:

Patch:
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
index 5c28e97d..7f9a624c 100644
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -229,7 +229,9 @@ def to_node(
         for source in scope.sources.values():
             if isinstance(source, Scope):
                 source = source.expression
-            node.downstream.append(Node(name=select.sql(), source=source, expression=source))
+            node.downstream.append(
+                Node(name=select.sql(comments=False), source=source, expression=source)
+            )
 
     # Find all columns that went into creating this one to list their lineage nodes.
     source_columns = set(find_all_in_scope(select, exp.Column))
@@ -278,7 +280,9 @@ def to_node(
             # it means this column's lineage is unknown. This can happen if the definition of a source used in a query
             # is not passed into the `sources` map.
             source = source or exp.Placeholder()
-            node.downstream.append(Node(name=c.sql(), source=source, expression=source))
+            node.downstream.append(
+                Node(name=c.sql(comments=False), source=source, expression=source)
+            )
 
     return node
 
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index b0a25167..51022bbe 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -5974,7 +5974,7 @@ class Parser(metaclass=_Parser):
         if kind in ("GLOBAL", "SESSION") and self._match_text_seq("TRANSACTION"):
             return self._parse_set_transaction(global_=kind == "GLOBAL")
 
-        left = self._parse_primary() or self._parse_id_var()
+        left = self._parse_primary() or self._parse_column()
         assignment_delimiter = self._match_texts(("=", "TO"))
 
         if not left or (self.SET_REQUIRES_ASSIGNMENT_DELIMITER and not assignment_delimiter):


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index dae83557..97b83ab9 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -10,14 +10,6 @@ class TestSnowflake(Validator):
     dialect = "snowflake"
 
     def test_snowflake(self):
-        self.validate_identity(
-            "MERGE INTO my_db AS ids USING (SELECT new_id FROM my_model WHERE NOT col IS NULL) AS new_ids ON ids.type = new_ids.type AND ids.source = new_ids.source WHEN NOT MATCHED THEN INSERT VALUES (new_ids.new_id)"
-        )
-        self.validate_identity("ALTER TABLE table1 CLUSTER BY (name DESC)")
-        self.validate_identity(
-            "INSERT OVERWRITE TABLE t SELECT 1", "INSERT OVERWRITE INTO t SELECT 1"
-        )
-        self.validate_identity("SELECT rename, replace")
         expr = parse_one("SELECT APPROX_TOP_K(C4, 3, 5) FROM t")
         expr.selects[0].assert_is(exp.AggFunc)
         self.assertEqual(expr.sql(dialect="snowflake"), "SELECT APPROX_TOP_K(C4, 3, 5) FROM t")
@@ -43,6 +35,9 @@ WHERE
   )""",
         )
 
+        self.validate_identity("ALTER TABLE table1 CLUSTER BY (name DESC)")
+        self.validate_identity("SELECT rename, replace")
+        self.validate_identity("ALTER TABLE table1 SET TAG foo.bar = 'baz'")
         self.validate_identity("SELECT TIMEADD(HOUR, 2, CAST('09:05:03' AS TIME))")
         self.validate_identity("SELECT CAST(OBJECT_CONSTRUCT('a', 1) AS MAP(VARCHAR, INT))")
         self.validate_identity("SELECT CAST(OBJECT_CONSTRUCT('a', 1) AS OBJECT(a CHAR NOT NULL))")
@@ -95,6 +90,12 @@ WHERE
         self.validate_identity("SELECT CONVERT_TIMEZONE('UTC', 'America/Los_Angeles', col)")
         self.validate_identity("ALTER TABLE a SWAP WITH b")
         self.validate_identity("SELECT MATCH_CONDITION")
+        self.validate_identity(
+            "MERGE INTO my_db AS ids USING (SELECT new_id FROM my_model WHERE NOT col IS NULL) AS new_ids ON ids.type = new_ids.type AND ids.source = new_ids.source WHEN NOT MATCHED THEN INSERT VALUES (new_ids.new_id)"
+        )
+        self.validate_identity(
+            "INSERT OVERWRITE TABLE t SELECT 1", "INSERT OVERWRITE INTO t SELECT 1"
+        )
         self.validate_identity(
             'DESCRIBE TABLE "SNOWFLAKE_SAMPLE_DATA"."TPCDS_SF100TCL"."WEB_SITE" type=stage'
         )
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index 3e17f956..036f1467 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -487,3 +487,11 @@ class TestLineage(unittest.TestCase):
         downstream = node.downstream[0]
         self.assertEqual(downstream.name, "z.a")
         self.assertEqual(downstream.source.sql(), "SELECT y.a AS a, y.b AS b, y.c AS c FROM y AS y")
+
+    def test_node_name_doesnt_contain_comment(self) -> None:
+        sql = "SELECT * FROM (SELECT x /* c */ FROM t1) AS t2"
+        node = lineage("x", sql)
+
+        self.assertEqual(len(node.downstream), 1)
+        self.assertEqual(len(node.downstream[0].downstream), 1)
+        self.assertEqual(node.downstream[0].downstream[0].name, "t1.x")
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 6bcdb643..1b9639a7 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -503,7 +503,7 @@ class TestParser(unittest.TestCase):
 
         self.assertIsInstance(set_item, exp.SetItem)
         self.assertIsInstance(set_item.this, exp.EQ)
-        self.assertIsInstance(set_item.this.this, exp.Identifier)
+        self.assertIsInstance(set_item.this.this, exp.Column)
         self.assertIsInstance(set_item.this.expression, exp.Literal)
 
         self.assertEqual(set_item.args.get("kind"), "SESSION")