From a950e069aef442fac7e8697018d969d98c767214 Mon Sep 17 00:00:00 2001 From: Michele Riva Date: Mon, 16 Feb 2026 16:54:46 -0800 Subject: [PATCH 1/2] feat: regenerate read api key --- src/collection_manager/sides/operation/op.rs | 10 +++++ .../sides/read/collection.rs | 28 ++++++++---- src/collection_manager/sides/read/mod.rs | 45 ++++++++++++++----- .../sides/write/collection.rs | 8 ++++ .../sides/write/collections.rs | 28 +++++++++++- src/collection_manager/sides/write/mod.rs | 20 +++++++++ src/tests/mod.rs | 1 + src/web_server/api/collection/admin.rs | 15 +++++++ 8 files changed, 135 insertions(+), 20 deletions(-) diff --git a/src/collection_manager/sides/operation/op.rs b/src/collection_manager/sides/operation/op.rs index ddc0d7d8..9972514d 100644 --- a/src/collection_manager/sides/operation/op.rs +++ b/src/collection_manager/sides/operation/op.rs @@ -253,6 +253,13 @@ pub enum CollectionWriteOperation { UpdateMcpDescription { mcp_description: Option, }, + UpdateReadApiKey { + #[serde( + deserialize_with = "deserialize_api_key", + serialize_with = "serialize_api_key" + )] + read_api_key: ApiKey, + }, PinRule(PinRuleOperation), Shelf(ShelfOperation), DocumentStorage(DocumentStorageWriteOperation), @@ -395,6 +402,9 @@ impl WriteOperation { _, CollectionWriteOperation::UpdateMcpDescription { .. }, ) => "update_mcp_description", + WriteOperation::Collection(_, CollectionWriteOperation::UpdateReadApiKey { .. }) => { + "update_read_api_key" + } WriteOperation::Collection( _, CollectionWriteOperation::IndexWriteOperation(_, IndexWriteOperation::Index { .. }), diff --git a/src/collection_manager/sides/read/collection.rs b/src/collection_manager/sides/read/collection.rs index de702367..5a6d4c6f 100644 --- a/src/collection_manager/sides/read/collection.rs +++ b/src/collection_manager/sides/read/collection.rs @@ -125,7 +125,7 @@ pub struct CollectionReader { default_locale: Locale, deleted: bool, - read_api_key: ApiKey, + read_api_key: OramaAsyncLock, write_api_key: Option, context: ReadSideContext, offload_config: OffloadFieldConfig, @@ -196,7 +196,7 @@ impl CollectionReader { default_locale, deleted: false, - read_api_key, + read_api_key: OramaAsyncLock::new("collection_read_api_key", read_api_key), write_api_key, context, @@ -361,7 +361,7 @@ impl CollectionReader { default_locale: dump.default_locale, deleted: false, - read_api_key: dump.read_api_key, + read_api_key: OramaAsyncLock::new("collection_read_api_key", dump.read_api_key), write_api_key: dump.write_api_key, context, @@ -546,7 +546,7 @@ impl CollectionReader { description: self.description.clone(), mcp_description: self.mcp_description.read("commit").await.clone(), default_locale: self.default_locale, - read_api_key: self.read_api_key, + read_api_key: **self.read_api_key.read("commit").await, write_api_key: self.write_api_key, index_ids, temp_index_ids, @@ -672,16 +672,16 @@ impl CollectionReader { .load(std::sync::atomic::Ordering::Relaxed) } - #[inline] #[allow(clippy::result_large_err)] - pub fn check_read_api_key( + pub async fn check_read_api_key( &self, api_key: &ReadApiKey, master_api_key: Option, ) -> Result<(), ReadError> { + let read_api_key = self.read_api_key.read("check_read_api_key").await; match api_key { ReadApiKey::ApiKey(api_key) => { - if *api_key == self.read_api_key { + if *api_key == **read_api_key { return Ok(()); } if let Some(write_api_key) = self.write_api_key { @@ -697,7 +697,7 @@ impl CollectionReader { } ReadApiKey::Claims(claims) => { // For JWT claims, verify the orak matches this collection's read API key - if claims.orak == self.read_api_key { + if claims.orak == **read_api_key { return Ok(()); } } @@ -760,6 +760,15 @@ impl CollectionReader { Ok(()) } + /// Updates the read API key for this collection. + pub async fn update_read_api_key(&self, new_key: ApiKey) -> Result<()> { + let mut read_api_key_lock = self.read_api_key.write("update_read_api_key").await; + **read_api_key_lock = new_key; + drop(read_api_key_lock); + + Ok(()) + } + pub async fn nlp_search( &self, read_side: State>, @@ -1107,6 +1116,9 @@ impl CollectionReader { CollectionWriteOperation::UpdateMcpDescription { mcp_description } => { self.update_mcp_description(mcp_description).await?; } + CollectionWriteOperation::UpdateReadApiKey { read_api_key } => { + self.update_read_api_key(read_api_key).await?; + } CollectionWriteOperation::PinRule(op) => { println!("Applying pin rule operation: {op:?}"); let mut pin_rules_lock = self.pin_rules_reader.write("update_pin_rule").await; diff --git a/src/collection_manager/sides/read/mod.rs b/src/collection_manager/sides/read/mod.rs index ca95d30d..51dd902a 100644 --- a/src/collection_manager/sides/read/mod.rs +++ b/src/collection_manager/sides/read/mod.rs @@ -409,7 +409,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; collection.stats(req).await } @@ -428,7 +430,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; collection.batch_get_documents(doc_id_strs).await } @@ -444,7 +448,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; let fields = collection.get_filterable_fields(with_keys).await?; @@ -571,7 +577,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; // Extract extra claims from JWT token if present, otherwise use None for plain API key let claims: Option> = match read_api_key { @@ -628,7 +636,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(&read_api_key, self.master_api_key)?; + collection + .check_read_api_key(&read_api_key, self.master_api_key) + .await?; let collection_stats = self .collection_stats( @@ -665,7 +675,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; let collection_stats = self .collection_stats( @@ -697,7 +709,9 @@ impl ReadSide { None => return Err(ReadError::NotFound(collection_id)), Some(collection) => collection, }; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; Ok(collection) } @@ -784,7 +798,9 @@ impl ReadSide { .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key) + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await } pub fn is_gpu_overloaded(&self) -> bool { @@ -851,7 +867,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; let known_prompt: KnownPrompts = system_prompt_id .as_str() @@ -874,7 +892,9 @@ impl ReadSide { .get_collection(collection_id) .await .ok_or_else(|| ReadError::NotFound(collection_id))?; - collection.check_read_api_key(read_api_key, self.master_api_key)?; + collection + .check_read_api_key(read_api_key, self.master_api_key) + .await?; match self .training_sets @@ -899,7 +919,10 @@ impl ReadSide { None => return Some(Err(ReadError::NotFound(collection_id))), }; - if let Err(e) = collection.check_read_api_key(read_api_key, self.master_api_key) { + if let Err(e) = collection + .check_read_api_key(read_api_key, self.master_api_key) + .await + { return Some(Err(e)); } diff --git a/src/collection_manager/sides/write/collection.rs b/src/collection_manager/sides/write/collection.rs index 2e86efe9..76fbc52a 100644 --- a/src/collection_manager/sides/write/collection.rs +++ b/src/collection_manager/sides/write/collection.rs @@ -733,6 +733,14 @@ impl CollectionWriter { Ok(()) } + /// Generates a new random read API key, updates it in-memory, and returns it. + pub fn regenerate_read_api_key(&mut self) -> ApiKey { + let new_key = ApiKey::try_new(cuid2::create_id()) + .expect("cuid2 IDs are always valid API keys (under 64 chars)"); + self.read_api_key = new_key; + new_key + } + pub async fn as_dto(&self) -> DescribeCollectionResponse { let mut indexes_desc = vec![]; let mut document_count = 0_usize; diff --git a/src/collection_manager/sides/write/collections.rs b/src/collection_manager/sides/write/collections.rs index 620c2d34..7966f025 100644 --- a/src/collection_manager/sides/write/collections.rs +++ b/src/collection_manager/sides/write/collections.rs @@ -17,7 +17,7 @@ use crate::lock::{OramaAsyncLock, OramaAsyncLockReadGuard}; use crate::metrics::commit::COMMIT_CALCULATION_TIME; use crate::metrics::CollectionCommitLabels; use crate::python::embeddings::Model; -use crate::types::CollectionId; +use crate::types::{ApiKey, CollectionId}; use crate::types::{CreateCollection, DescribeCollectionResponse, LanguageDTO}; use oramacore_lib::fs::{create_if_not_exists, BufferedFile}; use oramacore_lib::nlp::locales::Locale; @@ -183,6 +183,32 @@ impl CollectionsWriter { Ok(()) } + /// Regenerates the read API key for a collection and sends the update to the read side. + pub async fn regenerate_read_api_key( + &self, + collection_id: CollectionId, + sender: OperationSender, + ) -> Result { + let mut collections = self.collections.write("regenerate_read_api_key").await; + let collection = collections + .get_mut(&collection_id) + .ok_or(WriteError::CollectionNotFound(collection_id))?; + + let new_key = collection.regenerate_read_api_key(); + + sender + .send(WriteOperation::Collection( + collection_id, + CollectionWriteOperation::UpdateReadApiKey { + read_api_key: new_key, + }, + )) + .await + .context("Cannot send update read API key operation")?; + + Ok(new_key) + } + pub async fn list(&self) -> Vec { let collections = self.collections.read("list").await; diff --git a/src/collection_manager/sides/write/mod.rs b/src/collection_manager/sides/write/mod.rs index 70f21e46..fa2d611a 100644 --- a/src/collection_manager/sides/write/mod.rs +++ b/src/collection_manager/sides/write/mod.rs @@ -433,6 +433,26 @@ impl WriteSide { res } + /// Regenerates the read API key for a collection. Requires write access. + pub async fn regenerate_read_api_key( + &self, + write_api_key: WriteApiKey, + collection_id: CollectionId, + ) -> Result { + // Verify the collection exists and we have write access + let _collection = self.get_collection(collection_id, write_api_key).await?; + drop(_collection); + + self.write_operation_counter.fetch_add(1, Ordering::Relaxed); + let res = self + .collections + .regenerate_read_api_key(collection_id, self.op_sender.clone()) + .await; + self.write_operation_counter.fetch_sub(1, Ordering::Relaxed); + + res + } + pub async fn create_index( &self, write_api_key: WriteApiKey, diff --git a/src/tests/mod.rs b/src/tests/mod.rs index 98ae96e1..765ffa68 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -34,6 +34,7 @@ mod omc_test; mod openai_chat; mod pin_rules; mod quick_fulltext_benchmark; +mod regenerate_read_api_key; mod replace_doc_on_insert; mod replace_index; mod shelves; diff --git a/src/web_server/api/collection/admin.rs b/src/web_server/api/collection/admin.rs index d0f1743f..153911b9 100644 --- a/src/web_server/api/collection/admin.rs +++ b/src/web_server/api/collection/admin.rs @@ -54,6 +54,10 @@ pub fn apis(write_side: Arc) -> Router { "/v1/collections/{collection_id}/replace-index", post(replace_index), ) + .route( + "/v1/collections/{collection_id}/regenerate-read-api-key", + post(regenerate_read_api_key), + ) .with_state(write_side) } @@ -203,6 +207,17 @@ async fn create_temp_index( .map(|_| Json(json!({ "message": "temp collection created" }))) } +async fn regenerate_read_api_key( + Path(collection_id): Path, + write_side: State>, + write_api_key: WriteApiKey, +) -> impl IntoResponse { + write_side + .regenerate_read_api_key(write_api_key, collection_id) + .await + .map(|new_key| Json(json!({ "read_api_key": new_key.expose() }))) +} + async fn replace_index( Path(collection_id): Path, write_side: State>, From 2083a0db1b3028214a80851241208be82dbd6545 Mon Sep 17 00:00:00 2001 From: Michele Riva Date: Mon, 16 Feb 2026 17:57:04 -0800 Subject: [PATCH 2/2] fixe tests --- src/tests/regenerate_read_api_key.rs | 133 +++++++++++++++++++++++++++ 1 file changed, 133 insertions(+) create mode 100644 src/tests/regenerate_read_api_key.rs diff --git a/src/tests/regenerate_read_api_key.rs b/src/tests/regenerate_read_api_key.rs new file mode 100644 index 00000000..c232b291 --- /dev/null +++ b/src/tests/regenerate_read_api_key.rs @@ -0,0 +1,133 @@ +use anyhow::bail; +use futures::FutureExt; +use serde_json::json; + +use crate::tests::utils::{init_log, wait_for, TestContext}; +use crate::types::{DocumentList, ReadApiKey}; + +#[tokio::test(flavor = "multi_thread")] +async fn test_regenerate_read_api_key() { + init_log(); + + let test_context = TestContext::new().await; + let collection_client = test_context.create_collection().await.unwrap(); + let index_client = collection_client.create_index().await.unwrap(); + + // Insert test documents + let documents: DocumentList = json!([ + { "id": "1", "text": "The quick brown fox" }, + { "id": "2", "text": "jumps over the lazy dog" }, + ]) + .try_into() + .unwrap(); + index_client.insert_documents(documents).await.unwrap(); + + // Verify search works with original read key + let result = collection_client + .search(json!({ "term": "fox" }).try_into().unwrap()) + .await + .unwrap(); + assert_eq!( + result.count, 1, + "Search should find one document with the original key" + ); + + // Save the original read API key + let original_read_api_key = collection_client.read_api_key.clone(); + + // Regenerate the read API key + let new_read_api_key = test_context + .writer + .regenerate_read_api_key( + collection_client.write_api_key, + collection_client.collection_id, + ) + .await + .expect("Should regenerate read API key successfully"); + + // Wait for the read side to pick up the new key + let new_read_api_key_for_wait = ReadApiKey::from_api_key(new_read_api_key); + let collection_id = collection_client.collection_id; + wait_for(&test_context, |ctx| { + let reader = ctx.reader.clone(); + let new_key = new_read_api_key_for_wait.clone(); + async move { + // Try using the new key - if it works, the update has propagated + let stats = reader + .collection_stats( + &new_key, + collection_id, + crate::types::CollectionStatsRequest { with_keys: false }, + ) + .await; + match stats { + Ok(_) => Ok(()), + Err(_) => bail!("New read API key not yet accepted by reader"), + } + } + .boxed() + }) + .await + .expect("Reader should accept the new read API key"); + + // Verify the old key no longer works + let old_key_result = test_context + .reader + .collection_stats( + &original_read_api_key, + collection_id, + crate::types::CollectionStatsRequest { with_keys: false }, + ) + .await; + assert!( + old_key_result.is_err(), + "Old read API key should no longer work after regeneration" + ); + + // Verify search works with the new key + let new_collection_client = test_context + .get_test_collection_client( + collection_client.collection_id, + collection_client.write_api_key, + new_read_api_key_for_wait.clone(), + ) + .unwrap(); + + let result = new_collection_client + .search(json!({ "term": "fox" }).try_into().unwrap()) + .await + .unwrap(); + assert_eq!(result.count, 1, "Search should work with the new key"); + + // Commit and reload to verify persistence + test_context.commit_all().await.unwrap(); + let test_context = test_context.reload().await; + + // After reload, the new key should still work + let result = test_context + .reader + .collection_stats( + &new_read_api_key_for_wait, + collection_id, + crate::types::CollectionStatsRequest { with_keys: false }, + ) + .await; + assert!( + result.is_ok(), + "New read API key should persist across restarts" + ); + + // After reload, the old key should still not work + let old_key_result = test_context + .reader + .collection_stats( + &original_read_api_key, + collection_id, + crate::types::CollectionStatsRequest { with_keys: false }, + ) + .await; + assert!( + old_key_result.is_err(), + "Old read API key should still not work after restart" + ); +}