Skip to content

Commit

Permalink
fix(query): add_hours function may panic if the argument is too big
Browse files Browse the repository at this point in the history
  • Loading branch information
@TCeason
TCeason committed Nov 25, 2024
1 parent cf150d9 commit 33805b5
Show file tree
Hide file tree
Showing 3 changed files with 154 additions and 27 deletions.
7 changes: 4 additions & 3 deletions src/query/expression/src/utils/date_helper.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -293,7 +293,7 @@ impl EvalDaysImpl {
}

pub fn eval_timestamp(date: i64, delta: impl AsPrimitive<i64>) -> i64 {
let mut value = date.wrapping_add(delta.as_() * MICROSECS_PER_DAY);
let mut value = date.wrapping_add(delta.as_().wrapping_mul(MICROSECS_PER_DAY));
clamp_timestamp(&mut value);
value
}
Expand All @@ -311,12 +311,13 @@ pub struct EvalTimesImpl;
impl EvalTimesImpl {
pub fn eval_date(date: i32, delta: impl AsPrimitive<i64>, factor: i64) -> i32 {
clamp_date(
(date as i64 * MICROSECS_PER_DAY).wrapping_add(delta.as_() * factor * MICROS_PER_SEC),
(date as i64 * MICROSECS_PER_DAY)
.wrapping_add(delta.as_().wrapping_mul(factor * MICROS_PER_SEC)),
)
}

pub fn eval_timestamp(us: i64, delta: impl AsPrimitive<i64>, factor: i64) -> i64 {
let mut ts = us.wrapping_add(delta.as_() * factor * MICROS_PER_SEC);
let mut ts = us.wrapping_add(delta.as_().wrapping_mul(factor * MICROS_PER_SEC));
clamp_timestamp(&mut ts);
ts
}
Expand Down
158 changes: 134 additions & 24 deletions src/query/functions/src/scalars/datetime.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -872,44 +872,74 @@ macro_rules! impl_register_arith_functions {

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
match EvalYearsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) {
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
match EvalYearsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) {
Ok(t) => builder.push(t),
Err(e) => {
ctx.set_error(builder.len(), e);
builder.push(0);
},
}

}

Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}

}

}),
);
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
concat!($op, "_years"),

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|ts, delta, builder, ctx| {
match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) {
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) {
Ok(t) => builder.push(t),
Err(e) => {
ctx.set_error(builder.len(), e);
builder.push(0);
},
}
}
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
},
),
);

registry.register_passthrough_nullable_2_arg::<DateType, Int64Type, DateType, _, _>(
concat!($op, "_quarters"),

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta} * 3) {
let delta_i32: Result<i32, _> = (delta*3).try_into();
match delta_i32 {
Ok(delta_i32) => {
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) {
Ok(t) => builder.push(t),
Err(e) => {
ctx.set_error(builder.len(), e);
builder.push(0);
},
}
}

Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
}),
);
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
Expand All @@ -918,13 +948,23 @@ macro_rules! impl_register_arith_functions {
|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|ts, delta, builder, ctx| {
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta} * 3) {
let delta_i32: Result<i32, _> = (delta*3).try_into();
match delta_i32 {
Ok(delta_i32) => {
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) {
Ok(t) => builder.push(t),
Err(e) => {
ctx.set_error(builder.len(), e);
builder.push(0);
},
}
}

Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
},
),
);
Expand All @@ -934,13 +974,22 @@ macro_rules! impl_register_arith_functions {

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) {
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) {
Ok(t) => builder.push(t),
Err(e) => {
ctx.set_error(builder.len(), e);
builder.push(0);
},
}
}
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
}),
);
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
Expand All @@ -949,13 +998,23 @@ macro_rules! impl_register_arith_functions {
|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|ts, delta, builder, ctx| {
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) {
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) {
Ok(t) => builder.push(t),
Err(e) => {
ctx.set_error(builder.len(), e);
builder.push(0);
},
}
}
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}

},
),
);
Expand All @@ -964,17 +1023,35 @@ macro_rules! impl_register_arith_functions {
concat!($op, "_days"),

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, _| {
builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta}))
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {

let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta_i32}));
}
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
}),
);
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
concat!($op, "_days"),

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|ts, delta, builder, _| {
builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta}))
|ts, delta, builder, ctx| {
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta_i32})); }
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
},
),
);
Expand All @@ -983,19 +1060,35 @@ macro_rules! impl_register_arith_functions {
concat!($op, "_weeks"),

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, _| {
vectorize_with_builder_2_arg::<DateType, Int64Type, DateType>(|date, delta, builder, ctx| {
let delta = 7 * delta;
builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta}))
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta_i32})) }
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
}),
);
registry.register_passthrough_nullable_2_arg::<TimestampType, Int64Type, TimestampType, _, _>(
concat!($op, "_weeks"),

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|ts, delta, builder, _| {
|ts, delta, builder, ctx| {
let delta = 7 * delta;
builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta}))
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta_i32})) }
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
},
),
);
Expand All @@ -1005,13 +1098,21 @@ macro_rules! impl_register_arith_functions {

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<DateType, Int64Type, TimestampType>(
|ts, delta, builder, _| {
let val = (ts as i64) * 24 * 3600 * MICROS_PER_SEC;
|ts, delta, builder, ctx| {
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
let val = (ts as i64) * 24 * 3600 * MICROS_PER_SEC;
builder.push(EvalTimesImpl::eval_timestamp(
val,
$signed_wrapper!{delta},
$signed_wrapper!{delta_i32},
FACTOR_HOUR,
));
)); }
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
},
),
);
Expand All @@ -1020,12 +1121,21 @@ macro_rules! impl_register_arith_functions {

|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<TimestampType, Int64Type, TimestampType>(
|ts, delta, builder, _| {
builder.push(EvalTimesImpl::eval_timestamp(
|ts, delta, builder, ctx| {
let delta_i32: Result<i32, _> = delta.try_into();
match delta_i32 {
Ok(delta_i32) => {
builder.push(EvalTimesImpl::eval_timestamp(
ts,
$signed_wrapper!{delta},
$signed_wrapper!{delta_i32},
FACTOR_HOUR,
));
));
}
Err(_) => {
ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta));
builder.push(0);
}
}
},
),
);
Expand Down
16 changes: 16 additions & 0 deletions tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1466,3 +1466,19 @@ query T
select TRY_TO_TIMESTAMP(1, 0), TRY_TO_TIMESTAMP(1, null);
----
1970-01-01 00:00:01.000000 NULL

statement error 1006
SELECT add_hours(to_date(710455), 1512263452497496403);

statement error 1006
SELECT add_hours(to_timestamp(710455), 1512263452497496403);

query T
SELECT add_hours(to_date(710455), 2147483647);
----
1000-01-01 00:00:00.000000

query T
SELECT add_hours(to_timestamp(710455), 2147483647);
----
1000-01-01 00:00:00.000000

0 comments on commit 33805b5

Please sign in to comment.