nidx-binding 6.9.6.post677__tar.gz → 6.9.7.post680__tar.gz

{nidx_binding-6.9.6.post677 → nidx_binding-6.9.7.post680}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nidx_binding
-Version: 6.9.6.post677
+Version: 6.9.7.post680
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{nidx_binding-6.9.6.post677 → nidx_binding-6.9.7.post680}/nidx_protos/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "pdm.backend"
 
 [project]
 name = "nidx_protos"
-version = "6.9.6.post677"
+version = "6.9.7.post680"
 license = "AGPL-3.0-or-later"
 description = "Protobuf definitions for nucliadb/nidx"
 authors = [{ name = "Nuclia", email = "nucliadb@nuclia.com" }]

{nidx_binding-6.9.6.post677 → nidx_binding-6.9.7.post680}/pyproject.toml

@@ -13,7 +13,7 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
-version = "6.9.6.post677"
+version = "6.9.7.post680"
 
 [project.urls]
 Homepage = "https://nuclia.com"

{nidx_binding-6.9.6.post677 → nidx_binding-6.9.7.post680}/src/searcher/grpc.rs

@@ -39,6 +39,9 @@ use super::streams;
 use super::{index_cache::IndexCache, shard_search, shard_suggest};
 use tracing::*;
 
+/// When this header is set, we only try to serve the request from the local node
+const HEADER_LOCAL_ONLY: &str = "nidx-local-only";
+
 #[derive(Clone)]
 struct TelemetryInterceptor;
 #[cfg(feature = "telemetry")]
@@ -101,7 +104,12 @@ macro_rules! shard_request {
         $(, STREAM => $stream_type:ty)?
     } => {{
         let mut errors = Vec::new();
-        let nodes = $self.shard_selector.nodes_for_shard(&$shard_id);
+        let local_only = $request.metadata().contains_key(HEADER_LOCAL_ONLY);
+        let nodes = if local_only {
+            vec![SearcherNode::This]
+        } else {
+            $self.shard_selector.nodes_for_shard(&$shard_id)
+        };
         for node in &nodes {
             debug!(?node, ?$shard_id, concat!("Attempting ", $name));
             let result = match node {
@@ -113,7 +121,9 @@ macro_rules! shard_request {
                 SearcherNode::Remote(hostname) => {
                     match $self.get_client(hostname).await {
                         Ok(mut client) => {
-                            client.$remote_method(Request::new($request.clone()))
+                            let mut request = Request::new($request.get_ref().clone());
+                            request.metadata_mut().insert(HEADER_LOCAL_ONLY, "1".parse().unwrap());
+                            client.$remote_method(request)
                                 .await
                                 .map_err(|e| match e.code() {
                                     tonic::Code::NotFound => NidxError::NotFound,
@@ -146,48 +156,48 @@ macro_rules! shard_request {
 #[tonic::async_trait]
 impl NidxSearcher for SearchServer {
     async fn search(&self, request: Request<SearchRequest>) -> Result<Response<SearchResponse>> {
-        let request = request.into_inner();
-        let shard_id = uuid::Uuid::parse_str(&request.shard).map_err(NidxError::from)?;
+        let message = request.get_ref();
+        let shard_id = uuid::Uuid::parse_str(&message.shard).map_err(NidxError::from)?;
         shard_request! {
             "search",
             self,
             request,
             shard_id,
-            LOCAL => shard_search::search(Arc::clone(&self.index_cache), request.clone()).await,
+            LOCAL => shard_search::search(Arc::clone(&self.index_cache), message.clone()).await,
             REMOTE => search
         }
     }
 
     async fn suggest(&self, request: Request<SuggestRequest>) -> Result<Response<SuggestResponse>> {
-        let request = request.into_inner();
-        let shard_id = uuid::Uuid::parse_str(&request.shard).map_err(NidxError::from)?;
+        let message = request.get_ref();
+        let shard_id = uuid::Uuid::parse_str(&message.shard).map_err(NidxError::from)?;
         shard_request! {
             "suggest",
             self,
             request,
             shard_id,
-            LOCAL => shard_suggest::suggest(Arc::clone(&self.index_cache), request.clone()).await,
+            LOCAL => shard_suggest::suggest(Arc::clone(&self.index_cache), message.clone()).await,
             REMOTE => suggest
         }
     }
 
     async fn graph_search(&self, request: Request<GraphSearchRequest>) -> Result<Response<GraphSearchResponse>> {
-        let request = request.into_inner();
-        let shard_id = uuid::Uuid::parse_str(&request.shard).map_err(NidxError::from)?;
+        let message = request.get_ref();
+        let shard_id = uuid::Uuid::parse_str(&message.shard).map_err(NidxError::from)?;
         shard_request! {
             "graph_search",
             self,
             request,
             shard_id,
-            LOCAL => shard_search::graph_search(Arc::clone(&self.index_cache), request.clone()).await,
+            LOCAL => shard_search::graph_search(Arc::clone(&self.index_cache), message.clone()).await,
             REMOTE => graph_search
         }
     }
 
     type ParagraphsStream = Pin<Box<dyn Stream<Item = Result<ParagraphItem, Status>> + Send>>;
     async fn paragraphs(&self, request: Request<StreamRequest>) -> Result<Response<Self::ParagraphsStream>> {
-        let request = request.into_inner();
-        let Some(shard) = &request.shard_id.as_ref() else {
+        let message = request.get_ref();
+        let Some(shard) = &message.shard_id else {
             return Err(NidxError::invalid("Uuid is required").into());
         };
         let shard_id = uuid::Uuid::parse_str(&shard.id).map_err(NidxError::from)?;
@@ -196,7 +206,7 @@ impl NidxSearcher for SearchServer {
             self,
             request,
             shard_id,
-            LOCAL => streams::paragraph_iterator(Arc::clone(&self.index_cache), request.clone()).await,
+            LOCAL => streams::paragraph_iterator(Arc::clone(&self.index_cache), message.clone()).await,
             REMOTE => paragraphs,
             STREAM => Self::ParagraphsStream
         }
@@ -204,8 +214,8 @@ impl NidxSearcher for SearchServer {
 
     type DocumentsStream = Pin<Box<dyn Stream<Item = Result<DocumentItem, Status>> + Send>>;
     async fn documents(&self, request: Request<StreamRequest>) -> Result<Response<Self::DocumentsStream>> {
-        let request = request.into_inner();
-        let Some(shard) = &request.shard_id.as_ref() else {
+        let message = request.get_ref();
+        let Some(shard) = &message.shard_id else {
             return Err(NidxError::invalid("Uuid is required").into());
         };
         let shard_id = uuid::Uuid::parse_str(&shard.id).map_err(NidxError::from)?;
@@ -214,7 +224,7 @@ impl NidxSearcher for SearchServer {
             self,
             request,
             shard_id,
-            LOCAL => streams::document_iterator(Arc::clone(&self.index_cache), request.clone()).await,
+            LOCAL => streams::document_iterator(Arc::clone(&self.index_cache), message.clone()).await,
             REMOTE => documents,
             STREAM => Self::DocumentsStream
         }

{nidx_binding-6.9.6.post677 → nidx_binding-6.9.7.post680}/tests/test_searcher_cluster.rs

@@ -34,6 +34,7 @@ use std::{sync::Arc, time::Duration};
 use tempfile::tempdir;
 use tokio_util::sync::CancellationToken;
 use tonic::Request;
+use uuid::Uuid;
 
 struct ManualListNodes(Arc<Mutex<Vec<String>>>, usize);
 impl ListNodes for ManualListNodes {
@@ -311,3 +312,71 @@ async fn test_search_cluster_shard_distribution(pool: PgPool) -> anyhow::Result<
 
     Ok(())
 }
+
+#[sqlx::test]
+async fn test_search_cluster_shards_not_accessible(pool: PgPool) -> anyhow::Result<()> {
+    // Request a shard that does not exist, we should get a NotFound error
+    // This test makes sure we don't get the searchers stuck in a loop asking each other
+    // for the shard.
+    let fixture = NidxFixture::new(pool).await?;
+
+    // Create a cluster with 3 nodes and 2 replicas for each shard
+    let nodes = Arc::new(Mutex::new(Vec::new()));
+    let mut searchers = Vec::new();
+    for i in 0..3 {
+        let searcher_server = GrpcServer::new("localhost:0").await?;
+        let searcher_port = searcher_server.port()?;
+        nodes.lock().unwrap().push(format!("localhost:{searcher_port}"));
+        let list_nodes = ManualListNodes(nodes.clone(), i);
+        let work_dir = tempdir()?;
+        let searcher = SyncedSearcher::new(fixture.settings.metadata.clone(), work_dir.path());
+        let arc_list_nodes = Arc::new(list_nodes);
+        let searcher_api = SearchServer::new(searcher.index_cache(), ShardSelector::new(arc_list_nodes.clone(), 2));
+        searchers.push(format!("localhost:{searcher_port}"));
+        let settings_copy = fixture.settings.clone();
+        let shutdown = CancellationToken::new();
+        tokio::task::spawn(searcher_server.serve(searcher_api.into_router(), shutdown.clone()));
+        tokio::task::spawn(async move {
+            searcher
+                .run(
+                    settings_copy.storage.as_ref().unwrap().object_store.clone(),
+                    settings_copy.searcher.clone().unwrap_or_default(),
+                    shutdown.clone(),
+                    ShardSelector::new(arc_list_nodes, 2),
+                    None,
+                    None,
+                )
+                .await
+        });
+    }
+    tokio::time::sleep(Duration::from_secs(1)).await;
+
+    // Same behaviour from all nodes
+    let fake_uuid = Uuid::new_v4();
+    for searcher in &searchers {
+        let mut request = Request::new(SearchRequest {
+            shard: fake_uuid.to_string(),
+            result_per_page: 20,
+            body: "Hola".into(),
+            paragraph: true,
+            ..Default::default()
+        });
+        request.set_timeout(Duration::from_secs(1));
+
+        let response = NidxSearcherClient::connect(format!("http://{searcher}"))
+            .await?
+            .search(request)
+            .await;
+
+        let Err(err) = response else {
+            panic!("Expected error response")
+        };
+        assert_eq!(
+            err.code(),
+            tonic::Code::NotFound,
+            "Expected shard NotFound. Cancelled error means timeout which means something went wrong"
+        );
+    }
+
+    Ok(())
+}