rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-10-03 02:25:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: serialize struct to json in postgres (#9170)

Browse Source 
* 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 <sunning@greptime.com>

* 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 <sunning@greptime.com>

---------

Signed-off-by: Ning Sun <sunning@greptime.com>
This commit is contained in:
Ning Sun
2026-09-23 01:55:46 +00:00
committed by GitHub
parent f3eb8e6c72
commit 0f118bc5ba
6 changed files with 711 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/datatypes/src/data_type.rs
+3 -1
View File
@@ -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(_)
src/datatypes/src/types/struct_type.rs
+19
View File
@@ -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<Self> {
let struct_fields = fields
.iter()
.map(|field| {
Ok(StructField::new(
field.name().clone(),
ConcreteDataType::try_from(field.data_type())?,
field.is_nullable(),
))
})
.collect::<Result<Vec<_>>>()?;
Ok(StructType {
fields: Arc::new(struct_fields),
})
}
pub fn fields(&self) -> Arc<Vec<StructField>> {
self.fields.clone()
}
src/datatypes/src/value.rs
+94 -5
View File
@@ -1152,12 +1152,13 @@ impl TryFrom<ScalarValue> 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<ScalarValue> 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<ScalarValue> for Value {
| ScalarValue::LargeListView(_)
| ScalarValue::Union(_, _, _)
| ScalarValue::Float16(_)
| ScalarValue::Utf8View(_)
| ScalarValue::BinaryView(_)
| ScalarValue::Map(_)
| ScalarValue::Date64(_)
| ScalarValue::RunEndEncoded(_, _, _) => {
@@ -1259,6 +1266,21 @@ impl TryFrom<ScalarValue> 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<Value> {
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<ValueRef<'_>> 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]
src/servers/src/postgres/types.rs
+308 -7
View File
@@ -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<S: Encoder>(
_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<S: Encoder>(
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<S, B>
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::<i32>();
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<Type> {
&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<DataRow> = 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<Option<String>> = 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<DataRow> = 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<Option<String>> = 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 = [
tests-integration/src/test_util.rs
+55
View File
@@ -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<Box<dyn Server>>) {
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::<SocketAddr>().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()
tests-integration/tests/sql.rs
+232 -1
View File
@@ -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,