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/common/function/src/system.rs

@@ -22,7 +22,9 @@ mod version;
 use std::sync::Arc;
 
 use build::BuildFunction;
-use database::{CurrentSchemaFunction, DatabaseFunction, SessionUserFunction};
+use database::{
+    CurrentSchemaFunction, DatabaseFunction, ReadPreferenceFunction, SessionUserFunction,
+};
 use pg_catalog::PGCatalogFunction;
 use procedure_state::ProcedureStateFunction;
 use timezone::TimezoneFunction;
@@ -39,6 +41,7 @@ impl SystemFunction {
         registry.register(Arc::new(CurrentSchemaFunction));
         registry.register(Arc::new(DatabaseFunction));
         registry.register(Arc::new(SessionUserFunction));
+        registry.register(Arc::new(ReadPreferenceFunction));
         registry.register(Arc::new(TimezoneFunction));
         registry.register_async(Arc::new(ProcedureStateFunction));
         PGCatalogFunction::register(registry);

src/common/function/src/system/database.rs

@@ -30,9 +30,12 @@ pub struct DatabaseFunction;
 pub struct CurrentSchemaFunction;
 pub struct SessionUserFunction;
 
+pub struct ReadPreferenceFunction;
+
 const DATABASE_FUNCTION_NAME: &str = "database";
 const CURRENT_SCHEMA_FUNCTION_NAME: &str = "current_schema";
 const SESSION_USER_FUNCTION_NAME: &str = "session_user";
+const READ_PREFERENCE_FUNCTION_NAME: &str = "read_preference";
 
 impl Function for DatabaseFunction {
     fn name(&self) -> &str {
@@ -94,6 +97,26 @@ impl Function for SessionUserFunction {
     }
 }
 
+impl Function for ReadPreferenceFunction {
+    fn name(&self) -> &str {
+        READ_PREFERENCE_FUNCTION_NAME
+    }
+
+    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
+        Ok(ConcreteDataType::string_datatype())
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::nullary(Volatility::Immutable)
+    }
+
+    fn eval(&self, func_ctx: &FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
+        let read_preference = func_ctx.query_ctx.read_preference();
+
+        Ok(Arc::new(StringVector::from_slice(&[read_preference.as_ref()])) as _)
+    }
+}
+
 impl fmt::Display for DatabaseFunction {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "DATABASE")
@@ -112,6 +135,12 @@ impl fmt::Display for SessionUserFunction {
     }
 }
 
+impl fmt::Display for ReadPreferenceFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "READ_PREFERENCE")
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

src/common/session/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "common-session"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+strum.workspace = true

src/common/session/src/lib.rs

@@ -0,0 +1,45 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use strum::{AsRefStr, Display, EnumString};
+
+/// Defines the read preference for frontend route operations,
+/// determining whether to read from the region leader or follower.
+#[derive(Debug, Clone, Copy, Default, EnumString, Display, AsRefStr, PartialEq, Eq)]
+pub enum ReadPreference {
+    #[default]
+    // Reads all operations from the region leader. This is the default mode.
+    #[strum(serialize = "leader", to_string = "LEADER")]
+    Leader,
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use crate::ReadPreference;
+
+    #[test]
+    fn test_read_preference() {
+        assert_eq!(ReadPreference::Leader.to_string(), "LEADER");
+
+        let read_preference = ReadPreference::from_str("LEADER").unwrap();
+        assert_eq!(read_preference, ReadPreference::Leader);
+
+        let read_preference = ReadPreference::from_str("leader").unwrap();
+        assert_eq!(read_preference, ReadPreference::Leader);
+
+        ReadPreference::from_str("follower").unwrap_err();
+    }
+}