Merging upstream version 4.0.1.

Signed-off-by: Daniel Baumann <daniel@debian.org>

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"""
@@ -76,7 +97,6 @@ class ProtocolSafeCursor(psycopg.Cursor):
 
 
 class PGExecute:
-
     # The boolean argument to the current_schemas function indicates whether
     # implicit schemas, e.g. pg_catalog
     search_path_query = """
@@ -180,7 +200,6 @@ class PGExecute:
         dsn=None,
         **kwargs,
     ):
-
         conn_params = self._conn_params.copy()
 
         new_params = {
@@ -203,7 +222,11 @@ class PGExecute:
 
         conn_params.update({k: v for k, v in new_params.items() if v})
 
-        conn_info = make_conninfo(**conn_params)
+        if "dsn" in conn_params:
+            other_params = {k: v for k, v in conn_params.items() if k != "dsn"}
+            conn_info = make_conninfo(conn_params["dsn"], **other_params)
+        else:
+            conn_info = make_conninfo(**conn_params)
         conn = psycopg.connect(conn_info)
         conn.cursor_factory = ProtocolSafeCursor
 
@@ -309,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:
@@ -470,7 +492,7 @@ class PGExecute:
             return (
                 psycopg.sql.SQL(template)
                 .format(
-                    name=psycopg.sql.Identifier(f"{result.nspname}.{result.relname}"),
+                    name=psycopg.sql.Identifier(result.nspname, result.relname),
                     stmt=psycopg.sql.SQL(result.viewdef),
                 )
                 .as_string(self.conn)