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