fixed #1403, improved comment handling (#1404)

* fixed #1403, improved comment handling

* black + hooks + changelog

---------

Co-authored-by: Irina Truong <i.chernyavska@gmail.com>

changelog.rst

@@ -31,6 +31,7 @@ Bug fixes:
 * Allow specifying an `alias_map_file` in the config that will use
   predetermined table aliases instead of generating aliases programmatically on
   the fly
+* Fixed SQL error when there is a comment on the first line: ([issue 1403](https://github.com/dbcli/pgcli/issues/1403))
 
 Internal:
 ---------

pgcli/pgexecute.py

@@ -1,7 +1,7 @@
 import logging
 import traceback
 from collections import namedtuple
-
+import re
 import pgspecial as special
 import psycopg
 import psycopg.sql
@@ -17,6 +17,27 @@ ViewDef = namedtuple(
 )
 
 
+# we added this funcion to strip beginning comments
+# because sqlparse didn't handle tem well.  It won't be needed if sqlparse
+# does parsing of this situation better
+
+
+def remove_beginning_comments(command):
+    # Regular expression pattern to match comments
+    pattern = r"^(/\*.*?\*/|--.*?)(?:\n|$)"
+
+    # Find and remove all comments from the beginning
+    cleaned_command = command
+    comments = []
+    match = re.match(pattern, cleaned_command, re.DOTALL)
+    while match:
+        comments.append(match.group())
+        cleaned_command = cleaned_command[len(match.group()) :].lstrip()
+        match = re.match(pattern, cleaned_command, re.DOTALL)
+
+    return [cleaned_command, comments]
+
+
 def register_typecasters(connection):
     """Casts date and timestamp values to string, resolves issues with out-of-range
     dates (e.g. BC) which psycopg can't handle"""
@@ -311,21 +332,20 @@ class PGExecute:
         # sql parse doesn't split on a comment first + special
         # so we're going to do it
 
-        sqltemp = []
+        removed_comments = []
         sqlarr = []
+        cleaned_command = ""
 
-        if statement.startswith("--"):
-            sqltemp = statement.split("\n")
-            sqlarr.append(sqltemp[0])
-            for i in sqlparse.split(sqltemp[1]):
-                sqlarr.append(i)
-        elif statement.startswith("/*"):
-            sqltemp = statement.split("*/")
-            sqltemp[0] = sqltemp[0] + "*/"
-            for i in sqlparse.split(sqltemp[1]):
-                sqlarr.append(i)
-        else:
-            sqlarr = sqlparse.split(statement)
+        # could skip if statement doesn't match ^-- or ^/*
+        cleaned_command, removed_comments = remove_beginning_comments(statement)
+
+        sqlarr = sqlparse.split(cleaned_command)
+
+        # now re-add the beginning comments if there are any, so that they show up in
+        # log files etc when running these commands
+
+        if len(removed_comments) > 0:
+            sqlarr = removed_comments + sqlarr
 
         # run each sql query
         for sql in sqlarr:

tests/test_pgexecute.py

@@ -304,9 +304,7 @@ def test_execute_from_commented_file_that_executes_another_file(
 
 @dbtest
 def test_execute_commented_first_line_and_special(executor, pgspecial, tmpdir):
-    # https://github.com/dbcli/pgcli/issues/1362
-
-    # just some base caes that should work also
+    # just some base cases that should work also
     statement = "--comment\nselect now();"
     result = run(executor, statement, pgspecial=pgspecial)
     assert result != None
@@ -317,23 +315,43 @@ def test_execute_commented_first_line_and_special(executor, pgspecial, tmpdir):
     assert result != None
     assert result[1].find("now") >= 0
 
-    statement = "/*comment\ncomment line2*/\nselect now();"
-    result = run(executor, statement, pgspecial=pgspecial)
-    assert result != None
-    assert result[1].find("now") >= 0
-
+    # https://github.com/dbcli/pgcli/issues/1362
     statement = "--comment\n\\h"
     result = run(executor, statement, pgspecial=pgspecial)
     assert result != None
     assert result[1].find("ALTER") >= 0
     assert result[1].find("ABORT") >= 0
 
+    statement = "--comment1\n--comment2\n\\h"
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[1].find("ALTER") >= 0
+    assert result[1].find("ABORT") >= 0
+
     statement = "/*comment*/\n\h;"
     result = run(executor, statement, pgspecial=pgspecial)
     assert result != None
     assert result[1].find("ALTER") >= 0
     assert result[1].find("ABORT") >= 0
 
+    statement = """/*comment1
+    comment2*/
+    \h"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[1].find("ALTER") >= 0
+    assert result[1].find("ABORT") >= 0
+
+    statement = """/*comment1
+    comment2*/
+    /*comment 3
+    comment4*/
+    \\h"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[1].find("ALTER") >= 0
+    assert result[1].find("ABORT") >= 0
+
     statement = "    /*comment*/\n\h;"
     result = run(executor, statement, pgspecial=pgspecial)
     assert result != None
@@ -352,6 +370,126 @@ def test_execute_commented_first_line_and_special(executor, pgspecial, tmpdir):
     assert result[1].find("ALTER") >= 0
     assert result[1].find("ABORT") >= 0
 
+    statement = """\\h /*comment4 */"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    print(result)
+    assert result != None
+    assert result[0].find("No help") >= 0
+
+    # TODO: we probably don't want to do this but sqlparse is not parsing things well
+    # we relly want it to find help but right now, sqlparse isn't dropping the /*comment*/
+    # style comments after command
+
+    statement = """/*comment1*/
+    \h
+    /*comment4 */"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[0].find("No help") >= 0
+
+    # TODO: same for this one
+    statement = """/*comment1
+    comment3
+    comment2*/
+    \\h
+    /*comment4
+    comment5
+    comment6*/"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[0].find("No help") >= 0
+
+
+@dbtest
+def test_execute_commented_first_line_and_normal(executor, pgspecial, tmpdir):
+    # https://github.com/dbcli/pgcli/issues/1403
+
+    # just some base cases that should work also
+    statement = "--comment\nselect now();"
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[1].find("now") >= 0
+
+    statement = "/*comment*/\nselect now();"
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[1].find("now") >= 0
+
+    # this simulates the original error (1403) without having to add/drop tables
+    # since it was just an error on reading input files and not the actual
+    # command itself
+
+    # test that the statement works
+    statement = """VALUES (1, 'one'), (2, 'two'), (3, 'three');"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # test the statement with a \n in the middle
+    statement = """VALUES (1, 'one'),\n (2, 'two'), (3, 'three');"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # test the statement with a newline in the middle
+    statement = """VALUES (1, 'one'),
+     (2, 'two'), (3, 'three');"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # now add a single comment line
+    statement = """--comment\nVALUES (1, 'one'), (2, 'two'), (3, 'three');"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # doing without special char \n
+    statement = """--comment
+    VALUES (1,'one'),
+    (2, 'two'), (3, 'three');"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # two comment lines
+    statement = """--comment\n--comment2\nVALUES (1,'one'), (2, 'two'), (3, 'three');"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # doing without special char \n
+    statement = """--comment
+    --comment2
+    VALUES (1,'one'), (2, 'two'), (3, 'three');
+    """
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # multiline comment + newline in middle of the statement
+    statement = """/*comment
+comment2
+comment3*/
+VALUES (1,'one'),
+(2, 'two'), (3, 'three');"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
+    # multiline comment + newline in middle of the statement
+    # + comments after the statement
+    statement = """/*comment
+comment2
+comment3*/
+VALUES (1,'one'),
+(2, 'two'), (3, 'three');
+--comment4
+--comment5"""
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[5].find("three") >= 0
+
 
 @dbtest
 def test_multiple_queries_same_line(executor):