diff options
| -rw-r--r-- | changelog.rst | 1 | ||||
| -rw-r--r-- | pgcli/pgexecute.py | 48 | ||||
| -rw-r--r-- | tests/test_pgexecute.py | 150 |
3 files changed, 179 insertions, 20 deletions
diff --git a/changelog.rst b/changelog.rst index 94aa040c..5825675a 100644 --- a/changelog.rst +++ b/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: --------- diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 2addb738..497d6816 100644 --- a/pgcli/pgexecute.py +++ b/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: diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py index b6ee0abe..636795ba 100644 --- a/tests/test_pgexecute.py +++ b/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,12 +315,14 @@ 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();" + # https://github.com/dbcli/pgcli/issues/1362 + statement = "--comment\n\\h" result = run(executor, statement, pgspecial=pgspecial) assert result != None - assert result[1].find("now") >= 0 + assert result[1].find("ALTER") >= 0 + assert result[1].find("ABORT") >= 0 - statement = "--comment\n\\h" + statement = "--comment1\n--comment2\n\\h" result = run(executor, statement, pgspecial=pgspecial) assert result != None assert result[1].find("ALTER") >= 0 @@ -334,6 +334,24 @@ def test_execute_commented_first_line_and_special(executor, pgspecial, tmpdir): 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): |