fixed comments with special see Issue #1336

2 files changed, 89 insertions, 2 deletions

diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 758a3159..85a52c0d 100644
--- a/pgcli/pgexecute.py
+++ b/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()
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 7fef4658..05313c72 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -2,6 +2,7 @@ from textwrap import dedent
 
 import psycopg
 import pytest
+import re
 from unittest.mock import patch, MagicMock
 from pgspecial.main import PGSpecial, NO_QUERY
 from utils import run, dbtest, requires_json, requires_jsonb
@@ -281,6 +282,73 @@ def test_execute_from_file_io_error(os, executor, pgspecial):
     assert success == False
     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/1336
+
+    # 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):