diff options
Diffstat (limited to 'pgcli/pgexecute.py')
| -rw-r--r-- | pgcli/pgexecute.py | 58 |
1 files changed, 40 insertions, 18 deletions
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 8f2968d..497d681 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""" @@ -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) |