use std::sync::Arc; use std::time::{Duration, Instant}; use api::grpc::qdrant::points_server::Points; use api::grpc::qdrant::{ ClearPayloadPoints, CountPoints, CountResponse, CreateFieldIndexCollection, DeleteFieldIndexCollection, DeletePayloadPoints, DeletePointVectors, DeletePoints, DiscoverBatchPoints, DiscoverBatchResponse, DiscoverPoints, DiscoverResponse, FacetCounts, FacetResponse, GetPoints, GetResponse, PointsOperationResponse, QueryBatchPoints, QueryBatchResponse, QueryGroupsResponse, QueryPointGroups, QueryPoints, QueryResponse, RecommendBatchPoints, RecommendBatchResponse, RecommendGroupsResponse, RecommendPointGroups, RecommendPoints, RecommendResponse, ScrollPoints, ScrollResponse, SearchBatchPoints, SearchBatchResponse, SearchGroupsResponse, SearchMatrixOffsets, SearchMatrixOffsetsResponse, SearchMatrixPairs, SearchMatrixPairsResponse, SearchMatrixPoints, SearchPointGroups, SearchPoints, SearchResponse, SetPayloadPoints, UpdateBatchPoints, UpdateBatchResponse, UpdatePointVectors, UpsertPoints, }; use collection::operations::types::CoreSearchRequest; use collection::operations::verification::new_unchecked_verification_pass; use common::counter::hardware_accumulator::HwMeasurementAcc; use storage::content_manager::toc::request_hw_counter::RequestHwCounter; use storage::dispatcher::Dispatcher; use tonic::{Request, Response, Status}; use super::points_common::{ delete_vectors, discover, discover_batch, facet, query, query_batch, query_groups, recommend_groups, scroll, search_groups, search_points_matrix, update_batch, update_vectors, }; use super::validate; use crate::settings::ServiceConfig; use crate::tonic::api::points_common::{ clear_payload, convert_shard_selector_for_read, core_search_batch, count, create_field_index, delete, delete_field_index, delete_payload, get, overwrite_payload, recommend, recommend_batch, search, set_payload, upsert, }; use crate::tonic::auth::extract_access; use crate::tonic::verification::StrictModeCheckedTocProvider; pub struct PointsService { dispatcher: Arc, service_config: ServiceConfig, } impl PointsService { pub fn new(dispatcher: Arc, service_config: ServiceConfig) -> Self { Self { dispatcher, service_config, } } fn get_request_collection_hw_usage_counter(&self, collection_name: String) -> RequestHwCounter { let counter = HwMeasurementAcc::new_with_drain( &self.dispatcher.get_collection_hw_metrics(collection_name), ); RequestHwCounter::new(counter, self.service_config.hardware_reporting(), false) } } #[tonic::async_trait] impl Points for PointsService { async fn upsert( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; // Nothing to verify here. let pass = new_unchecked_verification_pass(); let access = extract_access(&mut request); upsert( self.dispatcher.toc(&access, &pass).clone(), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn delete( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); delete( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn get(&self, mut request: Request) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); get( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, access, ) .await } async fn update_vectors( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; // Nothing to verify here. let pass = new_unchecked_verification_pass(); let access = extract_access(&mut request); update_vectors( self.dispatcher.toc(&access, &pass).clone(), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn delete_vectors( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); delete_vectors( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn set_payload( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); set_payload( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn overwrite_payload( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); overwrite_payload( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn delete_payload( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); delete_payload( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn clear_payload( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); clear_payload( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn update_batch( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); update_batch(&self.dispatcher, request.into_inner(), None, None, access).await } async fn create_field_index( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); create_field_index( self.dispatcher.clone(), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn delete_field_index( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); delete_field_index( self.dispatcher.clone(), request.into_inner(), None, None, access, ) .await .map(|resp| resp.map(Into::into)) } async fn search( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = search( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, access, hw_metrics, ) .await?; Ok(res) } async fn search_batch( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let SearchBatchPoints { collection_name, search_points, read_consistency, timeout, } = request.into_inner(); let timeout = timeout.map(Duration::from_secs); let mut requests = Vec::new(); for mut search_point in search_points { let shard_key = search_point.shard_key_selector.take(); let shard_selector = convert_shard_selector_for_read(None, shard_key); let core_search_request = CoreSearchRequest::try_from(search_point)?; requests.push((core_search_request, shard_selector)); } let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name.clone()); let res = core_search_batch( StrictModeCheckedTocProvider::new(&self.dispatcher), &collection_name, requests, read_consistency, access, timeout, hw_metrics, ) .await?; Ok(res) } async fn search_groups( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = search_groups( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, access, hw_metrics, ) .await?; Ok(res) } async fn scroll( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); scroll( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, access, ) .await } async fn recommend( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = recommend( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), access, hw_metrics, ) .await?; Ok(res) } async fn recommend_batch( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let RecommendBatchPoints { collection_name, recommend_points, read_consistency, timeout, } = request.into_inner(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name.clone()); let res = recommend_batch( StrictModeCheckedTocProvider::new(&self.dispatcher), &collection_name, recommend_points, read_consistency, access, timeout.map(Duration::from_secs), hw_metrics, ) .await?; Ok(res) } async fn recommend_groups( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = recommend_groups( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), access, hw_metrics, ) .await?; Ok(res) } async fn discover( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = discover( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), access, hw_metrics, ) .await?; Ok(res) } async fn discover_batch( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let DiscoverBatchPoints { collection_name, discover_points, read_consistency, timeout, } = request.into_inner(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name.clone()); let res = discover_batch( StrictModeCheckedTocProvider::new(&self.dispatcher), &collection_name, discover_points, read_consistency, access, timeout.map(Duration::from_secs), hw_metrics, ) .await?; Ok(res) } async fn count( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = count( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, &access, hw_metrics, ) .await?; Ok(res) } async fn query( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = query( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, access, hw_metrics, ) .await?; Ok(res) } async fn query_batch( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let request = request.into_inner(); let QueryBatchPoints { collection_name, query_points, read_consistency, timeout, } = request; let timeout = timeout.map(Duration::from_secs); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name.clone()); let res = query_batch( StrictModeCheckedTocProvider::new(&self.dispatcher), &collection_name, query_points, read_consistency, access, timeout, hw_metrics, ) .await?; Ok(res) } async fn query_groups( &self, mut request: Request, ) -> Result, Status> { let access = extract_access(&mut request); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let res = query_groups( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), None, access, hw_metrics, ) .await?; Ok(res) } async fn facet( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); facet( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), access, ) .await } async fn search_matrix_pairs( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let timing = Instant::now(); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let search_matrix_response = search_points_matrix( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), access, hw_metrics.get_counter(), ) .await?; let pairs_response = SearchMatrixPairsResponse { result: Some(SearchMatrixPairs::from(search_matrix_response)), time: timing.elapsed().as_secs_f64(), usage: hw_metrics.to_grpc_api(), }; Ok(Response::new(pairs_response)) } async fn search_matrix_offsets( &self, mut request: Request, ) -> Result, Status> { validate(request.get_ref())?; let access = extract_access(&mut request); let timing = Instant::now(); let collection_name = request.get_ref().collection_name.clone(); let hw_metrics = self.get_request_collection_hw_usage_counter(collection_name); let search_matrix_response = search_points_matrix( StrictModeCheckedTocProvider::new(&self.dispatcher), request.into_inner(), access, hw_metrics.get_counter(), ) .await?; let offsets_response = SearchMatrixOffsetsResponse { result: Some(SearchMatrixOffsets::from(search_matrix_response)), time: timing.elapsed().as_secs_f64(), usage: hw_metrics.to_grpc_api(), }; Ok(Response::new(offsets_response)) } }