From 783b586de93ffa210d6c34fd7cc266220a57a4d2 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Tue, 19 Mar 2024 16:57:51 +0100
Subject: [PATCH] [bucket-id-prefix] CLI: allow manipulating buckets by
 prefixes of their full IDs

---
 src/garage/admin/bucket.rs | 40 ++++++++++++++---------------------
 src/model/helper/bucket.rs | 43 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/src/garage/admin/bucket.rs b/src/garage/admin/bucket.rs
index 803b55bd..ac43e122 100644
--- a/src/garage/admin/bucket.rs
+++ b/src/garage/admin/bucket.rs
@@ -54,9 +54,8 @@ impl AdminRpcHandler {
 		let bucket_id = self
 			.garage
 			.bucket_helper()
-			.resolve_global_bucket_name(&query.name)
-			.await?
-			.ok_or_bad_request("Bucket not found")?;
+			.admin_get_existing_matching_bucket(&query.name)
+			.await?;
 
 		let bucket = self
 			.garage
@@ -157,9 +156,8 @@ impl AdminRpcHandler {
 
 		let bucket_id = helper
 			.bucket()
-			.resolve_global_bucket_name(&query.name)
-			.await?
-			.ok_or_bad_request("Bucket not found")?;
+			.admin_get_existing_matching_bucket(&query.name)
+			.await?;
 
 		// Get the alias, but keep in minde here the bucket name
 		// given in parameter can also be directly the bucket's ID.
@@ -235,9 +233,8 @@ impl AdminRpcHandler {
 
 		let bucket_id = helper
 			.bucket()
-			.resolve_global_bucket_name(&query.existing_bucket)
-			.await?
-			.ok_or_bad_request("Bucket not found")?;
+			.admin_get_existing_matching_bucket(&query.existing_bucket)
+			.await?;
 
 		if let Some(key_pattern) = &query.local {
 			let key = helper.key().get_existing_matching_key(key_pattern).await?;
@@ -307,9 +304,8 @@ impl AdminRpcHandler {
 
 		let bucket_id = helper
 			.bucket()
-			.resolve_global_bucket_name(&query.bucket)
-			.await?
-			.ok_or_bad_request("Bucket not found")?;
+			.admin_get_existing_matching_bucket(&query.bucket)
+			.await?;
 		let key = helper
 			.key()
 			.get_existing_matching_key(&query.key_pattern)
@@ -343,9 +339,8 @@ impl AdminRpcHandler {
 
 		let bucket_id = helper
 			.bucket()
-			.resolve_global_bucket_name(&query.bucket)
-			.await?
-			.ok_or_bad_request("Bucket not found")?;
+			.admin_get_existing_matching_bucket(&query.bucket)
+			.await?;
 		let key = helper
 			.key()
 			.get_existing_matching_key(&query.key_pattern)
@@ -378,9 +373,8 @@ impl AdminRpcHandler {
 		let bucket_id = self
 			.garage
 			.bucket_helper()
-			.resolve_global_bucket_name(&query.bucket)
-			.await?
-			.ok_or_bad_request("Bucket not found")?;
+			.admin_get_existing_matching_bucket(&query.bucket)
+			.await?;
 
 		let mut bucket = self
 			.garage
@@ -420,9 +414,8 @@ impl AdminRpcHandler {
 		let bucket_id = self
 			.garage
 			.bucket_helper()
-			.resolve_global_bucket_name(&query.bucket)
-			.await?
-			.ok_or_bad_request("Bucket not found")?;
+			.admin_get_existing_matching_bucket(&query.bucket)
+			.await?;
 
 		let mut bucket = self
 			.garage
@@ -479,9 +472,8 @@ impl AdminRpcHandler {
 			bucket_ids.push(
 				self.garage
 					.bucket_helper()
-					.resolve_global_bucket_name(b)
-					.await?
-					.ok_or_bad_request(format!("Bucket not found: {}", b))?,
+					.admin_get_existing_matching_bucket(b)
+					.await?,
 			);
 		}
 
diff --git a/src/model/helper/bucket.rs b/src/model/helper/bucket.rs
index f4e669c3..4ae9122f 100644
--- a/src/model/helper/bucket.rs
+++ b/src/model/helper/bucket.rs
@@ -67,6 +67,49 @@ impl<'a> BucketHelper<'a> {
 		}
 	}
 
+	/// Find a bucket by its global alias or a prefix of its uuid
+	pub async fn admin_get_existing_matching_bucket(
+		&self,
+		pattern: &String,
+	) -> Result<Uuid, Error> {
+		if let Some(uuid) = self.resolve_global_bucket_name(pattern).await? {
+			return Ok(uuid);
+		} else if pattern.len() >= 2 {
+			let hexdec = pattern
+				.get(..pattern.len() & !1)
+				.and_then(|x| hex::decode(x).ok());
+			if let Some(hex) = hexdec {
+				let mut start = [0u8; 32];
+				start
+					.as_mut_slice()
+					.get_mut(..hex.len())
+					.ok_or_bad_request("invalid length")?
+					.copy_from_slice(&hex);
+				let mut candidates = self
+					.0
+					.bucket_table
+					.get_range(
+						&EmptyKey,
+						Some(start.into()),
+						Some(DeletedFilter::NotDeleted),
+						10,
+						EnumerationOrder::Forward,
+					)
+					.await?
+					.into_iter()
+					.collect::<Vec<_>>();
+				candidates.retain(|x| hex::encode(x.id).starts_with(pattern));
+				if candidates.len() == 1 {
+					return Ok(candidates.into_iter().next().unwrap().id);
+				}
+			}
+		}
+		Err(Error::BadRequest(format!(
+			"Bucket not found / several matching buckets: {}",
+			pattern
+		)))
+	}
+
 	/// Returns a Bucket if it is present in bucket table,
 	/// even if it is in deleted state. Querying a non-existing
 	/// bucket ID returns an internal error.