Merge pull request #1363 from ERYoung11/fix_comments_plus_special

Fix comments plus special - Issue #1362

changelog.rst

@@ -5,6 +5,7 @@ Bug fixes:
 ----------
 
 * Fix exception when retrieving password from keyring ([issue 1338](https://github.com/dbcli/pgcli/issues/1338)).
+* Fix using comments with special commands ([issue 1362](https://github.com/dbcli/pgcli/issues/1362)).
 
 Internal:
 ---------

pgcli/pgexecute.py

@@ -306,8 +306,27 @@ class PGExecute:
         if not statement:  # Empty string
             yield None, None, None, None, statement, False, False
 
-        # Split the sql into separate queries and run each one.
-        for sql in sqlparse.split(statement):
+        # sql parse doesn't split on a comment first + special
+        # so we're going to do it
+
+        sqltemp = []
+        sqlarr = []
+
+        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)
+
+        # run each sql query
+        for sql in sqlarr:
             # Remove spaces, eol and semi-colons.
             sql = sql.rstrip(";")
             sql = sqlparse.format(sql, strip_comments=False).strip()

tests/test_pgexecute.py

@@ -282,6 +282,77 @@ def test_execute_from_file_io_error(os, executor, pgspecial):
     assert is_special == True
 
 
+@dbtest
+def test_execute_from_commented_file_that_executes_another_file(
+    executor, pgspecial, tmpdir
+):
+    # https://github.com/dbcli/pgcli/issues/1336
+    sqlfile1 = tmpdir.join("test01.sql")
+    sqlfile1.write("-- asdf \n\\h")
+    sqlfile2 = tmpdir.join("test00.sql")
+    sqlfile2.write("--An useless comment;\nselect now();\n-- another useless comment")
+
+    rcfile = str(tmpdir.join("rcfile"))
+    print(rcfile)
+    cli = PGCli(pgexecute=executor, pgclirc_file=rcfile)
+    assert cli != None
+    statement = "--comment\n\\h"
+    result = run(executor, statement, pgspecial=cli.pgspecial)
+    assert result != None
+    assert result[0].find("ALTER TABLE")
+
+
+@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
+    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
+
+    statement = "/*comment\ncomment line2*/\nselect now();"
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[1].find("now") >= 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 = "/*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 = "    /*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 = "/*comment\ncomment line2*/\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\ncomment line2*/\n\h;"
+    result = run(executor, statement, pgspecial=pgspecial)
+    assert result != None
+    assert result[1].find("ALTER") >= 0
+    assert result[1].find("ABORT") >= 0
+
+
 @dbtest
 def test_multiple_queries_same_line(executor):
     result = run(executor, "select 'foo'; select 'bar'")