From d6b2d1dfb8e4dcab043f8cca0ba94d2946458e6e Mon Sep 17 00:00:00 2001 From: JohnsonLee <53596783+J0HN50N133@users.noreply.github.com> Date: Mon, 1 Apr 2024 15:31:36 +0800 Subject: [PATCH] feat: Support outputting various date styles for postgresql (#3602) * test: add integration_test for datetime style * feat: support various datestyle for postgres * doc: rewrite the comment about merge_datestyle_value * test: add more test to illustrate valid datestyle input --- src/operator/src/statement.rs | 101 +---- src/operator/src/statement/set.rs | 179 +++++++++ src/query/src/sql.rs | 4 + src/servers/src/postgres/types.rs | 13 +- src/servers/src/postgres/types/datetime.rs | 406 +++++++++++++++++++++ src/session/src/context.rs | 12 +- src/session/src/session_config.rs | 113 ++++++ tests-integration/tests/sql.rs | 228 +++++++++++- 8 files changed, 957 insertions(+), 99 deletions(-) create mode 100644 src/operator/src/statement/set.rs create mode 100644 src/servers/src/postgres/types/datetime.rs diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index d8c0d262fa..7d66ac2775 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -18,6 +18,7 @@ mod copy_table_to; mod ddl; mod describe; mod dml; +mod set; mod show; mod tql; @@ -33,28 +34,27 @@ use common_meta::table_name::TableName; use common_query::Output; use common_telemetry::tracing; use common_time::range::TimestampRange; -use common_time::{Timestamp, Timezone}; +use common_time::Timestamp; use partition::manager::{PartitionRuleManager, PartitionRuleManagerRef}; use query::parser::QueryStatement; use query::plan::LogicalPlan; use query::QueryEngineRef; use session::context::QueryContextRef; -use session::session_config::PGByteaOutputValue; use session::table_name::table_idents_to_full_name; -use snafu::{ensure, OptionExt, ResultExt}; +use snafu::{OptionExt, ResultExt}; use sql::statements::copy::{CopyDatabase, CopyDatabaseArgument, CopyTable, CopyTableArgument}; -use sql::statements::set_variables::SetVariables; use sql::statements::statement::Statement; use sql::statements::OptionMap; use sql::util::format_raw_object_name; -use sqlparser::ast::{Expr, Ident, ObjectName, Value}; +use sqlparser::ast::ObjectName; use table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest}; use table::table_reference::TableReference; use table::TableRef; +use self::set::{set_bytea_output, set_datestyle, set_timezone, validate_client_encoding}; use crate::error::{ - self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidConfigValueSnafu, - InvalidSqlSnafu, NotSupportedSnafu, PlanStatementSnafu, Result, TableNotFoundSnafu, + self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu, + PlanStatementSnafu, Result, TableNotFoundSnafu, }; use crate::insert::InserterRef; use crate::statement::copy_database::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY}; @@ -215,18 +215,12 @@ impl StatementExecutor { match var_name.as_str() { "TIMEZONE" | "TIME_ZONE" => set_timezone(set_var.value, query_ctx)?, - // Some postgresql client app may submit a "SET bytea_output" stmt upon connection. - // However, currently we lack the support for it (tracked in https://github.com/GreptimeTeam/greptimedb/issues/3438), - // so we just ignore it here instead of returning an error to break the connection. - // Since the "bytea_output" only determines the output format of binary values, - // it won't cause much trouble if we do so. "BYTEA_OUTPUT" => set_bytea_output(set_var.value, query_ctx)?, // Same as "bytea_output", we just ignore it here. // Not harmful since it only relates to how date is viewed in client app's output. // The tracked issue is https://github.com/GreptimeTeam/greptimedb/issues/3442. - // TODO(#3442): Remove this temporary workaround after the feature is implemented. - "DATESTYLE" => (), + "DATESTYLE" => set_datestyle(set_var.value, query_ctx)?, "CLIENT_ENCODING" => validate_client_encoding(set_var)?, _ => { @@ -283,85 +277,6 @@ impl StatementExecutor { } } -fn validate_client_encoding(set: SetVariables) -> Result<()> { - let Some((encoding, [])) = set.value.split_first() else { - return InvalidSqlSnafu { - err_msg: "must provide one and only one client encoding value", - } - .fail(); - }; - let encoding = match encoding { - Expr::Value(Value::SingleQuotedString(x)) - | Expr::Identifier(Ident { - value: x, - quote_style: _, - }) => x.to_uppercase(), - _ => { - return InvalidSqlSnafu { - err_msg: format!("client encoding must be a string, actual: {:?}", encoding), - } - .fail(); - } - }; - // For the sake of simplicity, we only support "UTF8" ("UNICODE" is the alias for it, - // see https://www.postgresql.org/docs/current/multibyte.html#MULTIBYTE-CHARSET-SUPPORTED). - // "UTF8" is universal and sufficient for almost all cases. - // GreptimeDB itself is always using "UTF8" as the internal encoding. - ensure!( - encoding == "UTF8" || encoding == "UNICODE", - NotSupportedSnafu { - feat: format!("client encoding of '{}'", encoding) - } - ); - Ok(()) -} - -fn set_timezone(exprs: Vec, ctx: QueryContextRef) -> Result<()> { - let tz_expr = exprs.first().context(NotSupportedSnafu { - feat: "No timezone find in set variable statement", - })?; - match tz_expr { - Expr::Value(Value::SingleQuotedString(tz)) | Expr::Value(Value::DoubleQuotedString(tz)) => { - match Timezone::from_tz_string(tz.as_str()) { - Ok(timezone) => ctx.set_timezone(timezone), - Err(_) => { - return NotSupportedSnafu { - feat: format!("Invalid timezone expr {} in set variable statement", tz), - } - .fail() - } - } - Ok(()) - } - expr => NotSupportedSnafu { - feat: format!( - "Unsupported timezone expr {} in set variable statement", - expr - ), - } - .fail(), - } -} - -fn set_bytea_output(exprs: Vec, ctx: QueryContextRef) -> Result<()> { - let Some((var_value, [])) = exprs.split_first() else { - return (NotSupportedSnafu { - feat: "Set variable value must have one and only one value for bytea_output", - }) - .fail(); - }; - let Expr::Value(value) = var_value else { - return (NotSupportedSnafu { - feat: "Set variable value must be a value", - }) - .fail(); - }; - ctx.configuration_parameter().set_postgres_bytea_output( - PGByteaOutputValue::try_from(value.clone()).context(InvalidConfigValueSnafu)?, - ); - Ok(()) -} - fn to_copy_table_request(stmt: CopyTable, query_ctx: QueryContextRef) -> Result { let direction = match stmt { CopyTable::To(_) => CopyDirection::Export, diff --git a/src/operator/src/statement/set.rs b/src/operator/src/statement/set.rs new file mode 100644 index 0000000000..6436f136d9 --- /dev/null +++ b/src/operator/src/statement/set.rs @@ -0,0 +1,179 @@ +// 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 common_time::Timezone; +use session::context::QueryContextRef; +use session::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle}; +use snafu::{ensure, OptionExt, ResultExt}; +use sql::ast::{Expr, Ident, Value}; +use sql::statements::set_variables::SetVariables; + +use crate::error::{InvalidConfigValueSnafu, InvalidSqlSnafu, NotSupportedSnafu, Result}; + +pub fn set_timezone(exprs: Vec, ctx: QueryContextRef) -> Result<()> { + let tz_expr = exprs.first().context(NotSupportedSnafu { + feat: "No timezone find in set variable statement", + })?; + match tz_expr { + Expr::Value(Value::SingleQuotedString(tz)) | Expr::Value(Value::DoubleQuotedString(tz)) => { + match Timezone::from_tz_string(tz.as_str()) { + Ok(timezone) => ctx.set_timezone(timezone), + Err(_) => { + return NotSupportedSnafu { + feat: format!("Invalid timezone expr {} in set variable statement", tz), + } + .fail() + } + } + Ok(()) + } + expr => NotSupportedSnafu { + feat: format!( + "Unsupported timezone expr {} in set variable statement", + expr + ), + } + .fail(), + } +} + +pub fn set_bytea_output(exprs: Vec, ctx: QueryContextRef) -> Result<()> { + let Some((var_value, [])) = exprs.split_first() else { + return (NotSupportedSnafu { + feat: "Set variable value must have one and only one value for bytea_output", + }) + .fail(); + }; + let Expr::Value(value) = var_value else { + return (NotSupportedSnafu { + feat: "Set variable value must be a value", + }) + .fail(); + }; + ctx.configuration_parameter().set_postgres_bytea_output( + PGByteaOutputValue::try_from(value.clone()).context(InvalidConfigValueSnafu)?, + ); + Ok(()) +} + +pub fn validate_client_encoding(set: SetVariables) -> Result<()> { + let Some((encoding, [])) = set.value.split_first() else { + return InvalidSqlSnafu { + err_msg: "must provide one and only one client encoding value", + } + .fail(); + }; + let encoding = match encoding { + Expr::Value(Value::SingleQuotedString(x)) + | Expr::Identifier(Ident { + value: x, + quote_style: _, + }) => x.to_uppercase(), + _ => { + return InvalidSqlSnafu { + err_msg: format!("client encoding must be a string, actual: {:?}", encoding), + } + .fail(); + } + }; + // For the sake of simplicity, we only support "UTF8" ("UNICODE" is the alias for it, + // see https://www.postgresql.org/docs/current/multibyte.html#MULTIBYTE-CHARSET-SUPPORTED). + // "UTF8" is universal and sufficient for almost all cases. + // GreptimeDB itself is always using "UTF8" as the internal encoding. + ensure!( + encoding == "UTF8" || encoding == "UNICODE", + NotSupportedSnafu { + feat: format!("client encoding of '{}'", encoding) + } + ); + Ok(()) +} + +// if one of original value and new value is none, return the other one +// returns new values only when it equals to original one else return error. +// This is only used for handling datestyle +fn merge_datestyle_value(value: Option, new_value: Option) -> Result> +where + T: PartialEq, +{ + match (&value, &new_value) { + (None, _) => Ok(new_value), + (_, None) => Ok(value), + (Some(v1), Some(v2)) if v1 == v2 => Ok(new_value), + _ => InvalidSqlSnafu { + err_msg: "Conflicting \"datestyle\" specifications.", + } + .fail(), + } +} + +fn try_parse_datestyle(expr: &Expr) -> Result<(Option, Option)> { + enum ParsedDateStyle { + Order(PGDateOrder), + Style(PGDateTimeStyle), + } + fn try_parse_str(s: &str) -> Result { + PGDateTimeStyle::try_from(s) + .map_or_else( + |_| PGDateOrder::try_from(s).map(ParsedDateStyle::Order), + |style| Ok(ParsedDateStyle::Style(style)), + ) + .context(InvalidConfigValueSnafu) + } + match expr { + Expr::Identifier(Ident { + value: s, + quote_style: _, + }) + | Expr::Value(Value::SingleQuotedString(s)) + | Expr::Value(Value::DoubleQuotedString(s)) => { + s.split(',') + .map(|s| s.trim()) + .try_fold((None, None), |(style, order), s| match try_parse_str(s)? { + ParsedDateStyle::Order(o) => { + Ok((style, merge_datestyle_value(order, Some(o))?)) + } + ParsedDateStyle::Style(s) => { + Ok((merge_datestyle_value(style, Some(s))?, order)) + } + }) + } + _ => NotSupportedSnafu { + feat: "Not supported expression for datestyle", + } + .fail(), + } +} + +pub fn set_datestyle(exprs: Vec, ctx: QueryContextRef) -> Result<()> { + // ORDER, + // STYLE, + // ORDER,ORDER + // ORDER,STYLE + // STYLE,ORDER + let (style, order) = exprs + .iter() + .try_fold((None, None), |(style, order), expr| { + let (new_style, new_order) = try_parse_datestyle(expr)?; + Ok(( + merge_datestyle_value(style, new_style)?, + merge_datestyle_value(order, new_order)?, + )) + })?; + + let (old_style, older_order) = *ctx.configuration_parameter().pg_datetime_style(); + ctx.configuration_parameter() + .set_pg_datetime_style(style.unwrap_or(old_style), order.unwrap_or(older_order)); + Ok(()) +} diff --git a/src/query/src/sql.rs b/src/query/src/sql.rs index 99ca99f54a..b3a0c4806e 100644 --- a/src/query/src/sql.rs +++ b/src/query/src/sql.rs @@ -499,6 +499,10 @@ 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(), + "DATESTYLE" => { + let (style, order) = *query_ctx.configuration_parameter().pg_datetime_style(); + format!("{}, {}", style, order) + } _ => return UnsupportedVariableSnafu { name: variable }.fail(), }; let schema = Arc::new(Schema::new(vec![ColumnSchema::new( diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 01351f8541..269eac2a95 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -12,7 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub mod bytea; +mod bytea; +mod datetime; mod interval; use std::collections::HashMap; @@ -33,6 +34,7 @@ use session::context::QueryContextRef; use session::session_config::PGByteaOutputValue; use self::bytea::{EscapeOutputBytea, HexOutputBytea}; +use self::datetime::{StylingDate, StylingDateTime}; use self::interval::PgInterval; use crate::error::{self, Error, Result}; use crate::SqlPlan; @@ -82,7 +84,8 @@ pub(super) fn encode_value( } Value::Date(v) => { if let Some(date) = v.to_chrono_date() { - builder.encode_field(&date) + let (style, order) = *query_ctx.configuration_parameter().pg_datetime_style(); + builder.encode_field(&StylingDate(&date, style, order)) } else { Err(PgWireError::ApiError(Box::new(Error::Internal { err_msg: format!("Failed to convert date to postgres type {v:?}",), @@ -91,7 +94,8 @@ pub(super) fn encode_value( } Value::DateTime(v) => { if let Some(datetime) = v.to_chrono_datetime() { - builder.encode_field(&datetime) + let (style, order) = *query_ctx.configuration_parameter().pg_datetime_style(); + builder.encode_field(&StylingDateTime(&datetime, style, order)) } else { Err(PgWireError::ApiError(Box::new(Error::Internal { err_msg: format!("Failed to convert date to postgres type {v:?}",), @@ -100,7 +104,8 @@ pub(super) fn encode_value( } Value::Timestamp(v) => { if let Some(datetime) = v.to_chrono_datetime() { - builder.encode_field(&datetime) + let (style, order) = *query_ctx.configuration_parameter().pg_datetime_style(); + builder.encode_field(&StylingDateTime(&datetime, style, order)) } else { Err(PgWireError::ApiError(Box::new(Error::Internal { err_msg: format!("Failed to convert date to postgres type {v:?}",), diff --git a/src/servers/src/postgres/types/datetime.rs b/src/servers/src/postgres/types/datetime.rs new file mode 100644 index 0000000000..fc324c047b --- /dev/null +++ b/src/servers/src/postgres/types/datetime.rs @@ -0,0 +1,406 @@ +// 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 bytes::BufMut; +use chrono::{NaiveDate, NaiveDateTime}; +use pgwire::types::ToSqlText; +use postgres_types::{IsNull, ToSql, Type}; +use session::session_config::{PGDateOrder, PGDateTimeStyle}; + +#[derive(Debug)] +pub struct StylingDate<'a>(pub &'a NaiveDate, pub PGDateTimeStyle, pub PGDateOrder); + +#[derive(Debug)] +pub struct StylingDateTime<'a>(pub &'a NaiveDateTime, pub PGDateTimeStyle, pub PGDateOrder); + +fn date_format_string(style: PGDateTimeStyle, order: PGDateOrder) -> &'static str { + match style { + PGDateTimeStyle::ISO => "%Y-%m-%d", + PGDateTimeStyle::German => "%d.%m.%Y", + PGDateTimeStyle::Postgres => match order { + PGDateOrder::MDY | PGDateOrder::YMD => "%m-%d-%Y", + PGDateOrder::DMY => "%d-%m-%Y", + }, + PGDateTimeStyle::SQL => match order { + PGDateOrder::MDY | PGDateOrder::YMD => "%m/%d/%Y", + PGDateOrder::DMY => "%d/%m/%Y", + }, + } +} + +fn datetime_format_string(style: PGDateTimeStyle, order: PGDateOrder) -> &'static str { + match style { + PGDateTimeStyle::ISO => "%Y-%m-%d %H:%M:%S%.6f", + PGDateTimeStyle::German => "%d.%m.%Y %H:%M:%S%.6f", + PGDateTimeStyle::Postgres => match order { + PGDateOrder::MDY | PGDateOrder::YMD => "%a %b %d %H:%M:%S%.6f %Y", + PGDateOrder::DMY => "%a %d %b %H:%M:%S%.6f %Y", + }, + PGDateTimeStyle::SQL => match order { + PGDateOrder::MDY | PGDateOrder::YMD => "%m/%d/%Y %H:%M:%S%.6f", + PGDateOrder::DMY => "%d/%m/%Y %H:%M:%S%.6f", + }, + } +} +impl ToSqlText for StylingDate<'_> { + fn to_sql_text( + &self, + ty: &Type, + out: &mut bytes::BytesMut, + ) -> std::result::Result> + where + Self: Sized, + { + match *ty { + Type::DATE => { + let fmt = self + .0 + .format(date_format_string(self.1, self.2)) + .to_string(); + out.put_slice(fmt.as_bytes()); + } + _ => { + self.0.to_sql_text(ty, out)?; + } + } + Ok(IsNull::No) + } +} + +impl ToSqlText for StylingDateTime<'_> { + fn to_sql_text( + &self, + ty: &Type, + out: &mut bytes::BytesMut, + ) -> Result> + where + Self: Sized, + { + match *ty { + Type::TIMESTAMP => { + let fmt = self + .0 + .format(datetime_format_string(self.1, self.2)) + .to_string(); + out.put_slice(fmt.as_bytes()); + } + Type::DATE => { + let fmt = self + .0 + .format(date_format_string(self.1, self.2)) + .to_string(); + out.put_slice(fmt.as_bytes()); + } + _ => { + self.0.to_sql_text(ty, out)?; + } + } + Ok(IsNull::No) + } +} + +macro_rules! delegate_to_sql { + ($delegator:ident, $delegatee:ident) => { + impl ToSql for $delegator<'_> { + fn to_sql( + &self, + ty: &Type, + out: &mut bytes::BytesMut, + ) -> std::result::Result> { + self.0.to_sql(ty, out) + } + + fn accepts(ty: &Type) -> bool { + <$delegatee as ToSql>::accepts(ty) + } + + fn to_sql_checked( + &self, + ty: &Type, + out: &mut bytes::BytesMut, + ) -> Result> { + self.0.to_sql_checked(ty, out) + } + } + }; +} + +delegate_to_sql!(StylingDate, NaiveDate); +delegate_to_sql!(StylingDateTime, NaiveDateTime); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_styling_date() { + let naive_date = NaiveDate::from_ymd_opt(1997, 12, 17).unwrap(); + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::ISO, PGDateOrder::MDY); + let expected = "1997-12-17"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::ISO, PGDateOrder::YMD); + let expected = "1997-12-17"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::ISO, PGDateOrder::DMY); + let expected = "1997-12-17"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::German, PGDateOrder::MDY); + let expected = "17.12.1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::German, PGDateOrder::YMD); + let expected = "17.12.1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::German, PGDateOrder::DMY); + let expected = "17.12.1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = + StylingDate(&naive_date, PGDateTimeStyle::Postgres, PGDateOrder::MDY); + let expected = "12-17-1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = + StylingDate(&naive_date, PGDateTimeStyle::Postgres, PGDateOrder::YMD); + let expected = "12-17-1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = + StylingDate(&naive_date, PGDateTimeStyle::Postgres, PGDateOrder::DMY); + let expected = "17-12-1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::SQL, PGDateOrder::MDY); + let expected = "12/17/1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::SQL, PGDateOrder::YMD); + let expected = "12/17/1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_date = StylingDate(&naive_date, PGDateTimeStyle::SQL, PGDateOrder::DMY); + let expected = "17/12/1997"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_date.to_sql_text(&Type::DATE, &mut out).unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + } + + #[test] + fn test_styling_datetime() { + let input = + NaiveDateTime::parse_from_str("2021-09-01 12:34:56.789012", "%Y-%m-%d %H:%M:%S%.f") + .unwrap(); + + { + let styling_datetime = StylingDateTime(&input, PGDateTimeStyle::ISO, PGDateOrder::MDY); + let expected = "2021-09-01 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = StylingDateTime(&input, PGDateTimeStyle::ISO, PGDateOrder::YMD); + let expected = "2021-09-01 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = StylingDateTime(&input, PGDateTimeStyle::ISO, PGDateOrder::DMY); + let expected = "2021-09-01 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = + StylingDateTime(&input, PGDateTimeStyle::German, PGDateOrder::MDY); + let expected = "01.09.2021 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = + StylingDateTime(&input, PGDateTimeStyle::German, PGDateOrder::YMD); + let expected = "01.09.2021 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = + StylingDateTime(&input, PGDateTimeStyle::German, PGDateOrder::DMY); + let expected = "01.09.2021 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = + StylingDateTime(&input, PGDateTimeStyle::Postgres, PGDateOrder::MDY); + let expected = "Wed Sep 01 12:34:56.789012 2021"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = + StylingDateTime(&input, PGDateTimeStyle::Postgres, PGDateOrder::YMD); + let expected = "Wed Sep 01 12:34:56.789012 2021"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = + StylingDateTime(&input, PGDateTimeStyle::Postgres, PGDateOrder::DMY); + let expected = "Wed 01 Sep 12:34:56.789012 2021"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = StylingDateTime(&input, PGDateTimeStyle::SQL, PGDateOrder::MDY); + let expected = "09/01/2021 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = StylingDateTime(&input, PGDateTimeStyle::SQL, PGDateOrder::YMD); + let expected = "09/01/2021 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + + { + let styling_datetime = StylingDateTime(&input, PGDateTimeStyle::SQL, PGDateOrder::DMY); + let expected = "01/09/2021 12:34:56.789012"; + let mut out = bytes::BytesMut::new(); + let is_null = styling_datetime + .to_sql_text(&Type::TIMESTAMP, &mut out) + .unwrap(); + assert!(matches!(is_null, IsNull::No)); + assert_eq!(out, expected.as_bytes()); + } + } +} diff --git a/src/session/src/context.rs b/src/session/src/context.rs index 181d46d87f..25091e1cf5 100644 --- a/src/session/src/context.rs +++ b/src/session/src/context.rs @@ -27,7 +27,7 @@ use common_time::Timezone; use derive_builder::Builder; use sql::dialect::{Dialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect}; -use crate::session_config::PGByteaOutputValue; +use crate::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle}; use crate::SessionRef; pub type QueryContextRef = Arc; @@ -282,12 +282,14 @@ impl Display for Channel { #[derive(Default, Debug)] pub struct ConfigurationVariables { postgres_bytea_output: ArcSwap, + pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>, } impl Clone for ConfigurationVariables { fn clone(&self) -> Self { Self { postgres_bytea_output: ArcSwap::new(self.postgres_bytea_output.load().clone()), + pg_datestyle_format: ArcSwap::new(self.pg_datestyle_format.load().clone()), } } } @@ -304,6 +306,14 @@ impl ConfigurationVariables { pub fn postgres_bytea_output(&self) -> Arc { self.postgres_bytea_output.load().clone() } + + pub fn pg_datetime_style(&self) -> Arc<(PGDateTimeStyle, PGDateOrder)> { + self.pg_datestyle_format.load().clone() + } + + pub fn set_pg_datetime_style(&self, style: PGDateTimeStyle, order: PGDateOrder) { + self.pg_datestyle_format.swap(Arc::new((style, order))); + } } #[cfg(test)] diff --git a/src/session/src/session_config.rs b/src/session/src/session_config.rs index aad50e70c1..2f6bdbbf40 100644 --- a/src/session/src/session_config.rs +++ b/src/session/src/session_config.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::fmt::Display; + use common_macro::stack_trace_debug; use snafu::{Location, Snafu}; use sql::ast::Value; @@ -62,3 +64,114 @@ impl TryFrom for PGByteaOutputValue { } } } + +// Refers to: https://www.postgresql.org/docs/current/runtime-config-client.html#GUC-DATESTYLE +#[derive(Default, PartialEq, Eq, Clone, Copy, Debug)] +pub enum PGDateOrder { + #[default] + MDY, + DMY, + YMD, +} + +impl Display for PGDateOrder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PGDateOrder::MDY => write!(f, "MDY"), + PGDateOrder::DMY => write!(f, "DMY"), + PGDateOrder::YMD => write!(f, "YMD"), + } + } +} + +impl TryFrom<&str> for PGDateOrder { + type Error = Error; + + fn try_from(s: &str) -> Result { + match s.to_uppercase().as_str() { + "US" | "NONEURO" | "NONEUROPEAN" | "MDY" => Ok(PGDateOrder::MDY), + "EUROPEAN" | "EURO" | "DMY" => Ok(PGDateOrder::DMY), + "YMD" => Ok(PGDateOrder::YMD), + _ => InvalidConfigValueSnafu { + name: "DateStyle", + value: s, + hint: format!("Unrecognized key word: {}", s), + } + .fail(), + } + } +} +impl TryFrom<&Value> for PGDateOrder { + type Error = Error; + + fn try_from(value: &Value) -> Result { + match value { + Value::DoubleQuotedString(s) | Value::SingleQuotedString(s) => { + Self::try_from(s.as_str()) + } + _ => InvalidConfigValueSnafu { + name: "DateStyle", + value: value.to_string(), + hint: format!("Unrecognized key word: {}", value), + } + .fail(), + } + } +} + +#[derive(Default, PartialEq, Eq, Clone, Copy, Debug)] +pub enum PGDateTimeStyle { + #[default] + ISO, + SQL, + Postgres, + German, +} + +impl Display for PGDateTimeStyle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PGDateTimeStyle::ISO => write!(f, "ISO"), + PGDateTimeStyle::SQL => write!(f, "SQL"), + PGDateTimeStyle::Postgres => write!(f, "Postgres"), + PGDateTimeStyle::German => write!(f, "German"), + } + } +} + +impl TryFrom<&str> for PGDateTimeStyle { + type Error = Error; + + fn try_from(s: &str) -> Result { + match s.to_uppercase().as_str() { + "ISO" => Ok(PGDateTimeStyle::ISO), + "SQL" => Ok(PGDateTimeStyle::SQL), + "POSTGRES" => Ok(PGDateTimeStyle::Postgres), + "GERMAN" => Ok(PGDateTimeStyle::German), + _ => InvalidConfigValueSnafu { + name: "DateStyle", + value: s, + hint: format!("Unrecognized key word: {}", s), + } + .fail(), + } + } +} + +impl TryFrom<&Value> for PGDateTimeStyle { + type Error = Error; + + fn try_from(value: &Value) -> Result { + match value { + Value::DoubleQuotedString(s) | Value::SingleQuotedString(s) => { + Self::try_from(s.as_str()) + } + _ => InvalidConfigValueSnafu { + name: "DateStyle", + value: value.to_string(), + hint: format!("Unrecognized key word: {}", value), + } + .fail(), + } + } +} diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 4628bc372f..0bf238b8a1 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + use auth::user_provider_from_option; use chrono::{DateTime, NaiveDate, NaiveDateTime, SecondsFormat, Utc}; use sqlx::mysql::{MySqlConnection, MySqlDatabaseError, MySqlPoolOptions}; @@ -21,7 +23,7 @@ use tests_integration::test_util::{ setup_mysql_server, setup_mysql_server_with_user_provider, setup_pg_server, setup_pg_server_with_user_provider, StorageType, }; -use tokio_postgres::{NoTls, SimpleQueryMessage}; +use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; #[macro_export] macro_rules! sql_test { @@ -61,6 +63,7 @@ macro_rules! sql_tests { test_postgres_crud, test_postgres_timezone, test_postgres_bytea, + test_postgres_datestyle, test_postgres_parameter_inference, test_mysql_prepare_stmt_insert_timestamp, ); @@ -479,6 +482,229 @@ pub async fn test_postgres_bytea(store_type: StorageType) { let _ = fe_pg_server.shutdown().await; guard.remove_all().await; } + +pub async fn test_postgres_datestyle(store_type: StorageType) { + let (addr, mut guard, fe_pg_server) = setup_pg_server(store_type, "various datestyle").await; + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + + tokio::spawn(async move { + connection.await.unwrap(); + }); + + let validate_datestyle = |client: Client, datestyle: &str, is_valid: bool| { + let datestyle = datestyle.to_string(); + async move { + assert_eq!( + client + .simple_query(format!("SET DATESTYLE={}", datestyle).as_str()) + .await + .is_ok(), + is_valid + ); + client + } + }; + + // style followed by order is valid + let client = validate_datestyle(client, "'ISO,MDY'", true).await; + + // Mix of string and ident is valid + let client = validate_datestyle(client, "'ISO',MDY", true).await; + + // list of string that didn't corrupt is valid + let client = validate_datestyle(client, "'ISO,MDY','ISO,MDY'", true).await; + + // corrupted style + let client = validate_datestyle(client, "'ISO,German'", false).await; + + // corrupted order + let client = validate_datestyle(client, "'ISO,DMY','ISO,MDY'", false).await; + + // as long as the value is not corrupted, it's valid + let client = validate_datestyle(client, "ISO,ISO,ISO,ISO,ISO,MDY,MDY,MDY,MDY", true).await; + + let _ = client + .simple_query("CREATE TABLE ts_test(ts TIMESTAMP TIME INDEX)") + .await + .expect("CREATE TABLE ts_test ERROR"); + let _ = client + .simple_query("CREATE TABLE date_test(d date, ts TIMESTAMP TIME INDEX)") + .await + .expect("CREATE TABLE date_test ERROR"); + + let _ = client + .simple_query("CREATE TABLE dt_test(dt datetime, ts TIMESTAMP TIME INDEX)") + .await + .expect("CREATE TABLE dt_test ERROR"); + + let _ = client + .simple_query("INSERT INTO ts_test VALUES('1997-12-17 07:37:16.123')") + .await + .expect("INSERT INTO ts_test ERROR"); + + let _ = client + .simple_query("INSERT INTO date_test VALUES('1997-12-17', '1997-12-17 07:37:16.123')") + .await + .expect("INSERT INTO date_test ERROR"); + + let _ = client + .simple_query( + "INSERT INTO dt_test VALUES('1997-12-17 07:37:16.123', '1997-12-17 07:37:16.123')", + ) + .await + .expect("INSERT INTO dt_test ERROR"); + + let get_row = |mess: Vec| -> String { + match &mess[0] { + SimpleQueryMessage::Row(row) => row.get(0).unwrap().to_string(), + _ => unreachable!(), + } + }; + + let date = "DATE"; + let datetime = "TIMESTAMP"; + let timestamp = "TIMESTAMP"; + + let iso = "ISO"; + let sql = "SQL"; + let postgres = "Postgres"; + let german = "German"; + + let expected_set: HashMap<&str, HashMap<&str, HashMap<&str, &str>>> = HashMap::from([ + ( + date, + HashMap::from([ + ( + iso, + HashMap::from([ + ("MDY", "1997-12-17"), + ("DMY", "1997-12-17"), + ("YMD", "1997-12-17"), + ]), + ), + ( + sql, + HashMap::from([ + ("MDY", "12/17/1997"), + ("DMY", "17/12/1997"), + ("YMD", "12/17/1997"), + ]), + ), + ( + postgres, + HashMap::from([ + ("MDY", "12-17-1997"), + ("DMY", "17-12-1997"), + ("YMD", "12-17-1997"), + ]), + ), + ( + german, + HashMap::from([ + ("MDY", "17.12.1997"), + ("DMY", "17.12.1997"), + ("YMD", "17.12.1997"), + ]), + ), + ]), + ), + ( + timestamp, + HashMap::from([ + ( + iso, + HashMap::from([ + ("MDY", "1997-12-17 07:37:16.123000"), + ("DMY", "1997-12-17 07:37:16.123000"), + ("YMD", "1997-12-17 07:37:16.123000"), + ]), + ), + ( + sql, + HashMap::from([ + ("MDY", "12/17/1997 07:37:16.123000"), + ("DMY", "17/12/1997 07:37:16.123000"), + ("YMD", "12/17/1997 07:37:16.123000"), + ]), + ), + ( + postgres, + HashMap::from([ + ("MDY", "Wed Dec 17 07:37:16.123000 1997"), + ("DMY", "Wed 17 Dec 07:37:16.123000 1997"), + ("YMD", "Wed Dec 17 07:37:16.123000 1997"), + ]), + ), + ( + german, + HashMap::from([ + ("MDY", "17.12.1997 07:37:16.123000"), + ("DMY", "17.12.1997 07:37:16.123000"), + ("YMD", "17.12.1997 07:37:16.123000"), + ]), + ), + ]), + ), + ]); + + let get_expected = |ty: &str, style: &str, order: &str| { + expected_set + .get(ty) + .and_then(|m| m.get(style)) + .and_then(|m2| m2.get(order)) + .unwrap() + .to_string() + }; + + for style in ["ISO", "SQL", "Postgres", "German"] { + for order in ["MDY", "DMY", "YMD"] { + let _ = client + .simple_query(&format!("SET DATESTYLE='{}', '{}'", style, order)) + .await + .expect("SET DATESTYLE ERROR"); + + let r = client.simple_query("SELECT ts FROM ts_test").await.unwrap(); + let ts = get_row(r); + assert_eq!( + ts, + get_expected(timestamp, style, order), + "style: {}, order: {}", + style, + order + ); + + let r = client + .simple_query("SELECT d FROM date_test") + .await + .unwrap(); + let d = get_row(r); + assert_eq!( + d, + get_expected(date, style, order), + "style: {}, order: {}", + style, + order + ); + + let r = client.simple_query("SELECT dt FROM dt_test").await.unwrap(); + let dt = get_row(r); + assert_eq!( + dt, + get_expected(datetime, style, order), + "style: {}, order: {}", + style, + order + ); + } + } + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_timezone(store_type: StorageType) { let (addr, mut guard, fe_pg_server) = setup_pg_server(store_type, "sql_inference").await;