Skip to content

Commit 7f2efe9

Browse files
allevoallevo
authored
Merge pull request #322 from oramasearch/feat/regenerate-api-key
feat: regenerate read api key
2 parents 7d38204 + 2083a0d commit 7f2efe9

File tree

9 files changed

+268
-20
lines changed

9 files changed

+268
-20
lines changed

‎src/collection_manager/sides/operation/op.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -253,6 +253,13 @@ pub enum CollectionWriteOperation {
253253
UpdateMcpDescription {
254254
mcp_description: Option<String>,
255255
},
256+
UpdateReadApiKey {
257+
#[serde(
258+
deserialize_with = "deserialize_api_key",
259+
serialize_with = "serialize_api_key"
260+
)]
261+
read_api_key: ApiKey,
262+
},
256263
PinRule(PinRuleOperation<DocumentId>),
257264
Shelf(ShelfOperation<DocumentId>),
258265
DocumentStorage(DocumentStorageWriteOperation),
@@ -395,6 +402,9 @@ impl WriteOperation {
395402
_,
396403
CollectionWriteOperation::UpdateMcpDescription { .. },
397404
) => "update_mcp_description",
405+
WriteOperation::Collection(_, CollectionWriteOperation::UpdateReadApiKey { .. }) => {
406+
"update_read_api_key"
407+
}
398408
WriteOperation::Collection(
399409
_,
400410
CollectionWriteOperation::IndexWriteOperation(_, IndexWriteOperation::Index { .. }),

‎src/collection_manager/sides/read/collection.rs‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ pub struct CollectionReader {
125125
default_locale: Locale,
126126
deleted: bool,
127127

128-
read_api_key: ApiKey,
128+
read_api_key: OramaAsyncLock<ApiKey>,
129129
write_api_key: Option<ApiKey>,
130130
context: ReadSideContext,
131131
offload_config: OffloadFieldConfig,
@@ -196,7 +196,7 @@ impl CollectionReader {
196196
default_locale,
197197
deleted: false,
198198

199-
read_api_key,
199+
read_api_key: OramaAsyncLock::new("collection_read_api_key", read_api_key),
200200
write_api_key,
201201

202202
context,
@@ -361,7 +361,7 @@ impl CollectionReader {
361361
default_locale: dump.default_locale,
362362
deleted: false,
363363

364-
read_api_key: dump.read_api_key,
364+
read_api_key: OramaAsyncLock::new("collection_read_api_key", dump.read_api_key),
365365
write_api_key: dump.write_api_key,
366366

367367
context,
@@ -546,7 +546,7 @@ impl CollectionReader {
546546
description: self.description.clone(),
547547
mcp_description: self.mcp_description.read("commit").await.clone(),
548548
default_locale: self.default_locale,
549-
read_api_key: self.read_api_key,
549+
read_api_key: **self.read_api_key.read("commit").await,
550550
write_api_key: self.write_api_key,
551551
index_ids,
552552
temp_index_ids,
@@ -672,16 +672,16 @@ impl CollectionReader {
672672
.load(std::sync::atomic::Ordering::Relaxed)
673673
}
674674

675-
#[inline]
676675
#[allow(clippy::result_large_err)]
677-
pub fn check_read_api_key(
676+
pub async fn check_read_api_key(
678677
&self,
679678
api_key: &ReadApiKey,
680679
master_api_key: Option<ApiKey>,
681680
) -> Result<(), ReadError> {
681+
let read_api_key = self.read_api_key.read("check_read_api_key").await;
682682
match api_key {
683683
ReadApiKey::ApiKey(api_key) => {
684-
if *api_key == self.read_api_key {
684+
if *api_key == **read_api_key {
685685
return Ok(());
686686
}
687687
if let Some(write_api_key) = self.write_api_key {
@@ -697,7 +697,7 @@ impl CollectionReader {
697697
}
698698
ReadApiKey::Claims(claims) => {
699699
// For JWT claims, verify the orak matches this collection's read API key
700-
if claims.orak == self.read_api_key {
700+
if claims.orak == **read_api_key {
701701
return Ok(());
702702
}
703703
}
@@ -760,6 +760,15 @@ impl CollectionReader {
760760
Ok(())
761761
}
762762

763+
/// Updates the read API key for this collection.
764+
pub async fn update_read_api_key(&self, new_key: ApiKey) -> Result<()> {
765+
let mut read_api_key_lock = self.read_api_key.write("update_read_api_key").await;
766+
**read_api_key_lock = new_key;
767+
drop(read_api_key_lock);
768+
769+
Ok(())
770+
}
771+
763772
pub async fn nlp_search(
764773
&self,
765774
read_side: State<Arc<ReadSide>>,
@@ -1107,6 +1116,9 @@ impl CollectionReader {
11071116
CollectionWriteOperation::UpdateMcpDescription { mcp_description } => {
11081117
self.update_mcp_description(mcp_description).await?;
11091118
}
1119+
CollectionWriteOperation::UpdateReadApiKey { read_api_key } => {
1120+
self.update_read_api_key(read_api_key).await?;
1121+
}
11101122
CollectionWriteOperation::PinRule(op) => {
11111123
println!("Applying pin rule operation: {op:?}");
11121124
let mut pin_rules_lock = self.pin_rules_reader.write("update_pin_rule").await;

‎src/collection_manager/sides/read/mod.rs‎

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -409,7 +409,9 @@ impl ReadSide {
409409
.get_collection(collection_id)
410410
.await
411411
.ok_or_else(|| ReadError::NotFound(collection_id))?;
412-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
412+
collection
413+
.check_read_api_key(read_api_key, self.master_api_key)
414+
.await?;
413415

414416
collection.stats(req).await
415417
}
@@ -428,7 +430,9 @@ impl ReadSide {
428430
.get_collection(collection_id)
429431
.await
430432
.ok_or_else(|| ReadError::NotFound(collection_id))?;
431-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
433+
collection
434+
.check_read_api_key(read_api_key, self.master_api_key)
435+
.await?;
432436

433437
collection.batch_get_documents(doc_id_strs).await
434438
}
@@ -444,7 +448,9 @@ impl ReadSide {
444448
.get_collection(collection_id)
445449
.await
446450
.ok_or_else(|| ReadError::NotFound(collection_id))?;
447-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
451+
collection
452+
.check_read_api_key(read_api_key, self.master_api_key)
453+
.await?;
448454

449455
let fields = collection.get_filterable_fields(with_keys).await?;
450456

@@ -571,7 +577,9 @@ impl ReadSide {
571577
.get_collection(collection_id)
572578
.await
573579
.ok_or_else(|| ReadError::NotFound(collection_id))?;
574-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
580+
collection
581+
.check_read_api_key(read_api_key, self.master_api_key)
582+
.await?;
575583

576584
// Extract extra claims from JWT token if present, otherwise use None for plain API key
577585
let claims: Option<HashMap<String, Value>> = match read_api_key {
@@ -628,7 +636,9 @@ impl ReadSide {
628636
.get_collection(collection_id)
629637
.await
630638
.ok_or_else(|| ReadError::NotFound(collection_id))?;
631-
collection.check_read_api_key(&read_api_key, self.master_api_key)?;
639+
collection
640+
.check_read_api_key(&read_api_key, self.master_api_key)
641+
.await?;
632642

633643
let collection_stats = self
634644
.collection_stats(
@@ -665,7 +675,9 @@ impl ReadSide {
665675
.get_collection(collection_id)
666676
.await
667677
.ok_or_else(|| ReadError::NotFound(collection_id))?;
668-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
678+
collection
679+
.check_read_api_key(read_api_key, self.master_api_key)
680+
.await?;
669681

670682
let collection_stats = self
671683
.collection_stats(
@@ -697,7 +709,9 @@ impl ReadSide {
697709
None => return Err(ReadError::NotFound(collection_id)),
698710
Some(collection) => collection,
699711
};
700-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
712+
collection
713+
.check_read_api_key(read_api_key, self.master_api_key)
714+
.await?;
701715

702716
Ok(collection)
703717
}
@@ -784,7 +798,9 @@ impl ReadSide {
784798
.await
785799
.ok_or_else(|| ReadError::NotFound(collection_id))?;
786800

787-
collection.check_read_api_key(read_api_key, self.master_api_key)
801+
collection
802+
.check_read_api_key(read_api_key, self.master_api_key)
803+
.await
788804
}
789805

790806
pub fn is_gpu_overloaded(&self) -> bool {
@@ -851,7 +867,9 @@ impl ReadSide {
851867
.get_collection(collection_id)
852868
.await
853869
.ok_or_else(|| ReadError::NotFound(collection_id))?;
854-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
870+
collection
871+
.check_read_api_key(read_api_key, self.master_api_key)
872+
.await?;
855873

856874
let known_prompt: KnownPrompts = system_prompt_id
857875
.as_str()
@@ -874,7 +892,9 @@ impl ReadSide {
874892
.get_collection(collection_id)
875893
.await
876894
.ok_or_else(|| ReadError::NotFound(collection_id))?;
877-
collection.check_read_api_key(read_api_key, self.master_api_key)?;
895+
collection
896+
.check_read_api_key(read_api_key, self.master_api_key)
897+
.await?;
878898

879899
match self
880900
.training_sets
@@ -899,7 +919,10 @@ impl ReadSide {
899919
None => return Some(Err(ReadError::NotFound(collection_id))),
900920
};
901921

902-
if let Err(e) = collection.check_read_api_key(read_api_key, self.master_api_key) {
922+
if let Err(e) = collection
923+
.check_read_api_key(read_api_key, self.master_api_key)
924+
.await
925+
{
903926
return Some(Err(e));
904927
}
905928

‎src/collection_manager/sides/write/collection.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -733,6 +733,14 @@ impl CollectionWriter {
733733
Ok(())
734734
}
735735

736+
/// Generates a new random read API key, updates it in-memory, and returns it.
737+
pub fn regenerate_read_api_key(&mut self) -> ApiKey {
738+
let new_key = ApiKey::try_new(cuid2::create_id())
739+
.expect("cuid2 IDs are always valid API keys (under 64 chars)");
740+
self.read_api_key = new_key;
741+
new_key
742+
}
743+
736744
pub async fn as_dto(&self) -> DescribeCollectionResponse {
737745
let mut indexes_desc = vec![];
738746
let mut document_count = 0_usize;

‎src/collection_manager/sides/write/collections.rs‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use crate::lock::{OramaAsyncLock, OramaAsyncLockReadGuard};
1717
use crate::metrics::commit::COMMIT_CALCULATION_TIME;
1818
use crate::metrics::CollectionCommitLabels;
1919
use crate::python::embeddings::Model;
20-
use crate::types::CollectionId;
20+
use crate::types::{ApiKey, CollectionId};
2121
use crate::types::{CreateCollection, DescribeCollectionResponse, LanguageDTO};
2222
use oramacore_lib::fs::{create_if_not_exists, BufferedFile};
2323
use oramacore_lib::nlp::locales::Locale;
@@ -183,6 +183,32 @@ impl CollectionsWriter {
183183
Ok(())
184184
}
185185

186+
/// Regenerates the read API key for a collection and sends the update to the read side.
187+
pub async fn regenerate_read_api_key(
188+
&self,
189+
collection_id: CollectionId,
190+
sender: OperationSender,
191+
) -> Result<ApiKey, WriteError> {
192+
let mut collections = self.collections.write("regenerate_read_api_key").await;
193+
let collection = collections
194+
.get_mut(&collection_id)
195+
.ok_or(WriteError::CollectionNotFound(collection_id))?;
196+
197+
let new_key = collection.regenerate_read_api_key();
198+
199+
sender
200+
.send(WriteOperation::Collection(
201+
collection_id,
202+
CollectionWriteOperation::UpdateReadApiKey {
203+
read_api_key: new_key,
204+
},
205+
))
206+
.await
207+
.context("Cannot send update read API key operation")?;
208+
209+
Ok(new_key)
210+
}
211+
186212
pub async fn list(&self) -> Vec<DescribeCollectionResponse> {
187213
let collections = self.collections.read("list").await;
188214

‎src/collection_manager/sides/write/mod.rs‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -433,6 +433,26 @@ impl WriteSide {
433433
res
434434
}
435435

436+
/// Regenerates the read API key for a collection. Requires write access.
437+
pub async fn regenerate_read_api_key(
438+
&self,
439+
write_api_key: WriteApiKey,
440+
collection_id: CollectionId,
441+
) -> Result<ApiKey, WriteError> {
442+
// Verify the collection exists and we have write access
443+
let _collection = self.get_collection(collection_id, write_api_key).await?;
444+
drop(_collection);
445+
446+
self.write_operation_counter.fetch_add(1, Ordering::Relaxed);
447+
let res = self
448+
.collections
449+
.regenerate_read_api_key(collection_id, self.op_sender.clone())
450+
.await;
451+
self.write_operation_counter.fetch_sub(1, Ordering::Relaxed);
452+
453+
res
454+
}
455+
436456
pub async fn create_index(
437457
&self,
438458
write_api_key: WriteApiKey,

‎src/tests/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ mod omc_test;
3434
mod openai_chat;
3535
mod pin_rules;
3636
mod quick_fulltext_benchmark;
37+
mod regenerate_read_api_key;
3738
mod replace_doc_on_insert;
3839
mod replace_index;
3940
mod shelves;

0 commit comments

Comments
 (0)