feat: Support printing postgresql's bytea data type in its "hex" and "escape" format

src/operator/src/error.rs

@@ -20,9 +20,9 @@ use common_error::status_code::StatusCode;
 use common_macro::stack_trace_debug;
 use datafusion::parquet;
 use datatypes::arrow::error::ArrowError;
-use datatypes::value::Value;
 use servers::define_into_tonic_status;
 use snafu::{Location, Snafu};
+use sql::ast::Value;
 
 #[derive(Snafu)]
 #[snafu(visibility(pub))]
@@ -528,6 +528,12 @@ pub enum Error {
 
     #[snafu(display("Invalid partition rule: {}", reason))]
     InvalidPartitionRule { reason: String, location: Location },
+
+    #[snafu(display("Invalid configuration value."))]
+    InvalidConfigValue {
+        source: session::session_config::Error,
+        location: Location,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -536,6 +542,7 @@ impl ErrorExt for Error {
     fn status_code(&self) -> StatusCode {
         match self {
             Error::InvalidSql { .. }
+            | Error::InvalidConfigValue { .. }
             | Error::InvalidInsertRequest { .. }
             | Error::InvalidDeleteRequest { .. }
             | Error::IllegalPrimaryKeysDef { .. }

src/operator/src/statement.rs

@@ -39,6 +39,7 @@ 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 sql::statements::copy::{CopyDatabase, CopyDatabaseArgument, CopyTable, CopyTableArgument};
@@ -52,8 +53,8 @@ use table::table_reference::TableReference;
 use table::TableRef;
 
 use crate::error::{
-    self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu,
-    PlanStatementSnafu, Result, TableNotFoundSnafu,
+    self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidConfigValueSnafu,
+    InvalidSqlSnafu, NotSupportedSnafu, PlanStatementSnafu, Result, TableNotFoundSnafu,
 };
 use crate::insert::InserterRef;
 use crate::statement::copy_database::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY};
@@ -219,8 +220,7 @@ impl StatementExecutor {
                     // 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.
-                    // TODO(#3438): Remove this temporary workaround after the feature is implemented.
-                    "BYTEA_OUTPUT" => (),
+                    "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.
@@ -339,6 +339,25 @@ fn set_timezone(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
     }
 }
 
+fn set_bytea_output(exprs: Vec<Expr>, 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<CopyTableRequest> {
     let direction = match stmt {
         CopyTable::To(_) => CopyDirection::Export,

src/servers/Cargo.toml

@@ -41,6 +41,7 @@ common-recordbatch.workspace = true
 common-runtime.workspace = true
 common-telemetry.workspace = true
 common-time.workspace = true
+dashmap.workspace = true
 datafusion.workspace = true
 datafusion-common.workspace = true
 datatypes.workspace = true

src/servers/src/mysql/federated.rs

@@ -310,7 +310,7 @@ mod test {
 
     #[test]
     fn test_check() {
-        let session = Arc::new(Session::new(None, Channel::Mysql));
+        let session = Arc::new(Session::new(None, Channel::Mysql, Default::default()));
         let query = "select 1";
         let result = check(query, QueryContext::arc(), session.clone());
         assert!(result.is_none());
@@ -320,7 +320,7 @@ mod test {
         assert!(output.is_none());
 
         fn test(query: &str, expected: &str) {
-            let session = Arc::new(Session::new(None, Channel::Mysql));
+            let session = Arc::new(Session::new(None, Channel::Mysql, Default::default()));
             let output = check(query, QueryContext::arc(), session.clone());
             match output.unwrap().data {
                 OutputData::RecordBatches(r) => {

src/servers/src/mysql/handler.rs

@@ -85,7 +85,11 @@ impl MysqlInstanceShim {
         MysqlInstanceShim {
             query_handler,
             salt: scramble,
-            session: Arc::new(Session::new(Some(client_addr), Channel::Mysql)),
+            session: Arc::new(Session::new(
+                Some(client_addr),
+                Channel::Mysql,
+                Default::default(),
+            )),
             user_provider,
             prepared_stmts: Default::default(),
             prepared_stmts_counter: AtomicU32::new(1),

src/servers/src/postgres.rs

@@ -88,7 +88,7 @@ pub(crate) struct MakePostgresServerHandler {
 
 impl MakePostgresServerHandler {
     fn make(&self, addr: Option<SocketAddr>) -> PostgresServerHandler {
-        let session = Arc::new(Session::new(addr, Channel::Postgres));
+        let session = Arc::new(Session::new(addr, Channel::Postgres, Default::default()));
         PostgresServerHandler {
             query_handler: self.query_handler.clone(),
             login_verifier: PgLoginVerifier::new(self.user_provider.clone()),

src/servers/src/postgres/handler.rs

@@ -31,6 +31,7 @@ use pgwire::api::stmt::{QueryParser, StoredStatement};
 use pgwire::api::{ClientInfo, Type};
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 use query::query_engine::DescribeResult;
+use session::context::QueryContextRef;
 use session::Session;
 use sql::dialect::PostgreSqlDialect;
 use sql::parser::{ParseOptions, ParserContext};
@@ -63,7 +64,7 @@ impl SimpleQueryHandler for PostgresServerHandler {
         let mut results = Vec::with_capacity(outputs.len());
 
         for output in outputs {
-            let resp = output_to_query_response(output, &Format::UnifiedText)?;
+            let resp = output_to_query_response(query_ctx.clone(), output, &Format::UnifiedText)?;
             results.push(resp);
         }
 
@@ -72,6 +73,7 @@ impl SimpleQueryHandler for PostgresServerHandler {
 }
 
 fn output_to_query_response<'a>(
+    query_ctx: QueryContextRef,
     output: Result<Output>,
     field_format: &Format,
 ) -> PgWireResult<Response<'a>> {
@@ -82,11 +84,16 @@ fn output_to_query_response<'a>(
             }
             OutputData::Stream(record_stream) => {
                 let schema = record_stream.schema();
-                recordbatches_to_query_response(record_stream, schema, field_format)
+                recordbatches_to_query_response(query_ctx, record_stream, schema, field_format)
             }
             OutputData::RecordBatches(recordbatches) => {
                 let schema = recordbatches.schema();
-                recordbatches_to_query_response(recordbatches.as_stream(), schema, field_format)
+                recordbatches_to_query_response(
+                    query_ctx,
+                    recordbatches.as_stream(),
+                    schema,
+                    field_format,
+                )
             }
         },
         Err(e) => Ok(Response::Error(Box::new(ErrorInfo::new(
@@ -98,6 +105,7 @@ fn output_to_query_response<'a>(
 }
 
 fn recordbatches_to_query_response<'a, S>(
+    query_ctx: QueryContextRef,
     recordbatches_stream: S,
     schema: SchemaRef,
     field_format: &Format,
@@ -125,7 +133,7 @@ where
             row.and_then(|row| {
                 let mut encoder = DataRowEncoder::new(pg_schema_ref.clone());
                 for value in row.iter() {
-                    encode_value(value, &mut encoder)?;
+                    encode_value(&query_ctx, value, &mut encoder)?;
                 }
                 encoder.finish()
             })
@@ -224,7 +232,9 @@ impl ExtendedQueryHandler for PostgresServerHandler {
             let plan = plan
                 .replace_params_with_values(parameters_to_scalar_values(plan, portal)?.as_ref())
                 .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
-            self.query_handler.do_exec_plan(plan, query_ctx).await
+            self.query_handler
+                .do_exec_plan(plan, query_ctx.clone())
+                .await
         } else {
             // manually replace variables in prepared statement when no
             // logical_plan is generated. This happens when logical plan is not
@@ -234,10 +244,13 @@ impl ExtendedQueryHandler for PostgresServerHandler {
                 sql = sql.replace(&format!("${}", i + 1), &parameter_to_string(portal, i)?);
             }
 
-            self.query_handler.do_query(&sql, query_ctx).await.remove(0)
+            self.query_handler
+                .do_query(&sql, query_ctx.clone())
+                .await
+                .remove(0)
         };
 
-        output_to_query_response(output, &portal.result_column_format)
+        output_to_query_response(query_ctx, output, &portal.result_column_format)
     }
 
     async fn do_describe_statement<C>(

src/servers/src/postgres/types.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod bytea;
 mod interval;
 
 use std::collections::HashMap;
@@ -28,7 +29,10 @@ use pgwire::api::results::{DataRowEncoder, FieldInfo};
 use pgwire::api::Type;
 use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
 use query::plan::LogicalPlan;
+use session::context::QueryContextRef;
+use session::session_config::PGByteaOutputValue;
 
+use self::bytea::{EscapeOutputBytea, HexOutputBytea};
 use self::interval::PgInterval;
 use crate::error::{self, Error, Result};
 use crate::SqlPlan;
@@ -50,7 +54,11 @@ pub(super) fn schema_to_pg(origin: &Schema, field_formats: &Format) -> Result<Ve
         .collect::<Result<Vec<FieldInfo>>>()
 }
 
-pub(super) fn encode_value(value: &Value, builder: &mut DataRowEncoder) -> PgWireResult<()> {
+pub(super) fn encode_value(
+    query_ctx: &QueryContextRef,
+    value: &Value,
+    builder: &mut DataRowEncoder,
+) -> PgWireResult<()> {
     match value {
         Value::Null => builder.encode_field(&None::<&i8>),
         Value::Boolean(v) => builder.encode_field(v),
@@ -65,7 +73,13 @@ pub(super) fn encode_value(value: &Value, builder: &mut DataRowEncoder) -> PgWir
         Value::Float32(v) => builder.encode_field(&v.0),
         Value::Float64(v) => builder.encode_field(&v.0),
         Value::String(v) => builder.encode_field(&v.as_utf8()),
-        Value::Binary(v) => builder.encode_field(&v.deref()),
+        Value::Binary(v) => {
+            let bytea_output = query_ctx.configuration_parameter().postgres_bytea_output();
+            match *bytea_output {
+                PGByteaOutputValue::ESCAPE => builder.encode_field(&EscapeOutputBytea(v.deref())),
+                PGByteaOutputValue::HEX => builder.encode_field(&HexOutputBytea(v.deref())),
+            }
+        }
         Value::Date(v) => {
             if let Some(date) = v.to_chrono_date() {
                 builder.encode_field(&date)
@@ -563,6 +577,7 @@ mod test {
     use datatypes::value::ListValue;
     use pgwire::api::results::{FieldFormat, FieldInfo};
     use pgwire::api::Type;
+    use session::context::QueryContextBuilder;
 
     use super::*;
 
@@ -784,12 +799,16 @@ mod test {
             Value::Timestamp(1000001i64.into()),
             Value::Interval(1000001i128.into()),
         ];
+        let query_context = QueryContextBuilder::default()
+            .configuration_parameter(Default::default())
+            .build();
         let mut builder = DataRowEncoder::new(Arc::new(schema));
         for i in values.iter() {
-            encode_value(i, &mut builder).unwrap();
+            encode_value(&query_context, i, &mut builder).unwrap();
         }
 
         let err = encode_value(
+            &query_context,
             &Value::List(ListValue::new(
                 Some(Box::default()),
                 ConcreteDataType::int16_datatype(),