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:

Resolving patch:
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
--- 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
 


Test patch:
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -487,3 +487,11 @@ def test_trim(self) -> None:
         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")