SERVER-69034 Use multiple tenants in shard merge passthroughs

2 files changed, 34 insertions, 36 deletions

diff --git a/buildscripts/resmokelib/testing/fixtures/shard_merge.py b/buildscripts/resmokelib/testing/fixtures/shard_merge.py
index 60044f5007e..6c88b3066b5 100644
--- a/buildscripts/resmokelib/testing/fixtures/shard_merge.py
+++ b/buildscripts/resmokelib/testing/fixtures/shard_merge.py
@@ -1,9 +1,7 @@
 """Fixture with multiple replica sets for executing JSTests against."""
 
 import os.path
-
 import buildscripts.resmokelib.testing.fixtures.interface as interface
-from buildscripts.resmokelib.testing.fixtures.fixturelib import FixtureLib
 
 
 class ShardMergeFixture(interface.MultiClusterFixture):  # pylint: disable=too-many-instance-attributes
diff --git a/buildscripts/resmokelib/testing/hooks/shard_merge.py b/buildscripts/resmokelib/testing/hooks/shard_merge.py
index 9b6216b2cdc..c8b932bd151 100644
--- a/buildscripts/resmokelib/testing/hooks/shard_merge.py
+++ b/buildscripts/resmokelib/testing/hooks/shard_merge.py
@@ -8,9 +8,9 @@ import time
 import uuid
 
 import bson
-import pymongo.errors
-
+from bson.binary import Binary, UUID_SUBTYPE
 from bson.objectid import ObjectId
+import pymongo.errors
 
 from buildscripts.resmokelib import errors
 from buildscripts.resmokelib.testing.fixtures import interface as fixture_interface
@@ -19,7 +19,7 @@ from buildscripts.resmokelib.testing.hooks import dbhash_tenant_migration
 from buildscripts.resmokelib.testing.hooks import interface
 
 
-class ContinuousShardMerge(interface.Hook):  # pylint: disable=too-many-instance-attributes
+class ContinuousShardMerge(interface.Hook):
     """Starts a shard merge thread at the beginning of each test."""
 
     DESCRIPTION = ("Continuous shard merges")
@@ -190,12 +190,12 @@ def get_primary(rs, logger, max_tries=5):  # noqa: D205,D400
 
 class _ShardMergeOptions:  # pylint:disable=too-many-instance-attributes
     def __init__(  # pylint: disable=too-many-arguments
-            self, donor_rs, recipient_rs, tenant_id, read_preference, logger, donor_rs_index,
+            self, donor_rs, recipient_rs, tenant_ids, read_preference, logger, donor_rs_index,
             recipient_rs_index):
         self.donor_rs = donor_rs
         self.recipient_rs = recipient_rs
         self.migration_id = uuid.uuid4()
-        self.tenant_id = tenant_id
+        self.tenant_ids = tenant_ids
         self.read_preference = read_preference
         self.logger = logger
         self.donor_rs_index = donor_rs_index
@@ -228,7 +228,7 @@ class _ShardMergeOptions:  # pylint:disable=too-many-instance-attributes
     def __str__(self):
         opts = {
             "donor": self.get_donor_name(), "recipient": self.get_recipient_name(),
-            "migration_id": self.migration_id, "tenant_id": self.tenant_id,
+            "migration_id": self.migration_id, "tenant_ids": self.tenant_ids,
             "read_preference": self.read_preference
         }
         return str(opts)
@@ -253,8 +253,9 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
         self.daemon = True
         self.logger = logger
         self._shard_merge_fixture = shard_merge_fixture
-        # TODO SERVER-69034 : replace tenantId with tenantIds
-        self._tenant_id = shell_options["global_vars"]["TestData"]["tenantId"]
+        self._tenant_ids = []
+        for tenant_id in shell_options["global_vars"]["TestData"]["tenantIds"]:
+            self._tenant_ids.append(ObjectId(tenant_id))
         self._auth_options = shell_options["global_vars"]["TestData"]["authOptions"]
         self._test = None
         self._test_report = test_report
@@ -375,7 +376,7 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
         donor_rs = self._shard_merge_fixture.get_replset(donor_rs_index)
         recipient_rs = self._shard_merge_fixture.get_replset(recipient_rs_index)
         read_preference = {"mode": "primary"}
-        return _ShardMergeOptions(donor_rs, recipient_rs, self._tenant_id, read_preference,
+        return _ShardMergeOptions(donor_rs, recipient_rs, self._tenant_ids, read_preference,
                                   self.logger, donor_rs_index, recipient_rs_index)
 
     def _create_client(self, node):
@@ -474,21 +475,19 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
         """
         cmd_obj = {
             "donorStartMigration":
-                1,
-            "migrationId":
-                bson.Binary(migration_opts.migration_id.bytes, 4),
+                1, "migrationId":
+                    Binary(migration_opts.migration_id.bytes, UUID_SUBTYPE),
             "recipientConnectionString":
-                migration_opts.recipient_rs.get_driver_connection_url(),
-            "readPreference":
-                migration_opts.read_preference,
-            "donorCertificateForRecipient":
-                get_certificate_and_private_key("jstests/libs/tenant_migration_donor.pem"),
+                migration_opts.recipient_rs.get_driver_connection_url(), "readPreference":
+                    migration_opts.read_preference, "donorCertificateForRecipient":
+                        get_certificate_and_private_key("jstests/libs/tenant_migration_donor.pem"),
             "recipientCertificateForDonor":
                 get_certificate_and_private_key("jstests/libs/tenant_migration_recipient.pem"),
             "protocol":
-                "shard merge",
-            "tenantIds": [ObjectId(migration_opts.tenant_id)],
+                "shard merge", "tenantIds":
+                    migration_opts.tenant_ids
         }
+
         donor_primary = migration_opts.get_donor_primary()
         # TODO(SERVER-68643) We no longer need to override the failpoint once milestone 3 is done
         # for shard merge.
@@ -512,8 +511,7 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
                 # Keep polling the migration state until the migration completes.
                 donor_primary_client = self._create_client(donor_primary)
                 res = donor_primary_client.admin.command(
-                    cmd_obj,
-                    bson.codec_options.CodecOptions(uuid_representation=bson.binary.UUID_SUBTYPE))
+                    cmd_obj, bson.codec_options.CodecOptions(uuid_representation=UUID_SUBTYPE))
             except (pymongo.errors.AutoReconnect, pymongo.errors.NotPrimaryError):
                 donor_primary = migration_opts.get_donor_primary()
                 self.logger.info(
@@ -546,8 +544,8 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
         self.logger.info("Forgetting shard merge: %s.", str(migration_opts))
 
         cmd_obj = {
-            "donorForgetMigration": 1, "migrationId": bson.Binary(migration_opts.migration_id.bytes,
-                                                                  4)
+            "donorForgetMigration": 1, "migrationId": Binary(migration_opts.migration_id.bytes,
+                                                             UUID_SUBTYPE)
         }
         donor_primary = migration_opts.get_donor_primary()
 
@@ -559,8 +557,7 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
             try:
                 donor_primary_client = self._create_client(donor_primary)
                 donor_primary_client.admin.command(
-                    cmd_obj,
-                    bson.codec_options.CodecOptions(uuid_representation=bson.binary.UUID_SUBTYPE))
+                    cmd_obj, bson.codec_options.CodecOptions(uuid_representation=UUID_SUBTYPE))
                 return
             except (pymongo.errors.AutoReconnect, pymongo.errors.NotPrimaryError):
                 donor_primary = migration_opts.get_donor_primary()
@@ -601,8 +598,9 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
                     try:
                         donor_node_client = self._create_client(donor_node)
                         res = donor_node_client.config.command({
-                            "count": "tenantMigrationDonors",
-                            "query": {"_id": bson.Binary(migration_opts.migration_id.bytes, 4)}
+                            "count": "tenantMigrationDonors", "query": {
+                                "_id": Binary(migration_opts.migration_id.bytes, UUID_SUBTYPE)
+                            }
                         })
                         if res["n"] == 0:
                             break
@@ -629,8 +627,9 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
                     try:
                         recipient_node_client = self._create_client(recipient_node)
                         res = recipient_node_client.config.command({
-                            "count": "tenantMigrationRecipients",
-                            "query": {"_id": bson.Binary(migration_opts.migration_id.bytes, 4)}
+                            "count": "tenantMigrationRecipients", "query": {
+                                "_id": Binary(migration_opts.migration_id.bytes, UUID_SUBTYPE)
+                            }
                         })
                         if res["n"] == 0:
                             break
@@ -666,7 +665,7 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
             try:
                 donor_primary_client = self._create_client(donor_primary)
                 doc = donor_primary_client["local"]["rerouted"].find_one(
-                    {"_id": bson.Binary(migration_opts.migration_id.bytes, 4)})
+                    {"_id": Binary(migration_opts.migration_id.bytes, UUID_SUBTYPE)})
                 if doc is not None:
                     return
             except (pymongo.errors.AutoReconnect, pymongo.errors.NotPrimaryError):
@@ -699,8 +698,9 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
                 res = primary_client.admin.command({"listDatabases": 1})
                 for database in res["databases"]:
                     db_name = database["name"]
-                    if db_name.startswith(self._tenant_id + "_"):
-                        primary_client.drop_database(db_name)
+                    for tenant_id in self._tenant_ids:
+                        if db_name.startswith(f"{str(tenant_id)}_"):
+                            primary_client.drop_database(db_name)
                 return
             # We retry on all write concern errors because we assume the only reason waiting for
             # write concern should fail is because of a failover.
@@ -713,8 +713,8 @@ class _ShardMergeThread(threading.Thread):  # pylint: disable=too-many-instance-
                 continue
             except pymongo.errors.PyMongoError:
                 self.logger.exception(
-                    "Error dropping databases for tenant id '%s' on primary on" +
-                    " port %d of replica set '%s' to be garbage collection.", self._tenant_id,
+                    "Error dropping databases for tenants %s on primary on" +
+                    " port %d of replica set '%s' to be garbage collection.", str(self._tenant_ids),
                     primary.port, rs.replset_name)
                 raise