Getting the underlying exception in `copy()`

[bluetech]

Problem

When an exception is raised from the user python code during copy(), the traceback looks like this:

Traceback (most recent call last):
    <snipped>
requests.exceptions.ConnectionError: ('Connection aborted.', RemoteDisconnected('Remote end closed connection without response'))

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "venv/lib/python3.11/site-packages/psycopg/cursor.py", line 900, in copy
    yield copy
  File "myfile.py", line 86, in myfunc
    raise MyException() from e
MyException

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "myotherfile.py", line 1251, in myotherfunc
    with cursor.copy(sql.SQL('''
  File "/usr/lib/python3.11/contextlib.py", line 155, in __exit__
    self.gen.throw(typ, value, traceback)
  File "venv/lib/python3.11/site-packages/psycopg/cursor.py", line 902, in copy
    raise ex.with_traceback(None)
psycopg.errors.QueryCanceled: COPY from stdin failed: error from Python: MyException - 

If I want to inspect the underlying exception, I currently need to check __context__. But __context__ should be reserved for unintended bugs in error handling code. This makes it hard to distinguish real underlying errors from exceptions unintentionally raised in psycopg's exception handling code.

Suggested solution

Change to explicit chaining, that is raise QueryCanceled() from underlying_exc instead of raise without from in the exception handling context.

This will change the traceback message to say "The above exception was the direct cause of the following exception". Also will change __context__ to None and make __cause__ have the exception instead.

There are linters that can help with detecting this more holistically, for example flake8-bugbear B904, though I didn't check if it will detect this case.

[dvarrazzo]

The code path I believe you refer to is the one where Python, on exception, sends a client-side error message to the server. The server closes the copy operation and sends back a QueryCanceled error with the message "proposed" by Python, but with a different traceback and error class. If you think the current implementation can be improved, a Pull Request is more than welcome. I apologize if I have misunderstood something in your report: I am, and will be, out of home without computer access for a few days.

[dvarrazzo]

I have taken a look, and actually the best option here is to just ignore the QueryCanceled and re-raise the original exception.

In a block such as:

with cur.copy("COPY ... FROM STDIN"):
    1 / 0  # unexpected exception

the ZeroDivisionError is the most important thing a developer would want to see reported. What psycopg does is to send the server a message to stop the COPY operation with error. This result in the server cancelling the operation and returning a QueryCanceled error message, that, so far, became the exception raised.

However, the QueryCanceled is part of the normal flow by now, not an unexpected exception. We do copy the exception string, which is repeated back as the QueryCanceled message (the "error from Python: {exc}" message) but we have lost the type, location, traceback of the message.

With the simple change proposed in #609, the Python exception is now re-raised, after we verify that we have received a QueryCanceled. This results in that the test script:

# test_593.py

import psycopg

def main() -> None:
    conn = psycopg.connect()
    cur = conn.cursor()
    cur.execute("create table if not exists copy_in (id serial primary key)")
    with cur.copy("copy copy_in from stdin (format text)"):
        1 / 0

if __name__ == "__main__":
    main()

instead of raising the chained exception reported here (the chaining is an accident: pytest swallows it so I didn't notice):

python test_593.py 
Traceback (most recent call last):
  File "/home/piro/dev/psycopg3/psycopg/build/__editable__.psycopg-3.1.10.dev1-py3-none-any/psycopg/cursor.py", line 900, in copy
    yield copy
  File "/home/piro/dev/psycopg3/test_593.py", line 10, in main
    1 / 0
ZeroDivisionError: division by zero

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/piro/dev/psycopg3/test_593.py", line 13, in <module>
    main()
  File "/home/piro/dev/psycopg3/test_593.py", line 9, in main
    with cur.copy("copy copy_in from stdin (format text)"):
  File "/usr/lib/python3.10/contextlib.py", line 153, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/home/piro/dev/psycopg3/psycopg/build/__editable__.psycopg-3.1.10.dev1-py3-none-any/psycopg/cursor.py", line 902, in copy
    raise ex.with_traceback(None)
psycopg.errors.QueryCanceled: COPY from stdin failed: error from Python: ZeroDivisionError - division by zero
CONTEXT:  COPY copy_in, line 1

to raise just the Python exception:

$ python test_593.py 
Traceback (most recent call last):
  File "/home/piro/dev/psycopg3/test_593.py", line 13, in <module>
    main()
  File "/home/piro/dev/psycopg3/test_593.py", line 10, in main
    1 / 0
ZeroDivisionError: division by zero

Does it sound good?

[dlax]

Isn't there any value, for the user, to be informed that the query got cancelled?
If so, the proposed solution (OP) seems interesting to me; meaning the following change (not a BC by the way):

diff --git a/psycopg/psycopg/copy.py b/psycopg/psycopg/copy.py
index 26a2d9e9..25ddef38 100644
--- a/psycopg/psycopg/copy.py
+++ b/psycopg/psycopg/copy.py
@@ -318,7 +318,13 @@ class Copy(BaseCopy["Connection[Any]"]):
             data = self.formatter.end()
             if data:
                 self._write(data)
-            self.writer.finish(exc)
+            if exc:
+                try:
+                    self.writer.finish(exc)
+                except e.QueryCanceled as qce:
+                    raise qce from exc
+            else:
+                self.writer.finish(exc)
             self._finished = True
         else:

which behaves as:

$ python test_593.py 
Traceback (most recent call last):
  File "/home/denis/src/psycopg3/psycopg/build/__editable__.psycopg-3.2.0.dev1-py3-none-any/psycopg/cursor.py", line 884, in copy
    yield copy
  File "/home/denis/src/psycopg3/test_593.py", line 10, in main
    1 / 0
    ~~^~~
ZeroDivisionError: division by zero

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/denis/src/psycopg3/test_593.py", line 13, in <module>
    main()
  File "/home/denis/src/psycopg3/test_593.py", line 9, in main
    with cur.copy("copy copy_in from stdin (format text)"):
  File "/usr/lib/python3.11/contextlib.py", line 155, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/home/denis/src/psycopg3/psycopg/build/__editable__.psycopg-3.2.0.dev1-py3-none-any/psycopg/cursor.py", line 886, in copy
    raise ex.with_traceback(None)
psycopg.errors.QueryCanceled: COPY from stdin failed: error from Python: ZeroDivisionError - division by zero
CONTEXT:  COPY copy_in, line 1

---

More generally, it's not clear to me if/why a QueryCanceled would always be raised; unless I'm missing something the libpq is not explicit in this respect?

[dvarrazzo]

I think QueryCanceled should be always raised, because it comes as an effect of using PQputCopyEnd() with a non-null error message. It should be more reliable than throwing a cancel to the server and expecting a QueryCanceled back, which is something we do in normal querying:

psycopg/psycopg/psycopg/connection.py

Lines 958 to 967 in 1f56776

	except KeyboardInterrupt:
	# On Ctrl-C, try to cancel the query in the server, otherwise
	# the connection will remain stuck in ACTIVE state.
	c = self.pgconn.get_cancel()
	c.cancel()
	try:
	waiting.wait(gen, self.pgconn.socket, timeout=timeout)
	except e.QueryCanceled:
	pass # as expected
	raise

and in COPY TO:

psycopg/psycopg/psycopg/copy.py

Lines 193 to 202 in 1f56776

	# Throw a cancel to the server, then consume the rest of the copy data
	# (which might or might not have been already transferred entirely to
	# the client, so we won't necessary see the exception associated with
	# canceling).
	self.connection.cancel()
	try:
	while (yield from self._read_gen()):
	pass
	except e.QueryCanceled:
	pass

In both these cases we consider the original exception more important than the QueryCanceled.

For me, giving the possibility to write the following code makes sense:

import time
import psycopg

conn = psycopg.connect()
cur = conn.cursor()
cur.execute("create table if not exists copy_in (id serial primary key)")
try:
    with cur.copy("copy copy_in from stdin (format text)"):
        time.sleep(10)
except KeyboardInterrupt:
    print("ctrl-c")

this only works if the main exception raised is the KeyboardInterrupt. But, if we chain it to the QueryCanceled error, the except doesn't work anymore. I think an user would often want to write error handlers for "control exceptions" or exceptions that they know Python may raise during copy (this is what the OP is also trying to catch, I understand). I doubt they would be interested in catching a QueryCanceled itself - at most they might want to catch any OperationalError the database might throw at them.

If any, modifying your example, I would raise exc from qce, which makes the above except work and, if uncaught, results in (after getting rid of the traceback leading to wait):

$ python test_593.py 
^C
  File "/home/piro/dev/psycopg3/psycopg/build/__editable__.psycopg-3.1.10.dev1-py3-none-any/psycopg/copy.py", line 326, in finish
    raise ex.with_traceback(None)
psycopg.errors.QueryCanceled: COPY from stdin failed: error from Python: KeyboardInterrupt - 
CONTEXT:  COPY copy_in, line 1

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/piro/dev/psycopg3/test_593.py", line 17, in <module>
    main()
  File "/home/piro/dev/psycopg3/test_593.py", line 12, in main
    time.sleep(10)
KeyboardInterrupt

I still think that the detail of the QueryCanceled is not necessary, and, if it is, we should show it consistently in normal querying and in COPY TO as well.

[dlax]

Ok; makes sense. Thank you for those details!

[dvarrazzo]

Ok, merging! Thank you very much for your review, @dlax.

@bluetech: after this change is merged, the exception in your example will be raised simply as MyException and can be caught as such outside the copy block. I assume that, in a transition period before 3.1.10 is released, you might want to catch both the exception, although you will still need the clumsy workaround you described if you get a QueryCanceled.

[bluetech]

@dvarrazzo Sorry I haven't replied earlier, I haven't found the time for a proper answer.

The merged solution is good. I think it's pretty intuitive that an exception raised within the with would simply propagate through the with; I'd say that's the expected behavior unless said otherwise. In our use cases we have no need for the exception to be wrapped in a QueryCanceled.

Thank you for working on it!