SNOW-733718 Fix remove_comments when comments are at end of line

src/snowflake/connector/util_text.py

@@ -146,6 +146,8 @@ def split_statements(
                     if not remove_comments:
                         # keep the comment
                         statement.append((line[col:], False))
+                    else:
+                        statement.append(("\n", True))
                     col = len_line + 1
                     col0 = col
                 elif line[col:].startswith("/*") and not line[col0:].startswith(

test/unit/test_split_statement.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 from io import StringIO
+from textwrap import dedent
 
 import pytest
 
@@ -499,6 +500,31 @@ def test_comments_with_semicolon():
             next(itr)
 
 
+@pytest.mark.skipif(split_statements is None, reason="No split_statements is available")
+def test_remove_comments_end_of_line():
+    with StringIO(
+        dedent(
+            """\
+           select L_ORDERKEY,
+           L_SHIPINSTRUCT-- this is a comment
+           from lineitem
+           where L_ORDERKEY=5795896006;
+           """
+        )
+    ) as f:
+        itr = split_statements(f, remove_comments=True)
+        expected_sql = dedent(
+            """\
+        select L_ORDERKEY,
+        L_SHIPINSTRUCT
+        from lineitem
+        where L_ORDERKEY=5795896006;"""
+        )
+        assert next(itr) == (expected_sql, False)
+        with pytest.raises(StopIteration):
+            next(itr)
+
+
 @pytest.mark.skipif(split_statements is None, reason="No split_statements is available")
 def test_comment_in_values():
     """SNOW-51297: SnowSQL -o remove_comments=True breaks the query."""