test(crdb): add random copy test

Tests are currently skipped because of
https://github.com/cockroachdb/cockroach/issues/82792

diff --git a/tests/crdb/test_copy.py b/tests/crdb/test_copy.py
index 1dd7d32a61613eb2df1ed65613ecf4f61d5a9599..2c20f2fdec1d2305225a68031a480cb8a7815e4e 100644 (file)
--- a/tests/crdb/test_copy.py
+++ b/tests/crdb/test_copy.py
@@ -1,12 +1,14 @@
+import gc
 import pytest
 import string
 from random import randrange, choice
 
+from psycopg import sql, errors as e
 from psycopg.pq import Format
-from psycopg import errors as e
+from psycopg.adapt import PyFormat
 from psycopg.types.numeric import Int4
 
-from ..utils import eur
+from ..utils import eur, gc_collect
 from ..test_copy import sample_text, sample_binary  # noqa
 from ..test_copy import ensure_table, sample_records
 from ..test_copy import sample_tabledef as sample_tabledef_pg
@@ -171,8 +173,49 @@ from copy_in group by 1, 2, 3
     assert data == [(True, True, 1, 256)]
 
 
-def copyopt(format):
-    return "with binary" if format == Format.BINARY else ""
+@pytest.mark.slow
+@pytest.mark.parametrize(
+    "fmt, set_types",
+    [(Format.TEXT, True), (Format.TEXT, False), (Format.BINARY, True)],
+)
+@pytest.mark.crdb_skip("copy array")
+def test_copy_from_leaks(conn_cls, dsn, faker, fmt, set_types):
+    faker.format = PyFormat.from_pq(fmt)
+    faker.choose_schema(ncols=20)
+    faker.make_records(20)
+
+    def work():
+        with conn_cls.connect(dsn) as conn:
+            with conn.cursor(binary=fmt) as cur:
+                cur.execute(faker.drop_stmt)
+                cur.execute(faker.create_stmt)
+
+                stmt = sql.SQL("copy {} ({}) from stdin {}").format(
+                    faker.table_name,
+                    sql.SQL(", ").join(faker.fields_names),
+                    sql.SQL("with binary" if fmt else ""),
+                )
+                with cur.copy(stmt) as copy:
+                    if set_types:
+                        copy.set_types(faker.types_names)
+                    for row in faker.records:
+                        copy.write_row(row)
+
+                cur.execute(faker.select_stmt)
+                recs = cur.fetchall()
+
+                for got, want in zip(recs, faker.records):
+                    faker.assert_record(got, want)
+
+    gc_collect()
+    n = []
+    for i in range(3):
+        work()
+        gc_collect()
+        n.append(len(gc.get_objects()))
+
+    assert n[0] == n[1] == n[2], f"objects leaked: {n[1] - n[0]}, {n[2] - n[1]}"
 
 
-# TODOCRDB: random tests
+def copyopt(format):
+    return "with binary" if format == Format.BINARY else ""

diff --git a/tests/crdb/test_copy_async.py b/tests/crdb/test_copy_async.py
index b156266b97e979a531dfa4cf7109ff476936f1ff..464277c369b12b5575ec311d621d247f4933cf18 100644 (file)
--- a/tests/crdb/test_copy_async.py
+++ b/tests/crdb/test_copy_async.py
@@ -1,12 +1,14 @@
+import gc
 import pytest
 import string
 from random import randrange, choice
 
 from psycopg.pq import Format
-from psycopg import errors as e
+from psycopg import sql, errors as e
+from psycopg.adapt import PyFormat
 from psycopg.types.numeric import Int4
 
-from ..utils import eur
+from ..utils import eur, gc_collect
 from ..test_copy import sample_text, sample_binary  # noqa
 from ..test_copy import sample_records
 from ..test_copy_async import ensure_table
@@ -177,4 +179,45 @@ from copy_in group by 1, 2, 3
     assert data == [(True, True, 1, 256)]
 
 
-# TODOCRDB: random tests
+@pytest.mark.slow
+@pytest.mark.parametrize(
+    "fmt, set_types",
+    [(Format.TEXT, True), (Format.TEXT, False), (Format.BINARY, True)],
+)
+@pytest.mark.crdb_skip("copy array")
+async def test_copy_from_leaks(aconn_cls, dsn, faker, fmt, set_types):
+    faker.format = PyFormat.from_pq(fmt)
+    faker.choose_schema(ncols=20)
+    faker.make_records(20)
+
+    async def work():
+        async with await aconn_cls.connect(dsn) as conn:
+            async with conn.cursor(binary=fmt) as cur:
+                await cur.execute(faker.drop_stmt)
+                await cur.execute(faker.create_stmt)
+
+                stmt = sql.SQL("copy {} ({}) from stdin {}").format(
+                    faker.table_name,
+                    sql.SQL(", ").join(faker.fields_names),
+                    sql.SQL("with binary" if fmt else ""),
+                )
+                async with cur.copy(stmt) as copy:
+                    if set_types:
+                        copy.set_types(faker.types_names)
+                    for row in faker.records:
+                        await copy.write_row(row)
+
+                await cur.execute(faker.select_stmt)
+                recs = await cur.fetchall()
+
+                for got, want in zip(recs, faker.records):
+                    faker.assert_record(got, want)
+
+    gc_collect()
+    n = []
+    for i in range(3):
+        await work()
+        gc_collect()
+        n.append(len(gc.get_objects()))
+
+    assert n[0] == n[1] == n[2], f"objects leaked: {n[1] - n[0]}, {n[2] - n[1]}"

diff --git a/tests/fix_crdb.py b/tests/fix_crdb.py
index 0700395c906e8d23f278182cb89d1cd0bb4920bd..5a73d005c8c6b8a54c482004958f5e478573e7ee 100644 (file)
--- a/tests/fix_crdb.py
+++ b/tests/fix_crdb.py
@@ -95,6 +95,7 @@ _crdb_reasons = {
     "cast adds tz": 51692,
     "cidr": 18846,
     "composite": 27792,
+    "copy array": 82792,
     "copy canceled": 81559,
     "copy": 41608,
     "cursor with hold": 77101,