Improve metadata bwc test for Logical Replication

Author: jeeminso

tests/bwc/test_rolling_upgrade.py

@@ -1,8 +1,9 @@
+import time
 import unittest
 from crate.client import connect
 from crate.client.exceptions import ProgrammingError
 
-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
@@ -30,23 +31,20 @@
 )
 
 ROLLING_UPGRADES_V5 = (
-    UpgradePath('5.0.x', '5.1.x'),
-    UpgradePath('5.1.x', '5.2.x'),
-    UpgradePath('5.2.x', '5.3.x'),
-    UpgradePath('5.3.x', '5.4.x'),
-    UpgradePath('5.4.x', '5.5.x'),
-    UpgradePath('5.5.x', '5.6.x'),
-    UpgradePath('5.6.x', '5.7.x'),
-    UpgradePath('5.7.x', '5.8.x'),
-    UpgradePath('5.8.x', '5.9.x'),
-    UpgradePath('5.9.x', '5.10.x'),
-    UpgradePath('5.10.x', '5.10'),
-    UpgradePath('5.10', 'latest-nightly'),
+    UpgradePath('5.10', 'branch:master'),
 )
 
 
 class RollingUpgradeTest(NodeProvider, unittest.TestCase):
 
+    def _num_docs(self, cursor, schema, table):
+        cursor.execute("select sum(num_docs) from sys.shards where schema_name = ? and table_name = ?", (schema, table))
+        return cursor.fetchall()[0][0]
+
+    def _assert_num_docs(self, cursor, schema, table, expected_count):
+        count = self._num_docs(cursor, schema, table)
+        self.assertEqual(expected_count, count)
+
     def test_rolling_upgrade_4_to_5(self):
         print("")  # force newline for first print
         for path in ROLLING_UPGRADES_V4:
@@ -88,6 +86,10 @@ def _test_rolling_upgrade(self, path, nodes):
         }
         cluster = self._new_cluster(path.from_version, nodes, settings=settings)
         cluster.start()
+        replica_cluster = None
+        if int(path.from_version.split('.')[0]) >= 5 and int(path.from_version.split('.')[1]) >= 10:
+            replica_cluster = self._new_cluster(path.from_version, 1, settings=settings, explicit_discovery=False)
+            replica_cluster.start()
         with connect(cluster.node().http_url, error_trace=True) as conn:
             c = conn.cursor()
             c.execute("create user arthur with (password = 'secret')")
@@ -152,6 +154,24 @@ def _test_rolling_upgrade(self, path, nodes):
             # Add the shards of the new partition primaries
             expected_active_shards += shards
 
+            # Set up tables for logical replications
+            if int(path.from_version.split('.')[0]) >= 5 and int(path.from_version.split('.')[1]) >= 10:
+                c.execute("create table doc.x (a int) clustered into 1 shards with (number_of_replicas=0)")
+                expected_active_shards += 1
+                c.execute("create publication p for table doc.x")
+                with connect(replica_cluster.node().http_url, error_trace=True) as replica_conn:
+                    rc = replica_conn.cursor()
+                    transport_port = cluster.node().addresses.transport.port
+                    replica_transport_port = replica_cluster.node().addresses.transport.port
+                    assert 4300 <= transport_port <= 4310 and 4300 <= replica_transport_port <= 4310
+                    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:{transport_port}?user=crate&sslmode=sniff' publication p")
+                    assert_busy(lambda: self._assert_num_docs(rc, "doc", "x", 0))
+                c.execute(f"create subscription s connection 'crate://localhost:{replica_transport_port}?user=crate&sslmode=sniff' publication rp")
+                assert_busy(lambda: self._assert_num_docs(c, "doc", "rx", 0))
+                expected_active_shards += 1
+
         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.
             # Otherwise upgraded node simply requests N-1 columns from old version with N columns and it always works.
@@ -282,6 +302,25 @@ def _test_rolling_upgrade(self, path, nodes):
                         c.execute("select version['created'] from information_schema.table_partitions where table_name = 't3' and values['a'] = ?", [idx])
                         self.assertEqual(c.fetchall(), [[partition_version]])
 
+                # Ensure logical replications works
+                if int(path.from_version.split('.')[0]) >= 5 and int(path.from_version.split('.')[1]) >= 10:
+                    with connect(replica_cluster.node().http_url, error_trace=True) as replica_conn:
+                        rc = replica_conn.cursor()
+
+                        # Cannot drop replicated tables
+                        with self.assertRaises(ProgrammingError):
+                            rc.execute("drop table doc.x")
+                            c.execute("drop table doc.rx")
+
+                        count = self._num_docs(rc, "doc", "x")
+                        count2 = self._num_docs(c, "doc", "rx")
+
+                        c.execute("insert into doc.x values (1)")
+                        rc.execute("insert into doc.rx values (1)")
+
+                        assert_busy(lambda: self._assert_num_docs(rc, "doc", "x", count + 1))
+                        assert_busy(lambda: self._assert_num_docs(c, "doc", "rx", count2 + 1))
+
         # Finally validate that all shards (primaries and replicas) of all partitions are started
         # and writes into the partitioned table while upgrading were successful
         with connect(cluster.node().http_url, error_trace=True) as conn: