diff --git a/AUTHORS b/AUTHORS index 691396045..3f3b56db4 100644 --- a/AUTHORS +++ b/AUTHORS @@ -155,6 +155,7 @@ Contributors: * Diego * Chris (ChrisJr404) * Pieter Ouwerkerk (pouwerkerk) + * Anand Hegde (anandghegde) Creator: -------- diff --git a/changelog.rst b/changelog.rst index ee5d1fcf5..1226693fe 100644 --- a/changelog.rst +++ b/changelog.rst @@ -39,6 +39,11 @@ Bug fixes: ``sslmode``, everything) and silently fell back to a local socket connection as the OS user. The database argument is now kept, like psql; only when no database is given at all does the listing connect to ``postgres``. +* Fix ``COPY ... TO STDOUT`` and ``COPY ... FROM STDIN`` (instead of ``\copy``) + leaving the connection stuck: every later query failed with "another command + is already in progress" and quitting asked about an ongoing transaction. + The COPY is now ended cleanly and the error suggests ``\copy`` instead + ([issue 1505](https://github.com/dbcli/pgcli/issues/1505)). 4.6.0 (2026-08-26) ================== diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 1cac82eb4..0fa045409 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -471,7 +471,16 @@ def handle_notices(n): return title, None, None, res.command_status.decode() cur = self.conn.cursor() - cur.execute(split_sql) + try: + cur.execute(split_sql) + except psycopg.ProgrammingError as e: + # psycopg rejects COPY ... TO STDOUT / FROM STDIN only after the + # server has started the COPY, which leaves the connection busy: + # every later query fails and pgcli thinks a transaction is ongoing. + if self.conn.info.transaction_status != psycopg.pq.TransactionStatus.ACTIVE: + raise + self._end_copy() + raise psycopg.ProgrammingError("COPY to STDOUT or from STDIN is not supported, use \\copy instead") from e # cur.description will be None for operations that do not return # rows. @@ -485,6 +494,25 @@ def handle_notices(n): _logger.debug("No rows in result.") return title, None, None, cur.statusmessage + def _end_copy(self): + """Finish a COPY that was started by ``cursor.execute()``.""" + pgconn = self.conn.pgconn + nonblocking = pgconn.nonblocking + # Blocking mode, so libpq waits for the server instead of us polling. + pgconn.nonblocking = 0 + try: + try: + # COPY FROM STDIN: abort it without sending any data. + pgconn.put_copy_end(b"use \\copy instead") + except psycopg.OperationalError: + # COPY TO STDOUT: discard the rows sent by the server. + while pgconn.get_copy_data(0)[0] > 0: + pass + while pgconn.get_result() is not None: + pass + finally: + pgconn.nonblocking = nonblocking + def search_path(self): """Returns the current search path as a list of schema names""" diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py index c5fcaa2cd..6393b59dc 100644 --- a/tests/test_pgexecute.py +++ b/tests/test_pgexecute.py @@ -644,6 +644,63 @@ def test_on_error_stop(executor, exception_formatter): assert len(result) == 2 +@dbtest +@pytest.mark.parametrize( + "sql", + [ + "copy (select 1) to stdout", + "copy (select generate_series(1, 10000)) to stdout csv", + "copy test from stdin", + ], +) +def test_copy_stdout_stdin_does_not_break_connection(executor, exception_formatter, sql): + run(executor, "create table test(a text)") + try: + result = run(executor, sql, exception_formatter=exception_formatter) + assert "use \\copy instead" in result[-1] + assert not executor.valid_transaction() + assert run(executor, "select 42", join=True) == dedent( + """\ + +----------+ + | ?column? | + |----------| + | 42 | + +----------+ + SELECT 1""" + ) + finally: + # A connection stuck in COPY keeps its locks, which would block the + # fixture from dropping the tables. + executor.conn.close() + + +@dbtest +def test_copy_stdout_keeps_open_transaction(executor, exception_formatter): + try: + run(executor, "begin") + run(executor, "copy (select 1) to stdout", exception_formatter=exception_formatter) + assert executor.conn.info.transaction_status == psycopg.pq.TransactionStatus.INTRANS + run(executor, "rollback") + assert executor.conn.info.transaction_status == psycopg.pq.TransactionStatus.IDLE + finally: + executor.conn.close() + + +@dbtest +@pytest.mark.parametrize( + ("sql", "error"), + [ + ("select from", psycopg.errors.SyntaxError), + ("select * from no_such_table", psycopg.errors.UndefinedTable), + ], +) +def test_other_programming_errors_are_not_rewritten(executor, sql, error): + with pytest.raises(error) as excinfo: + list(executor.run(sql)) + assert "\\copy" not in str(excinfo.value) + assert run(executor, "select 1", join=True).endswith("SELECT 1") + + # @dbtest # def test_unicode_notices(executor): # sql = "DO language plpgsql $$ BEGIN RAISE NOTICE '有人更改'; END $$;"