diff --git a/AUTHORS b/AUTHORS index d0b4a5cee..ea45b833b 100644 --- a/AUTHORS +++ b/AUTHORS @@ -158,6 +158,7 @@ Contributors: * Chris (ChrisJr404) * Pieter Ouwerkerk (pouwerkerk) * VXNCXNX + * Anand Hegde (anandghegde) * Wei-Lun Li (b84330808) * jackwalkerlabs (Codex-assisted) diff --git a/changelog.rst b/changelog.rst index fc362092b..a2695d292 100644 --- a/changelog.rst +++ b/changelog.rst @@ -10,6 +10,13 @@ Features: Bug fixes: ---------- +* 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)). Outside a + transaction, a ``COPY ... TO STDOUT`` is cancelled first, so the error does + not wait for the whole result to be read. * Display database errors from ``\ef`` and ``\ev`` without exiting the interactive session, allowing recovery from an aborted transaction. * Fix the development setup instructions to activate the ``.venv`` directory diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py index 1cac82eb4..e7fabdf29 100644 --- a/pgcli/pgexecute.py +++ b/pgcli/pgexecute.py @@ -471,7 +471,18 @@ def handle_notices(n): return title, None, None, res.command_status.decode() cur = self.conn.cursor() - cur.execute(split_sql) + # Read before execute(): while a COPY runs the status is ACTIVE either way. + outside_transaction = self.conn.info.transaction_status == psycopg.pq.TransactionStatus.IDLE + 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(cancel=outside_transaction) + 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 +496,33 @@ def handle_notices(n): _logger.debug("No rows in result.") return title, None, None, cur.statusmessage + def _end_copy(self, cancel): + """Finish a COPY that was started by ``cursor.execute()``. + + ``cancel`` asks the server to stop a COPY TO STDOUT instead of sending + all of its rows. It must be false inside a transaction: the cancel would + abort the transaction, and the work done in it would be lost. + """ + 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. Without a + # cancel this takes as long as reading the whole result. + if cancel: + self.conn.cancel() + 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..966c9a103 100644 --- a/tests/test_pgexecute.py +++ b/tests/test_pgexecute.py @@ -644,6 +644,96 @@ 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 +def test_copy_stdout_is_cancelled_outside_a_transaction(executor, exception_formatter, monkeypatch): + # Without the cancel, the error waits until every row has been read and discarded. + cancels = [] + cancel = executor.conn.cancel + monkeypatch.setattr(executor.conn, "cancel", lambda: cancels.append(1) or cancel()) + try: + result = run(executor, "copy (select generate_series(1, 1000000)) to stdout", exception_formatter=exception_formatter) + assert "use \\copy instead" in result[-1] + assert cancels == [1] + assert run(executor, "select 1", join=True).endswith("SELECT 1") + finally: + executor.conn.close() + + +@dbtest +def test_copy_stdout_in_a_transaction_keeps_its_work(executor, exception_formatter, monkeypatch): + # A cancel would abort the transaction, so the commit would roll back the insert. + cancels = [] + cancel = executor.conn.cancel + monkeypatch.setattr(executor.conn, "cancel", lambda: cancels.append(1) or cancel()) + try: + run(executor, "create table test(a text)") + run(executor, "begin") + run(executor, "insert into test values ('kept')") + run(executor, "copy (select generate_series(1, 1000)) to stdout", exception_formatter=exception_formatter) + assert cancels == [] + run(executor, "commit") + assert run(executor, "select a from test", join=True).endswith("SELECT 1") + 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 $$;"