diff --git a/Cargo.lock b/Cargo.lock index e7cd25c366..b70498a449 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8675,6 +8675,7 @@ dependencies = [ "futures", "futures-util", "humantime", + "itertools 0.14.0", "jsonb", "lazy_static", "meta-client", diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 42c831ebc1..82ddb12e20 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -47,6 +47,7 @@ file-engine.workspace = true futures.workspace = true futures-util.workspace = true humantime.workspace = true +itertools.workspace = true jsonb.workspace = true lazy_static.workspace = true meta-client.workspace = true diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 5dd39681b6..e4abf941c0 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -46,12 +46,13 @@ use common_meta::key::{TableMetadataManager, TableMetadataManagerRef}; use common_meta::kv_backend::KvBackendRef; use common_meta::procedure_executor::ProcedureExecutorRef; use common_query::Output; -use common_telemetry::tracing; +use common_telemetry::{debug, tracing}; use common_time::Timestamp; use common_time::range::TimestampRange; use datafusion_expr::LogicalPlan; use datatypes::prelude::ConcreteDataType; use humantime::format_duration; +use itertools::Itertools; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use query::QueryEngineRef; use query::parser::QueryStatement; @@ -452,6 +453,13 @@ impl StatementExecutor { fn set_variables(&self, set_var: SetVariables, query_ctx: QueryContextRef) -> Result { let var_name = set_var.variable.to_string().to_uppercase(); + debug!( + "Trying to set {}={} for session: {} ", + var_name, + set_var.value.iter().map(|e| e.to_string()).join(", "), + query_ctx.conn_info() + ); + match var_name.as_str() { "READ_PREFERENCE" => set_read_preference(set_var.value, query_ctx)?, diff --git a/src/servers/src/postgres/auth_handler.rs b/src/servers/src/postgres/auth_handler.rs index 14450289ba..286c39da64 100644 --- a/src/servers/src/postgres/auth_handler.rs +++ b/src/servers/src/postgres/auth_handler.rs @@ -19,6 +19,7 @@ use ::auth::{Identity, Password, UserInfoRef, UserProviderRef, userinfo_by_name} use async_trait::async_trait; use common_catalog::parse_catalog_and_schema_from_db_string; use common_error::ext::ErrorExt; +use common_time::Timezone; use futures::{Sink, SinkExt}; use pgwire::api::auth::StartupHandler; use pgwire::api::{ClientInfo, PgWireConnectionState, auth}; @@ -171,6 +172,23 @@ impl StartupHandler for PostgresServerHandlerInner { } } + // try to set TimeZone + if let Some(tz) = client.metadata().get("TimeZone") { + match Timezone::from_tz_string(tz) { + Ok(tz) => self.session.set_timezone(tz), + Err(_) => { + send_error( + client, + PgErrorCode::Ec22023 + .to_err_info(format!("Invalid TimeZone: {}", tz)), + ) + .await?; + + return Ok(()); + } + } + } + if self.login_verifier.user_provider.is_some() { client.set_state(PgWireConnectionState::AuthenticationInProgress); client