Add tenant lock commands and fdcli tests

fdbcli/TenantCommands.actor.cpp

@@ -652,40 +652,91 @@ ACTOR Future<bool> tenantRenameCommand(Reference<IDatabase> db, std::vector<Stri
 }
 
 ACTOR Future<bool> tenantLockCommand(Reference<IDatabase> db, std::vector<StringRef> tokens) {
-	state KeyRef name;
 	state UID uid;
 	state Reference<ITransaction> tr;
-	if (tokens.size() < 3 || tokens.size() > 4) {
+	state StringRef name;
-		fmt::print("Usage: tenant lock <NAME> [UID]\n\n");
+	state Key nameKey;
-		fmt::print("Locks a tenant with a given UID. If no UID is passed, fdbcli will\n");
+	state TenantLockState desiredLockState;
-		fmt::print("generate one. UID has to be a 16-byte number represented in hex.\n");
+	state int uidIdx;
+	if (tokens[1] == "lock"_sr && (tokens.size() < 3 || tokens.size() > 5)) {
+		fmt::print("Usage: tenant lock <NAME> [w|rw] [UID]\n\n");
+		fmt::print("Locks a tenant with for read-write or read-only with a given UID.\n");
+		fmt::print("generate one. By default a read-write lock is created.\n");
+		fmt::print(" If no UID is passed, fdbcli will UID has to be a 16-byte number represented in hex.");
+		return false;
+	} else if (tokens[1] == "unlock"_sr && tokens.size() != 4) {
+		fmt::print("Usage: tenant unlock <NAME> <UID>\n\n");
 		return false;
 	}
 	name = tokens[2];
-	if (tokens.size() > 3) {
+	nameKey = tenantMapSpecialKeyRange.begin.withSuffix(name);
-		auto uidStr = tokens[3].toString();
+	if (tokens[1] == "unlock"_sr) {
+		uidIdx = 3;
+		desiredLockState = TenantLockState::UNLOCKED;
+	} else {
+		uidIdx = 4;
+		if (tokens.size() > 3) {
+			if (tokens[3] == "w"_sr) {
+				desiredLockState = TenantLockState::READ_ONLY;
+			} else if (tokens[3] == "rw"_sr) {
+				desiredLockState = TenantLockState::LOCKED;
+			} else {
+				fmt::print(stderr, "Invalid lock type \"{}\"\n", tokens[3]);
+				return false;
+			}
+		} else {
+			desiredLockState = TenantLockState::LOCKED;
+		}
+	}
+	if (tokens.size() > uidIdx) {
 		try {
-			uid = UID::fromStringThrowsOnFailure(uidStr);
+			uid = UID::fromStringThrowsOnFailure(tokens[uidIdx].toString());
 		} catch (Error& e) {
 			ASSERT(e.code() == error_code_operation_failed);
-			fmt::print(stderr, "Couldn't not parse {} as a valid UID", uidStr);
+			fmt::print(stderr, "Couldn't not parse {} as a valid UID", tokens[uidIdx].toString());
 		}
 	} else {
+		ASSERT(desiredLockState != TenantLockState::UNLOCKED);
 		uid = deterministicRandom()->randomUniqueID();
 	}
 	tr = db->createTransaction();
 	loop {
 		try {
-			// TenantAPI::changeLockState(tr, )
+			tr->setOption(FDBTransactionOptions::READ_SYSTEM_KEYS);
+			ClusterType clusterType = wait(TenantAPI::getClusterType(tr));
+			if (clusterType == ClusterType::METACLUSTER_MANAGEMENT) {
+				fmt::print(stderr, "Locking a cluster through a management cluster not yet supported\n");
+				return false;
+			}
+			auto f = tr->get(nameKey);
+			Optional<Value> entry = wait(safeThreadFutureToFuture(f));
+			if (!entry.present()) {
+				fmt::print(stderr, "Tenant \"{}\" does not exist\n", name);
+				return false;
+			}
+			auto tenantId = getTenantId(entry.get());
+			wait(TenantAPI::changeLockState(tr.getPtr(), tenantId, desiredLockState, uid));
+			wait(safeThreadFutureToFuture(tr->commit()));
+			if (desiredLockState != TenantLockState::UNLOCKED) {
+				fmt::print("Locked tenant {} with UID {}", name.toString(), uid.toString());
+			} else {
+				fmt::print("Unlocked tenant {}", name.toString());
+			}
 			return true;
 		} catch (Error& e) {
+			if (e.code() == error_code_tenant_locked) {
+				if (desiredLockState == TenantLockState::UNLOCKED) {
+					fmt::print(stderr, "Wrong lock UID\n");
+				} else {
+					fmt::print(stderr, "Tenant locked with a different UID\n");
+				}
+				return false;
+			}
 			wait(safeThreadFutureToFuture(tr->onError(e)));
 		}
 	}
 }
 
-ACTOR Future<bool> tenantUnlockCommand(Reference<IDatabase> db, std::vector<StringRef> tokens) {}
-
 // tenant command
 Future<bool> tenantCommand(Reference<IDatabase> db, std::vector<StringRef> tokens) {
 	if (tokens.size() == 1) {
@@ -708,7 +759,7 @@ Future<bool> tenantCommand(Reference<IDatabase> db, std::vector<StringRef> token
 	} else if (tokencmp(tokens[1], "lock")) {
 		return tenantLockCommand(db, tokens);
 	} else if (tokencmp(tokens[1], "unlock")) {
-		return tenantUnlockCommand(db, tokens);
+		return tenantLockCommand(db, tokens);
 	} else {
 		printUsage(tokens[0]);
 		return true;

fdbcli/tests/fdbcli_tests.py

@@ -810,6 +810,59 @@ def tenant_list(logger):
     output = run_fdbcli_command_and_get_error('tenant list a b state=14z')
     assert output == 'ERROR: unrecognized tenant state(s) `14z\'.'
 
+
+@enable_logging()
+def tenant_lock(logger):
+    logger.debug('Create tenant')
+    setup_tenants(['tenant'])
+
+    logger.debug('Write test key')
+    run_fdbcli_command('usetenant tenant; writemode on; set foo bar')
+    logger.debug('Lock tenant in read-only mode')
+    output = run_fdbcli_command('tenant lock tenant w')
+    output = output.strip()
+    logger.debug('output: {}'.format(output))
+    assert output.startswith("Locked tenant tenant with UID ")
+    uid_str = output.strip("Locked tenant tenant with UID ")
+    assert len(uid_str) <= 32  # could be smaller than 32 if the first 4 bits are 0
+
+    logger.debug('Verify tenant is readable')
+    output = run_fdbcli_command('usetenant tenant; get foo').strip()
+    logger.debug('output: {}'.format(output))
+    lines = output.split('\n')
+    assert lines[-1] == "`foo' is `bar'"
+
+    logger.debug('Verify tenant is NOT writeable')
+    output = run_fdbcli_command_and_get_error('usetenant tenant; writemode on; set foo bar2').strip()
+    logger.debug('output: {}'.format(output))
+    assert output == 'ERROR: Tenant is locked (2144)'
+
+    logger.debug('Unlock tenant')
+    output = run_fdbcli_command('tenant unlock tenant {}'.format(uid_str))
+    logger.debug('output: {}'.format(output.strip()))
+    assert output.strip() == 'Unlocked tenant tenant'
+
+    logger.debug('Lock tenant in rw mode')
+    output = run_fdbcli_command('tenant lock tenant rw {}'.format(uid_str)).strip()
+    logger.debug('output: {}'.format(output))
+    assert output == 'Locked tenant tenant with UID {}'.format(uid_str)
+
+    logger.debug('Verify tenant is NOT readable')
+    output = run_fdbcli_command_and_get_error('usetenant tenant; get foo').strip()
+    logger.debug('output: {}'.format(output))
+    assert output == 'ERROR: Tenant is locked (2144)'
+
+    logger.debug('Verify tenant is NOT writeable')
+    output = run_fdbcli_command_and_get_error('usetenant tenant; writemode on; set foo bar2').strip()
+    logger.debug('output: {}'.format(output))
+    assert output == 'ERROR: Tenant is locked (2144)'
+
+    logger.debug('Unlock tenant')
+    output = run_fdbcli_command('tenant unlock tenant {}'.format(uid_str))
+    logger.debug('output: {}'.format(output.strip()))
+    assert output.strip() == 'Unlocked tenant tenant'
+
+
 @enable_logging()
 def tenant_get(logger):
     setup_tenants(['tenant', 'tenant2 tenant_group=tenant_group2'])
@@ -1043,6 +1096,7 @@ def tenants():
     run_tenant_test(tenant_old_commands)
     run_tenant_test(tenant_group_list)
     run_tenant_test(tenant_group_get)
+    run_tenant_test(tenant_lock)
 
 def integer_options():
     process = subprocess.Popen(command_template[:-1], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=fdbcli_env)