test(crdb): add tests to test CHANGEFEED together with stream()

diff --git a/tests/crdb/test_cursor.py b/tests/crdb/test_cursor.py
new file mode 100644 (file)
index 0000000..991b084
--- /dev/null
+++ b/tests/crdb/test_cursor.py
@@ -0,0 +1,65 @@
+import json
+import threading
+from uuid import uuid4
+from queue import Queue
+from typing import Any
+
+import pytest
+from psycopg import pq, errors as e
+from psycopg.rows import namedtuple_row
+
+pytestmark = pytest.mark.crdb
+
+
+@pytest.fixture
+def testfeed(svcconn):
+    name = f"test_feed_{str(uuid4()).replace('-', '')}"
+    svcconn.execute("set cluster setting kv.rangefeed.enabled to true")
+    svcconn.execute(f"create table {name} (id serial primary key, data text)")
+    yield name
+    svcconn.execute(f"drop table {name}")
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("fmt_out", pq.Format)
+def test_changefeed(conn_cls, dsn, conn, testfeed, fmt_out):
+    conn.autocommit = True
+    q: "Queue[Any]" = Queue()
+
+    def worker():
+        try:
+            with conn_cls.connect(dsn, autocommit=True) as conn:
+                cur = conn.cursor(binary=fmt_out, row_factory=namedtuple_row)
+                try:
+                    for row in cur.stream(f"experimental changefeed for {testfeed}"):
+                        q.put(row)
+                except e.QueryCanceled:
+                    assert conn.info.transaction_status == conn.TransactionStatus.IDLE
+                    q.put(None)
+        except Exception as ex:
+            q.put(ex)
+
+    t = threading.Thread(target=worker)
+    t.start()
+
+    cur = conn.cursor()
+    cur.execute(f"insert into {testfeed} (data) values ('hello') returning id")
+    (key,) = cur.fetchone()
+    row = q.get(timeout=1)
+    assert row.table == testfeed
+    assert json.loads(row.key) == [key]
+    assert json.loads(row.value)["after"] == {"id": key, "data": "hello"}
+
+    cur.execute(f"delete from {testfeed} where id = %s", [key])
+    row = q.get(timeout=1)
+    assert row.table == testfeed
+    assert json.loads(row.key) == [key]
+    assert json.loads(row.value)["after"] is None
+
+    cur.execute("select query_id from [show statements] where query !~ 'show'")
+    (qid,) = cur.fetchone()
+    cur.execute("cancel query %s", [qid])
+    assert cur.statusmessage == "CANCEL QUERIES 1"
+
+    assert q.get(timeout=1) is None
+    t.join()

diff --git a/tests/crdb/test_cursor_async.py b/tests/crdb/test_cursor_async.py
new file mode 100644 (file)
index 0000000..229295d
--- /dev/null
+++ b/tests/crdb/test_cursor_async.py
@@ -0,0 +1,61 @@
+import json
+import asyncio
+from typing import Any
+from asyncio.queues import Queue
+
+import pytest
+from psycopg import pq, errors as e
+from psycopg.rows import namedtuple_row
+from psycopg._compat import create_task
+
+from .test_cursor import testfeed
+
+testfeed  # fixture
+
+pytestmark = [pytest.mark.crdb, pytest.mark.asyncio]
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("fmt_out", pq.Format)
+async def test_changefeed(aconn_cls, dsn, aconn, testfeed, fmt_out):
+    await aconn.set_autocommit(True)
+    q: "Queue[Any]" = Queue()
+
+    async def worker():
+        try:
+            async with await aconn_cls.connect(dsn, autocommit=True) as conn:
+                cur = conn.cursor(binary=fmt_out, row_factory=namedtuple_row)
+                try:
+                    async for row in cur.stream(
+                        f"experimental changefeed for {testfeed}"
+                    ):
+                        q.put_nowait(row)
+                except e.QueryCanceled:
+                    assert conn.info.transaction_status == conn.TransactionStatus.IDLE
+                    q.put_nowait(None)
+        except Exception as ex:
+            q.put_nowait(ex)
+
+    t = create_task(worker())
+
+    cur = aconn.cursor()
+    await cur.execute(f"insert into {testfeed} (data) values ('hello') returning id")
+    (key,) = await cur.fetchone()
+    row = await asyncio.wait_for(q.get(), 1.0)
+    assert row.table == testfeed
+    assert json.loads(row.key) == [key]
+    assert json.loads(row.value)["after"] == {"id": key, "data": "hello"}
+
+    await cur.execute(f"delete from {testfeed} where id = %s", [key])
+    row = await asyncio.wait_for(q.get(), 1.0)
+    assert row.table == testfeed
+    assert json.loads(row.key) == [key]
+    assert json.loads(row.value)["after"] is None
+
+    await cur.execute("select query_id from [show statements] where query !~ 'show'")
+    (qid,) = await cur.fetchone()
+    await cur.execute("cancel query %s", [qid])
+    assert cur.statusmessage == "CANCEL QUERIES 1"
+
+    assert await asyncio.wait_for(q.get(), 1.0) is None
+    await asyncio.gather(t)