1use std::borrow::Cow;
17use std::cmp;
18use std::fmt;
19
20use aws_lc_rs::hmac as aws_hmac;
21use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
22use fallible_iterator::FallibleIterator;
23use hmac::{Hmac, KeyInit, Mac};
24use itertools::Itertools;
25use md5::Md5;
26use mz_expr_derive::sqlfunc;
27use mz_ore::cast::{CastFrom, ReinterpretCast};
28use mz_pgtz::timezone::TimezoneSpec;
29use mz_repr::ReprColumnType;
30use mz_repr::adt::array::{Array, ArrayDimension, ArrayDimensions, InvalidArrayError};
31use mz_repr::adt::mz_acl_item::{AclItem, AclMode, MzAclItem};
32use mz_repr::adt::range::{InvalidRangeError, Range, RangeBound, parse_range_bound_flags};
33use mz_repr::adt::system::Oid;
34use mz_repr::adt::timestamp::CheckedTimestamp;
35use mz_repr::role_id::RoleId;
36use mz_repr::{
37 ColumnName, Datum, DatumList, DatumMap, FromDatum, InputDatumType, OptionalArg,
38 OutputDatumType, Row, RowArena, SqlColumnType, SqlScalarType, Variadic,
39};
40use serde::{Deserialize, Serialize};
41
42use crate::func::{
43 CaseLiteral, array_create_scalar, build_regex, check_build_fits_budget,
44 check_datums_fit_budget, date_bin, max_string_func_result_bytes, parse_timezone,
45 regexp_match_static, regexp_replace_parse_flags, regexp_split_to_array_re, stringify_datum,
46 timezone_time,
47};
48use crate::{Eval, EvalError, MirScalarExpr};
49use mz_repr::adt::date::Date;
50use mz_repr::adt::interval::Interval;
51use mz_repr::adt::jsonb::JsonbRef;
52
53#[derive(
54 Ord,
55 PartialOrd,
56 Clone,
57 Debug,
58 Eq,
59 PartialEq,
60 Serialize,
61 Deserialize,
62 Hash
63)]
64pub struct And;
65
66impl fmt::Display for And {
67 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
68 f.write_str("AND")
69 }
70}
71
72impl LazyVariadicFunc for And {
73 fn eval<'a>(
74 &'a self,
75 datums: &[Datum<'a>],
76 temp_storage: &'a RowArena,
77 exprs: &'a [impl Eval],
78 ) -> Result<Datum<'a>, EvalError> {
79 let mut null = false;
81 let mut err = None;
82 for expr in exprs {
83 match expr.eval(datums, temp_storage) {
84 Ok(Datum::False) => return Ok(Datum::False), Ok(Datum::True) => {}
86 Ok(Datum::Null) => null = true,
88 Err(this_err) => err = std::cmp::max(err.take(), Some(this_err)),
89 _ => unreachable!(),
90 }
91 }
92 match (err, null) {
93 (Some(err), _) => Err(err),
94 (None, true) => Ok(Datum::Null),
95 (None, false) => Ok(Datum::True),
96 }
97 }
98
99 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
100 let in_nullable = input_types.iter().any(|t| t.nullable);
101 SqlScalarType::Bool.nullable(in_nullable)
102 }
103
104 fn propagates_nulls(&self) -> bool {
105 false
106 }
107
108 fn introduces_nulls(&self) -> bool {
109 false
110 }
111
112 fn could_error(&self) -> bool {
113 false
114 }
115
116 fn is_monotone(&self) -> bool {
117 true
118 }
119
120 fn is_associative(&self) -> bool {
121 true
122 }
123
124 fn is_infix_op(&self) -> bool {
125 true
126 }
127}
128
129#[derive(
130 Ord,
131 PartialOrd,
132 Clone,
133 Debug,
134 Eq,
135 PartialEq,
136 Serialize,
137 Deserialize,
138 Hash
139)]
140pub struct ArrayCreate {
141 pub elem_type: SqlScalarType,
142}
143
144#[sqlfunc(
160 ArrayCreate,
161 output_type_expr = "match &self.elem_type { SqlScalarType::Array(_) => self.elem_type.clone().nullable(false), _ => SqlScalarType::Array(Box::new(self.elem_type.clone())).nullable(false) }"
162)]
163fn array_create<'a>(
164 &self,
165 datums: Variadic<Datum<'a>>,
166 temp_storage: &'a RowArena,
167) -> Result<Array<'a>, EvalError> {
168 match &self.elem_type {
169 SqlScalarType::Array(_) => array_create_multidim(&datums, temp_storage),
170 _ => array_create_scalar(&datums, temp_storage),
171 }
172}
173fn array_create_multidim<'a>(
174 datums: &[Datum<'a>],
175 temp_storage: &'a RowArena,
176) -> Result<Array<'a>, EvalError> {
177 let mut dim: Option<ArrayDimensions> = None;
178 for datum in datums {
179 let actual_dims = match datum {
180 Datum::Null => ArrayDimensions::default(),
181 Datum::Array(arr) => arr.dims(),
182 d => panic!("unexpected datum {d}"),
183 };
184 if let Some(expected) = &dim {
185 if actual_dims.ndims() != expected.ndims() {
186 let actual = actual_dims.ndims().into();
187 let expected = expected.ndims().into();
188 return Err(InvalidArrayError::WrongCardinality { actual, expected }.into());
190 }
191 if let Some((e, a)) = expected
192 .into_iter()
193 .zip_eq(actual_dims)
194 .find(|(e, a)| e != a)
195 {
196 let actual = a.length;
197 let expected = e.length;
198 return Err(InvalidArrayError::WrongCardinality { actual, expected }.into());
200 }
201 }
202 dim = Some(actual_dims);
203 }
204 if dim.as_ref().map_or(true, ArrayDimensions::is_empty) {
206 let datum = temp_storage.try_make_datum(|packer| packer.try_push_array(&[], &[]))?;
207 return Ok(datum.unwrap_array());
208 }
209
210 let mut dims = vec![ArrayDimension {
211 lower_bound: 1,
212 length: datums.len(),
213 }];
214 if let Some(d) = datums.first() {
215 dims.extend(d.unwrap_array().dims());
216 };
217 check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
218 let elements = datums
219 .iter()
220 .flat_map(|d| d.unwrap_array().elements().iter());
221 let datum =
222 temp_storage.try_make_datum(move |packer| packer.try_push_array(&dims, elements))?;
223 Ok(datum.unwrap_array())
224}
225
226#[derive(
227 Ord,
228 PartialOrd,
229 Clone,
230 Debug,
231 Eq,
232 PartialEq,
233 Serialize,
234 Deserialize,
235 Hash
236)]
237pub struct ArrayFill {
238 pub elem_type: SqlScalarType,
239}
240
241#[sqlfunc(
242 ArrayFill,
243 output_type_expr = "SqlScalarType::Array(Box::new(self.elem_type.clone())).nullable(false)"
244)]
245fn array_fill<'a>(
246 &self,
247 fill: Datum<'a>,
248 dims: Option<Array<'a>>,
249 lower_bounds: OptionalArg<Option<Array<'a>>>,
250 temp_storage: &'a RowArena,
251) -> Result<Array<'a>, EvalError> {
252 const MAX_SIZE: usize = (1 << 28) - 1;
253 const NULL_ARR_ERR: &str = "dimension array or low bound array";
254 const NULL_ELEM_ERR: &str = "dimension values";
255
256 if matches!(fill, Datum::Array(_)) {
257 return Err(EvalError::Unsupported {
258 feature: "array_fill with arrays".into(),
259 discussion_no: None,
260 });
261 }
262
263 let Some(arr) = dims else {
264 return Err(EvalError::MustNotBeNull(NULL_ARR_ERR.into()));
265 };
266
267 if arr.dims().ndims() > 1 {
269 return Err(EvalError::ArrayFillWrongArraySubscripts);
270 }
271
272 let dimensions = arr
273 .elements()
274 .iter()
275 .map(|d| match d {
276 Datum::Null => Err(EvalError::MustNotBeNull(NULL_ELEM_ERR.into())),
277 d => Ok(usize::cast_from(u32::reinterpret_cast(d.unwrap_int32()))),
278 })
279 .collect::<Result<Vec<_>, _>>()?;
280
281 let lower_bounds = match *lower_bounds {
282 Some(d) => {
283 let Some(arr) = d else {
284 return Err(EvalError::MustNotBeNull(NULL_ARR_ERR.into()));
285 };
286
287 if arr.dims().ndims() > 1 {
289 return Err(EvalError::ArrayFillWrongArraySubscripts);
290 }
291
292 arr.elements()
293 .iter()
294 .map(|l| match l {
295 Datum::Null => Err(EvalError::MustNotBeNull(NULL_ELEM_ERR.into())),
296 l => Ok(isize::cast_from(l.unwrap_int32())),
297 })
298 .collect::<Result<Vec<_>, _>>()?
299 }
300 None => {
301 vec![1isize; dimensions.len()]
302 }
303 };
304
305 if lower_bounds.len() != dimensions.len() {
306 return Err(EvalError::ArrayFillWrongArraySubscripts);
307 }
308
309 let fill_count: usize = dimensions
310 .iter()
311 .cloned()
312 .map(Some)
313 .reduce(|a, b| match (a, b) {
314 (Some(a), Some(b)) => a.checked_mul(b),
315 _ => None,
316 })
317 .flatten()
318 .ok_or(EvalError::MaxArraySizeExceeded(MAX_SIZE))?;
319
320 if matches!(
321 mz_repr::datum_size(&fill).checked_mul(fill_count),
322 None | Some(MAX_SIZE..)
323 ) {
324 return Err(EvalError::MaxArraySizeExceeded(MAX_SIZE));
325 }
326
327 let budget_remaining = temp_storage.budget_remaining();
334 let packed_size = mz_repr::datum_size(&fill).saturating_mul(fill_count);
335 let build_size = fill_count.saturating_mul(std::mem::size_of::<Datum<'a>>());
336 if packed_size > budget_remaining || build_size > budget_remaining {
337 return Err(EvalError::TempStorageBudgetExceeded);
338 }
339
340 let array_dimensions = if fill_count == 0 {
341 vec![ArrayDimension {
342 lower_bound: 1,
343 length: 0,
344 }]
345 } else {
346 dimensions
347 .into_iter()
348 .zip_eq(lower_bounds)
349 .map(|(length, lower_bound)| ArrayDimension {
350 lower_bound,
351 length,
352 })
353 .collect()
354 };
355
356 let datum = temp_storage.try_make_datum(|packer| {
357 packer.try_push_array(&array_dimensions, vec![fill; fill_count])
358 })?;
359 Ok(datum.unwrap_array())
360}
361
362#[derive(
363 Ord,
364 PartialOrd,
365 Clone,
366 Debug,
367 Eq,
368 PartialEq,
369 Serialize,
370 Deserialize,
371 Hash
372)]
373pub struct ArrayIndex {
374 pub offset: i64,
375}
376#[sqlfunc(ArrayIndex, sqlname = "array_index")]
377fn array_index<'a, T: FromDatum<'a>>(
378 &self,
379 array: Array<'a, T>,
380 indices: Variadic<i64>,
381) -> Option<T> {
382 mz_ore::soft_assert_no_log!(
383 self.offset == 0 || self.offset == 1,
384 "offset must be either 0 or 1"
385 );
386
387 let dims = array.dims();
388 if dims.len() != indices.len() {
389 return None;
391 }
392
393 let mut final_idx = 0;
394
395 for (d, idx) in dims.into_iter().zip_eq(indices.iter()) {
396 let idx = isize::cast_from(*idx + self.offset);
398
399 let (lower, upper) = d.dimension_bounds();
400
401 if !(lower..upper + 1).contains(&idx) {
404 return None;
405 }
406
407 final_idx *= d.length;
409
410 final_idx += usize::try_from(idx - d.lower_bound)
414 .expect("previous bounds check ensures physical index is at least 0");
415 }
416
417 array.elements().typed_iter().nth(final_idx)
418}
419
420#[sqlfunc]
421fn array_position<'a>(
422 array: Array<'a>,
423 search: Datum<'a>,
424 initial_pos: OptionalArg<Option<i32>>,
425) -> Result<Option<i32>, EvalError> {
426 if array.dims().len() > 1 {
427 return Err(EvalError::MultiDimensionalArraySearch);
428 }
429
430 let skip = match initial_pos.0 {
442 None => 0,
443 Some(None) => return Err(EvalError::MustNotBeNull("initial position".into())),
444 Some(Some(o)) => usize::try_from(o).unwrap_or(0).saturating_sub(1),
445 };
446
447 let Some(r) = array.elements().iter().skip(skip).position(|d| d == search) else {
448 return Ok(None);
449 };
450
451 let p = i32::try_from(r + skip + 1).expect("fewer than i32::MAX elements in array");
453 Ok(Some(p))
454}
455
456#[derive(
457 Ord,
458 PartialOrd,
459 Clone,
460 Debug,
461 Eq,
462 PartialEq,
463 Serialize,
464 Deserialize,
465 Hash
466)]
467pub struct ArrayToString {
468 pub elem_type: SqlScalarType,
469}
470
471#[sqlfunc]
472fn array_to_string<'a>(
473 &self,
474 array: Array<'a>,
475 delimiter: &str,
476 null_str_arg: OptionalArg<Option<&str>>,
477 temp_storage: &RowArena,
478) -> Result<String, EvalError> {
479 let null_str = null_str_arg.flatten();
482 let max_result_bytes = max_string_func_result_bytes(temp_storage);
486 let mut out = String::new();
487 for elem in array.elements().iter() {
488 if elem.is_null() {
489 if let Some(null_str) = null_str {
490 out.push_str(null_str);
491 out.push_str(delimiter);
492 }
493 } else {
494 stringify_datum(&mut out, elem, &self.elem_type)?;
495 out.push_str(delimiter);
496 }
497 if out.len() > max_result_bytes {
498 return Err(EvalError::LengthTooLarge);
499 }
500 }
501 if out.len() > 0 {
502 out.truncate(out.len() - delimiter.len());
504 }
505 Ok(out)
506}
507
508#[derive(
509 Ord,
510 PartialOrd,
511 Clone,
512 Debug,
513 Eq,
514 PartialEq,
515 Serialize,
516 Deserialize,
517 Hash
518)]
519pub struct Coalesce;
520
521impl fmt::Display for Coalesce {
522 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
523 f.write_str("coalesce")
524 }
525}
526
527impl LazyVariadicFunc for Coalesce {
528 fn eval<'a>(
529 &'a self,
530 datums: &[Datum<'a>],
531 temp_storage: &'a RowArena,
532 exprs: &'a [impl Eval],
533 ) -> Result<Datum<'a>, EvalError> {
534 for e in exprs {
535 let d = e.eval(datums, temp_storage)?;
536 if !d.is_null() {
537 return Ok(d);
538 }
539 }
540 Ok(Datum::Null)
541 }
542
543 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
544 let nullable = input_types.iter().all(|typ| typ.nullable);
548 SqlColumnType::union_many(input_types).nullable(nullable)
549 }
550
551 fn propagates_nulls(&self) -> bool {
552 false
553 }
554
555 fn introduces_nulls(&self) -> bool {
556 false
557 }
558
559 fn could_error(&self) -> bool {
560 false
561 }
562
563 fn is_monotone(&self) -> bool {
564 true
565 }
566
567 fn is_associative(&self) -> bool {
568 true
569 }
570}
571
572#[derive(
573 Ord,
574 PartialOrd,
575 Clone,
576 Debug,
577 Eq,
578 PartialEq,
579 Serialize,
580 Deserialize,
581 Hash
582)]
583pub struct RangeCreate {
584 pub elem_type: SqlScalarType,
585}
586
587impl fmt::Display for RangeCreate {
588 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
589 f.write_str(match &self.elem_type {
590 SqlScalarType::Int32 => "int4range",
591 SqlScalarType::Int64 => "int8range",
592 SqlScalarType::Date => "daterange",
593 SqlScalarType::Numeric { .. } => "numrange",
594 SqlScalarType::Timestamp { .. } => "tsrange",
595 SqlScalarType::TimestampTz { .. } => "tstzrange",
596 _ => unreachable!(),
597 })
598 }
599}
600
601#[sqlfunc(
602 RangeCreate,
603 skip_display = true,
604 output_type_expr = "SqlScalarType::Range { element_type: Box::new(self.elem_type.clone()) }.nullable(false)"
605)]
606fn range_create<'a>(
607 &self,
608 lower: Datum<'a>,
609 upper: Datum<'a>,
610 flags_datum: Datum<'a>,
611) -> Result<Range<Datum<'a>>, EvalError> {
612 let flags = match flags_datum {
613 Datum::Null => {
614 return Err(EvalError::InvalidRange(
615 InvalidRangeError::NullRangeBoundFlags,
616 ));
617 }
618 o => o.unwrap_str(),
619 };
620
621 let (lower_inclusive, upper_inclusive) = parse_range_bound_flags(flags)?;
622
623 let mut range = Range::new(Some((
624 RangeBound::new(lower, lower_inclusive),
625 RangeBound::new(upper, upper_inclusive),
626 )));
627
628 range.canonicalize()?;
629
630 Ok(range)
631}
632
633#[sqlfunc(sqlname = "datediff")]
634fn date_diff_date(unit_str: &str, a: Date, b: Date) -> Result<i64, EvalError> {
635 let unit = unit_str
636 .parse()
637 .map_err(|_| EvalError::InvalidDatePart(unit_str.into()))?;
638
639 let a_ts = CheckedTimestamp::try_from(NaiveDate::from(a).and_hms_opt(0, 0, 0).unwrap())?;
641 let b_ts = CheckedTimestamp::try_from(NaiveDate::from(b).and_hms_opt(0, 0, 0).unwrap())?;
642 let diff = b_ts.diff_as(&a_ts, unit)?;
643 Ok(diff)
644}
645
646#[sqlfunc(sqlname = "datediff")]
647fn date_diff_time(unit_str: &str, a: NaiveTime, b: NaiveTime) -> Result<i64, EvalError> {
648 let unit = unit_str
649 .parse()
650 .map_err(|_| EvalError::InvalidDatePart(unit_str.into()))?;
651
652 let a_ts =
654 CheckedTimestamp::try_from(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap().and_time(a))?;
655 let b_ts =
656 CheckedTimestamp::try_from(NaiveDate::from_ymd_opt(1970, 1, 1).unwrap().and_time(b))?;
657 let diff = b_ts.diff_as(&a_ts, unit)?;
658 Ok(diff)
659}
660
661#[sqlfunc(sqlname = "datediff")]
662fn date_diff_timestamp(
663 unit: &str,
664 a: CheckedTimestamp<NaiveDateTime>,
665 b: CheckedTimestamp<NaiveDateTime>,
666) -> Result<i64, EvalError> {
667 let unit = unit
668 .parse()
669 .map_err(|_| EvalError::InvalidDatePart(unit.into()))?;
670
671 let diff = b.diff_as(&a, unit)?;
672 Ok(diff)
673}
674
675#[sqlfunc(sqlname = "datediff")]
676fn date_diff_timestamp_tz(
677 unit: &str,
678 a: CheckedTimestamp<DateTime<Utc>>,
679 b: CheckedTimestamp<DateTime<Utc>>,
680) -> Result<i64, EvalError> {
681 let unit = unit
682 .parse()
683 .map_err(|_| EvalError::InvalidDatePart(unit.into()))?;
684
685 let diff = b.diff_as(&a, unit)?;
686 Ok(diff)
687}
688
689#[derive(
690 Ord,
691 PartialOrd,
692 Clone,
693 Debug,
694 Eq,
695 PartialEq,
696 Serialize,
697 Deserialize,
698 Hash
699)]
700pub struct ErrorIfNull;
701
702impl fmt::Display for ErrorIfNull {
703 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
704 f.write_str("error_if_null")
705 }
706}
707
708impl LazyVariadicFunc for ErrorIfNull {
709 fn eval<'a>(
710 &'a self,
711 datums: &[Datum<'a>],
712 temp_storage: &'a RowArena,
713 exprs: &'a [impl Eval],
714 ) -> Result<Datum<'a>, EvalError> {
715 let first = exprs[0].eval(datums, temp_storage)?;
716 match first {
717 Datum::Null => {
718 let err_msg = match exprs[1].eval(datums, temp_storage)? {
719 Datum::Null => {
720 return Err(EvalError::Internal(
721 "unexpected NULL in error side of error_if_null".into(),
722 ));
723 }
724 o => o.unwrap_str(),
725 };
726 Err(EvalError::IfNullError(err_msg.into()))
727 }
728 _ => Ok(first),
729 }
730 }
731
732 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
733 input_types[0].scalar_type.clone().nullable(false)
734 }
735
736 fn propagates_nulls(&self) -> bool {
737 false
738 }
739
740 fn introduces_nulls(&self) -> bool {
741 false
742 }
743}
744
745#[derive(
746 Ord,
747 PartialOrd,
748 Clone,
749 Debug,
750 Eq,
751 PartialEq,
752 Serialize,
753 Deserialize,
754 Hash
755)]
756pub struct Greatest;
757
758impl fmt::Display for Greatest {
759 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
760 f.write_str("greatest")
761 }
762}
763
764impl LazyVariadicFunc for Greatest {
765 fn eval<'a>(
766 &'a self,
767 datums: &[Datum<'a>],
768 temp_storage: &'a RowArena,
769 exprs: &'a [impl Eval],
770 ) -> Result<Datum<'a>, EvalError> {
771 let datums = fallible_iterator::convert(exprs.iter().map(|e| e.eval(datums, temp_storage)));
772 Ok(datums
773 .filter(|d| Ok(!d.is_null()))
774 .max()?
775 .unwrap_or(Datum::Null))
776 }
777
778 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
779 SqlColumnType::union_many(input_types)
780 }
781
782 fn propagates_nulls(&self) -> bool {
783 false
784 }
785
786 fn introduces_nulls(&self) -> bool {
787 false
788 }
789
790 fn could_error(&self) -> bool {
791 false
792 }
793
794 fn is_monotone(&self) -> bool {
795 true
796 }
797
798 fn is_associative(&self) -> bool {
799 true
800 }
801}
802
803#[sqlfunc(sqlname = "hmac")]
804fn hmac_string(to_digest: &str, key: &str, typ: &str) -> Result<Vec<u8>, EvalError> {
805 let to_digest = to_digest.as_bytes();
806 let key = key.as_bytes();
807 hmac_inner(to_digest, key, typ)
808}
809
810#[sqlfunc(sqlname = "hmac")]
811fn hmac_bytes(to_digest: &[u8], key: &[u8], typ: &str) -> Result<Vec<u8>, EvalError> {
812 hmac_inner(to_digest, key, typ)
813}
814
815pub fn hmac_inner(to_digest: &[u8], key: &[u8], typ: &str) -> Result<Vec<u8>, EvalError> {
816 match typ {
817 "md5" => {
818 let mut mac = Hmac::<Md5>::new_from_slice(key).expect("HMAC accepts any key size");
819 mac.update(to_digest);
820 Ok(mac.finalize().into_bytes().to_vec())
821 }
822 "sha1" => {
823 let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA1_FOR_LEGACY_USE_ONLY, key);
824 Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
825 }
826 "sha224" => {
827 let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA224, key);
828 Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
829 }
830 "sha256" => {
831 let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA256, key);
832 Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
833 }
834 "sha384" => {
835 let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA384, key);
836 Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
837 }
838 "sha512" => {
839 let k = aws_hmac::Key::new(aws_hmac::HMAC_SHA512, key);
840 Ok(aws_hmac::sign(&k, to_digest).as_ref().to_vec())
841 }
842 other => Err(EvalError::InvalidHashAlgorithm(other.into())),
843 }
844}
845
846#[sqlfunc]
847fn jsonb_build_array<'a>(
848 datums: Variadic<Datum<'a>>,
849 temp_storage: &'a RowArena,
850) -> Result<JsonbRef<'a>, EvalError> {
851 check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
852 let datum = temp_storage.make_datum(|packer| {
853 packer.push_list(datums.into_iter().map(|d| match d {
854 Datum::Null => Datum::JsonNull,
855 d => d,
856 }))
857 });
858 Ok(JsonbRef::from_datum(datum))
859}
860
861#[sqlfunc]
862fn jsonb_build_object<'a>(
863 mut kvs: Variadic<(Datum<'a>, Datum<'a>)>,
864 temp_storage: &'a RowArena,
865) -> Result<JsonbRef<'a>, EvalError> {
866 kvs.0.sort_by(|kv1, kv2| kv1.0.cmp(&kv2.0));
867 kvs.0.dedup_by(|kv1, kv2| kv1.0 == kv2.0);
868 check_datums_fit_budget(kvs.0.iter().flat_map(|(k, v)| [*k, *v]), temp_storage)?;
870 let datum = temp_storage.try_make_datum(|packer| {
871 packer.push_dict_with(|packer| {
872 for (k, v) in kvs {
873 if k.is_null() {
874 return Err(EvalError::KeyCannotBeNull);
875 }
876 let v = match v {
877 Datum::Null => Datum::JsonNull,
878 d => d,
879 };
880 packer.push(k);
881 packer.push(v);
882 }
883 Ok(())
884 })
885 })?;
886 Ok(JsonbRef::from_datum(datum))
887}
888
889#[derive(
890 Ord,
891 PartialOrd,
892 Clone,
893 Debug,
894 Eq,
895 PartialEq,
896 Serialize,
897 Deserialize,
898 Hash
899)]
900pub struct Least;
901
902impl fmt::Display for Least {
903 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
904 f.write_str("least")
905 }
906}
907
908impl LazyVariadicFunc for Least {
909 fn eval<'a>(
910 &'a self,
911 datums: &[Datum<'a>],
912 temp_storage: &'a RowArena,
913 exprs: &'a [impl Eval],
914 ) -> Result<Datum<'a>, EvalError> {
915 let datums = fallible_iterator::convert(exprs.iter().map(|e| e.eval(datums, temp_storage)));
916 Ok(datums
917 .filter(|d| Ok(!d.is_null()))
918 .min()?
919 .unwrap_or(Datum::Null))
920 }
921
922 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
923 SqlColumnType::union_many(input_types)
924 }
925
926 fn propagates_nulls(&self) -> bool {
927 false
928 }
929
930 fn introduces_nulls(&self) -> bool {
931 false
932 }
933
934 fn could_error(&self) -> bool {
935 false
936 }
937
938 fn is_monotone(&self) -> bool {
939 true
940 }
941
942 fn is_associative(&self) -> bool {
943 true
944 }
945}
946
947#[derive(
948 Ord,
949 PartialOrd,
950 Clone,
951 Debug,
952 Eq,
953 PartialEq,
954 Serialize,
955 Deserialize,
956 Hash
957)]
958pub struct ListCreate {
959 pub elem_type: SqlScalarType,
960}
961
962#[sqlfunc(
963 output_type_expr = "SqlScalarType::List { element_type: Box::new(self.elem_type.clone()), custom_id: None }.nullable(false)"
964)]
965fn list_create<'a>(
966 &self,
967 datums: Variadic<Datum<'a>>,
968 temp_storage: &'a RowArena,
969) -> Result<DatumList<'a>, EvalError> {
970 check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
971 Ok(temp_storage
972 .make_datum(|packer| packer.push_list(datums))
973 .unwrap_list())
974}
975
976#[derive(
977 Ord,
978 PartialOrd,
979 Clone,
980 Debug,
981 Eq,
982 PartialEq,
983 Serialize,
984 Deserialize,
985 Hash
986)]
987pub struct RecordCreate {
988 pub field_names: Vec<ColumnName>,
989}
990
991#[sqlfunc(
992 output_type_expr = "SqlScalarType::Record { fields: self.field_names.clone().into_iter().zip_eq(input_types.iter().cloned()).collect(), custom_id: None }.nullable(false)"
993)]
994fn record_create<'a>(
995 &self,
996 datums: Variadic<Datum<'a>>,
997 temp_storage: &'a RowArena,
998) -> Result<DatumList<'a>, EvalError> {
999 check_datums_fit_budget(datums.iter().copied(), temp_storage)?;
1000 Ok(temp_storage
1001 .make_datum(|packer| packer.push_list(datums.iter().copied()))
1002 .unwrap_list())
1003}
1004
1005#[sqlfunc(
1006 output_type_expr = "input_types[0].scalar_type.unwrap_list_nth_layer_type(input_types.len() - 1).clone().nullable(true)"
1007)]
1008#[allow(clippy::as_conversions)]
1010fn list_index<'a>(buf: DatumList<'a>, indices: Variadic<i64>) -> Datum<'a> {
1011 let mut buf = Datum::List(buf);
1012 for i in indices {
1013 if buf.is_null() {
1014 break;
1015 }
1016 if i < 1 {
1017 return Datum::Null;
1018 }
1019
1020 buf = match buf.unwrap_list().iter().nth(i as usize - 1) {
1021 Some(datum) => datum,
1022 None => return Datum::Null,
1023 }
1024 }
1025 buf
1026}
1027
1028#[sqlfunc(sqlname = "makeaclitem")]
1029fn make_acl_item(
1030 grantee_oid: u32,
1031 grantor_oid: u32,
1032 privileges: &str,
1033 is_grantable: bool,
1034) -> Result<AclItem, EvalError> {
1035 let grantee = Oid(grantee_oid);
1036 let grantor = Oid(grantor_oid);
1037 let acl_mode = AclMode::parse_multiple_privileges(privileges)
1038 .map_err(|e: anyhow::Error| EvalError::InvalidPrivileges(e.to_string().into()))?;
1039 if is_grantable {
1040 return Err(EvalError::Unsupported {
1041 feature: "GRANT OPTION".into(),
1042 discussion_no: None,
1043 });
1044 }
1045
1046 Ok(AclItem {
1047 grantee,
1048 grantor,
1049 acl_mode,
1050 })
1051}
1052
1053#[sqlfunc(sqlname = "make_mz_aclitem")]
1054fn make_mz_acl_item(
1055 grantee_str: &str,
1056 grantor_str: &str,
1057 privileges: &str,
1058) -> Result<MzAclItem, EvalError> {
1059 let grantee: RoleId = grantee_str
1060 .parse()
1061 .map_err(|e: anyhow::Error| EvalError::InvalidRoleId(e.to_string().into()))?;
1062 let grantor: RoleId = grantor_str
1063 .parse()
1064 .map_err(|e: anyhow::Error| EvalError::InvalidRoleId(e.to_string().into()))?;
1065 if grantor == RoleId::Public {
1066 return Err(EvalError::InvalidRoleId(
1067 "mz_aclitem grantor cannot be PUBLIC role".into(),
1068 ));
1069 }
1070 let acl_mode = AclMode::parse_multiple_privileges(privileges)
1071 .map_err(|e: anyhow::Error| EvalError::InvalidPrivileges(e.to_string().into()))?;
1072
1073 Ok(MzAclItem {
1074 grantee,
1075 grantor,
1076 acl_mode,
1077 })
1078}
1079
1080#[sqlfunc(sqlname = "makets")]
1081#[allow(clippy::as_conversions)]
1082fn make_timestamp(
1083 year: i64,
1084 month: i64,
1085 day: i64,
1086 hour: i64,
1087 minute: i64,
1088 second_float: f64,
1089) -> Result<Option<CheckedTimestamp<NaiveDateTime>>, EvalError> {
1090 if year == 0 {
1097 return Err(EvalError::DateOutOfRange);
1098 }
1099 let year = if year < 0 { year + 1 } else { year };
1100 let year: i32 = year.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1101 let month: u32 = month.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1102 let day: u32 = day.try_into().map_err(|_| EvalError::DateOutOfRange)?;
1103 let date = NaiveDate::from_ymd_opt(year, month, day).ok_or(EvalError::DateOutOfRange)?;
1104 if !(0..=24).contains(&hour)
1107 || !(0..=59).contains(&minute)
1108 || !(0.0..=60.0).contains(&second_float)
1109 || (hour == 24 && (minute > 0 || second_float > 0.0))
1110 {
1111 return Err(EvalError::TimestampOutOfRange);
1112 }
1113 let micros = (second_float * 1_000_000.0).round() as i64;
1115 let time = chrono::Duration::microseconds((hour * 60 + minute) * 60 * 1_000_000 + micros);
1116 let timestamp = date
1117 .and_hms_opt(0, 0, 0)
1118 .unwrap()
1119 .checked_add_signed(time)
1120 .ok_or(EvalError::TimestampOutOfRange)?;
1121 Ok(Some(timestamp.try_into()?))
1122}
1123
1124#[derive(
1125 Ord,
1126 PartialOrd,
1127 Clone,
1128 Debug,
1129 Eq,
1130 PartialEq,
1131 Serialize,
1132 Deserialize,
1133 Hash
1134)]
1135pub struct MapBuild {
1136 pub value_type: SqlScalarType,
1137}
1138
1139#[sqlfunc(
1140 output_type_expr = "SqlScalarType::Map { value_type: Box::new(self.value_type.clone()), custom_id: None }.nullable(false)"
1141)]
1142fn map_build<'a>(
1143 &self,
1144 datums: Variadic<(Option<&str>, Datum<'a>)>,
1145 temp_storage: &'a RowArena,
1146) -> Result<DatumMap<'a>, EvalError> {
1147 let map: std::collections::BTreeMap<&str, _> = datums
1149 .into_iter()
1150 .filter_map(|(k, v)| k.map(|k| (k, v)))
1151 .collect();
1152
1153 check_datums_fit_budget(
1157 map.iter().flat_map(|(k, v)| [Datum::String(k), *v]),
1158 temp_storage,
1159 )?;
1160 Ok(temp_storage
1161 .make_datum(|packer| packer.push_dict(map))
1162 .unwrap_map())
1163}
1164
1165#[derive(
1166 Ord,
1167 PartialOrd,
1168 Clone,
1169 Debug,
1170 Eq,
1171 PartialEq,
1172 Serialize,
1173 Deserialize,
1174 Hash
1175)]
1176pub struct Or;
1177
1178impl fmt::Display for Or {
1179 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1180 f.write_str("OR")
1181 }
1182}
1183
1184impl LazyVariadicFunc for Or {
1185 fn eval<'a>(
1186 &'a self,
1187 datums: &[Datum<'a>],
1188 temp_storage: &'a RowArena,
1189 exprs: &'a [impl Eval],
1190 ) -> Result<Datum<'a>, EvalError> {
1191 let mut null = false;
1193 let mut err = None;
1194 for expr in exprs {
1195 match expr.eval(datums, temp_storage) {
1196 Ok(Datum::False) => {}
1197 Ok(Datum::True) => return Ok(Datum::True), Ok(Datum::Null) => null = true,
1200 Err(this_err) => err = std::cmp::max(err.take(), Some(this_err)),
1201 _ => unreachable!(),
1202 }
1203 }
1204 match (err, null) {
1205 (Some(err), _) => Err(err),
1206 (None, true) => Ok(Datum::Null),
1207 (None, false) => Ok(Datum::False),
1208 }
1209 }
1210
1211 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
1212 let in_nullable = input_types.iter().any(|t| t.nullable);
1213 SqlScalarType::Bool.nullable(in_nullable)
1214 }
1215
1216 fn propagates_nulls(&self) -> bool {
1217 false
1218 }
1219
1220 fn introduces_nulls(&self) -> bool {
1221 false
1222 }
1223
1224 fn could_error(&self) -> bool {
1225 false
1226 }
1227
1228 fn is_monotone(&self) -> bool {
1229 true
1230 }
1231
1232 fn is_associative(&self) -> bool {
1233 true
1234 }
1235
1236 fn is_infix_op(&self) -> bool {
1237 true
1238 }
1239}
1240
1241#[sqlfunc(sqlname = "lpad")]
1242fn pad_leading(
1243 string: &str,
1244 raw_len: i32,
1245 pad: OptionalArg<&str>,
1246 temp_storage: &RowArena,
1247) -> Result<String, EvalError> {
1248 let len = match usize::try_from(raw_len) {
1249 Ok(len) => len,
1250 Err(_) => {
1251 return Err(EvalError::InvalidParameterValue(
1252 "length must be nonnegative".into(),
1253 ));
1254 }
1255 };
1256 let pad_string = pad.unwrap_or(" ");
1257
1258 let widest_pad = pad_string.chars().map(char::len_utf8).max().unwrap_or(0);
1264 let pad_chars = len.saturating_sub(string.chars().count());
1265 let max_result_bytes = string
1266 .len()
1267 .saturating_add(pad_chars.saturating_mul(widest_pad));
1268 if max_result_bytes > max_string_func_result_bytes(temp_storage) {
1269 return Err(EvalError::LengthTooLarge);
1270 }
1271
1272 let (end_char, end_char_byte_offset) = string
1273 .chars()
1274 .take(len)
1275 .fold((0, 0), |acc, char| (acc.0 + 1, acc.1 + char.len_utf8()));
1276
1277 let truncating = len == end_char;
1281 let mut buf = String::with_capacity(if truncating {
1282 end_char_byte_offset
1283 } else {
1284 max_result_bytes
1285 });
1286 if truncating {
1287 buf.push_str(&string[0..end_char_byte_offset]);
1288 } else {
1289 buf.extend(pad_string.chars().cycle().take(len - end_char));
1290 buf.push_str(string);
1291 }
1292
1293 Ok(buf)
1294}
1295
1296#[sqlfunc(
1297 output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(true)"
1298)]
1299fn regexp_match<'a>(
1300 haystack: &'a str,
1301 needle: &str,
1302 flags: OptionalArg<&str>,
1303 temp_storage: &'a RowArena,
1304) -> Result<Option<Array<'a>>, EvalError> {
1305 let flags = flags.unwrap_or("");
1306 let needle = build_regex(needle, flags)?;
1307 regexp_match_static(haystack, temp_storage, &needle)
1308}
1309
1310#[sqlfunc(
1311 output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)"
1312)]
1313fn regexp_split_to_array<'a>(
1314 text: &str,
1315 regexp_str: &str,
1316 flags: OptionalArg<&str>,
1317 temp_storage: &'a RowArena,
1318) -> Result<Array<'a>, EvalError> {
1319 let flags = flags.unwrap_or("");
1320 let regexp = build_regex(regexp_str, flags)?;
1321 regexp_split_to_array_re(text, ®exp, temp_storage)
1322}
1323
1324#[sqlfunc]
1325fn regexp_replace<'a>(
1326 source: &'a str,
1327 pattern: &str,
1328 replacement: &str,
1329 flags_opt: OptionalArg<&str>,
1330 temp_storage: &RowArena,
1331) -> Result<Cow<'a, str>, EvalError> {
1332 let flags = flags_opt.0.unwrap_or("");
1333 let (limit, flags) = regexp_replace_parse_flags(flags);
1334 let regexp = build_regex(pattern, &flags)?;
1335
1336 let dollars = replacement.matches('$').count();
1340 let max_result_bytes = |matches: usize| {
1341 let expansions = match matches {
1342 0 => 0,
1343 _ => dollars.saturating_mul(source.len()),
1344 };
1345 source
1346 .len()
1347 .saturating_add(matches.saturating_mul(replacement.len()))
1348 .saturating_add(expansions)
1349 };
1350 let most_matches = match limit {
1354 0 => source.len().saturating_add(1),
1355 limit => limit,
1356 };
1357 let ceiling = max_string_func_result_bytes(temp_storage);
1358 if max_result_bytes(most_matches) > ceiling {
1359 let matches = regexp.find_iter(source).take(most_matches).count();
1360 if max_result_bytes(matches) > ceiling {
1361 return Err(EvalError::LengthTooLarge);
1362 }
1363 }
1364
1365 Ok(regexp.replacen(source, limit, replacement))
1366}
1367
1368#[sqlfunc]
1369fn replace(text: &str, from: &str, to: &str, temp_storage: &RowArena) -> Result<String, EvalError> {
1370 let max_result_bytes = max_string_func_result_bytes(temp_storage);
1377 let possible_size = text.len() * to.len();
1378 if possible_size > max_result_bytes {
1379 let replacement_count = text.matches(from).count();
1380 let estimated_size = text.len() + replacement_count * (to.len().saturating_sub(from.len()));
1381 if estimated_size > max_result_bytes {
1382 return Err(EvalError::LengthTooLarge);
1383 }
1384 }
1385
1386 Ok(text.replace(from, to))
1387}
1388
1389#[sqlfunc(
1390 output_type_expr = "SqlScalarType::Array(Box::new(SqlScalarType::String)).nullable(false)",
1391 propagates_nulls = false
1392)]
1393fn string_to_array<'a>(
1394 string: &'a str,
1395 delimiter: Option<&'a str>,
1396 null_string: OptionalArg<Option<&'a str>>,
1397 temp_storage: &'a RowArena,
1398) -> Result<Array<'a>, EvalError> {
1399 if string.is_empty() {
1400 let mut row = Row::default();
1401 let mut packer = row.packer();
1402 packer.try_push_array(&[], std::iter::empty::<Datum>())?;
1403
1404 return Ok(temp_storage.push_unary_row(row).unwrap_array());
1405 }
1406
1407 let Some(delimiter) = delimiter else {
1408 let split_all_chars_delimiter = "";
1409 return string_to_array_impl(
1410 string,
1411 split_all_chars_delimiter,
1412 null_string.flatten(),
1413 temp_storage,
1414 );
1415 };
1416
1417 if delimiter.is_empty() {
1418 let mut row = Row::default();
1419 let mut packer = row.packer();
1420 let dims = &[ArrayDimension {
1421 lower_bound: 1,
1422 length: 1,
1423 }];
1424 match null_string.flatten() {
1425 Some(null_string) if null_string == string => {
1426 packer.try_push_array(dims, std::iter::once(Datum::Null))?;
1427 }
1428 _ => {
1429 packer.try_push_array(dims, vec![string].into_iter().map(Datum::String))?;
1430 }
1431 }
1432 Ok(temp_storage.push_unary_row(row).unwrap_array())
1433 } else {
1434 string_to_array_impl(string, delimiter, null_string.flatten(), temp_storage)
1435 }
1436}
1437
1438fn string_to_array_impl<'a>(
1439 string: &str,
1440 delimiter: &str,
1441 null_string: Option<&'a str>,
1442 temp_storage: &'a RowArena,
1443) -> Result<Array<'a>, EvalError> {
1444 check_build_fits_budget(
1447 || {
1448 if delimiter.is_empty() {
1449 string.split(delimiter).filter(|s| !s.is_empty()).count()
1450 } else {
1451 string.split(delimiter).count()
1452 }
1453 },
1454 std::mem::size_of::<&str>(),
1455 temp_storage,
1456 )?;
1457
1458 let mut row = Row::default();
1459 let mut packer = row.packer();
1460
1461 let result = string.split(delimiter);
1462 let found: Vec<&str> = if delimiter.is_empty() {
1463 result.filter(|s| !s.is_empty()).collect()
1464 } else {
1465 result.collect()
1466 };
1467 let array_dimensions = [ArrayDimension {
1468 lower_bound: 1,
1469 length: found.len(),
1470 }];
1471
1472 check_datums_fit_budget(found.iter().copied().map(Datum::String), temp_storage)?;
1476
1477 if let Some(null_string) = null_string {
1478 let found_datums = found.into_iter().map(|chunk| {
1479 if chunk.eq(null_string) {
1480 Datum::Null
1481 } else {
1482 Datum::String(chunk)
1483 }
1484 });
1485
1486 packer.try_push_array(&array_dimensions, found_datums)?;
1487 } else {
1488 packer.try_push_array(&array_dimensions, found.into_iter().map(Datum::String))?;
1489 }
1490
1491 Ok(temp_storage.push_unary_row(row).unwrap_array())
1492}
1493
1494#[sqlfunc]
1495fn substr<'a>(s: &'a str, start: i32, length: OptionalArg<i32>) -> Result<&'a str, EvalError> {
1496 let raw_start_idx = i64::from(start) - 1;
1497 let start_idx = match usize::try_from(cmp::max(raw_start_idx, 0)) {
1498 Ok(i) => i,
1499 Err(_) => {
1500 return Err(EvalError::InvalidParameterValue(
1501 format!(
1502 "substring starting index ({}) exceeds min/max position",
1503 raw_start_idx
1504 )
1505 .into(),
1506 ));
1507 }
1508 };
1509
1510 let mut char_indices = s.char_indices();
1511 let get_str_index = |(index, _char)| index;
1512
1513 let str_len = s.len();
1514 let start_char_idx = char_indices.nth(start_idx).map_or(str_len, get_str_index);
1515
1516 if let OptionalArg(Some(len)) = length {
1517 let end_idx = match i64::from(len) {
1518 e if e < 0 => {
1519 return Err(EvalError::InvalidParameterValue(
1520 "negative substring length not allowed".into(),
1521 ));
1522 }
1523 e if e == 0 || e + raw_start_idx < 1 => return Ok(""),
1524 e => {
1525 let e = cmp::min(raw_start_idx + e - 1, e - 1);
1526 match usize::try_from(e) {
1527 Ok(i) => i,
1528 Err(_) => {
1529 return Err(EvalError::InvalidParameterValue(
1530 format!("substring length ({}) exceeds max position", e).into(),
1531 ));
1532 }
1533 }
1534 }
1535 };
1536
1537 let end_char_idx = char_indices.nth(end_idx).map_or(str_len, get_str_index);
1538
1539 Ok(&s[start_char_idx..end_char_idx])
1540 } else {
1541 Ok(&s[start_char_idx..])
1542 }
1543}
1544
1545#[sqlfunc(sqlname = "split_string")]
1546fn split_part<'a>(string: &'a str, delimiter: &str, field: i32) -> Result<&'a str, EvalError> {
1547 let index = match usize::try_from(i64::from(field) - 1) {
1548 Ok(index) => index,
1549 Err(_) => {
1550 return Err(EvalError::InvalidParameterValue(
1551 "field position must be greater than zero".into(),
1552 ));
1553 }
1554 };
1555
1556 if delimiter.is_empty() {
1560 if index == 0 {
1561 return Ok(string);
1562 } else {
1563 return Ok("");
1564 }
1565 }
1566
1567 Ok(string.split(delimiter).nth(index).unwrap_or(""))
1570}
1571
1572#[sqlfunc(is_associative = true)]
1573fn concat(strs: Variadic<Option<&str>>, temp_storage: &RowArena) -> Result<String, EvalError> {
1574 let max_result_bytes = max_string_func_result_bytes(temp_storage);
1575 let mut total_size = 0;
1576 for s in &strs {
1577 if let Some(s) = s {
1578 total_size += s.len();
1579 if total_size > max_result_bytes {
1580 return Err(EvalError::LengthTooLarge);
1581 }
1582 }
1583 }
1584 let mut buf = String::with_capacity(total_size);
1585 for s in strs {
1586 if let Some(s) = s {
1587 buf.push_str(s);
1588 }
1589 }
1590 Ok(buf)
1591}
1592
1593#[sqlfunc]
1594fn concat_ws(
1595 ws: &str,
1596 rest: Variadic<Option<&str>>,
1597 temp_storage: &RowArena,
1598) -> Result<String, EvalError> {
1599 let max_result_bytes = max_string_func_result_bytes(temp_storage);
1600 let mut total_size = 0;
1601 for s in &rest {
1602 if let Some(s) = s {
1603 total_size += s.len();
1604 total_size += ws.len();
1605 if total_size > max_result_bytes {
1606 return Err(EvalError::LengthTooLarge);
1607 }
1608 }
1609 }
1610
1611 let buf = Itertools::join(&mut rest.into_iter().filter_map(|s| s), ws);
1612
1613 Ok(buf)
1614}
1615
1616#[sqlfunc]
1617fn translate(
1618 string: &str,
1619 from_str: &str,
1620 to_str: &str,
1621 temp_storage: &RowArena,
1622) -> Result<String, EvalError> {
1623 let widest = to_str.chars().map(char::len_utf8).max().unwrap_or(0);
1630 let max_result_bytes = match widest.saturating_sub(1) {
1631 0 => string.len(),
1634 widening => string
1635 .len()
1636 .saturating_add(string.chars().count().saturating_mul(widening)),
1637 };
1638 if max_result_bytes > max_string_func_result_bytes(temp_storage) {
1639 return Err(EvalError::LengthTooLarge);
1640 }
1641
1642 let from = from_str.chars().collect::<Vec<_>>();
1643 let to = to_str.chars().collect::<Vec<_>>();
1644
1645 Ok(string
1646 .chars()
1647 .filter_map(|c| match from.iter().position(|f| f == &c) {
1648 Some(idx) => to.get(idx).copied(),
1649 None => Some(c),
1650 })
1651 .collect())
1652}
1653
1654#[sqlfunc(output_type_expr = "input_types[0].scalar_type.clone().nullable(false)")]
1655#[allow(clippy::as_conversions)]
1657fn list_slice_linear<'a>(
1658 list: DatumList<'a>,
1659 first: (i64, i64),
1660 remainder: Variadic<(i64, i64)>,
1661 temp_storage: &'a RowArena,
1662) -> DatumList<'a> {
1663 let mut start_idx = 0;
1664 let mut total_length = usize::MAX;
1665
1666 for (start, end) in std::iter::once(first).chain(remainder) {
1667 let start = std::cmp::max(start, 1);
1668
1669 if start > end {
1671 start_idx = 0;
1672 total_length = 0;
1673 break;
1674 }
1675
1676 let start_inner = start as usize - 1;
1677 start_idx += start_inner;
1679
1680 let length_inner = (end - start) as usize + 1;
1682 total_length = std::cmp::min(length_inner, total_length - start_inner);
1683 }
1684
1685 let iter = list.iter().skip(start_idx).take(total_length);
1686
1687 temp_storage
1688 .make_datum(|row| {
1689 row.push_list_with(|row| {
1690 for d in iter {
1692 row.push(d);
1693 }
1694 });
1695 })
1696 .unwrap_list()
1697}
1698
1699#[sqlfunc(sqlname = "timestamp_bin")]
1700fn date_bin_timestamp(
1701 stride: Interval,
1702 source: CheckedTimestamp<NaiveDateTime>,
1703 origin: CheckedTimestamp<NaiveDateTime>,
1704) -> Result<CheckedTimestamp<NaiveDateTime>, EvalError> {
1705 date_bin(stride, source, origin)
1706}
1707
1708#[sqlfunc(sqlname = "timestamptz_bin")]
1709fn date_bin_timestamp_tz(
1710 stride: Interval,
1711 source: CheckedTimestamp<DateTime<Utc>>,
1712 origin: CheckedTimestamp<DateTime<Utc>>,
1713) -> Result<CheckedTimestamp<DateTime<Utc>>, EvalError> {
1714 date_bin(stride, source, origin)
1715}
1716
1717#[sqlfunc(sqlname = "timezonet")]
1718fn timezone_time_variadic(
1719 tz_str: &str,
1720 time: NaiveTime,
1721 wall_time: CheckedTimestamp<DateTime<Utc>>,
1722) -> Result<NaiveTime, EvalError> {
1723 parse_timezone(tz_str, TimezoneSpec::Posix)
1724 .map(|tz| timezone_time(tz, time, &wall_time.naive_utc()))
1725}
1726pub(crate) trait LazyVariadicFunc: fmt::Display {
1727 fn eval<'a>(
1728 &'a self,
1729 datums: &[Datum<'a>],
1730 temp_storage: &'a RowArena,
1731 exprs: &'a [impl Eval],
1732 ) -> Result<Datum<'a>, EvalError>;
1733
1734 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType;
1736
1737 fn propagates_nulls(&self) -> bool;
1739
1740 fn introduces_nulls(&self) -> bool;
1742
1743 fn could_error(&self) -> bool {
1745 true
1746 }
1747
1748 fn is_monotone(&self) -> bool {
1750 false
1751 }
1752
1753 fn is_associative(&self) -> bool {
1755 false
1756 }
1757
1758 fn is_infix_op(&self) -> bool {
1760 false
1761 }
1762}
1763
1764pub(crate) trait EagerVariadicFunc: fmt::Display {
1765 type Input<'a>: InputDatumType<'a, EvalError>;
1766 type Output<'a>: OutputDatumType<'a, EvalError>;
1767
1768 fn call<'a>(&self, input: Self::Input<'a>, temp_storage: &'a RowArena) -> Self::Output<'a>;
1769
1770 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType;
1771
1772 fn propagates_nulls(&self) -> bool {
1773 !Self::Input::nullable()
1774 }
1775
1776 fn introduces_nulls(&self) -> bool {
1777 Self::Output::nullable()
1778 }
1779
1780 fn could_error(&self) -> bool {
1781 Self::Output::fallible()
1782 }
1783
1784 fn is_monotone(&self) -> bool {
1785 false
1786 }
1787
1788 fn is_associative(&self) -> bool {
1789 false
1790 }
1791
1792 fn is_infix_op(&self) -> bool {
1793 false
1794 }
1795}
1796
1797impl<T: EagerVariadicFunc> LazyVariadicFunc for T {
1800 fn eval<'a>(
1801 &'a self,
1802 datums: &[Datum<'a>],
1803 temp_storage: &'a RowArena,
1804 exprs: &'a [impl Eval],
1805 ) -> Result<Datum<'a>, EvalError> {
1806 let mut datums = exprs.iter().map(|e| e.eval(datums, temp_storage));
1807 match T::Input::try_from_iter(&mut datums) {
1808 Ok(input) => self.call(input, temp_storage).into_result(temp_storage),
1809 Err(Ok(None)) => Err(EvalError::Internal("missing parameter".into())),
1810 Err(Ok(Some(datum))) if datum.is_null() => Ok(datum),
1811 Err(Ok(Some(_datum))) => {
1812 Err(EvalError::Internal("invalid input type".into()))
1814 }
1815 Err(Err(res)) => Err(res),
1816 }
1817 }
1818
1819 fn output_type(&self, input_types: &[SqlColumnType]) -> SqlColumnType {
1820 self.output_type(input_types)
1821 }
1822
1823 fn propagates_nulls(&self) -> bool {
1824 self.propagates_nulls()
1825 }
1826
1827 fn introduces_nulls(&self) -> bool {
1828 self.introduces_nulls()
1829 }
1830
1831 fn could_error(&self) -> bool {
1832 self.could_error()
1833 }
1834
1835 fn is_monotone(&self) -> bool {
1836 self.is_monotone()
1837 }
1838
1839 fn is_associative(&self) -> bool {
1840 self.is_associative()
1841 }
1842
1843 fn is_infix_op(&self) -> bool {
1844 self.is_infix_op()
1845 }
1846}
1847
1848derive_variadic! {
1849 Coalesce(Coalesce),
1850 Greatest(Greatest),
1851 Least(Least),
1852 Concat(Concat),
1853 ConcatWs(ConcatWs),
1854 MakeTimestamp(MakeTimestamp),
1855 PadLeading(PadLeading),
1856 Substr(Substr),
1857 Replace(Replace),
1858 JsonbBuildArray(JsonbBuildArray),
1859 JsonbBuildObject(JsonbBuildObject),
1860 MapBuild(MapBuild),
1861 ArrayCreate(ArrayCreate),
1862 ArrayToString(ArrayToString),
1863 ArrayIndex(ArrayIndex),
1864 ListCreate(ListCreate),
1865 RecordCreate(RecordCreate),
1866 ListIndex(ListIndex),
1867 ListSliceLinear(ListSliceLinear),
1868 SplitPart(SplitPart),
1869 RegexpMatch(RegexpMatch),
1870 HmacString(HmacString),
1871 HmacBytes(HmacBytes),
1872 ErrorIfNull(ErrorIfNull),
1873 DateBinTimestamp(DateBinTimestamp),
1874 DateBinTimestampTz(DateBinTimestampTz),
1875 DateDiffTimestamp(DateDiffTimestamp),
1876 DateDiffTimestampTz(DateDiffTimestampTz),
1877 DateDiffDate(DateDiffDate),
1878 DateDiffTime(DateDiffTime),
1879 And(And),
1880 Or(Or),
1881 RangeCreate(RangeCreate),
1882 MakeAclItem(MakeAclItem),
1883 MakeMzAclItem(MakeMzAclItem),
1884 Translate(Translate),
1885 ArrayPosition(ArrayPosition),
1886 ArrayFill(ArrayFill),
1887 StringToArray(StringToArray),
1888 TimezoneTimeVariadic(TimezoneTimeVariadic),
1889 RegexpSplitToArray(RegexpSplitToArray),
1890 RegexpReplace(RegexpReplace),
1891 CaseLiteral(CaseLiteral),
1892}
1893
1894impl VariadicFunc {
1895 pub fn switch_and_or(&self) -> Self {
1896 match self {
1897 VariadicFunc::And(_) => Or.into(),
1898 VariadicFunc::Or(_) => And.into(),
1899 _ => unreachable!(),
1900 }
1901 }
1902
1903 pub fn unit_of_and_or(&self) -> MirScalarExpr {
1906 match self {
1907 VariadicFunc::And(_) => MirScalarExpr::literal_true(),
1908 VariadicFunc::Or(_) => MirScalarExpr::literal_false(),
1909 _ => unreachable!(),
1910 }
1911 }
1912
1913 pub fn zero_of_and_or(&self) -> MirScalarExpr {
1915 match self {
1916 VariadicFunc::And(_) => MirScalarExpr::literal_false(),
1917 VariadicFunc::Or(_) => MirScalarExpr::literal_true(),
1918 _ => unreachable!(),
1919 }
1920 }
1921}