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/query/Cargo.toml

@@ -51,6 +51,7 @@ meter-core.workspace = true
 meter-macros.workspace = true
 object-store.workspace = true
 once_cell.workspace = true
+partition.workspace = true
 prometheus.workspace = true
 promql.workspace = true
 promql-parser.workspace = true

src/query/src/dist_plan/merge_scan.rs

@@ -198,6 +198,7 @@ impl MergeScanExec {
         let dbname = context.task_id().unwrap_or_default();
         let tracing_context = TracingContext::from_json(context.session_id().as_str());
         let current_channel = self.query_ctx.channel();
+        let read_preference = self.query_ctx.read_preference();
 
         let stream = Box::pin(stream!({
             // only report metrics once for each MergeScan
@@ -226,7 +227,7 @@ impl MergeScanExec {
                 };
                 let do_get_start = Instant::now();
                 let mut stream = region_query_handler
-                    .do_get(request)
+                    .do_get(read_preference, request)
                     .await
                     .map_err(|e| {
                         MERGE_SCAN_ERRORS_TOTAL.inc();

src/query/src/region_query.rs

@@ -15,14 +15,33 @@
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use common_meta::node_manager::NodeManagerRef;
 use common_query::request::QueryRequest;
 use common_recordbatch::SendableRecordBatchStream;
+use partition::manager::PartitionRuleManagerRef;
+use session::ReadPreference;
 
 use crate::error::Result;
 
+/// A factory to create a [`RegionQueryHandler`].
+pub trait RegionQueryHandlerFactory: Send + Sync {
+    /// Build a [`RegionQueryHandler`] with the given partition manager and node manager.
+    fn build(
+        &self,
+        partition_manager: PartitionRuleManagerRef,
+        node_manager: NodeManagerRef,
+    ) -> RegionQueryHandlerRef;
+}
+
+pub type RegionQueryHandlerFactoryRef = Arc<dyn RegionQueryHandlerFactory>;
+
 #[async_trait]
 pub trait RegionQueryHandler: Send + Sync {
-    async fn do_get(&self, request: QueryRequest) -> Result<SendableRecordBatchStream>;
+    async fn do_get(
+        &self,
+        read_preference: ReadPreference,
+        request: QueryRequest,
+    ) -> Result<SendableRecordBatchStream>;
 }
 
 pub type RegionQueryHandlerRef = Arc<dyn RegionQueryHandler>;

src/query/src/sql.rs

@@ -724,6 +724,7 @@ pub fn show_variable(stmt: ShowVariables, query_ctx: QueryContextRef) -> Result<
     let value = match variable.as_str() {
         "SYSTEM_TIME_ZONE" | "SYSTEM_TIMEZONE" => get_timezone(None).to_string(),
         "TIME_ZONE" | "TIMEZONE" => query_ctx.timezone().to_string(),
+        "READ_PREFERENCE" => query_ctx.read_preference().to_string(),
         "DATESTYLE" => {
             let (style, order) = *query_ctx.configuration_parameter().pg_datetime_style();
             format!("{}, {}", style, order)