Improve metadata bwc test for Logical Replication

Author: jeeminso

tests/bwc/test_rolling_upgrade.py

@@ -5,7 +5,7 @@
 from crate.client.exceptions import ProgrammingError
 from cr8.run_crate import CrateNode
 
-from crate.qa.tests import NodeProvider, insert_data, wait_for_active_shards, UpgradePath
+from crate.qa.tests import NodeProvider, insert_data, wait_for_active_shards, UpgradePath, assert_busy
 
 ROLLING_UPGRADES_V4 = (
     # 4.0.0 -> 4.0.1 -> 4.0.2 don't support rolling upgrades due to a bug
@@ -96,9 +96,17 @@ def _test_rolling_upgrade(self, path: UpgradePath, nodes: int):
         cluster = self._new_cluster(path.from_version, nodes, settings=settings)
         cluster.start()
         node = cluster.node()
+        remote_node = None
         with connect(node.http_url, error_trace=True) as conn:
             new_shards = init_data(conn, node.version, shards, replicas)
             expected_active_shards += new_shards
+            if node.version >= (5, 10, 0):
+                remote_cluster = self._new_cluster(path.from_version, 1, settings=settings, explicit_discovery=False)
+                remote_cluster.start()
+                remote_node = remote_cluster.node()
+                with connect(remote_node.http_url, error_trace=True) as remote_conn:
+                    new_shards = init_logical_replication_data(self, conn, remote_conn, node.addresses.transport.port, remote_node.addresses.transport.port, expected_active_shards)
+                    expected_active_shards += new_shards
 
         for idx, node in enumerate(cluster):
             # Enforce an old version node be a handler to make sure that an upgraded node can serve 'select *' from an old version node.
@@ -129,6 +137,9 @@ def _test_rolling_upgrade(self, path: UpgradePath, nodes: int):
                 c = conn.cursor()
                 new_shards = self._test_queries_on_new_node(idx, c, node, new_node, nodes, shards, expected_active_shards)
                 expected_active_shards += new_shards
+                if node.version >= (5, 10, 0):
+                    with connect(remote_node.http_url, error_trace=True) as remote_conn:
+                        test_logical_replication_queries(self, conn, remote_conn)
 
         # Finally validate that all shards (primaries and replicas) of all partitions are started
         # and writes into the partitioned table while upgrading were successful
@@ -328,3 +339,55 @@ def init_data(conn: Connection, version: tuple[int, int, int], shards: int, repl
     c.execute("INSERT INTO doc.parted (id, value) VALUES (1, 1)")
     new_shards += shards
     return new_shards
+
+
+def init_logical_replication_data(self, local_conn: Connection, remote_conn: Connection, local_transport_port:int, remote_transport_port: int, local_active_shards: int) -> int:
+    assert 4300 <= local_transport_port <= 4310 and 4300 <= remote_transport_port <= 4310
+
+    c = local_conn.cursor()
+    c.execute("create table doc.x (a int) clustered into 1 shards with (number_of_replicas=0)")
+    c.execute("create publication p for table doc.x")
+    
+    rc = remote_conn.cursor()
+    rc.execute("create table doc.rx (a int) clustered into 1 shards with (number_of_replicas=0)")
+    rc.execute("create publication rp for table doc.rx")
+
+    rc.execute(f"create subscription rs connection 'crate://localhost:{local_transport_port}?user=crate&sslmode=sniff' publication p")
+    c.execute(f"create subscription s connection 'crate://localhost:{remote_transport_port}?user=crate&sslmode=sniff' publication rp")
+
+    new_shards = 2 # 1 shard for doc.x and another 1 shard for doc.rx
+    wait_for_active_shards(rc, new_shards)
+    wait_for_active_shards(c, local_active_shards + new_shards)
+    assert_busy(lambda: self.assertEqual(num_docs_x(rc), 0))        
+    assert_busy(lambda: self.assertEqual(num_docs_rx(c), 0))
+
+    return new_shards
+
+
+def test_logical_replication_queries(self, local_conn: Connection, remote_conn: Connection):
+    c = local_conn.cursor()
+    rc = remote_conn.cursor()
+
+    # Cannot drop replicated tables
+    with self.assertRaises(ProgrammingError):
+        rc.execute("drop table doc.x")
+        c.execute("drop table doc.rx")
+
+    count = num_docs_x(rc)
+    count2 = num_docs_rx(c)
+
+    c.execute("insert into doc.x values (1)")
+    rc.execute("insert into doc.rx values (1)")
+
+    assert_busy(lambda: self.assertEqual(num_docs_x(rc), count + 1))
+    assert_busy(lambda: self.assertEqual(num_docs_rx(c), count2 + 1))
+
+
+def num_docs_x(cursor):
+    cursor.execute("select count(*) from doc.x")
+    return cursor.fetchall()[0][0]
+
+
+def num_docs_rx(cursor):
+    cursor.execute("select count(*) from doc.rx")
+    return cursor.fetchall()[0][0]