test(crdb): add tests to expose canceled copy problem

See https://github.com/cockroachdb/cockroach/issues/81559

diff --git a/tests/crdb/test_copy.py b/tests/crdb/test_copy.py
index 2c20f2fdec1d2305225a68031a480cb8a7815e4e..21eb78b602cb02d548a344d6f1927f0a37e2dd0d 100644 (file)
--- a/tests/crdb/test_copy.py
+++ b/tests/crdb/test_copy.py
@@ -155,6 +155,19 @@ def test_copy_in_records_binary(conn, format):
     assert data == [(None, "hello"), (None, "world")]
 
 
+@pytest.mark.crdb_skip("copy canceled")
+def test_copy_in_buffers_with_py_error(conn):
+    cur = conn.cursor()
+    ensure_table(cur, sample_tabledef)
+    with pytest.raises(e.QueryCanceled) as exc:
+        with cur.copy("copy copy_in from stdin") as copy:
+            copy.write(sample_text)
+            raise Exception("nuttengoggenio")
+
+    assert "nuttengoggenio" in str(exc.value)
+    assert conn.info.transaction_status == conn.TransactionStatus.INERROR
+
+
 def test_copy_in_allchars(conn):
     cur = conn.cursor()
     ensure_table(cur, "col1 int primary key, col2 int, data text")

diff --git a/tests/crdb/test_copy_async.py b/tests/crdb/test_copy_async.py
index 464277c369b12b5575ec311d621d247f4933cf18..7642b56354f58bafa1bd7d5279f028eeceaf11da 100644 (file)
--- a/tests/crdb/test_copy_async.py
+++ b/tests/crdb/test_copy_async.py
@@ -160,6 +160,19 @@ async def test_copy_in_records_binary(aconn, format):
     assert data == [(None, "hello"), (None, "world")]
 
 
+@pytest.mark.crdb_skip("copy canceled")
+async def test_copy_in_buffers_with_py_error(aconn):
+    cur = aconn.cursor()
+    await ensure_table(cur, sample_tabledef)
+    with pytest.raises(e.QueryCanceled) as exc:
+        async with cur.copy("copy copy_in from stdin") as copy:
+            await copy.write(sample_text)
+            raise Exception("nuttengoggenio")
+
+    assert "nuttengoggenio" in str(exc.value)
+    assert aconn.info.transaction_status == aconn.TransactionStatus.INERROR
+
+
 async def test_copy_in_allchars(aconn):
     cur = aconn.cursor()
     await ensure_table(cur, "col1 int primary key, col2 int, data text")