1mod error;
16
17use std::collections::HashMap;
18use std::pin::Pin;
19use std::sync::{Arc, LazyLock};
20use std::task::{Context, Poll};
21
22use arrow::array::{Array, AsArray};
23use arrow_pg::encoder::{Encoder, encode_value};
24use arrow_pg::list_encoder::encode_list;
25use arrow_schema::{DataType, TimeUnit};
26use bytes::BufMut;
27use chrono::{DateTime, FixedOffset, NaiveDate, NaiveDateTime};
28use common_recordbatch::error::Result as RecordBatchResult;
29use common_recordbatch::{RecordBatch, map_dictionary_to_values_data_type};
30use common_time::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
31use datafusion_common::ScalarValue;
32use datafusion_expr::LogicalPlan;
33use datafusion_pg_catalog::pg_catalog::PgCatalogStaticTables;
34use datafusion_pg_catalog::pg_catalog::oid_field::{self, OID_ALIAS_KEY};
35use datatypes::arrow::datatypes::DataType as ArrowDataType;
36use datatypes::json::JsonSettings;
37use datatypes::prelude::{ConcreteDataType, DataType as _, Value};
38use datatypes::schema::{Schema, SchemaRef};
39use datatypes::types::{Decimal128Type, IntervalType, TimestampType, jsonb_to_string};
40use datatypes::value::StructValue;
41use futures::Stream;
42use pg_interval::Interval as PgInterval;
43use pgwire::api::Type;
44use pgwire::api::portal::{Format, Portal};
45use pgwire::api::results::FieldInfo;
46use pgwire::error::{PgWireError, PgWireResult};
47use pgwire::types::ToSqlText;
48use pgwire::types::format::FormatOptions as PgFormatOptions;
49use postgres_types::{IsNull, ToSql};
50use query::planner::DfLogicalPlanner;
51use rust_decimal::Decimal;
52use rust_decimal::prelude::ToPrimitive;
53use session::context::QueryContextRef;
54use snafu::ResultExt;
55
56pub use self::error::{PgErrorCode, PgErrorSeverity};
57use crate::error::{self as server_error, InferParameterTypesSnafu, Result};
58use crate::postgres::handler::PgSqlPlan;
59use crate::postgres::utils::convert_err;
60
61pub(super) fn schema_to_pg(
62 origin: &Schema,
63 field_formats: &Format,
64 format_options: Option<Arc<PgFormatOptions>>,
65) -> Result<Vec<FieldInfo>> {
66 origin
67 .column_schemas()
68 .iter()
69 .enumerate()
70 .map(|(idx, col)| {
71 let pg_type = match pg_oid_alias_type(col) {
72 Some(pg_type) => pg_type,
73 None => type_gt_to_pg(&col.data_type)?,
74 };
75 let mut field_info = FieldInfo::new(
76 col.name.clone(),
77 None,
78 None,
79 pg_type,
80 field_formats.format_for(idx),
81 );
82 if let Some(format_options) = &format_options {
83 field_info = field_info.with_format_options(format_options.clone());
84 }
85 Ok(field_info)
86 })
87 .collect::<Result<Vec<FieldInfo>>>()
88}
89
90fn pg_oid_alias_type(column: &datatypes::schema::ColumnSchema) -> Option<Type> {
99 if !matches!(
100 &column.data_type,
101 ConcreteDataType::Int32(_) | ConcreteDataType::String(_)
102 ) {
103 return None;
104 }
105
106 pg_oid_alias_type_name(column.metadata().get(OID_ALIAS_KEY)?)
107}
108
109fn pg_oid_alias_type_name(alias: &str) -> Option<Type> {
110 match alias {
111 oid_field::kind::OID => Some(Type::OID),
112 oid_field::kind::REGPROC => Some(Type::REGPROC),
113 oid_field::kind::REGCLASS => Some(Type::REGCLASS),
114 oid_field::kind::REGTYPE => Some(Type::REGTYPE),
115 oid_field::kind::REGNAMESPACE => Some(Type::REGNAMESPACE),
116 "regprocedure" => Some(Type::REGPROCEDURE),
117 "regoper" => Some(Type::REGOPER),
118 "regoperator" => Some(Type::REGOPERATOR),
119 "regcollation" => Some(Type::REGCOLLATION),
120 "regconfig" => Some(Type::REGCONFIG),
121 "regdictionary" => Some(Type::REGDICTIONARY),
122 "regrole" => Some(Type::REGROLE),
123 _ => None,
124 }
125}
126
127static REGPROC_OIDS: LazyLock<std::result::Result<HashMap<String, Option<u32>>, String>> =
129 LazyLock::new(|| {
130 let tables = PgCatalogStaticTables::try_new()
131 .map_err(|e| format!("load static PostgreSQL catalog tables: {e}"))?;
132 let mut oids = HashMap::new();
133
134 for batch in tables.pg_proc.data() {
135 let names = batch
136 .column_by_name("proname")
137 .ok_or("pg_proc is missing proname")?;
138 let procedure_oids = batch
139 .column_by_name("oid")
140 .ok_or("pg_proc is missing oid")?;
141 if names.data_type() != &DataType::Utf8
142 || procedure_oids.data_type() != &DataType::Int32
143 {
144 return Err("pg_proc has unexpected proname or oid types".to_string());
145 }
146
147 let names = names.as_string::<i32>();
148 let procedure_oids = procedure_oids.as_primitive::<arrow::datatypes::Int32Type>();
149 for (name, oid) in names.iter().zip(procedure_oids.iter()) {
150 let (Some(name), Some(oid)) = (name, oid) else {
151 continue;
152 };
153 let oid = u32::try_from(oid)
154 .map_err(|_| format!("pg_proc contains negative oid {oid}"))?;
155
156 if oids.insert(name.to_string(), Some(oid)).is_some() {
157 oids.insert(name.to_string(), None);
158 }
159 }
160 }
161
162 Ok(oids)
163 });
164
165#[derive(Debug)]
166struct OidAliasError(String);
167
168impl std::fmt::Display for OidAliasError {
169 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
170 f.write_str(&self.0)
171 }
172}
173
174impl std::error::Error for OidAliasError {}
175
176fn resolve_oid_alias(value: &str, alias: &str) -> std::result::Result<u32, OidAliasError> {
177 if value == "-" {
178 return Ok(0);
179 }
180
181 if let Ok(oid) = value.parse() {
182 return Ok(oid);
183 }
184
185 if alias == oid_field::kind::REGPROC {
186 let oids = REGPROC_OIDS
187 .as_ref()
188 .map_err(|e| OidAliasError(e.clone()))?;
189 return match oids.get(value) {
190 Some(Some(oid)) => Ok(*oid),
191 Some(None) => Err(OidAliasError(format!("ambiguous regproc name: {value}"))),
192 None => Err(OidAliasError(format!("unknown regproc name: {value}"))),
193 };
194 }
195
196 Err(OidAliasError(format!(
197 "named oid aliases are only supported for regproc: {value}"
198 )))
199}
200
201#[derive(Debug)]
202struct OidAliasValue<'a> {
203 text: &'a str,
204 alias: &'a str,
205}
206
207impl ToSql for OidAliasValue<'_> {
208 fn to_sql(
209 &self,
210 _ty: &Type,
211 out: &mut bytes::BytesMut,
212 ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
213 let oid = resolve_oid_alias(self.text, self.alias)?;
214 out.put_u32(oid);
215 Ok(IsNull::No)
216 }
217
218 fn accepts(ty: &Type) -> bool {
219 pg_oid_alias_type_name(ty.name()).is_some()
220 }
221
222 postgres_types::to_sql_checked!();
223}
224
225impl ToSqlText for OidAliasValue<'_> {
226 fn to_sql_text(
227 &self,
228 _ty: &Type,
229 out: &mut bytes::BytesMut,
230 _format_options: &PgFormatOptions,
231 ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
232 out.put_slice(self.text.as_bytes());
233 Ok(IsNull::No)
234 }
235}
236
237fn encode_struct<S: Encoder>(
247 _query_ctx: &QueryContextRef,
248 struct_value: StructValue,
249 builder: &mut S,
250 pg_field: &FieldInfo,
251) -> PgWireResult<()> {
252 let encoding_setting = JsonSettings::default();
253 let json_value = encoding_setting
254 .decode(Value::Struct(struct_value))
255 .map_err(|e| PgWireError::ApiError(Box::new(e)))?;
256
257 builder.encode_field(&json_value, pg_field)
258}
259
260pub(crate) struct RecordBatchRowStream<S, B>
261where
262 S: Encoder,
263 B: Stream<Item = RecordBatchResult<RecordBatch>>,
264{
265 query_ctx: QueryContextRef,
266 pg_schema: Arc<Vec<FieldInfo>>,
267 schema: SchemaRef,
268 record_batches: Pin<Box<B>>,
269 encoder: S,
270}
271
272impl<S, B> Stream for RecordBatchRowStream<S, B>
273where
274 S: Encoder + Unpin,
275 B: Stream<Item = RecordBatchResult<RecordBatch>>,
276{
277 type Item = PgWireResult<Vec<S::Item>>;
278
279 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
280 match self.record_batches.as_mut().poll_next(cx) {
281 Poll::Ready(Some(Ok(batch))) => {
282 let record_batch = batch.into_df_record_batch();
283 let num_rows = record_batch.num_rows();
284
285 if num_rows == 0 {
286 return Poll::Ready(Some(Ok(vec![])));
287 }
288
289 let arrow_schema = record_batch.schema();
290 let query_ctx = self.query_ctx.clone();
291 let pg_schema = self.pg_schema.clone();
292 let schema = self.schema.clone();
293 let mut results = Vec::with_capacity(num_rows);
294
295 for i in 0..num_rows {
296 if let Err(e) = Self::encode_row(
297 &query_ctx,
298 &pg_schema,
299 &schema,
300 arrow_schema.as_ref(),
301 &mut self.encoder,
302 &record_batch,
303 i,
304 ) {
305 return Poll::Ready(Some(Err(e)));
306 }
307 results.push(self.encoder.take_row());
308 }
309
310 Poll::Ready(Some(Ok(results)))
311 }
312 Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(convert_err(e)))),
313 Poll::Ready(None) => Poll::Ready(None),
314 Poll::Pending => Poll::Pending,
315 }
316 }
317}
318
319impl<S, B> RecordBatchRowStream<S, B>
320where
321 S: Encoder,
322 B: Stream<Item = RecordBatchResult<RecordBatch>>,
323{
324 pub(crate) fn new(
325 query_ctx: QueryContextRef,
326 pg_schema: Arc<Vec<FieldInfo>>,
327 schema: SchemaRef,
328 record_batches: B,
329 encoder: S,
330 ) -> Self {
331 Self {
332 query_ctx,
333 pg_schema,
334 schema,
335 record_batches: Box::pin(record_batches),
336 encoder,
337 }
338 }
339
340 fn encode_row(
341 query_ctx: &QueryContextRef,
342 pg_schema: &Arc<Vec<FieldInfo>>,
343 schema: &SchemaRef,
344 arrow_schema: &arrow::datatypes::Schema,
345 encoder: &mut S,
346 record_batch: &arrow::record_batch::RecordBatch,
347 i: usize,
348 ) -> PgWireResult<()> {
349 for (j, column) in record_batch.columns().iter().enumerate() {
350 let pg_field = &pg_schema[j];
351
352 if column.is_null(i) {
353 encoder.encode_field(&None::<&i8>, pg_field)?;
354 continue;
355 }
356
357 match column.data_type() {
358 DataType::Binary | DataType::LargeBinary | DataType::BinaryView => {
360 if let ConcreteDataType::Json(_) = &schema.column_schemas()[j].data_type {
362 let v = datatypes::arrow_array::binary_array_value(column, i);
363 let s = jsonb_to_string(v).map_err(convert_err)?;
364 encoder.encode_field(&s, pg_field)?;
365 } else {
366 let arrow_field = arrow_schema.field(j);
368 encode_value(encoder, column, i, arrow_field, pg_field)?;
369 }
370 }
371
372 DataType::List(_) => {
373 let array = column.as_list::<i32>();
374 let items = array.value(i);
375
376 encode_list(encoder, items, pg_field)?;
377 }
378 DataType::Struct(_) => {
379 encode_struct(query_ctx, Default::default(), encoder, pg_field)?;
380 }
381 DataType::Utf8 => {
382 let arrow_field = arrow_schema.field(j);
383 if let Some(alias) = arrow_field
384 .metadata()
385 .get(OID_ALIAS_KEY)
386 .filter(|alias| pg_oid_alias_type_name(alias).is_some())
387 {
388 let value = OidAliasValue {
389 text: column.as_string::<i32>().value(i),
390 alias,
391 };
392 encoder.encode_field(&value, pg_field)?;
393 } else {
394 encode_value(encoder, column, i, arrow_field, pg_field)?;
395 }
396 }
397 _ => {
398 let arrow_field = arrow_schema.field(j);
400 encode_value(encoder, column, i, arrow_field, pg_field)?;
401 }
402 }
403 }
404 Ok(())
405 }
406}
407
408pub(super) fn type_gt_to_pg(origin: &ConcreteDataType) -> Result<Type> {
409 let logical_type = map_dictionary_to_values_data_type(origin);
410 let origin = &logical_type;
411 match origin {
412 &ConcreteDataType::Null(_) => Ok(Type::UNKNOWN),
413 &ConcreteDataType::Boolean(_) => Ok(Type::BOOL),
414 &ConcreteDataType::Int8(_) => Ok(Type::INT2),
415 &ConcreteDataType::Int16(_) | &ConcreteDataType::UInt8(_) => Ok(Type::INT2),
416 &ConcreteDataType::Int32(_) | &ConcreteDataType::UInt16(_) => Ok(Type::INT4),
417 &ConcreteDataType::Int64(_) | &ConcreteDataType::UInt32(_) => Ok(Type::INT8),
418 &ConcreteDataType::UInt64(_) => Ok(Type::NUMERIC),
419 &ConcreteDataType::Float32(_) => Ok(Type::FLOAT4),
420 &ConcreteDataType::Float64(_) => Ok(Type::FLOAT8),
421 &ConcreteDataType::Binary(_) | &ConcreteDataType::Vector(_) => Ok(Type::BYTEA),
422 &ConcreteDataType::String(_) => Ok(Type::VARCHAR),
423 &ConcreteDataType::Date(_) => Ok(Type::DATE),
424 &ConcreteDataType::Timestamp(_) => Ok(Type::TIMESTAMP),
425 &ConcreteDataType::Time(_) => Ok(Type::TIME),
426 &ConcreteDataType::Interval(_) => Ok(Type::INTERVAL),
427 &ConcreteDataType::Decimal128(_) => Ok(Type::NUMERIC),
428 &ConcreteDataType::Json(_) => Ok(Type::JSON),
429 ConcreteDataType::List(list) => match list.item_type() {
430 &ConcreteDataType::Null(_) => Ok(Type::TEXT_ARRAY),
431 &ConcreteDataType::Boolean(_) => Ok(Type::BOOL_ARRAY),
432 &ConcreteDataType::Int8(_) => Ok(Type::INT2_ARRAY),
433 &ConcreteDataType::Int16(_) | &ConcreteDataType::UInt8(_) => Ok(Type::INT2_ARRAY),
434 &ConcreteDataType::Int32(_) | &ConcreteDataType::UInt16(_) => Ok(Type::INT4_ARRAY),
435 &ConcreteDataType::Int64(_) | &ConcreteDataType::UInt32(_) => Ok(Type::INT8_ARRAY),
436 &ConcreteDataType::UInt64(_) => Ok(Type::NUMERIC_ARRAY),
437 &ConcreteDataType::Float32(_) => Ok(Type::FLOAT4_ARRAY),
438 &ConcreteDataType::Float64(_) => Ok(Type::FLOAT8_ARRAY),
439 &ConcreteDataType::Binary(_) => Ok(Type::BYTEA_ARRAY),
440 &ConcreteDataType::String(_) => Ok(Type::VARCHAR_ARRAY),
441 &ConcreteDataType::Date(_) => Ok(Type::DATE_ARRAY),
442 &ConcreteDataType::Timestamp(_) => Ok(Type::TIMESTAMP_ARRAY),
443 &ConcreteDataType::Time(_) => Ok(Type::TIME_ARRAY),
444 &ConcreteDataType::Interval(_) => Ok(Type::INTERVAL_ARRAY),
445 &ConcreteDataType::Decimal128(_) => Ok(Type::NUMERIC_ARRAY),
446 &ConcreteDataType::Json(_) => Ok(Type::JSON_ARRAY),
447 &ConcreteDataType::Duration(_) => Ok(Type::INTERVAL_ARRAY),
448 &ConcreteDataType::Struct(_) => Ok(Type::JSON_ARRAY),
449 &ConcreteDataType::Dictionary(_)
450 | &ConcreteDataType::Vector(_)
451 | &ConcreteDataType::List(_) => server_error::UnsupportedDataTypeSnafu {
452 data_type: origin,
453 reason: "not implemented",
454 }
455 .fail(),
456 },
457 &ConcreteDataType::Dictionary(_) => server_error::UnsupportedDataTypeSnafu {
458 data_type: origin,
459 reason: "not implemented",
460 }
461 .fail(),
462 &ConcreteDataType::Duration(_) => Ok(Type::INTERVAL),
463 &ConcreteDataType::Struct(_) => Ok(Type::JSON),
464 }
465}
466
467#[allow(dead_code)]
468pub(super) fn type_pg_to_gt(origin: &Type) -> Result<ConcreteDataType> {
469 match origin {
471 &Type::BOOL => Ok(ConcreteDataType::boolean_datatype()),
472 &Type::INT2 => Ok(ConcreteDataType::int16_datatype()),
473 &Type::INT4 => Ok(ConcreteDataType::int32_datatype()),
474 &Type::INT8 => Ok(ConcreteDataType::int64_datatype()),
475 &Type::NUMERIC => Ok(ConcreteDataType::uint64_datatype()),
476 &Type::VARCHAR | &Type::CHAR | &Type::TEXT => Ok(ConcreteDataType::string_datatype()),
477 &Type::TIMESTAMP | &Type::TIMESTAMPTZ => Ok(ConcreteDataType::timestamp_datatype(
478 common_time::timestamp::TimeUnit::Millisecond,
479 )),
480 &Type::DATE => Ok(ConcreteDataType::date_datatype()),
481 &Type::TIME => Ok(ConcreteDataType::timestamp_datatype(
482 common_time::timestamp::TimeUnit::Microsecond,
483 )),
484 &Type::INT2_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new(
485 ConcreteDataType::int16_datatype(),
486 ))),
487 &Type::INT4_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new(
488 ConcreteDataType::int32_datatype(),
489 ))),
490 &Type::INT8_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new(
491 ConcreteDataType::int64_datatype(),
492 ))),
493 &Type::NUMERIC_ARRAY => Ok(ConcreteDataType::list_datatype(Arc::new(
494 ConcreteDataType::uint64_datatype(),
495 ))),
496 &Type::VARCHAR_ARRAY | &Type::CHAR_ARRAY | &Type::TEXT_ARRAY => Ok(
497 ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::string_datatype())),
498 ),
499 _ => server_error::InternalSnafu {
500 err_msg: format!("unimplemented datatype {origin:?}"),
501 }
502 .fail(),
503 }
504}
505
506pub(super) fn invalid_parameter_error(msg: &str, detail: Option<String>) -> PgWireError {
507 let mut error_info = PgErrorCode::Ec22023.to_err_info(msg.to_string());
508 error_info.detail = detail;
509 PgWireError::UserError(Box::new(error_info))
510}
511
512fn to_timestamp_scalar_value<T>(
513 data: Option<T>,
514 unit: &TimestampType,
515 ctype: &ConcreteDataType,
516) -> PgWireResult<ScalarValue>
517where
518 T: Into<i64>,
519{
520 if let Some(n) = data {
521 Value::Timestamp(unit.create_timestamp(n.into()))
522 .try_to_scalar_value(ctype)
523 .map_err(convert_err)
524 } else {
525 Ok(ScalarValue::Null)
526 }
527}
528
529fn to_decimal_scalar_value(data: Option<Decimal>, ctype: &Decimal128Type) -> ScalarValue {
530 if let Some(data) = data {
531 let mut value = data;
532 value.rescale(ctype.scale() as u32);
533
534 ScalarValue::Decimal128(Some(value.mantissa()), ctype.precision(), ctype.scale())
535 } else {
536 ScalarValue::Decimal128(None, ctype.precision(), ctype.scale())
537 }
538}
539
540fn numeric_out_of_range_error(value: impl std::fmt::Display) -> PgWireError {
541 invalid_parameter_error(
542 "numeric_value_out_of_range",
543 Some(format!("value {} is out of range for target type", value)),
544 )
545}
546
547fn string_parameter_to_scalar_value(
548 data: Option<String>,
549 data_type: &ConcreteDataType,
550) -> Option<ScalarValue> {
551 match data_type {
552 ConcreteDataType::String(string_type) => {
553 if string_type.is_large() {
554 Some(ScalarValue::LargeUtf8(data))
555 } else {
556 Some(ScalarValue::Utf8(data))
557 }
558 }
559 ConcreteDataType::Dictionary(dictionary) => Some(ScalarValue::Dictionary(
560 Box::new(dictionary.key_type().as_arrow_type()),
561 Box::new(string_parameter_to_scalar_value(
562 data,
563 dictionary.value_type(),
564 )?),
565 )),
566 _ => None,
567 }
568}
569
570pub(super) fn parameters_to_scalar_values(
571 plan: &LogicalPlan,
572 portal: &Portal<PgSqlPlan>,
573) -> PgWireResult<Vec<ScalarValue>> {
574 let param_count = portal.parameter_len();
575 let mut results = Vec::with_capacity(param_count);
576
577 let client_param_types = &portal.statement.parameter_types;
578 let server_param_types = DfLogicalPlanner::get_inferred_parameter_types(plan)
579 .context(InferParameterTypesSnafu)
580 .map_err(convert_err)?
581 .into_iter()
582 .map(|(k, v)| (k, v.map(|v| ConcreteDataType::from_arrow_type(&v))))
583 .collect::<HashMap<_, _>>();
584
585 for idx in 0..param_count {
586 let server_type = server_param_types
587 .get(&format!("${}", idx + 1))
588 .and_then(|t| t.as_ref());
589
590 let client_type = if let Some(Some(client_given_type)) = client_param_types.get(idx) {
591 client_given_type.clone()
592 } else if let Some(server_provided_type) = &server_type {
593 type_gt_to_pg(server_provided_type).map_err(convert_err)?
594 } else {
595 return Err(invalid_parameter_error(
596 "unknown_parameter_type",
597 Some(format!(
598 "Cannot get type for parameter {}, try to provide a type using ${}::<type>",
599 idx, idx
600 )),
601 ));
602 };
603
604 let value = match &client_type {
605 &Type::VARCHAR | &Type::TEXT | &Type::CHAR => {
606 let data = portal.parameter::<String>(idx, &client_type)?;
607 if let Some(server_type) = &server_type {
608 string_parameter_to_scalar_value(data, server_type).ok_or_else(|| {
609 invalid_parameter_error(
610 "invalid_parameter_type",
611 Some(format!("Expected: {}, found: {}", server_type, client_type)),
612 )
613 })?
614 } else {
615 ScalarValue::Utf8(data)
616 }
617 }
618 &Type::BOOL => {
619 let data = portal.parameter::<bool>(idx, &client_type)?;
620 if let Some(server_type) = &server_type {
621 match server_type {
622 ConcreteDataType::Boolean(_) => ScalarValue::Boolean(data),
623 _ => {
624 return Err(invalid_parameter_error(
625 "invalid_parameter_type",
626 Some(format!("Expected: {}, found: {}", server_type, client_type)),
627 ));
628 }
629 }
630 } else {
631 ScalarValue::Boolean(data)
632 }
633 }
634 &Type::INT2 => {
635 let data = portal.parameter::<i16>(idx, &client_type)?;
636 if let Some(server_type) = &server_type {
637 match server_type {
638 ConcreteDataType::Int8(_) => ScalarValue::Int8(
639 data.map(|n| n.to_i8().ok_or_else(|| numeric_out_of_range_error(n)))
640 .transpose()?,
641 ),
642 ConcreteDataType::Int16(_) => ScalarValue::Int16(data),
643 ConcreteDataType::Int32(_) => ScalarValue::Int32(data.map(|n| n as i32)),
644 ConcreteDataType::Int64(_) => ScalarValue::Int64(data.map(|n| n as i64)),
645 ConcreteDataType::UInt8(_) => ScalarValue::UInt8(
646 data.map(|n| n.to_u8().ok_or_else(|| numeric_out_of_range_error(n)))
647 .transpose()?,
648 ),
649 ConcreteDataType::UInt16(_) => ScalarValue::UInt16(
650 data.map(|n| n.to_u16().ok_or_else(|| numeric_out_of_range_error(n)))
651 .transpose()?,
652 ),
653 ConcreteDataType::UInt32(_) => ScalarValue::UInt32(
654 data.map(|n| n.to_u32().ok_or_else(|| numeric_out_of_range_error(n)))
655 .transpose()?,
656 ),
657 ConcreteDataType::UInt64(_) => ScalarValue::UInt64(
658 data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)))
659 .transpose()?,
660 ),
661 ConcreteDataType::Timestamp(unit) => {
662 to_timestamp_scalar_value(data, unit, server_type)?
663 }
664 _ => {
665 return Err(invalid_parameter_error(
666 "invalid_parameter_type",
667 Some(format!("Expected: {}, found: {}", server_type, client_type)),
668 ));
669 }
670 }
671 } else {
672 ScalarValue::Int16(data)
673 }
674 }
675 &Type::INT4 => {
676 let data = portal.parameter::<i32>(idx, &client_type)?;
677 if let Some(server_type) = &server_type {
678 match server_type {
679 ConcreteDataType::Int8(_) => ScalarValue::Int8(
680 data.map(|n| n.to_i8().ok_or_else(|| numeric_out_of_range_error(n)))
681 .transpose()?,
682 ),
683 ConcreteDataType::Int16(_) => ScalarValue::Int16(
684 data.map(|n| n.to_i16().ok_or_else(|| numeric_out_of_range_error(n)))
685 .transpose()?,
686 ),
687 ConcreteDataType::Int32(_) => ScalarValue::Int32(data),
688 ConcreteDataType::Int64(_) => ScalarValue::Int64(data.map(|n| n as i64)),
689 ConcreteDataType::UInt8(_) => ScalarValue::UInt8(
690 data.map(|n| n.to_u8().ok_or_else(|| numeric_out_of_range_error(n)))
691 .transpose()?,
692 ),
693 ConcreteDataType::UInt16(_) => ScalarValue::UInt16(
694 data.map(|n| n.to_u16().ok_or_else(|| numeric_out_of_range_error(n)))
695 .transpose()?,
696 ),
697 ConcreteDataType::UInt32(_) => ScalarValue::UInt32(
698 data.map(|n| n.to_u32().ok_or_else(|| numeric_out_of_range_error(n)))
699 .transpose()?,
700 ),
701 ConcreteDataType::UInt64(_) => ScalarValue::UInt64(
702 data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)))
703 .transpose()?,
704 ),
705 ConcreteDataType::Timestamp(unit) => {
706 to_timestamp_scalar_value(data, unit, server_type)?
707 }
708 _ => {
709 return Err(invalid_parameter_error(
710 "invalid_parameter_type",
711 Some(format!("Expected: {}, found: {}", server_type, client_type)),
712 ));
713 }
714 }
715 } else {
716 ScalarValue::Int32(data)
717 }
718 }
719 &Type::INT8 => {
720 let data = portal.parameter::<i64>(idx, &client_type)?;
721 if let Some(server_type) = &server_type {
722 match server_type {
723 ConcreteDataType::Int8(_) => ScalarValue::Int8(
724 data.map(|n| n.to_i8().ok_or_else(|| numeric_out_of_range_error(n)))
725 .transpose()?,
726 ),
727 ConcreteDataType::Int16(_) => ScalarValue::Int16(
728 data.map(|n| n.to_i16().ok_or_else(|| numeric_out_of_range_error(n)))
729 .transpose()?,
730 ),
731 ConcreteDataType::Int32(_) => ScalarValue::Int32(
732 data.map(|n| n.to_i32().ok_or_else(|| numeric_out_of_range_error(n)))
733 .transpose()?,
734 ),
735 ConcreteDataType::Int64(_) => ScalarValue::Int64(data),
736 ConcreteDataType::UInt8(_) => ScalarValue::UInt8(
737 data.map(|n| n.to_u8().ok_or_else(|| numeric_out_of_range_error(n)))
738 .transpose()?,
739 ),
740 ConcreteDataType::UInt16(_) => ScalarValue::UInt16(
741 data.map(|n| n.to_u16().ok_or_else(|| numeric_out_of_range_error(n)))
742 .transpose()?,
743 ),
744 ConcreteDataType::UInt32(_) => ScalarValue::UInt32(
745 data.map(|n| n.to_u32().ok_or_else(|| numeric_out_of_range_error(n)))
746 .transpose()?,
747 ),
748 ConcreteDataType::UInt64(_) => ScalarValue::UInt64(
749 data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)))
750 .transpose()?,
751 ),
752 ConcreteDataType::Timestamp(unit) => {
753 to_timestamp_scalar_value(data, unit, server_type)?
754 }
755 _ => {
756 return Err(invalid_parameter_error(
757 "invalid_parameter_type",
758 Some(format!("Expected: {}, found: {}", server_type, client_type)),
759 ));
760 }
761 }
762 } else {
763 ScalarValue::Int64(data)
764 }
765 }
766 &Type::NUMERIC => {
767 let data = portal.parameter::<Decimal>(idx, &client_type)?;
768 match &server_type {
769 Some(ConcreteDataType::Decimal128(dt)) => to_decimal_scalar_value(data, dt),
770 Some(st @ ConcreteDataType::Timestamp(unit)) => {
771 to_timestamp_scalar_value(data.and_then(|n| n.to_i64()), unit, st)?
772 }
773 Some(ConcreteDataType::UInt64(_)) | None => ScalarValue::UInt64(
774 data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)))
775 .transpose()?,
776 ),
777 Some(st) => {
778 return Err(invalid_parameter_error(
779 "invalid_parameter_type",
780 Some(format!("Expected: {}, found: {}", st, client_type)),
781 ));
782 }
783 }
784 }
785 &Type::FLOAT4 => {
786 let data = portal.parameter::<f32>(idx, &client_type)?;
787 if let Some(server_type) = &server_type {
788 match server_type {
789 ConcreteDataType::Int8(_) => ScalarValue::Int8(
790 data.map(|n| n.to_i8().ok_or_else(|| numeric_out_of_range_error(n)))
791 .transpose()?,
792 ),
793 ConcreteDataType::Int16(_) => ScalarValue::Int16(
794 data.map(|n| n.to_i16().ok_or_else(|| numeric_out_of_range_error(n)))
795 .transpose()?,
796 ),
797 ConcreteDataType::Int32(_) => ScalarValue::Int32(
798 data.map(|n| n.to_i32().ok_or_else(|| numeric_out_of_range_error(n)))
799 .transpose()?,
800 ),
801 ConcreteDataType::Int64(_) => ScalarValue::Int64(
802 data.map(|n| n.to_i64().ok_or_else(|| numeric_out_of_range_error(n)))
803 .transpose()?,
804 ),
805 ConcreteDataType::UInt8(_) => ScalarValue::UInt8(
806 data.map(|n| n.to_u8().ok_or_else(|| numeric_out_of_range_error(n)))
807 .transpose()?,
808 ),
809 ConcreteDataType::UInt16(_) => ScalarValue::UInt16(
810 data.map(|n| n.to_u16().ok_or_else(|| numeric_out_of_range_error(n)))
811 .transpose()?,
812 ),
813 ConcreteDataType::UInt32(_) => ScalarValue::UInt32(
814 data.map(|n| n.to_u32().ok_or_else(|| numeric_out_of_range_error(n)))
815 .transpose()?,
816 ),
817 ConcreteDataType::UInt64(_) => ScalarValue::UInt64(
818 data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)))
819 .transpose()?,
820 ),
821 ConcreteDataType::Float32(_) => ScalarValue::Float32(data),
822 ConcreteDataType::Float64(_) => {
823 ScalarValue::Float64(data.map(|n| n as f64))
824 }
825 _ => {
826 return Err(invalid_parameter_error(
827 "invalid_parameter_type",
828 Some(format!("Expected: {}, found: {}", server_type, client_type)),
829 ));
830 }
831 }
832 } else {
833 ScalarValue::Float32(data)
834 }
835 }
836 &Type::FLOAT8 => {
837 let data = portal.parameter::<f64>(idx, &client_type)?;
838 if let Some(server_type) = &server_type {
839 match server_type {
840 ConcreteDataType::Int8(_) => ScalarValue::Int8(
841 data.map(|n| n.to_i8().ok_or_else(|| numeric_out_of_range_error(n)))
842 .transpose()?,
843 ),
844 ConcreteDataType::Int16(_) => ScalarValue::Int16(
845 data.map(|n| n.to_i16().ok_or_else(|| numeric_out_of_range_error(n)))
846 .transpose()?,
847 ),
848 ConcreteDataType::Int32(_) => ScalarValue::Int32(
849 data.map(|n| n.to_i32().ok_or_else(|| numeric_out_of_range_error(n)))
850 .transpose()?,
851 ),
852 ConcreteDataType::Int64(_) => ScalarValue::Int64(
853 data.map(|n| n.to_i64().ok_or_else(|| numeric_out_of_range_error(n)))
854 .transpose()?,
855 ),
856 ConcreteDataType::UInt8(_) => ScalarValue::UInt8(
857 data.map(|n| n.to_u8().ok_or_else(|| numeric_out_of_range_error(n)))
858 .transpose()?,
859 ),
860 ConcreteDataType::UInt16(_) => ScalarValue::UInt16(
861 data.map(|n| n.to_u16().ok_or_else(|| numeric_out_of_range_error(n)))
862 .transpose()?,
863 ),
864 ConcreteDataType::UInt32(_) => ScalarValue::UInt32(
865 data.map(|n| n.to_u32().ok_or_else(|| numeric_out_of_range_error(n)))
866 .transpose()?,
867 ),
868 ConcreteDataType::UInt64(_) => ScalarValue::UInt64(
869 data.map(|n| n.to_u64().ok_or_else(|| numeric_out_of_range_error(n)))
870 .transpose()?,
871 ),
872 ConcreteDataType::Float32(_) => ScalarValue::Float32(
873 data.map(|n| n.to_f32().ok_or_else(|| numeric_out_of_range_error(n)))
874 .transpose()?,
875 ),
876 ConcreteDataType::Float64(_) => ScalarValue::Float64(data),
877 _ => {
878 return Err(invalid_parameter_error(
879 "invalid_parameter_type",
880 Some(format!("Expected: {}, found: {}", server_type, client_type)),
881 ));
882 }
883 }
884 } else {
885 ScalarValue::Float64(data)
886 }
887 }
888 &Type::TIMESTAMP => {
889 let data = portal.parameter::<NaiveDateTime>(idx, &client_type)?;
890 if let Some(server_type) = &server_type {
891 match server_type {
892 ConcreteDataType::Timestamp(unit) => match *unit {
893 TimestampType::Second(_) => ScalarValue::TimestampSecond(
894 data.map(|ts| ts.and_utc().timestamp()),
895 None,
896 ),
897 TimestampType::Millisecond(_) => ScalarValue::TimestampMillisecond(
898 data.map(|ts| ts.and_utc().timestamp_millis()),
899 None,
900 ),
901 TimestampType::Microsecond(_) => ScalarValue::TimestampMicrosecond(
902 data.map(|ts| ts.and_utc().timestamp_micros()),
903 None,
904 ),
905 TimestampType::Nanosecond(_) => ScalarValue::TimestampNanosecond(
906 data.and_then(|ts| ts.and_utc().timestamp_nanos_opt()),
907 None,
908 ),
909 },
910 _ => {
911 return Err(invalid_parameter_error(
912 "invalid_parameter_type",
913 Some(format!("Expected: {}, found: {}", server_type, client_type)),
914 ));
915 }
916 }
917 } else {
918 ScalarValue::TimestampMillisecond(
919 data.map(|ts| ts.and_utc().timestamp_millis()),
920 None,
921 )
922 }
923 }
924 &Type::TIMESTAMPTZ => {
925 let data = portal.parameter::<DateTime<FixedOffset>>(idx, &client_type)?;
926 if let Some(server_type) = &server_type {
927 match server_type {
928 ConcreteDataType::Timestamp(unit) => match *unit {
929 TimestampType::Second(_) => {
930 ScalarValue::TimestampSecond(data.map(|ts| ts.timestamp()), None)
931 }
932 TimestampType::Millisecond(_) => ScalarValue::TimestampMillisecond(
933 data.map(|ts| ts.timestamp_millis()),
934 None,
935 ),
936 TimestampType::Microsecond(_) => ScalarValue::TimestampMicrosecond(
937 data.map(|ts| ts.timestamp_micros()),
938 None,
939 ),
940 TimestampType::Nanosecond(_) => ScalarValue::TimestampNanosecond(
941 data.and_then(|ts| ts.timestamp_nanos_opt()),
942 None,
943 ),
944 },
945 _ => {
946 return Err(invalid_parameter_error(
947 "invalid_parameter_type",
948 Some(format!("Expected: {}, found: {}", server_type, client_type)),
949 ));
950 }
951 }
952 } else {
953 ScalarValue::TimestampMillisecond(data.map(|ts| ts.timestamp_millis()), None)
954 }
955 }
956 &Type::DATE => {
957 let data = portal.parameter::<NaiveDate>(idx, &client_type)?;
958 if let Some(server_type) = &server_type {
959 match server_type {
960 ConcreteDataType::Date(_) => ScalarValue::Date32(
961 data.map(|d| (d - DateTime::UNIX_EPOCH.date_naive()).num_days() as i32),
962 ),
963 _ => {
964 return Err(invalid_parameter_error(
965 "invalid_parameter_type",
966 Some(format!("Expected: {}, found: {}", server_type, client_type)),
967 ));
968 }
969 }
970 } else {
971 ScalarValue::Date32(
972 data.map(|d| (d - DateTime::UNIX_EPOCH.date_naive()).num_days() as i32),
973 )
974 }
975 }
976 &Type::INTERVAL => {
977 let data = portal.parameter::<PgInterval>(idx, &client_type)?;
978 if let Some(server_type) = &server_type {
979 match server_type {
980 ConcreteDataType::Interval(IntervalType::YearMonth(_)) => {
981 ScalarValue::IntervalYearMonth(
982 data.map(|i| {
983 if i.days != 0 || i.microseconds != 0 {
984 Err(invalid_parameter_error(
985 "invalid_parameter_type",
986 Some(format!(
987 "Expected: {}, found: {}",
988 server_type, client_type
989 )),
990 ))
991 } else {
992 Ok(IntervalYearMonth::new(i.months).to_i32())
993 }
994 })
995 .transpose()?,
996 )
997 }
998 ConcreteDataType::Interval(IntervalType::DayTime(_)) => {
999 ScalarValue::IntervalDayTime(
1000 data.map(|i| {
1001 if i.months != 0 || i.microseconds % 1000 != 0 {
1002 Err(invalid_parameter_error(
1003 "invalid_parameter_type",
1004 Some(format!(
1005 "Expected: {}, found: {}",
1006 server_type, client_type
1007 )),
1008 ))
1009 } else {
1010 Ok(IntervalDayTime::new(
1011 i.days,
1012 (i.microseconds / 1000) as i32,
1013 )
1014 .into())
1015 }
1016 })
1017 .transpose()?,
1018 )
1019 }
1020 ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => {
1021 ScalarValue::IntervalMonthDayNano(data.map(|i| {
1022 IntervalMonthDayNano::new(
1023 i.months,
1024 i.days,
1025 i.microseconds * 1_000i64,
1026 )
1027 .into()
1028 }))
1029 }
1030 _ => {
1031 return Err(invalid_parameter_error(
1032 "invalid_parameter_type",
1033 Some(format!("Expected: {}, found: {}", server_type, client_type)),
1034 ));
1035 }
1036 }
1037 } else {
1038 ScalarValue::IntervalMonthDayNano(data.map(|i| {
1039 IntervalMonthDayNano::new(i.months, i.days, i.microseconds * 1_000i64)
1040 .into()
1041 }))
1042 }
1043 }
1044 &Type::BYTEA => {
1045 let data = portal.parameter::<Vec<u8>>(idx, &client_type)?;
1046 if let Some(server_type) = &server_type {
1047 match server_type {
1048 ConcreteDataType::String(t) => {
1049 let s = data.map(|d| String::from_utf8_lossy(&d).to_string());
1050 if t.is_large() {
1051 ScalarValue::LargeUtf8(s)
1052 } else {
1053 ScalarValue::Utf8(s)
1054 }
1055 }
1056 ConcreteDataType::Binary(_) => ScalarValue::Binary(data),
1057 _ => {
1058 return Err(invalid_parameter_error(
1059 "invalid_parameter_type",
1060 Some(format!("Expected: {}, found: {}", server_type, client_type)),
1061 ));
1062 }
1063 }
1064 } else {
1065 ScalarValue::Binary(data)
1066 }
1067 }
1068 &Type::JSONB => {
1069 let data = portal.parameter::<serde_json::Value>(idx, &client_type)?;
1070 if let Some(server_type) = &server_type {
1071 match server_type {
1072 ConcreteDataType::Binary(_) => {
1073 ScalarValue::Binary(data.map(|d| d.to_string().into_bytes()))
1074 }
1075 _ => {
1076 return Err(invalid_parameter_error(
1077 "invalid_parameter_type",
1078 Some(format!("Expected: {}, found: {}", server_type, client_type)),
1079 ));
1080 }
1081 }
1082 } else {
1083 ScalarValue::Binary(data.map(|d| d.to_string().into_bytes()))
1084 }
1085 }
1086 &Type::INT2_ARRAY => {
1087 let data = portal.parameter::<Vec<Option<i16>>>(idx, &client_type)?;
1088 if let Some(data) = data {
1089 let values = data.into_iter().map(|i| i.into()).collect::<Vec<_>>();
1090 ScalarValue::List(ScalarValue::new_list(&values, &ArrowDataType::Int16, true))
1091 } else {
1092 ScalarValue::Null
1093 }
1094 }
1095 &Type::INT4_ARRAY => {
1096 let data = portal.parameter::<Vec<Option<i32>>>(idx, &client_type)?;
1097 if let Some(data) = data {
1098 let values = data.into_iter().map(|i| i.into()).collect::<Vec<_>>();
1099 ScalarValue::List(ScalarValue::new_list(&values, &ArrowDataType::Int32, true))
1100 } else {
1101 ScalarValue::Null
1102 }
1103 }
1104 &Type::INT8_ARRAY => {
1105 let data = portal.parameter::<Vec<Option<i64>>>(idx, &client_type)?;
1106 if let Some(data) = data {
1107 let values = data.into_iter().map(|i| i.into()).collect::<Vec<_>>();
1108 ScalarValue::List(ScalarValue::new_list(&values, &ArrowDataType::Int64, true))
1109 } else {
1110 ScalarValue::Null
1111 }
1112 }
1113 &Type::NUMERIC_ARRAY => {
1114 let data = portal.parameter::<Vec<Option<Decimal>>>(idx, &client_type)?;
1115 if let Some(data) = data {
1116 let build_u64_list = |data: Vec<Option<Decimal>>| -> PgWireResult<ScalarValue> {
1117 let values = data
1118 .into_iter()
1119 .map(|n| {
1120 Ok(ScalarValue::UInt64(
1121 n.map(|n| {
1122 n.to_u64().ok_or_else(|| numeric_out_of_range_error(n))
1123 })
1124 .transpose()?,
1125 ))
1126 })
1127 .collect::<PgWireResult<Vec<_>>>()?;
1128 Ok(ScalarValue::List(ScalarValue::new_list(
1129 &values,
1130 &ArrowDataType::UInt64,
1131 true,
1132 )))
1133 };
1134 if let Some(server_type) = &server_type {
1135 match server_type {
1136 ConcreteDataType::List(list_type) => match list_type.item_type() {
1137 ConcreteDataType::UInt64(_) => build_u64_list(data)?,
1138 ConcreteDataType::Decimal128(dt) => {
1139 let values = data
1140 .into_iter()
1141 .map(|n| to_decimal_scalar_value(n, dt))
1142 .collect::<Vec<_>>();
1143 ScalarValue::List(ScalarValue::new_list(
1144 &values,
1145 &ArrowDataType::Decimal128(dt.precision(), dt.scale()),
1146 true,
1147 ))
1148 }
1149 _ => {
1150 return Err(invalid_parameter_error(
1152 "invalid_parameter_type",
1153 Some(format!(
1154 "Expected: {}, found: {}",
1155 list_type.item_type(),
1156 client_type
1157 )),
1158 ));
1159 }
1160 },
1161 _ => {
1162 return Err(invalid_parameter_error(
1164 "invalid_parameter_type",
1165 Some(format!(
1166 "Expected: {}, found: {}",
1167 server_type, client_type
1168 )),
1169 ));
1170 }
1171 }
1172 } else {
1173 build_u64_list(data)?
1175 }
1176 } else {
1177 ScalarValue::Null
1178 }
1179 }
1180 &Type::VARCHAR_ARRAY | &Type::TEXT_ARRAY | &Type::CHAR_ARRAY => {
1181 let data = portal.parameter::<Vec<Option<String>>>(idx, &client_type)?;
1182 if let Some(data) = data {
1183 let values = data.into_iter().map(|i| i.into()).collect::<Vec<_>>();
1184 ScalarValue::List(ScalarValue::new_list(&values, &ArrowDataType::Utf8, true))
1185 } else {
1186 ScalarValue::Null
1187 }
1188 }
1189 &Type::TIMESTAMP_ARRAY => {
1190 let data = portal.parameter::<Vec<Option<NaiveDateTime>>>(idx, &client_type)?;
1191 if let Some(data) = data {
1192 if let Some(ConcreteDataType::List(list_type)) = &server_type {
1193 match list_type.item_type() {
1194 ConcreteDataType::Timestamp(unit) => match *unit {
1195 TimestampType::Second(_) => {
1196 let values = data
1197 .into_iter()
1198 .map(|ts| {
1199 ScalarValue::TimestampSecond(
1200 ts.map(|ts| ts.and_utc().timestamp()),
1201 None,
1202 )
1203 })
1204 .collect::<Vec<_>>();
1205 ScalarValue::List(ScalarValue::new_list(
1206 &values,
1207 &ArrowDataType::Timestamp(TimeUnit::Second, None),
1208 true,
1209 ))
1210 }
1211 TimestampType::Millisecond(_) => {
1212 let values = data
1213 .into_iter()
1214 .map(|ts| {
1215 ScalarValue::TimestampMillisecond(
1216 ts.map(|ts| ts.and_utc().timestamp_millis()),
1217 None,
1218 )
1219 })
1220 .collect::<Vec<_>>();
1221 ScalarValue::List(ScalarValue::new_list(
1222 &values,
1223 &ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
1224 true,
1225 ))
1226 }
1227 TimestampType::Microsecond(_) => {
1228 let values = data
1229 .into_iter()
1230 .map(|ts| {
1231 ScalarValue::TimestampMicrosecond(
1232 ts.map(|ts| ts.and_utc().timestamp_micros()),
1233 None,
1234 )
1235 })
1236 .collect::<Vec<_>>();
1237 ScalarValue::List(ScalarValue::new_list(
1238 &values,
1239 &ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
1240 true,
1241 ))
1242 }
1243 TimestampType::Nanosecond(_) => {
1244 let values = data
1245 .into_iter()
1246 .map(|ts| match ts {
1247 None => {
1248 Ok(ScalarValue::TimestampNanosecond(None, None))
1249 }
1250 Some(ts) => ts
1251 .and_utc()
1252 .timestamp_nanos_opt()
1253 .map(|nanos| {
1254 ScalarValue::TimestampNanosecond(
1255 Some(nanos),
1256 None,
1257 )
1258 })
1259 .ok_or_else(|| numeric_out_of_range_error(ts)),
1260 })
1261 .collect::<PgWireResult<Vec<_>>>()?;
1262 ScalarValue::List(ScalarValue::new_list(
1263 &values,
1264 &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1265 true,
1266 ))
1267 }
1268 },
1269 _ => {
1270 return Err(invalid_parameter_error(
1271 "invalid_parameter_type",
1272 Some(format!(
1273 "Expected: {}, found: {}",
1274 list_type.item_type(),
1275 client_type
1276 )),
1277 ));
1278 }
1279 }
1280 } else {
1281 let values = data
1282 .into_iter()
1283 .map(|ts| {
1284 ScalarValue::TimestampMillisecond(
1285 ts.map(|ts| ts.and_utc().timestamp_millis()),
1286 None,
1287 )
1288 })
1289 .collect::<Vec<_>>();
1290 ScalarValue::List(ScalarValue::new_list(
1291 &values,
1292 &ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
1293 true,
1294 ))
1295 }
1296 } else {
1297 ScalarValue::Null
1298 }
1299 }
1300 &Type::TIMESTAMPTZ_ARRAY => {
1301 let data =
1302 portal.parameter::<Vec<Option<DateTime<FixedOffset>>>>(idx, &client_type)?;
1303 if let Some(data) = data {
1304 if let Some(ConcreteDataType::List(list_type)) = &server_type {
1305 match list_type.item_type() {
1306 ConcreteDataType::Timestamp(unit) => match *unit {
1307 TimestampType::Second(_) => {
1308 let values = data
1309 .into_iter()
1310 .map(|ts| {
1311 ScalarValue::TimestampSecond(
1312 ts.map(|ts| ts.timestamp()),
1313 None,
1314 )
1315 })
1316 .collect::<Vec<_>>();
1317 ScalarValue::List(ScalarValue::new_list(
1318 &values,
1319 &ArrowDataType::Timestamp(TimeUnit::Second, None),
1320 true,
1321 ))
1322 }
1323 TimestampType::Millisecond(_) => {
1324 let values = data
1325 .into_iter()
1326 .map(|ts| {
1327 ScalarValue::TimestampMillisecond(
1328 ts.map(|ts| ts.timestamp_millis()),
1329 None,
1330 )
1331 })
1332 .collect::<Vec<_>>();
1333 ScalarValue::List(ScalarValue::new_list(
1334 &values,
1335 &ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
1336 true,
1337 ))
1338 }
1339 TimestampType::Microsecond(_) => {
1340 let values = data
1341 .into_iter()
1342 .map(|ts| {
1343 ScalarValue::TimestampMicrosecond(
1344 ts.map(|ts| ts.timestamp_micros()),
1345 None,
1346 )
1347 })
1348 .collect::<Vec<_>>();
1349 ScalarValue::List(ScalarValue::new_list(
1350 &values,
1351 &ArrowDataType::Timestamp(TimeUnit::Microsecond, None),
1352 true,
1353 ))
1354 }
1355 TimestampType::Nanosecond(_) => {
1356 let values = data
1357 .into_iter()
1358 .map(|ts| {
1359 ScalarValue::TimestampNanosecond(
1360 ts.and_then(|ts| ts.timestamp_nanos_opt()),
1361 None,
1362 )
1363 })
1364 .collect::<Vec<_>>();
1365 ScalarValue::List(ScalarValue::new_list(
1366 &values,
1367 &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
1368 true,
1369 ))
1370 }
1371 },
1372 _ => {
1373 return Err(invalid_parameter_error(
1374 "invalid_parameter_type",
1375 Some(format!(
1376 "Expected: {}, found: {}",
1377 list_type.item_type(),
1378 client_type
1379 )),
1380 ));
1381 }
1382 }
1383 } else {
1384 let values = data
1385 .into_iter()
1386 .map(|ts| {
1387 ScalarValue::TimestampMillisecond(
1388 ts.map(|ts| ts.timestamp_millis()),
1389 None,
1390 )
1391 })
1392 .collect::<Vec<_>>();
1393 ScalarValue::List(ScalarValue::new_list(
1394 &values,
1395 &ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
1396 true,
1397 ))
1398 }
1399 } else {
1400 ScalarValue::Null
1401 }
1402 }
1403 _ => Err(invalid_parameter_error(
1404 "unsupported_parameter_value",
1405 Some(format!("Found type: {}", client_type)),
1406 ))?,
1407 };
1408
1409 results.push(value);
1410 }
1411
1412 Ok(results)
1413}
1414
1415pub(super) fn param_types_to_pg_types(
1416 param_types: &HashMap<String, Option<ConcreteDataType>>,
1417) -> Result<Vec<Type>> {
1418 let param_count = param_types.len();
1419 let mut types = Vec::with_capacity(param_count);
1420 for i in 0..param_count {
1421 if let Some(Some(param_type)) = param_types.get(&format!("${}", i + 1)) {
1422 let pg_type = type_gt_to_pg(param_type)?;
1423 types.push(pg_type);
1424 } else {
1425 types.push(Type::UNKNOWN);
1426 }
1427 }
1428 Ok(types)
1429}
1430
1431pub fn format_options_from_query_ctx(query_ctx: &QueryContextRef) -> Arc<PgFormatOptions> {
1432 let config = query_ctx.configuration_parameter();
1433 let (date_style, date_order) = *config.pg_datetime_style();
1434
1435 let mut format_options = PgFormatOptions::default();
1436 format_options.date_style = format!("{}, {}", date_style, date_order);
1437 format_options.interval_style = config.pg_intervalstyle_format().to_string();
1438 format_options.bytea_output = config.postgres_bytea_output().to_string();
1439 format_options.time_zone = query_ctx.timezone().to_string();
1440
1441 Arc::new(format_options)
1442}
1443
1444#[cfg(test)]
1445mod test {
1446 use std::str::FromStr;
1447 use std::sync::Arc;
1448
1449 use arrow::array::{
1450 Float64Builder, Int64Builder, ListBuilder, StringBuilder, TimestampSecondBuilder,
1451 };
1452 use arrow_schema::{Field, IntervalUnit};
1453 use bytes::Bytes;
1454 use datafusion_expr::expr::Placeholder;
1455 use datafusion_expr::{Expr, LogicalPlanBuilder};
1456 use datatypes::schema::{ColumnSchema, Schema};
1457 use datatypes::vectors::{
1458 BinaryVector, BooleanVector, DateVector, Float32Vector, Float64Vector, Int8Vector,
1459 Int16Vector, Int32Vector, Int64Vector, IntervalDayTimeVector, IntervalMonthDayNanoVector,
1460 IntervalYearMonthVector, ListVector, NullVector, StringVector, TimeSecondVector,
1461 TimestampSecondVector, UInt8Vector, UInt16Vector, UInt32Vector, UInt64Vector, VectorRef,
1462 };
1463 use futures::{StreamExt as FuturesStreamExt, stream};
1464 use pgwire::api::Type;
1465 use pgwire::api::portal::{Format, Portal};
1466 use pgwire::api::results::{
1467 CopyEncoder, CopyTextOptions, DataRowEncoder, FieldFormat, FieldInfo,
1468 };
1469 use pgwire::api::stmt::StoredStatement;
1470 use pgwire::messages::extendedquery::Bind;
1471 use session::context::QueryContextBuilder;
1472
1473 use super::*;
1474 use crate::SqlPlan;
1475 use crate::postgres::handler::PgSqlPlan;
1476
1477 #[test]
1478 fn test_null_array_maps_to_text_array() {
1479 let array = ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::null_datatype()));
1480 assert_eq!(Type::TEXT_ARRAY, type_gt_to_pg(&array).unwrap());
1481 }
1482
1483 #[test]
1484 fn test_schema_convert() {
1485 let column_schemas = vec![
1486 ColumnSchema::new("nulls", ConcreteDataType::null_datatype(), true),
1487 ColumnSchema::new("bools", ConcreteDataType::boolean_datatype(), true),
1488 ColumnSchema::new("int8s", ConcreteDataType::int8_datatype(), true),
1489 ColumnSchema::new("int16s", ConcreteDataType::int16_datatype(), true),
1490 ColumnSchema::new("int32s", ConcreteDataType::int32_datatype(), true),
1491 ColumnSchema::new("int64s", ConcreteDataType::int64_datatype(), true),
1492 ColumnSchema::new("uint8s", ConcreteDataType::uint8_datatype(), true),
1493 ColumnSchema::new("uint16s", ConcreteDataType::uint16_datatype(), true),
1494 ColumnSchema::new("uint32s", ConcreteDataType::uint32_datatype(), true),
1495 ColumnSchema::new("uint64s", ConcreteDataType::uint64_datatype(), true),
1496 ColumnSchema::new("float32s", ConcreteDataType::float32_datatype(), true),
1497 ColumnSchema::new("float64s", ConcreteDataType::float64_datatype(), true),
1498 ColumnSchema::new("binaries", ConcreteDataType::binary_datatype(), true),
1499 ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true),
1500 ColumnSchema::new(
1501 "timestamps",
1502 ConcreteDataType::timestamp_millisecond_datatype(),
1503 true,
1504 ),
1505 ColumnSchema::new("dates", ConcreteDataType::date_datatype(), true),
1506 ColumnSchema::new("times", ConcreteDataType::time_second_datatype(), true),
1507 ColumnSchema::new(
1508 "intervals",
1509 ConcreteDataType::interval_month_day_nano_datatype(),
1510 true,
1511 ),
1512 ];
1513 let pg_field_info = vec![
1514 FieldInfo::new("nulls".into(), None, None, Type::UNKNOWN, FieldFormat::Text),
1515 FieldInfo::new("bools".into(), None, None, Type::BOOL, FieldFormat::Text),
1516 FieldInfo::new("int8s".into(), None, None, Type::INT2, FieldFormat::Text),
1517 FieldInfo::new("int16s".into(), None, None, Type::INT2, FieldFormat::Text),
1518 FieldInfo::new("int32s".into(), None, None, Type::INT4, FieldFormat::Text),
1519 FieldInfo::new("int64s".into(), None, None, Type::INT8, FieldFormat::Text),
1520 FieldInfo::new("uint8s".into(), None, None, Type::INT2, FieldFormat::Text),
1521 FieldInfo::new("uint16s".into(), None, None, Type::INT4, FieldFormat::Text),
1522 FieldInfo::new("uint32s".into(), None, None, Type::INT8, FieldFormat::Text),
1523 FieldInfo::new(
1524 "uint64s".into(),
1525 None,
1526 None,
1527 Type::NUMERIC,
1528 FieldFormat::Text,
1529 ),
1530 FieldInfo::new(
1531 "float32s".into(),
1532 None,
1533 None,
1534 Type::FLOAT4,
1535 FieldFormat::Text,
1536 ),
1537 FieldInfo::new(
1538 "float64s".into(),
1539 None,
1540 None,
1541 Type::FLOAT8,
1542 FieldFormat::Text,
1543 ),
1544 FieldInfo::new(
1545 "binaries".into(),
1546 None,
1547 None,
1548 Type::BYTEA,
1549 FieldFormat::Text,
1550 ),
1551 FieldInfo::new(
1552 "strings".into(),
1553 None,
1554 None,
1555 Type::VARCHAR,
1556 FieldFormat::Text,
1557 ),
1558 FieldInfo::new(
1559 "timestamps".into(),
1560 None,
1561 None,
1562 Type::TIMESTAMP,
1563 FieldFormat::Text,
1564 ),
1565 FieldInfo::new("dates".into(), None, None, Type::DATE, FieldFormat::Text),
1566 FieldInfo::new("times".into(), None, None, Type::TIME, FieldFormat::Text),
1567 FieldInfo::new(
1568 "intervals".into(),
1569 None,
1570 None,
1571 Type::INTERVAL,
1572 FieldFormat::Text,
1573 ),
1574 ];
1575 let schema = Schema::new(column_schemas);
1576 let fs = schema_to_pg(&schema, &Format::UnifiedText, None).unwrap();
1577 assert_eq!(fs, pg_field_info);
1578 }
1579
1580 #[test]
1581 fn test_schema_convert_oid_alias_types() {
1582 let aliases = [
1583 (oid_field::kind::OID, Type::OID),
1584 (oid_field::kind::REGPROC, Type::REGPROC),
1585 ("regprocedure", Type::REGPROCEDURE),
1586 ("regoper", Type::REGOPER),
1587 ("regoperator", Type::REGOPERATOR),
1588 (oid_field::kind::REGCLASS, Type::REGCLASS),
1589 (oid_field::kind::REGTYPE, Type::REGTYPE),
1590 (oid_field::kind::REGNAMESPACE, Type::REGNAMESPACE),
1591 ("regrole", Type::REGROLE),
1592 ("regconfig", Type::REGCONFIG),
1593 ("regdictionary", Type::REGDICTIONARY),
1594 ("regcollation", Type::REGCOLLATION),
1595 ];
1596 let mut columns = Vec::new();
1597 let mut expected_oids = Vec::new();
1598 for (type_name, data_type, fallback_type) in [
1599 ("int32", ConcreteDataType::int32_datatype(), Type::INT4),
1600 ("utf8", ConcreteDataType::string_datatype(), Type::VARCHAR),
1601 ] {
1602 for (alias, pg_type) in &aliases {
1603 let mut column =
1604 ColumnSchema::new(format!("{type_name}_{alias}"), data_type.clone(), true);
1605 column
1606 .mut_metadata()
1607 .insert(OID_ALIAS_KEY.to_string(), alias.to_string());
1608 columns.push(column);
1609 expected_oids.push(pg_type.oid());
1610 }
1611
1612 columns.push(ColumnSchema::new(
1613 format!("{type_name}_untagged"),
1614 data_type.clone(),
1615 true,
1616 ));
1617 expected_oids.push(fallback_type.oid());
1618
1619 let mut unknown = ColumnSchema::new(format!("{type_name}_unknown"), data_type, true);
1620 unknown
1621 .mut_metadata()
1622 .insert(OID_ALIAS_KEY.to_string(), "unknown".to_string());
1623 columns.push(unknown);
1624 expected_oids.push(fallback_type.oid());
1625 }
1626
1627 let fields = schema_to_pg(&Schema::new(columns), &Format::UnifiedText, None).unwrap();
1628 let actual_oids = fields
1629 .iter()
1630 .map(|field| field.datatype().oid())
1631 .collect::<Vec<_>>();
1632
1633 assert_eq!(actual_oids, expected_oids);
1634 }
1635
1636 #[test]
1637 fn test_encode_text_format_data() {
1638 let pg_schema = vec![
1639 FieldInfo::new("nulls".into(), None, None, Type::UNKNOWN, FieldFormat::Text),
1640 FieldInfo::new("bools".into(), None, None, Type::BOOL, FieldFormat::Text),
1641 FieldInfo::new("uint8s".into(), None, None, Type::INT2, FieldFormat::Text),
1642 FieldInfo::new("uint16s".into(), None, None, Type::INT4, FieldFormat::Text),
1643 FieldInfo::new("uint32s".into(), None, None, Type::INT8, FieldFormat::Text),
1644 FieldInfo::new(
1645 "uint64s".into(),
1646 None,
1647 None,
1648 Type::NUMERIC,
1649 FieldFormat::Text,
1650 ),
1651 FieldInfo::new("int8s".into(), None, None, Type::INT2, FieldFormat::Text),
1652 FieldInfo::new("int16s".into(), None, None, Type::INT2, FieldFormat::Text),
1653 FieldInfo::new("int32s".into(), None, None, Type::INT4, FieldFormat::Text),
1654 FieldInfo::new("int64s".into(), None, None, Type::INT8, FieldFormat::Text),
1655 FieldInfo::new(
1656 "float32s".into(),
1657 None,
1658 None,
1659 Type::FLOAT4,
1660 FieldFormat::Text,
1661 ),
1662 FieldInfo::new(
1663 "float64s".into(),
1664 None,
1665 None,
1666 Type::FLOAT8,
1667 FieldFormat::Text,
1668 ),
1669 FieldInfo::new(
1670 "strings".into(),
1671 None,
1672 None,
1673 Type::VARCHAR,
1674 FieldFormat::Text,
1675 ),
1676 FieldInfo::new(
1677 "binaries".into(),
1678 None,
1679 None,
1680 Type::BYTEA,
1681 FieldFormat::Text,
1682 ),
1683 FieldInfo::new("dates".into(), None, None, Type::DATE, FieldFormat::Text),
1684 FieldInfo::new("times".into(), None, None, Type::TIME, FieldFormat::Text),
1685 FieldInfo::new(
1686 "timestamps".into(),
1687 None,
1688 None,
1689 Type::TIMESTAMP,
1690 FieldFormat::Text,
1691 ),
1692 FieldInfo::new(
1693 "interval_year_month".into(),
1694 None,
1695 None,
1696 Type::INTERVAL,
1697 FieldFormat::Text,
1698 ),
1699 FieldInfo::new(
1700 "interval_day_time".into(),
1701 None,
1702 None,
1703 Type::INTERVAL,
1704 FieldFormat::Text,
1705 ),
1706 FieldInfo::new(
1707 "interval_month_day_nano".into(),
1708 None,
1709 None,
1710 Type::INTERVAL,
1711 FieldFormat::Text,
1712 ),
1713 FieldInfo::new(
1714 "int_list".into(),
1715 None,
1716 None,
1717 Type::INT8_ARRAY,
1718 FieldFormat::Text,
1719 ),
1720 FieldInfo::new(
1721 "float_list".into(),
1722 None,
1723 None,
1724 Type::FLOAT8_ARRAY,
1725 FieldFormat::Text,
1726 ),
1727 FieldInfo::new(
1728 "string_list".into(),
1729 None,
1730 None,
1731 Type::VARCHAR_ARRAY,
1732 FieldFormat::Text,
1733 ),
1734 FieldInfo::new(
1735 "timestamp_list".into(),
1736 None,
1737 None,
1738 Type::TIMESTAMP_ARRAY,
1739 FieldFormat::Text,
1740 ),
1741 ];
1742
1743 let arrow_schema = arrow_schema::Schema::new(vec![
1744 Field::new("x", DataType::Null, true),
1745 Field::new("x", DataType::Boolean, true),
1746 Field::new("x", DataType::UInt8, true),
1747 Field::new("x", DataType::UInt16, true),
1748 Field::new("x", DataType::UInt32, true),
1749 Field::new("x", DataType::UInt64, true),
1750 Field::new("x", DataType::Int8, true),
1751 Field::new("x", DataType::Int16, true),
1752 Field::new("x", DataType::Int32, true),
1753 Field::new("x", DataType::Int64, true),
1754 Field::new("x", DataType::Float32, true),
1755 Field::new("x", DataType::Float64, true),
1756 Field::new("x", DataType::Utf8, true),
1757 Field::new("x", DataType::Binary, true),
1758 Field::new("x", DataType::Date32, true),
1759 Field::new("x", DataType::Time32(TimeUnit::Second), true),
1760 Field::new("x", DataType::Timestamp(TimeUnit::Second, None), true),
1761 Field::new("x", DataType::Interval(IntervalUnit::YearMonth), true),
1762 Field::new("x", DataType::Interval(IntervalUnit::DayTime), true),
1763 Field::new("x", DataType::Interval(IntervalUnit::MonthDayNano), true),
1764 Field::new(
1765 "x",
1766 DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
1767 true,
1768 ),
1769 Field::new(
1770 "x",
1771 DataType::List(Arc::new(Field::new("item", DataType::Float64, true))),
1772 true,
1773 ),
1774 Field::new(
1775 "x",
1776 DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
1777 true,
1778 ),
1779 Field::new(
1780 "x",
1781 DataType::List(Arc::new(Field::new(
1782 "item",
1783 DataType::Timestamp(TimeUnit::Second, None),
1784 true,
1785 ))),
1786 true,
1787 ),
1788 ]);
1789
1790 let mut builder = ListBuilder::new(Int64Builder::new());
1791 builder.append_value([Some(1i64), None, Some(2)]);
1792 builder.append_null();
1793 builder.append_value([Some(-1i64), None, Some(-2)]);
1794 let i64_list_array = builder.finish();
1795
1796 let mut builder = ListBuilder::new(Float64Builder::new());
1797 builder.append_value([Some(1.0f64), None, Some(2.0)]);
1798 builder.append_null();
1799 builder.append_value([Some(-1.0f64), None, Some(-2.0)]);
1800 let f64_list_array = builder.finish();
1801
1802 let mut builder = ListBuilder::new(StringBuilder::new());
1803 builder.append_value([Some("a"), None, Some("b")]);
1804 builder.append_null();
1805 builder.append_value([Some("c"), None, Some("d")]);
1806 let string_list_array = builder.finish();
1807
1808 let mut builder = ListBuilder::new(TimestampSecondBuilder::new());
1809 builder.append_value([Some(1i64), None, Some(2)]);
1810 builder.append_null();
1811 builder.append_value([Some(3i64), None, Some(4)]);
1812 let timestamp_list_array = builder.finish();
1813
1814 let values = vec![
1815 Arc::new(NullVector::new(3)) as VectorRef,
1816 Arc::new(BooleanVector::from(vec![Some(true), Some(false), None])),
1817 Arc::new(UInt8Vector::from(vec![Some(u8::MAX), Some(u8::MIN), None])),
1818 Arc::new(UInt16Vector::from(vec![
1819 Some(u16::MAX),
1820 Some(u16::MIN),
1821 None,
1822 ])),
1823 Arc::new(UInt32Vector::from(vec![
1824 Some(u32::MAX),
1825 Some(u32::MIN),
1826 None,
1827 ])),
1828 Arc::new(UInt64Vector::from(vec![
1829 Some(u64::MAX),
1830 Some(u64::MIN),
1831 None,
1832 ])),
1833 Arc::new(Int8Vector::from(vec![Some(i8::MAX), Some(i8::MIN), None])),
1834 Arc::new(Int16Vector::from(vec![
1835 Some(i16::MAX),
1836 Some(i16::MIN),
1837 None,
1838 ])),
1839 Arc::new(Int32Vector::from(vec![
1840 Some(i32::MAX),
1841 Some(i32::MIN),
1842 None,
1843 ])),
1844 Arc::new(Int64Vector::from(vec![
1845 Some(i64::MAX),
1846 Some(i64::MIN),
1847 None,
1848 ])),
1849 Arc::new(Float32Vector::from(vec![
1850 None,
1851 Some(f32::MAX),
1852 Some(f32::MIN),
1853 ])),
1854 Arc::new(Float64Vector::from(vec![
1855 None,
1856 Some(f64::MAX),
1857 Some(f64::MIN),
1858 ])),
1859 Arc::new(StringVector::from(vec![
1860 None,
1861 Some("hello"),
1862 Some("greptime"),
1863 ])),
1864 Arc::new(BinaryVector::from(vec![
1865 None,
1866 Some("hello".as_bytes().to_vec()),
1867 Some("world".as_bytes().to_vec()),
1868 ])),
1869 Arc::new(DateVector::from(vec![Some(1001), None, Some(1)])),
1870 Arc::new(TimeSecondVector::from(vec![Some(1001), None, Some(1)])),
1871 Arc::new(TimestampSecondVector::from(vec![
1872 Some(1000001),
1873 None,
1874 Some(1),
1875 ])),
1876 Arc::new(IntervalYearMonthVector::from(vec![Some(1), None, Some(2)])),
1877 Arc::new(IntervalDayTimeVector::from(vec![
1878 Some(arrow::datatypes::IntervalDayTime::new(1, 1)),
1879 None,
1880 Some(arrow::datatypes::IntervalDayTime::new(2, 2)),
1881 ])),
1882 Arc::new(IntervalMonthDayNanoVector::from(vec![
1883 Some(arrow::datatypes::IntervalMonthDayNano::new(1, 1, 10)),
1884 None,
1885 Some(arrow::datatypes::IntervalMonthDayNano::new(2, 2, 20)),
1886 ])),
1887 Arc::new(ListVector::from(i64_list_array)),
1888 Arc::new(ListVector::from(f64_list_array)),
1889 Arc::new(ListVector::from(string_list_array)),
1890 Arc::new(ListVector::from(timestamp_list_array)),
1891 ];
1892 let record_batch =
1893 RecordBatch::new(Arc::new(arrow_schema.try_into().unwrap()), values).unwrap();
1894
1895 let query_context = QueryContextBuilder::default()
1896 .configuration_parameter(Default::default())
1897 .build()
1898 .into();
1899 let schema = record_batch.schema.clone();
1900 let pg_schema_ref = Arc::new(pg_schema);
1901
1902 let encoder = DataRowEncoder::new(pg_schema_ref.clone());
1903
1904 let row_stream = RecordBatchRowStream::new(
1905 query_context,
1906 pg_schema_ref.clone(),
1907 schema,
1908 stream::once(async { Ok(record_batch) }),
1909 encoder,
1910 );
1911
1912 let rows: Vec<_> = futures::executor::block_on(
1913 row_stream
1914 .filter_map(|x: PgWireResult<_>| async move { x.ok() })
1915 .flat_map(stream::iter)
1916 .collect::<Vec<_>>(),
1917 );
1918 assert_eq!(rows.len(), 3);
1919 for row in rows {
1920 assert_eq!(row.field_count, pg_schema_ref.len() as i16);
1921 }
1922 }
1923
1924 #[test]
1925 fn test_encode_utf8_oid_alias_data() {
1926 let aliases = [
1927 ("regproc_binary", FieldFormat::Binary, Some("boolrecv")),
1928 ("regproc_text", FieldFormat::Text, Some("boolrecv")),
1929 ("regproc_unknown_text", FieldFormat::Text, Some("unknown")),
1930 ("regproc_ambiguous_text", FieldFormat::Text, Some("int4")),
1931 ("regtype_text", FieldFormat::Text, Some("int4recv")),
1932 ("regproc_zero", FieldFormat::Binary, Some("-")),
1933 ("regproc_null", FieldFormat::Binary, None),
1934 ];
1935 let mut columns = Vec::new();
1936 let mut values = Vec::new();
1937 let mut formats = Vec::new();
1938
1939 for (name, format, value) in aliases {
1940 let mut column = ColumnSchema::new(name, ConcreteDataType::string_datatype(), true);
1941 column.mut_metadata().insert(
1942 OID_ALIAS_KEY.to_string(),
1943 if name == "regtype_text" {
1944 oid_field::kind::REGTYPE.to_string()
1945 } else {
1946 oid_field::kind::REGPROC.to_string()
1947 },
1948 );
1949 columns.push(column);
1950 values.push(Arc::new(StringVector::from(vec![value])) as VectorRef);
1951 formats.push(format.value());
1952 }
1953
1954 columns.push(ColumnSchema::new(
1955 "varchar",
1956 ConcreteDataType::string_datatype(),
1957 false,
1958 ));
1959 values.push(Arc::new(StringVector::from(vec![Some("varchar")])) as VectorRef);
1960 formats.push(FieldFormat::Binary.value());
1961
1962 let schema = Arc::new(Schema::new(columns));
1963 let pg_schema =
1964 Arc::new(schema_to_pg(&schema, &Format::Individual(formats), None).unwrap());
1965 let record_batch = RecordBatch::new(schema.clone(), values).unwrap();
1966 let query_context = QueryContextBuilder::default()
1967 .configuration_parameter(Default::default())
1968 .build()
1969 .into();
1970 let row_stream = RecordBatchRowStream::new(
1971 query_context,
1972 pg_schema.clone(),
1973 schema,
1974 stream::once(async { Ok(record_batch) }),
1975 DataRowEncoder::new(pg_schema),
1976 );
1977
1978 let row = futures::executor::block_on(row_stream.into_future())
1979 .0
1980 .unwrap()
1981 .unwrap()
1982 .pop()
1983 .unwrap();
1984 assert_eq!(row.field_count, 8);
1985 assert_eq!(
1986 &row.data[..],
1987 [
1988 0, 0, 0, 4, 0, 0, 9, 132, 0, 0, 0, 8, b'b', b'o', b'o', b'l', b'r', b'e', b'c', b'v', 0, 0, 0, 7, b'u', b'n', b'k', b'n', b'o', b'w', b'n', 0, 0, 0, 4, b'i', b'n', b't', b'4', 0, 0, 0, 8, b'i', b'n', b't', b'4', b'r', b'e', b'c', b'v', 0, 0, 0, 4, 0, 0, 0, 0, 255, 255, 255, 255, 0, 0, 0, 7, b'v', b'a', b'r', b'c', b'h', b'a', b'r',
1996 ]
1997 );
1998 }
1999
2000 #[test]
2001 fn test_encode_utf8_oid_alias_numeric_data() {
2002 let aliases = [
2003 oid_field::kind::OID,
2004 oid_field::kind::REGPROC,
2005 "regprocedure",
2006 "regoper",
2007 "regoperator",
2008 oid_field::kind::REGCLASS,
2009 oid_field::kind::REGTYPE,
2010 oid_field::kind::REGNAMESPACE,
2011 "regrole",
2012 "regconfig",
2013 "regdictionary",
2014 "regcollation",
2015 ];
2016 let mut columns = Vec::new();
2017 let mut values = Vec::new();
2018 for alias in aliases {
2019 let mut column = ColumnSchema::new(alias, ConcreteDataType::string_datatype(), false);
2020 column
2021 .mut_metadata()
2022 .insert(OID_ALIAS_KEY.to_string(), alias.to_string());
2023 columns.push(column);
2024 values.push(Arc::new(StringVector::from(vec![Some("4294967295")])) as VectorRef);
2025 }
2026
2027 let schema = Arc::new(Schema::new(columns));
2028 let pg_schema = Arc::new(schema_to_pg(&schema, &Format::UnifiedBinary, None).unwrap());
2029 let record_batch = RecordBatch::new(schema.clone(), values).unwrap();
2030 let query_context = QueryContextBuilder::default()
2031 .configuration_parameter(Default::default())
2032 .build()
2033 .into();
2034 let row_stream = RecordBatchRowStream::new(
2035 query_context,
2036 pg_schema.clone(),
2037 schema,
2038 stream::once(async { Ok(record_batch) }),
2039 DataRowEncoder::new(pg_schema),
2040 );
2041
2042 let row = futures::executor::block_on(row_stream.into_future())
2043 .0
2044 .unwrap()
2045 .unwrap()
2046 .pop()
2047 .unwrap();
2048 assert_eq!(row.field_count, aliases.len() as i16);
2049 assert_eq!(
2050 &row.data[..],
2051 &[[0, 0, 0, 4, 255, 255, 255, 255]; 12].concat()
2052 );
2053 }
2054
2055 #[test]
2056 fn test_encode_utf8_oid_alias_binary_errors() {
2057 for value in ["unknown", "int4"] {
2058 let mut column =
2059 ColumnSchema::new("regproc", ConcreteDataType::string_datatype(), false);
2060 column.mut_metadata().insert(
2061 OID_ALIAS_KEY.to_string(),
2062 oid_field::kind::REGPROC.to_string(),
2063 );
2064 let schema = Arc::new(Schema::new(vec![column]));
2065 let pg_schema = Arc::new(schema_to_pg(&schema, &Format::UnifiedBinary, None).unwrap());
2066 let record_batch = RecordBatch::new(
2067 schema.clone(),
2068 vec![Arc::new(StringVector::from(vec![Some(value)])) as VectorRef],
2069 )
2070 .unwrap();
2071 let query_context = QueryContextBuilder::default()
2072 .configuration_parameter(Default::default())
2073 .build()
2074 .into();
2075 let row_stream = RecordBatchRowStream::new(
2076 query_context,
2077 pg_schema.clone(),
2078 schema,
2079 stream::once(async { Ok(record_batch) }),
2080 DataRowEncoder::new(pg_schema),
2081 );
2082
2083 assert!(
2084 futures::executor::block_on(row_stream.into_future())
2085 .0
2086 .unwrap()
2087 .is_err()
2088 );
2089 }
2090 }
2091
2092 #[test]
2093 fn test_copy_text_utf8_oid_alias_preserves_named_value() {
2094 let mut column = ColumnSchema::new("regproc", ConcreteDataType::string_datatype(), false);
2095 column.mut_metadata().insert(
2096 OID_ALIAS_KEY.to_string(),
2097 oid_field::kind::REGPROC.to_string(),
2098 );
2099 let schema = Arc::new(Schema::new(vec![column]));
2100 let pg_schema = Arc::new(schema_to_pg(&schema, &Format::UnifiedBinary, None).unwrap());
2101 let record_batch = RecordBatch::new(
2102 schema.clone(),
2103 vec![Arc::new(StringVector::from(vec![Some("unknown")])) as VectorRef],
2104 )
2105 .unwrap();
2106 let query_context = QueryContextBuilder::default()
2107 .configuration_parameter(Default::default())
2108 .build()
2109 .into();
2110 let row_stream = RecordBatchRowStream::new(
2111 query_context,
2112 pg_schema.clone(),
2113 schema,
2114 stream::once(async { Ok(record_batch) }),
2115 CopyEncoder::new_text(pg_schema, CopyTextOptions::default()),
2116 );
2117
2118 let row = futures::executor::block_on(row_stream.into_future())
2119 .0
2120 .unwrap()
2121 .unwrap()
2122 .pop()
2123 .unwrap();
2124 assert_eq!(row.data.as_ref(), b"unknown\n");
2125 }
2126
2127 #[test]
2128 fn test_invalid_parameter() {
2129 let msg = "invalid_parameter_count";
2131 let error = invalid_parameter_error(msg, None);
2132 if let PgWireError::UserError(value) = error {
2133 assert_eq!("ERROR", value.severity);
2134 assert_eq!("22023", value.code);
2135 assert_eq!(msg, value.message);
2136 } else {
2137 panic!("test_invalid_parameter failed");
2138 }
2139 }
2140
2141 #[test]
2142 fn test_to_decimal_scalar_value() {
2143 let dt = Decimal128Type::new(18, 4);
2144
2145 let d = Decimal::from_str("12345.6789").unwrap();
2146 assert_eq!(d.mantissa(), 123456789i128);
2147 let scalar = to_decimal_scalar_value(Some(d), &dt);
2148 assert_eq!(scalar, ScalarValue::Decimal128(Some(123456789), 18, 4));
2149
2150 let d = Decimal::from_str("100.5").unwrap();
2151 assert_eq!(d.mantissa(), 1005);
2152 let scalar = to_decimal_scalar_value(Some(d), &dt);
2153 assert_eq!(scalar, ScalarValue::Decimal128(Some(1005000), 18, 4));
2154
2155 let d = Decimal::from_str("-9876.5432").unwrap();
2156 let scalar = to_decimal_scalar_value(Some(d), &dt);
2157 assert_eq!(scalar, ScalarValue::Decimal128(Some(-98765432), 18, 4));
2158
2159 let scalar = to_decimal_scalar_value(None, &dt);
2160 assert_eq!(scalar, ScalarValue::Decimal128(None, 18, 4));
2161 }
2162
2163 fn s(v: &str) -> Option<String> {
2164 Some(v.to_string())
2165 }
2166
2167 fn typed_param(id: &str, dt: DataType) -> Expr {
2168 Expr::Placeholder(Placeholder::new_with_field(
2169 id.to_string(),
2170 Some(Arc::new(arrow_schema::Field::new(id, dt, true))),
2171 ))
2172 }
2173
2174 fn build_plan_with_params(params: Vec<(&str, DataType)>) -> LogicalPlan {
2175 let exprs: Vec<Expr> = params
2176 .into_iter()
2177 .map(|(id, dt)| typed_param(id, dt))
2178 .collect();
2179 LogicalPlanBuilder::empty(true)
2180 .project(exprs)
2181 .unwrap()
2182 .build()
2183 .unwrap()
2184 }
2185
2186 fn make_portal(
2187 client_param_types: Vec<Option<Type>>,
2188 param_data: Vec<Option<String>>,
2189 ) -> Portal<PgSqlPlan> {
2190 let bind = Bind::new(
2191 None,
2192 None,
2193 vec![],
2194 param_data
2195 .into_iter()
2196 .map(|opt| opt.map(Bytes::from))
2197 .collect(),
2198 vec![],
2199 );
2200 let statement = Arc::new(StoredStatement::new(
2201 String::new(),
2202 PgSqlPlan {
2203 plan: SqlPlan::Empty,
2204 copy_to_stdout_format: None,
2205 },
2206 client_param_types,
2207 ));
2208 Portal::try_new(&bind, statement).unwrap()
2209 }
2210
2211 #[test]
2212 fn test_dictionary_string_parameter() {
2213 let dictionary_type =
2214 DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8));
2215 let plan = build_plan_with_params(vec![("$1", dictionary_type)]);
2216 let portal = make_portal(vec![None], vec![s("host-a")]);
2217
2218 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2219 assert_eq!(
2220 vec![ScalarValue::Dictionary(
2221 Box::new(DataType::UInt32),
2222 Box::new(ScalarValue::Utf8(s("host-a"))),
2223 )],
2224 values
2225 );
2226
2227 let param_types = HashMap::from([(
2228 "$1".to_string(),
2229 Some(ConcreteDataType::dictionary_datatype(
2230 ConcreteDataType::uint32_datatype(),
2231 ConcreteDataType::string_datatype(),
2232 )),
2233 )]);
2234 assert_eq!(
2235 vec![Type::VARCHAR],
2236 param_types_to_pg_types(¶m_types).unwrap()
2237 );
2238 }
2239
2240 #[test]
2241 fn test_int2_coerce_in_range() {
2242 let plan = build_plan_with_params(vec![
2243 ("$1", DataType::Int8),
2244 ("$2", DataType::Int16),
2245 ("$3", DataType::Int32),
2246 ("$4", DataType::Int64),
2247 ("$5", DataType::UInt8),
2248 ("$6", DataType::UInt16),
2249 ("$7", DataType::UInt32),
2250 ("$8", DataType::UInt64),
2251 ]);
2252 let portal = make_portal(
2253 vec![
2254 Some(Type::INT2),
2255 Some(Type::INT2),
2256 Some(Type::INT2),
2257 Some(Type::INT2),
2258 Some(Type::INT2),
2259 Some(Type::INT2),
2260 Some(Type::INT2),
2261 Some(Type::INT2),
2262 ],
2263 vec![
2264 s("100"),
2265 s("100"),
2266 s("100"),
2267 s("100"),
2268 s("100"),
2269 s("100"),
2270 s("100"),
2271 s("100"),
2272 ],
2273 );
2274
2275 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2276 assert_eq!(values[0], ScalarValue::Int8(Some(100)));
2277 assert_eq!(values[1], ScalarValue::Int16(Some(100)));
2278 assert_eq!(values[2], ScalarValue::Int32(Some(100)));
2279 assert_eq!(values[3], ScalarValue::Int64(Some(100)));
2280 assert_eq!(values[4], ScalarValue::UInt8(Some(100)));
2281 assert_eq!(values[5], ScalarValue::UInt16(Some(100)));
2282 assert_eq!(values[6], ScalarValue::UInt32(Some(100)));
2283 assert_eq!(values[7], ScalarValue::UInt64(Some(100)));
2284 }
2285
2286 #[test]
2287 fn test_int2_coerce_out_of_range() {
2288 let plan = build_plan_with_params(vec![("$1", DataType::Int8)]);
2289 let portal = make_portal(vec![Some(Type::INT2)], vec![s("200")]);
2290 let result = parameters_to_scalar_values(&plan, &portal);
2291 assert!(result.is_err());
2292 }
2293
2294 #[test]
2295 fn test_int2_coerce_negative_to_unsigned_out_of_range() {
2296 let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]);
2297 let portal = make_portal(vec![Some(Type::INT2)], vec![s("-1")]);
2298 let result = parameters_to_scalar_values(&plan, &portal);
2299 assert!(result.is_err());
2300 }
2301
2302 #[test]
2303 fn test_int4_coerce_in_range() {
2304 let plan = build_plan_with_params(vec![
2305 ("$1", DataType::Int8),
2306 ("$2", DataType::Int16),
2307 ("$3", DataType::Int32),
2308 ("$4", DataType::Int64),
2309 ("$5", DataType::UInt8),
2310 ("$6", DataType::UInt16),
2311 ("$7", DataType::UInt32),
2312 ("$8", DataType::UInt64),
2313 ]);
2314 let portal = make_portal(
2315 vec![
2316 Some(Type::INT4),
2317 Some(Type::INT4),
2318 Some(Type::INT4),
2319 Some(Type::INT4),
2320 Some(Type::INT4),
2321 Some(Type::INT4),
2322 Some(Type::INT4),
2323 Some(Type::INT4),
2324 ],
2325 vec![
2326 s("100"),
2327 s("1000"),
2328 s("100000"),
2329 s("100000"),
2330 s("200"),
2331 s("1000"),
2332 s("100000"),
2333 s("100000"),
2334 ],
2335 );
2336
2337 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2338 assert_eq!(values[0], ScalarValue::Int8(Some(100)));
2339 assert_eq!(values[1], ScalarValue::Int16(Some(1000)));
2340 assert_eq!(values[2], ScalarValue::Int32(Some(100000)));
2341 assert_eq!(values[3], ScalarValue::Int64(Some(100000)));
2342 assert_eq!(values[4], ScalarValue::UInt8(Some(200)));
2343 assert_eq!(values[5], ScalarValue::UInt16(Some(1000)));
2344 assert_eq!(values[6], ScalarValue::UInt32(Some(100000)));
2345 assert_eq!(values[7], ScalarValue::UInt64(Some(100000)));
2346 }
2347
2348 #[test]
2349 fn test_int4_coerce_out_of_range() {
2350 let plan = build_plan_with_params(vec![("$1", DataType::Int8)]);
2351 let portal = make_portal(vec![Some(Type::INT4)], vec![s("200")]);
2352 let result = parameters_to_scalar_values(&plan, &portal);
2353 assert!(result.is_err());
2354 }
2355
2356 #[test]
2357 fn test_int4_coerce_i32_max_to_i16_out_of_range() {
2358 let plan = build_plan_with_params(vec![("$1", DataType::Int16)]);
2359 let portal = make_portal(vec![Some(Type::INT4)], vec![Some(i32::MAX.to_string())]);
2360 let result = parameters_to_scalar_values(&plan, &portal);
2361 assert!(result.is_err());
2362 }
2363
2364 #[test]
2365 fn test_int8_coerce_in_range() {
2366 let plan = build_plan_with_params(vec![
2367 ("$1", DataType::Int8),
2368 ("$2", DataType::Int16),
2369 ("$3", DataType::Int32),
2370 ("$4", DataType::Int64),
2371 ("$5", DataType::UInt8),
2372 ("$6", DataType::UInt16),
2373 ("$7", DataType::UInt32),
2374 ("$8", DataType::UInt64),
2375 ]);
2376 let portal = make_portal(
2377 vec![
2378 Some(Type::INT8),
2379 Some(Type::INT8),
2380 Some(Type::INT8),
2381 Some(Type::INT8),
2382 Some(Type::INT8),
2383 Some(Type::INT8),
2384 Some(Type::INT8),
2385 Some(Type::INT8),
2386 ],
2387 vec![
2388 s("100"),
2389 s("1000"),
2390 s("100000"),
2391 s("100000"),
2392 s("200"),
2393 s("1000"),
2394 s("3000000000"),
2395 s("3000000000"),
2396 ],
2397 );
2398
2399 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2400 assert_eq!(values[0], ScalarValue::Int8(Some(100)));
2401 assert_eq!(values[1], ScalarValue::Int16(Some(1000)));
2402 assert_eq!(values[2], ScalarValue::Int32(Some(100000)));
2403 assert_eq!(values[3], ScalarValue::Int64(Some(100000)));
2404 assert_eq!(values[4], ScalarValue::UInt8(Some(200)));
2405 assert_eq!(values[5], ScalarValue::UInt16(Some(1000)));
2406 assert_eq!(values[6], ScalarValue::UInt32(Some(3000000000)));
2407 assert_eq!(values[7], ScalarValue::UInt64(Some(3000000000)));
2408 }
2409
2410 #[test]
2411 fn test_int8_coerce_out_of_range() {
2412 let plan = build_plan_with_params(vec![("$1", DataType::Int32)]);
2413 let portal = make_portal(
2414 vec![Some(Type::INT8)],
2415 vec![Some((i32::MAX as i64 + 1).to_string())],
2416 );
2417 let result = parameters_to_scalar_values(&plan, &portal);
2418 assert!(result.is_err());
2419 }
2420
2421 #[test]
2422 fn test_int8_coerce_negative_to_unsigned_out_of_range() {
2423 let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]);
2424 let portal = make_portal(vec![Some(Type::INT8)], vec![s("-1")]);
2425 let result = parameters_to_scalar_values(&plan, &portal);
2426 assert!(result.is_err());
2427 }
2428
2429 #[test]
2430 fn test_float4_coerce_in_range() {
2431 let plan =
2432 build_plan_with_params(vec![("$1", DataType::Float32), ("$2", DataType::Float64)]);
2433 let portal = make_portal(
2434 vec![Some(Type::FLOAT4), Some(Type::FLOAT4)],
2435 vec![s("1.5"), s("2.5")],
2436 );
2437
2438 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2439 assert_eq!(values[0], ScalarValue::Float32(Some(1.5)));
2440 assert_eq!(values[1], ScalarValue::Float64(Some(2.5)));
2441 }
2442
2443 #[test]
2444 fn test_float4_coerce_to_int_in_range() {
2445 let plan = build_plan_with_params(vec![
2446 ("$1", DataType::Int8),
2447 ("$2", DataType::Int32),
2448 ("$3", DataType::UInt64),
2449 ]);
2450 let portal = make_portal(
2451 vec![Some(Type::FLOAT4), Some(Type::FLOAT4), Some(Type::FLOAT4)],
2452 vec![s("100"), s("1000"), s("200")],
2453 );
2454
2455 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2456 assert_eq!(values[0], ScalarValue::Int8(Some(100)));
2457 assert_eq!(values[1], ScalarValue::Int32(Some(1000)));
2458 assert_eq!(values[2], ScalarValue::UInt64(Some(200)));
2459 }
2460
2461 #[test]
2462 fn test_float4_coerce_to_int_out_of_range() {
2463 let plan = build_plan_with_params(vec![("$1", DataType::Int8)]);
2464 let portal = make_portal(vec![Some(Type::FLOAT4)], vec![s("200")]);
2465 let result = parameters_to_scalar_values(&plan, &portal);
2466 assert!(result.is_err());
2467 }
2468
2469 #[test]
2470 fn test_float8_coerce_in_range() {
2471 let plan =
2472 build_plan_with_params(vec![("$1", DataType::Float32), ("$2", DataType::Float64)]);
2473 let portal = make_portal(
2474 vec![Some(Type::FLOAT8), Some(Type::FLOAT8)],
2475 vec![s("1.5"), s("2.5")],
2476 );
2477
2478 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2479 assert_eq!(values[0], ScalarValue::Float32(Some(1.5)));
2480 assert_eq!(values[1], ScalarValue::Float64(Some(2.5)));
2481 }
2482
2483 #[test]
2484 fn test_float8_coerce_to_int_in_range() {
2485 let plan = build_plan_with_params(vec![
2486 ("$1", DataType::Int8),
2487 ("$2", DataType::Int64),
2488 ("$3", DataType::UInt64),
2489 ]);
2490 let portal = make_portal(
2491 vec![Some(Type::FLOAT8), Some(Type::FLOAT8), Some(Type::FLOAT8)],
2492 vec![s("100"), s("1000000"), s("200")],
2493 );
2494
2495 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2496 assert_eq!(values[0], ScalarValue::Int8(Some(100)));
2497 assert_eq!(values[1], ScalarValue::Int64(Some(1000000)));
2498 assert_eq!(values[2], ScalarValue::UInt64(Some(200)));
2499 }
2500
2501 #[test]
2502 fn test_float8_coerce_to_int_out_of_range() {
2503 let plan = build_plan_with_params(vec![("$1", DataType::Int8)]);
2504 let portal = make_portal(vec![Some(Type::FLOAT8)], vec![s("200")]);
2505 let result = parameters_to_scalar_values(&plan, &portal);
2506 assert!(result.is_err());
2507 }
2508
2509 #[test]
2510 fn test_float8_coerce_negative_to_unsigned_out_of_range() {
2511 let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]);
2512 let portal = make_portal(vec![Some(Type::FLOAT8)], vec![s("-1")]);
2513 let result = parameters_to_scalar_values(&plan, &portal);
2514 assert!(result.is_err());
2515 }
2516
2517 #[test]
2518 fn test_null_parameter() {
2519 let plan = build_plan_with_params(vec![("$1", DataType::Int8)]);
2520 let portal = make_portal(vec![Some(Type::INT2)], vec![None]);
2521
2522 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2523 assert_eq!(values[0], ScalarValue::Int8(None));
2524 }
2525
2526 fn numeric_uint64_array_plan() -> LogicalPlan {
2527 build_plan_with_params(vec![(
2528 "$1",
2529 DataType::List(Arc::new(Field::new("item", DataType::UInt64, true))),
2530 )])
2531 }
2532
2533 fn assert_numeric_out_of_range(result: PgWireResult<Vec<ScalarValue>>) {
2534 match result.unwrap_err() {
2535 PgWireError::UserError(error) => {
2536 assert_eq!("22023", error.code);
2537 assert_eq!("numeric_value_out_of_range", error.message);
2538 }
2539 error => panic!("expected numeric out-of-range error, got {error:?}"),
2540 }
2541 }
2542
2543 fn qbs_timestamp_nanosecond_array_plan() -> LogicalPlan {
2544 build_plan_with_params(vec![(
2545 "$1",
2546 DataType::List(Arc::new(Field::new(
2547 "item",
2548 DataType::Timestamp(TimeUnit::Nanosecond, None),
2549 true,
2550 ))),
2551 )])
2552 }
2553
2554 #[test]
2555 fn test_qbs_pg_timestamp_nanosecond_array_preserves_null_slot() {
2556 let portal = make_portal(
2557 vec![Some(Type::TIMESTAMP_ARRAY)],
2558 vec![s(
2559 r#"{"2024-01-01 00:00:00.000001",NULL,"2024-01-01 00:00:00.000003"}"#,
2560 )],
2561 );
2562
2563 let values =
2564 parameters_to_scalar_values(&qbs_timestamp_nanosecond_array_plan(), &portal).unwrap();
2565 let expected = ScalarValue::List(ScalarValue::new_list(
2566 &[
2567 ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_001_000), None),
2568 ScalarValue::TimestampNanosecond(None, None),
2569 ScalarValue::TimestampNanosecond(Some(1_704_067_200_000_003_000), None),
2570 ],
2571 &ArrowDataType::Timestamp(TimeUnit::Nanosecond, None),
2572 true,
2573 ));
2574 assert_eq!(expected, values[0]);
2575 }
2576
2577 #[test]
2578 fn test_qbs_pg_timestamp_nanosecond_array_out_of_range_rejected() {
2579 let portal = make_portal(
2580 vec![Some(Type::TIMESTAMP_ARRAY)],
2581 vec![s(r#"{"3000-01-01 00:00:00"}"#)],
2582 );
2583
2584 assert_numeric_out_of_range(parameters_to_scalar_values(
2585 &qbs_timestamp_nanosecond_array_plan(),
2586 &portal,
2587 ));
2588 }
2589
2590 #[test]
2591 fn test_numeric_uint64_scalar_negative_rejected() {
2592 let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]);
2593 let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s("-1")]);
2594
2595 assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal));
2596 }
2597
2598 #[test]
2599 fn test_numeric_uint64_scalar_above_u64_max_rejected() {
2600 let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]);
2601 let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s("18446744073709551616")]);
2602
2603 assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal));
2604 }
2605
2606 #[test]
2607 fn test_numeric_uint64_scalar_null_preserved() {
2608 let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]);
2609 let portal = make_portal(vec![Some(Type::NUMERIC)], vec![None]);
2610
2611 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2612 assert_eq!(ScalarValue::UInt64(None), values[0]);
2613 }
2614
2615 #[test]
2616 fn test_numeric_uint64_scalar_u64_max_preserved() {
2617 let plan = build_plan_with_params(vec![("$1", DataType::UInt64)]);
2618 let portal = make_portal(vec![Some(Type::NUMERIC)], vec![s(&u64::MAX.to_string())]);
2619
2620 let values = parameters_to_scalar_values(&plan, &portal).unwrap();
2621 assert_eq!(ScalarValue::UInt64(Some(u64::MAX)), values[0]);
2622 }
2623
2624 #[test]
2625 fn test_numeric_uint64_array_outer_null_preserved() {
2626 let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![None]);
2627
2628 let values = parameters_to_scalar_values(&numeric_uint64_array_plan(), &portal).unwrap();
2629 assert_eq!(ScalarValue::Null, values[0]);
2630 }
2631
2632 #[test]
2633 fn test_numeric_uint64_array_preserves_values_and_null_slots() {
2634 let portal = make_portal(
2635 vec![Some(Type::NUMERIC_ARRAY)],
2636 vec![s("{42,NULL,18446744073709551615}")],
2637 );
2638
2639 let values = parameters_to_scalar_values(&numeric_uint64_array_plan(), &portal).unwrap();
2640 let expected = ScalarValue::List(ScalarValue::new_list(
2641 &[
2642 ScalarValue::UInt64(Some(42)),
2643 ScalarValue::UInt64(None),
2644 ScalarValue::UInt64(Some(u64::MAX)),
2645 ],
2646 &ArrowDataType::UInt64,
2647 true,
2648 ));
2649 assert_eq!(expected, values[0]);
2650 }
2651
2652 #[test]
2653 fn test_numeric_uint64_array_invalid_after_valid_and_null_prefix_rejected() {
2654 let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![s("{42,NULL,-1}")]);
2655
2656 assert_numeric_out_of_range(parameters_to_scalar_values(
2657 &numeric_uint64_array_plan(),
2658 &portal,
2659 ));
2660 }
2661
2662 #[test]
2663 fn test_numeric_uint64_array_above_u64_max_rejected() {
2664 let portal = make_portal(
2665 vec![Some(Type::NUMERIC_ARRAY)],
2666 vec![s("{18446744073709551616}")],
2667 );
2668
2669 assert_numeric_out_of_range(parameters_to_scalar_values(
2670 &numeric_uint64_array_plan(),
2671 &portal,
2672 ));
2673 }
2674
2675 #[test]
2676 fn test_numeric_uint64_uninferred_array_invalid_value_rejected() {
2677 let plan = LogicalPlanBuilder::empty(true).build().unwrap();
2678 let portal = make_portal(vec![Some(Type::NUMERIC_ARRAY)], vec![s("{-1}")]);
2679
2680 assert_numeric_out_of_range(parameters_to_scalar_values(&plan, &portal));
2681 }
2682}