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/operator/src/statement.rs

@@ -56,7 +56,7 @@ use query::parser::QueryStatement;
 use query::QueryEngineRef;
 use session::context::{Channel, QueryContextRef};
 use session::table_name::table_idents_to_full_name;
-use set::set_query_timeout;
+use set::{set_query_timeout, set_read_preference};
 use snafu::{ensure, OptionExt, ResultExt};
 use sql::statements::copy::{
     CopyDatabase, CopyDatabaseArgument, CopyQueryToArgument, CopyTable, CopyTableArgument,
@@ -379,6 +379,8 @@ impl StatementExecutor {
     fn set_variables(&self, set_var: SetVariables, query_ctx: QueryContextRef) -> Result<Output> {
         let var_name = set_var.variable.to_string().to_uppercase();
         match var_name.as_str() {
+            "READ_PREFERENCE" => set_read_preference(set_var.value, query_ctx)?,
+
             "TIMEZONE" | "TIME_ZONE" => set_timezone(set_var.value, query_ctx)?,
 
             "BYTEA_OUTPUT" => set_bytea_output(set_var.value, query_ctx)?,

src/operator/src/statement/set.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::str::FromStr;
 use std::time::Duration;
 
 use common_time::Timezone;
@@ -20,6 +21,7 @@ use regex::Regex;
 use session::context::Channel::Postgres;
 use session::context::QueryContextRef;
 use session::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle};
+use session::ReadPreference;
 use snafu::{ensure, OptionExt, ResultExt};
 use sql::ast::{Expr, Ident, Value};
 use sql::statements::set_variables::SetVariables;
@@ -35,6 +37,38 @@ lazy_static! {
     static ref PG_TIME_INPUT_REGEX: Regex = Regex::new(r"^(\d+)(ms|s|min|h|d)$").unwrap();
 }
 
+pub fn set_read_preference(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
+    let read_preference_expr = exprs.first().context(NotSupportedSnafu {
+        feat: "No read preference find in set variable statement",
+    })?;
+
+    match read_preference_expr {
+        Expr::Value(Value::SingleQuotedString(expr))
+        | Expr::Value(Value::DoubleQuotedString(expr)) => {
+            match ReadPreference::from_str(expr.as_str().to_lowercase().as_str()) {
+                Ok(read_preference) => ctx.set_read_preference(read_preference),
+                Err(_) => {
+                    return NotSupportedSnafu {
+                        feat: format!(
+                            "Invalid read preference expr {} in set variable statement",
+                            expr,
+                        ),
+                    }
+                    .fail()
+                }
+            }
+            Ok(())
+        }
+        expr => NotSupportedSnafu {
+            feat: format!(
+                "Unsupported read preference expr {} in set variable statement",
+                expr
+            ),
+        }
+        .fail(),
+    }
+}
+
 pub fn set_timezone(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
     let tz_expr = exprs.first().context(NotSupportedSnafu {
         feat: "No timezone find in set variable statement",