diff --git a/crates/index-scheduler/src/lib.rs b/crates/index-scheduler/src/lib.rs index 6ad7a8397..797ecc134 100644 --- a/crates/index-scheduler/src/lib.rs +++ b/crates/index-scheduler/src/lib.rs @@ -64,7 +64,7 @@ use meilisearch_types::milli::vector::{ }; use meilisearch_types::milli::{self, Index}; use meilisearch_types::task_view::TaskView; -use meilisearch_types::tasks::{KindWithContent, Task}; +use meilisearch_types::tasks::{KindWithContent, Task, TaskNetwork}; use meilisearch_types::webhooks::{Webhook, WebhooksDumpView, WebhooksView}; use milli::vector::db::IndexEmbeddingConfig; use processing::ProcessingTasks; @@ -666,6 +666,16 @@ impl IndexScheduler { self.queue.get_task_ids_from_authorized_indexes(&rtxn, query, filters, &processing) } + pub fn set_task_network(&self, task_id: TaskId, network: TaskNetwork) -> Result<()> { + let mut wtxn = self.env.write_txn()?; + let mut task = + self.queue.tasks.get_task(&wtxn, task_id)?.ok_or(Error::TaskNotFound(task_id))?; + task.network = Some(network); + self.queue.tasks.all_tasks.put(&mut wtxn, &task_id, &task)?; + wtxn.commit()?; + Ok(()) + } + /// Return the batches matching the query from the user's point of view along /// with the total number of batches matching the query, ignoring from and limit. ///