From 0f118bc5ba3ad591824562f3ede4b8fd536fb86d Mon Sep 17 00:00:00 2001 From: Ning Sun Date: Wed, 23 Sep 2026 01:55:46 +0000 Subject: [PATCH] fix: serialize struct to json in postgres (#9170) * feat: serialize struct to json in postgres * fix: support view scalars and preserve null structs in scalar-to-value conversion Address PR review: - Utf8View/BinaryView ScalarValues now convert like their non-view forms instead of failing row extraction for struct columns - a null struct scalar converts to Value::Null so a null struct inside a list stays null in the serialized JSON Signed-off-by: Ning Sun * fix: return errors instead of panics for unsupported arrow field types Struct-typed query results with arrow field types greptimedb cannot represent (e.g. Decimal256) used to panic during schema conversion and row extraction, dropping the client connection. They now surface as query errors: - ConcreteDataType::try_from builds struct types fallibly via the new StructType::try_from_arrow_fields - Value::try_from(ScalarValue::Struct) uses the same fallible path - new try_value_from_array converts an arrow element to Value with error propagation, used by the postgres struct encoding Signed-off-by: Ning Sun --------- Signed-off-by: Ning Sun --- src/datatypes/src/data_type.rs | 4 +- src/datatypes/src/types/struct_type.rs | 19 ++ src/datatypes/src/value.rs | 99 +++++++- src/servers/src/postgres/types.rs | 315 ++++++++++++++++++++++++- tests-integration/src/test_util.rs | 55 +++++ tests-integration/tests/sql.rs | 233 +++++++++++++++++- 6 files changed, 711 insertions(+), 14 deletions(-) diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs index f2ed05534d6..8ad388c2ff7 100644 --- a/src/datatypes/src/data_type.rs +++ b/src/datatypes/src/data_type.rs @@ -492,7 +492,9 @@ impl TryFrom<&ArrowDataType> for ConcreteDataType { ArrowDataType::Decimal128(precision, scale) => { ConcreteDataType::decimal128_datatype(*precision, *scale) } - ArrowDataType::Struct(fields) => ConcreteDataType::Struct(StructType::from(fields)), + ArrowDataType::Struct(fields) => { + ConcreteDataType::Struct(StructType::try_from_arrow_fields(fields)?) + } ArrowDataType::Float16 | ArrowDataType::Date64 | ArrowDataType::FixedSizeBinary(_) diff --git a/src/datatypes/src/types/struct_type.rs b/src/datatypes/src/types/struct_type.rs index ffd1d1c0b25..fdbfd4886b9 100644 --- a/src/datatypes/src/types/struct_type.rs +++ b/src/datatypes/src/types/struct_type.rs @@ -19,6 +19,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Field}; use arrow_schema::Fields; use serde::{Deserialize, Serialize}; +use crate::error::Result; use crate::prelude::{ConcreteDataType, DataType, LogicalTypeId}; use crate::value::Value; use crate::vectors::StructVectorBuilder; @@ -100,6 +101,24 @@ impl StructType { } } + /// Like the `From<&Fields>` impl, but returns an error for arrow field + /// types greptimedb cannot represent instead of panicking. + pub fn try_from_arrow_fields(fields: &Fields) -> Result { + let struct_fields = fields + .iter() + .map(|field| { + Ok(StructField::new( + field.name().clone(), + ConcreteDataType::try_from(field.data_type())?, + field.is_nullable(), + )) + }) + .collect::>>()?; + Ok(StructType { + fields: Arc::new(struct_fields), + }) + } + pub fn fields(&self) -> Arc> { self.fields.clone() } diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 56a551484f2..da1df882aba 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -1152,12 +1152,13 @@ impl TryFrom for Value { ScalarValue::UInt16(u) => Value::from(u), ScalarValue::UInt32(u) => Value::from(u), ScalarValue::UInt64(u) => Value::from(u), - ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) => { + ScalarValue::Utf8(s) | ScalarValue::LargeUtf8(s) | ScalarValue::Utf8View(s) => { Value::from(s.map(StringBytes::from)) } ScalarValue::Binary(b) | ScalarValue::LargeBinary(b) - | ScalarValue::FixedSizeBinary(_, b) => Value::from(b.map(Bytes::from)), + | ScalarValue::FixedSizeBinary(_, b) + | ScalarValue::BinaryView(b) => Value::from(b.map(Bytes::from)), ScalarValue::List(array) => { // this is for item type let datatype = ConcreteDataType::try_from(&array.value_type())?; @@ -1221,7 +1222,15 @@ impl TryFrom for Value { .map(|v| Value::Decimal128(Decimal128::new(v, p, s))) .unwrap_or(Value::Null), ScalarValue::Struct(struct_array) => { - let struct_type = StructType::from(struct_array.fields()); + // A struct scalar carries a single element; a null (or empty) + // element means the struct itself is null, for example a null + // struct inside a list. + if struct_array.is_empty() || struct_array.is_null(0) { + return Ok(Value::Null); + } + // Build the struct type fallibly: an unrepresentable arrow + // field type must surface as an error, not a panic. + let struct_type = StructType::try_from_arrow_fields(struct_array.fields())?; let items = struct_array .columns() .iter() @@ -1244,8 +1253,6 @@ impl TryFrom for Value { | ScalarValue::LargeListView(_) | ScalarValue::Union(_, _, _) | ScalarValue::Float16(_) - | ScalarValue::Utf8View(_) - | ScalarValue::BinaryView(_) | ScalarValue::Map(_) | ScalarValue::Date64(_) | ScalarValue::RunEndEncoded(_, _, _) => { @@ -1259,6 +1266,21 @@ impl TryFrom for Value { } } +/// Converts the value at `index` of an arrow array into a [`Value`]. +/// +/// Unlike the unchecked conversion in [`crate::vectors::StructVector`]'s `get`, +/// this propagates conversion errors (for example arrow types greptimedb +/// cannot represent, such as `Decimal256`) instead of panicking, so callers +/// can surface a proper error to the user. +pub fn try_value_from_array(array: &dyn Array, index: usize) -> Result { + if array.is_null(index) { + return Ok(Value::Null); + } + let scalar = + ScalarValue::try_from_array(array, index).context(ConvertArrowArrayToScalarsSnafu)?; + Value::try_from(scalar) +} + impl From> for Value { fn from(value: ValueRef<'_>) -> Self { match value { @@ -2157,6 +2179,73 @@ pub(crate) mod tests { Value::Struct(struct_value), scalar_struct_value.try_into().unwrap() ); + + // view-typed strings and binaries convert like their non-view forms + assert_eq!( + Value::String("abc".into()), + ScalarValue::Utf8View(Some("abc".into())) + .try_into() + .unwrap() + ); + assert_eq!(Value::Null, ScalarValue::Utf8View(None).try_into().unwrap()); + assert_eq!( + Value::Binary(Bytes::from(vec![1, 2])), + ScalarValue::BinaryView(Some(vec![1, 2])) + .try_into() + .unwrap() + ); + assert_eq!( + Value::Null, + ScalarValue::BinaryView(None).try_into().unwrap() + ); + + // a null struct scalar is a null value, not a struct of null fields + let null_struct = ScalarStructBuilder::new_null(build_struct_type().as_arrow_fields()); + assert_eq!(Value::Null, null_struct.try_into().unwrap()); + } + + #[test] + fn test_try_value_from_array() { + use arrow_array::Int32Array; + use datafusion_common::arrow::datatypes::i256; + + let array = Int32Array::from(vec![Some(1), None]); + assert_eq!(Value::Int32(1), try_value_from_array(&array, 0).unwrap()); + assert_eq!(Value::Null, try_value_from_array(&array, 1).unwrap()); + + // a struct with representable fields converts + let supported = ScalarStructBuilder::new() + .with_name_and_scalar("i", ScalarValue::Int32(Some(7))) + .build() + .unwrap(); + let ScalarValue::Struct(array) = supported else { + unreachable!(); + }; + assert_eq!( + Value::Struct(StructValue::new( + vec![Value::Int32(7)], + StructType::new(std::sync::Arc::new(vec![StructField::new( + "i", + ConcreteDataType::int32_datatype(), + false, + )])) + )), + try_value_from_array(array.as_ref(), 0).unwrap() + ); + + // an unsupported arrow field type errors instead of panicking + let unsupported = ScalarStructBuilder::new() + .with_name_and_scalar( + "d", + ScalarValue::Decimal256(Some(i256::from_i128(1)), 38, 10), + ) + .build() + .unwrap(); + let ScalarValue::Struct(array) = unsupported else { + unreachable!(); + }; + let error = try_value_from_array(array.as_ref(), 0).unwrap_err(); + assert!(error.to_string().contains("Unsupported arrow data type")); } #[test] diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 033a1c665e2..c83debd39a1 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -19,7 +19,7 @@ use std::pin::Pin; use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; -use arrow::array::{Array, AsArray}; +use arrow::array::{Array, ArrayRef, AsArray}; use arrow_pg::encoder::{Encoder, encode_value}; use arrow_pg::list_encoder::encode_list; use arrow_schema::{DataType, TimeUnit}; @@ -37,7 +37,7 @@ use datatypes::json::JsonSettings; use datatypes::prelude::{ConcreteDataType, DataType as _, Value}; use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::{Decimal128Type, IntervalType, TimestampType, jsonb_to_string}; -use datatypes::value::StructValue; +use datatypes::value::try_value_from_array; use futures::Stream; use pg_interval::Interval as PgInterval; use pgwire::api::Type; @@ -234,7 +234,8 @@ impl ToSqlText for OidAliasValue<'_> { } } -/// this function will encode greptime's `StructValue` into PostgreSQL jsonb type +/// this function will encode greptime's structured values (`StructValue`, or a +/// list of them) into PostgreSQL jsonb type /// /// Note that greptimedb has different types of StructValue for storing json data, /// based on policy defined in `JsonSettings`. But here the `StructValue` @@ -245,18 +246,35 @@ impl ToSqlText for OidAliasValue<'_> { /// arrays: element in array must be the same type fn encode_struct( _query_ctx: &QueryContextRef, - struct_value: StructValue, + value: Value, builder: &mut S, pg_field: &FieldInfo, ) -> PgWireResult<()> { let encoding_setting = JsonSettings::default(); let json_value = encoding_setting - .decode(Value::Struct(struct_value)) + .decode(value) .map_err(|e| PgWireError::ApiError(Box::new(e)))?; builder.encode_field(&json_value, pg_field) } +/// Encodes one row of a structured column (`Struct`, or `List` of `Struct`) as +/// a PostgreSQL `json` value. +/// +/// Returns an error for arrow field types greptimedb cannot represent +/// (e.g. `Decimal256`) instead of panicking, so the client receives a proper +/// error response. +fn encode_structured_json( + query_ctx: &QueryContextRef, + column: &ArrayRef, + i: usize, + encoder: &mut S, + pg_field: &FieldInfo, +) -> PgWireResult<()> { + let value = try_value_from_array(column.as_ref(), i).map_err(convert_err)?; + encode_struct(query_ctx, value, encoder, pg_field) +} + pub(crate) struct RecordBatchRowStream where S: Encoder, @@ -369,6 +387,11 @@ where } } + DataType::List(list_field) + if matches!(list_field.data_type(), DataType::Struct(_)) => + { + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; + } DataType::List(_) => { let array = column.as_list::(); let items = array.value(i); @@ -376,7 +399,7 @@ where encode_list(encoder, items, pg_field)?; } DataType::Struct(_) => { - encode_struct(query_ctx, Default::default(), encoder, pg_field)?; + encode_structured_json(query_ctx, column, i, encoder, pg_field)?; } DataType::Utf8 => { let arrow_field = arrow_schema.field(j); @@ -445,7 +468,9 @@ pub(super) fn type_gt_to_pg(origin: &ConcreteDataType) -> Result { &ConcreteDataType::Decimal128(_) => Ok(Type::NUMERIC_ARRAY), &ConcreteDataType::Json(_) => Ok(Type::JSON_ARRAY), &ConcreteDataType::Duration(_) => Ok(Type::INTERVAL_ARRAY), - &ConcreteDataType::Struct(_) => Ok(Type::JSON_ARRAY), + // a list of structs is encoded as a single JSON array, matching the + // struct-to-json mapping above + &ConcreteDataType::Struct(_) => Ok(Type::JSON), &ConcreteDataType::Dictionary(_) | &ConcreteDataType::Vector(_) | &ConcreteDataType::List(_) => server_error::UnsupportedDataTypeSnafu { @@ -1921,6 +1946,282 @@ mod test { } } + #[test] + fn test_encode_list_of_struct_as_json() { + use datatypes::types::{StructField, StructType}; + use datatypes::value::{ListValue, StructValue}; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + // struct type: {a: int32, b: float64} + let fields = vec![ + StructField::new("a".to_string(), ConcreteDataType::int32_datatype(), true), + StructField::new("b".to_string(), ConcreteDataType::float64_datatype(), true), + ]; + let struct_type = + ConcreteDataType::struct_datatype(StructType::new(Arc::new(fields.clone()))); + let list_type = ConcreteDataType::list_datatype(Arc::new(struct_type.clone())); + + let make_struct = |a: i32, b: f64| { + Value::Struct(StructValue::new( + vec![Value::Int32(a), Value::Float64(b.into())], + StructType::new(Arc::new(fields.clone())), + )) + }; + let rows = vec![ + Value::List(ListValue::new( + vec![make_struct(1, 1.5), Value::Null, make_struct(3, 3.5)], + Arc::new(struct_type.clone()), + )), + Value::List(ListValue::new(vec![], Arc::new(struct_type.clone()))), + Value::Null, + ]; + let vector = Helper::try_from_row_into_vector(&rows, &list_type).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "l".to_string(), + list_type.clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + // The list-of-struct column is declared as a single json value. + let pg_schema = Arc::new(vec![FieldInfo::new( + "l".into(), + None, + None, + type_gt_to_pg(&list_type).unwrap(), + FieldFormat::Text, + )]); + assert_eq!(pg_schema[0].datatype(), &Type::JSON); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL list row stays NULL. + assert_eq!(decoded[2], None); + + // A list of structs is a JSON array; null elements stay null. + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!([ + {"a": 1, "b": 1.5}, + null, + {"a": 3, "b": 3.5}, + ]) + ); + + // An empty list is an empty JSON array. + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!(json, serde_json::json!([])); + } + + #[test] + fn test_encode_struct_value_as_json() { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use api::helper::pb_value_to_value_ref; + use api::v1::Value as PbValue; + use common_query::native_histogram::{ + encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, + }; + use datatypes::vectors::Helper; + use pgwire::messages::data::DataRow; + + fn histogram_value(histogram: &Histogram) -> Value { + let value_data = encode_native_histogram(histogram).unwrap(); + let pb_value = PbValue { + value_data: Some(value_data), + }; + let column_schema = native_histogram_column_schema().unwrap(); + Value::from(pb_value_to_value_ref( + &pb_value, + column_schema.datatype_extension.as_ref(), + )) + } + + // One integer-count histogram and one float-count histogram, mirroring + // what prometheus remote-write v2 stores. The expected JSON matches the + // HTTP query output for the same payloads. + let int_histogram = Histogram { + schema: 1, + sum: 10.0, + zero_threshold: 0.001, + reset_hint: 2, + start_timestamp: 1500, + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + count: Some(Count::CountInt(8)), + zero_count: Some(ZeroCount::ZeroCountInt(1)), + ..Default::default() + }; + let float_histogram = Histogram { + schema: 2, + sum: 20.0, + zero_threshold: 0.002, + reset_hint: 3, + start_timestamp: 2500, + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + count: Some(Count::CountFloat(6.0)), + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + ..Default::default() + }; + + let rows = vec![ + histogram_value(&int_histogram), + histogram_value(&float_histogram), + Value::Null, + ]; + let vector = + Helper::try_from_row_into_vector(&rows, native_histogram_value_type()).unwrap(); + + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "greptime_native_histogram", + native_histogram_value_type().clone(), + true, + )])); + let record_batch = RecordBatch::new(schema.clone(), vec![vector]).unwrap(); + + let pg_schema = Arc::new(vec![FieldInfo::new( + "greptime_native_histogram".into(), + None, + None, + Type::JSON, + FieldFormat::Text, + )]); + + let query_context = QueryContextBuilder::default() + .configuration_parameter(Default::default()) + .build() + .into(); + let encoder = DataRowEncoder::new(pg_schema.clone()); + let row_stream = RecordBatchRowStream::new( + query_context, + pg_schema.clone(), + schema, + stream::once(async { Ok(record_batch) }), + encoder, + ); + + let data_rows: Vec = futures::executor::block_on( + row_stream + .filter_map(|x: PgWireResult<_>| async move { x.ok() }) + .flat_map(stream::iter) + .collect(), + ); + assert_eq!(data_rows.len(), 3); + + // Decode text-format fields: [i32 BE length][bytes], -1 for NULL. + let decoded: Vec> = data_rows + .iter() + .map(|row| { + assert_eq!(row.field_count, 1); + let len = i32::from_be_bytes(row.data[0..4].try_into().unwrap()); + if len < 0 { + return None; + } + Some(String::from_utf8(row.data[4..4 + len as usize].to_vec()).unwrap()) + }) + .collect(); + + // A NULL struct row stays NULL. + assert_eq!(decoded[2], None); + + let json: serde_json::Value = serde_json::from_str(decoded[0].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let json: serde_json::Value = serde_json::from_str(decoded[1].as_deref().unwrap()).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + } + #[test] fn test_encode_utf8_oid_alias_data() { let aliases = [ diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index fd02db2e5a1..37a1740ccd1 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1264,6 +1264,61 @@ pub async fn setup_pg_server_with_user_provider( (instance.guard, Arc::new(pg_server)) } +/// Sets up a standalone instance with both a Prometheus remote-write HTTP app +/// (native histograms enabled) and a Postgres server attached, so native +/// histogram data written via remote-write can be queried over the Postgres +/// protocol. +pub async fn setup_pg_server_with_prom_native_histogram( + store_type: StorageType, + name: &str, +) -> (TestGuard, Router, Arc>) { + unsafe { + std::env::set_var("TZ", "UTC"); + } + + let instance = setup_standalone_instance(name, store_type).await; + + // Prometheus remote-write HTTP app with native histograms enabled. + let http_server = build_test_prom_server(instance.fe_instance().clone(), false, true) + .with_greptime_config_options(instance.opts.datanode_options().to_toml().unwrap()) + .build(); + let app = http_server.build(http_server.make_app()).unwrap(); + + // Postgres server on the same instance. + let runtime = RuntimeBuilder::default() + .worker_threads(2) + .thread_name("pg-runtime") + .build() + .unwrap(); + + let fe_pg_addr = format!("127.0.0.1:{}", ports::get_port()); + let opts = PostgresOptions { + addr: fe_pg_addr.clone(), + ..Default::default() + }; + let tls_server_config = Arc::new( + ReloadableTlsServerConfig::try_new(opts.tls.clone()) + .expect("Failed to load certificates and keys"), + ); + + let mut pg_server = Box::new(PostgresServer::new( + instance.fe_instance().clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + )); + + pg_server + .start(fe_pg_addr.parse::().unwrap()) + .await + .unwrap(); + + (instance.guard, app, Arc::new(pg_server)) +} + pub(crate) async fn prepare_another_catalog_and_schema(instance: &Instance) { let catalog_manager = instance .catalog_manager() diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 6c58ab0246c..e3b66fc4447 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -31,7 +31,8 @@ use sqlx::{Connection, Executor, Row}; use tests_integration::test_util::{ StorageType, setup_mysql_server, setup_mysql_server_with_slow_query_threshold, setup_mysql_server_with_user_provider, setup_pg_server, - setup_pg_server_with_slow_query_threshold, setup_pg_server_with_user_provider, + setup_pg_server_with_prom_native_histogram, setup_pg_server_with_slow_query_threshold, + setup_pg_server_with_user_provider, }; use tokio_postgres::{Client, NoTls, SimpleQueryMessage}; @@ -94,6 +95,8 @@ macro_rules! sql_tests { test_mysql_federated_prepare_stmt, test_mysql_prepare_tql_and_show, test_postgres_extended_query_row_returning_statements, + test_postgres_native_histogram, + test_postgres_struct_types, test_declare_fetch_close_cursor, test_alter_update_on, ); @@ -879,6 +882,234 @@ pub async fn test_postgres_bytea(store_type: StorageType) { guard.remove_all().await; } +pub async fn test_postgres_native_histogram(store_type: StorageType) { + use api::greptime_proto::io::prometheus::write::v2::histogram::{Count, ZeroCount}; + use api::greptime_proto::io::prometheus::write::v2::{BucketSpan, Histogram}; + use axum::http::StatusCode; + use prost::Message; + use servers::http::test_helpers::TestClient; + use servers::prom_remote_write::v2::test_util as remote_write_v2; + use servers::prom_store; + + let (mut guard, app, fe_pg_server) = + setup_pg_server_with_prom_native_histogram(store_type, "test_postgres_native_histogram") + .await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + let http_client = TestClient::new(app).await; + + // Ingest one integer-count and one float-count native histogram via + // prometheus remote-write v2. + let write_request = remote_write_v2::request_with_labels_and_histograms( + vec![ + (prom_store::METRIC_NAME_LABEL, "pg_native_histogram_seconds"), + ("job", "api"), + ("instance", "localhost:9090"), + ], + vec![ + Histogram { + count: Some(Count::CountInt(8)), + sum: 10.0, + schema: 1, + zero_threshold: 0.001, + zero_count: Some(ZeroCount::ZeroCountInt(1)), + negative_spans: vec![BucketSpan { + offset: -2, + length: 1, + }], + negative_deltas: vec![1], + positive_spans: vec![BucketSpan { + offset: 0, + length: 3, + }], + positive_deltas: vec![1, 2, -1], + reset_hint: 2, + timestamp: 3000, + start_timestamp: 1500, + ..Default::default() + }, + Histogram { + count: Some(Count::CountFloat(6.0)), + sum: 20.0, + schema: 2, + zero_threshold: 0.002, + zero_count: Some(ZeroCount::ZeroCountFloat(0.5)), + positive_spans: vec![BucketSpan { + offset: 3, + length: 2, + }], + positive_counts: vec![2.0, 3.5], + reset_hint: 3, + timestamp: 4000, + start_timestamp: 2500, + ..Default::default() + }, + ], + ); + let compressed = prom_store::snappy_compress(&write_request.encode_to_vec()).unwrap(); + let res = http_client + .post("/v1/prometheus/write") + .header("Content-Encoding", "snappy") + .header( + "Content-Type", + "application/x-protobuf;proto=io.prometheus.write.v2.Request", + ) + .body(compressed) + .send() + .await; + assert_eq!(res.status(), StatusCode::NO_CONTENT); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + let rows = client + .simple_query( + "select greptime_native_histogram from pg_native_histogram_seconds order by greptime_timestamp", + ) + .await + .unwrap(); + let jsons: Vec<&str> = rows + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0), + _ => None, + }) + .collect(); + assert_eq!(jsons.len(), 2); + + // The struct column is encoded as a JSON object, matching the HTTP query + // output for the same histograms. Before the fix every row was encoded as + // an empty struct (`{}`). + let histogram: serde_json::Value = serde_json::from_str(jsons[0]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": null, + "count_i64": 8, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [1], + "negative_span_lengths": [1], + "negative_span_offsets": [-2], + "positive_buckets_f64": [], + "positive_buckets_i64": [1, 3, 2], + "positive_span_lengths": [3], + "positive_span_offsets": [0], + "reset_hint": 2, + "schema": 1, + "start_timestamp": 1500, + "sum": 10.0, + "zero_count_f64": null, + "zero_count_i64": 1, + "zero_threshold": 0.001 + }) + ); + + let histogram: serde_json::Value = serde_json::from_str(jsons[1]).unwrap(); + assert_eq!( + histogram, + serde_json::json!({ + "count_f64": 6.0, + "count_i64": null, + "custom_values": [], + "negative_buckets_f64": [], + "negative_buckets_i64": [], + "negative_span_lengths": [], + "negative_span_offsets": [], + "positive_buckets_f64": [2.0, 3.5], + "positive_buckets_i64": [], + "positive_span_lengths": [2], + "positive_span_offsets": [3], + "reset_hint": 3, + "schema": 2, + "start_timestamp": 2500, + "sum": 20.0, + "zero_count_f64": 0.5, + "zero_count_i64": null, + "zero_threshold": 0.002 + }) + ); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + +pub async fn test_postgres_struct_types(store_type: StorageType) { + let (mut guard, fe_pg_server) = setup_pg_server(store_type, "test_postgres_struct_types").await; + let addr = fe_pg_server.bind_addr().unwrap().to_string(); + + let (client, connection) = tokio_postgres::connect(&format!("postgres://{addr}/public"), NoTls) + .await + .unwrap(); + let (tx, rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + connection.await.unwrap(); + tx.send(()).unwrap(); + }); + + async fn query_one(client: &Client, sql: &str) -> String { + let messages = client.simple_query(sql).await.unwrap(); + messages + .iter() + .filter_map(|message| match message { + SimpleQueryMessage::Row(row) => row.get(0).map(|v| v.to_string()), + _ => None, + }) + .next() + .unwrap() + } + + // Struct and list-of-struct columns are serialized as JSON. View-typed + // struct fields must not fail row extraction, and a null struct inside a + // list stays null instead of becoming a struct of null fields. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + // '0102' casts to its UTF-8 bytes when interpreted as BinaryView. + let row = query_one(&client, "SELECT struct(arrow_cast('0102', 'BinaryView'))").await; + assert_eq!(row, "{\"c0\":[48,49,48,50]}"); + + let row = query_one(&client, "SELECT struct([struct(1), NULL])").await; + assert_eq!(row, "{\"c0\":[{\"c0\":1},null]}"); + + // Top-level control: the same null element in a bare list stays null too. + let row = query_one(&client, "SELECT [struct(1), NULL]").await; + assert_eq!(row, "[{\"c0\":1},null]"); + + // Unsupported arrow field types surface as query errors instead of + // dropping the connection. + let error = client + .simple_query("SELECT struct(arrow_cast('1', 'Decimal256(38, 10)'))") + .await + .unwrap_err(); + let message = match error.as_db_error() { + Some(db_error) => db_error.message().to_string(), + None => error.to_string(), + }; + assert!( + message.contains("Unsupported arrow data type"), + "unexpected error message: {message}" + ); + + // The connection stays usable after the error. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + + drop(client); + rx.await.unwrap(); + + let _ = fe_pg_server.shutdown().await; + guard.remove_all().await; +} + pub async fn test_postgres_slow_query(store_type: StorageType) { let (mut guard, fe_pg_server) = setup_pg_server_with_slow_query_threshold( store_type,