feat: introduce read preference (#5783)

* feat: introduce read preference

* feat: introduce `RegionQueryHandlerFactory`

* feat: extract ReadPreference from http header

* test: add more tests

* chore: apply suggestions from CR

* chore: apply suggestions from CR

src/frontend/src/error.rs

@@ -19,7 +19,9 @@ use common_error::define_into_tonic_status;
 use common_error::ext::{BoxedError, ErrorExt};
 use common_error::status_code::StatusCode;
 use common_macro::stack_trace_debug;
+use session::ReadPreference;
 use snafu::{Location, Snafu};
+use store_api::storage::RegionId;
 
 #[derive(Snafu)]
 #[snafu(visibility(pub))]
@@ -140,9 +142,14 @@ pub enum Error {
         location: Location,
     },
 
-    #[snafu(display("Failed to find table route for table id {}", table_id))]
-    FindTableRoute {
-        table_id: u32,
+    #[snafu(display(
+        "Failed to find region peer for region id {}, read preference: {}",
+        region_id,
+        read_preference
+    ))]
+    FindRegionPeer {
+        region_id: RegionId,
+        read_preference: ReadPreference,
         #[snafu(implicit)]
         location: Location,
         source: partition::error::Error,
@@ -410,7 +417,7 @@ impl ErrorExt for Error {
             Error::External { source, .. } | Error::InitPlugin { source, .. } => {
                 source.status_code()
             }
-            Error::FindTableRoute { source, .. } => source.status_code(),
+            Error::FindRegionPeer { source, .. } => source.status_code(),
 
             Error::TableOperation { source, .. } => source.status_code(),
 

src/frontend/src/instance/builder.rs

@@ -33,6 +33,7 @@ use operator::statement::{StatementExecutor, StatementExecutorRef};
 use operator::table::TableMutationOperator;
 use partition::manager::PartitionRuleManager;
 use pipeline::pipeline_operator::PipelineOperator;
+use query::region_query::RegionQueryHandlerFactoryRef;
 use query::stats::StatementStatistics;
 use query::QueryEngineFactory;
 use snafu::OptionExt;
@@ -114,7 +115,11 @@ impl FrontendBuilder {
             .unwrap_or_else(|| Arc::new(DummyCacheInvalidator));
 
         let region_query_handler =
-            FrontendRegionQueryHandler::arc(partition_manager.clone(), node_manager.clone());
+            if let Some(factory) = plugins.get::<RegionQueryHandlerFactoryRef>() {
+                factory.build(partition_manager.clone(), node_manager.clone())
+            } else {
+                FrontendRegionQueryHandler::arc(partition_manager.clone(), node_manager.clone())
+            };
 
         let table_flownode_cache =
             self.layered_cache_registry

src/frontend/src/instance/region_query.rs

@@ -22,9 +22,10 @@ use common_recordbatch::SendableRecordBatchStream;
 use partition::manager::PartitionRuleManagerRef;
 use query::error::{RegionQuerySnafu, Result as QueryResult};
 use query::region_query::RegionQueryHandler;
+use session::ReadPreference;
 use snafu::ResultExt;
 
-use crate::error::{FindTableRouteSnafu, RequestQuerySnafu, Result};
+use crate::error::{FindRegionPeerSnafu, RequestQuerySnafu, Result};
 
 pub(crate) struct FrontendRegionQueryHandler {
     partition_manager: PartitionRuleManagerRef,
@@ -45,8 +46,12 @@ impl FrontendRegionQueryHandler {
 
 #[async_trait]
 impl RegionQueryHandler for FrontendRegionQueryHandler {
-    async fn do_get(&self, request: QueryRequest) -> QueryResult<SendableRecordBatchStream> {
-        self.do_get_inner(request)
+    async fn do_get(
+        &self,
+        read_preference: ReadPreference,
+        request: QueryRequest,
+    ) -> QueryResult<SendableRecordBatchStream> {
+        self.do_get_inner(read_preference, request)
             .await
             .map_err(BoxedError::new)
             .context(RegionQuerySnafu)
@@ -54,15 +59,20 @@ impl RegionQueryHandler for FrontendRegionQueryHandler {
 }
 
 impl FrontendRegionQueryHandler {
-    async fn do_get_inner(&self, request: QueryRequest) -> Result<SendableRecordBatchStream> {
+    async fn do_get_inner(
+        &self,
+        read_preference: ReadPreference,
+        request: QueryRequest,
+    ) -> Result<SendableRecordBatchStream> {
         let region_id = request.region_id;
 
         let peer = &self
             .partition_manager
             .find_region_leader(region_id)
             .await
-            .context(FindTableRouteSnafu {
-                table_id: region_id.table_id(),
+            .context(FindRegionPeerSnafu {
+                region_id,
+                read_preference,
             })?;
 
         let client = self.node_manager.datanode(peer).await;