reifydb_routine/procedure/clock/
advance.rs1use std::sync::LazyLock;
5
6use reifydb_core::value::column::columns::Columns;
7use reifydb_runtime::context::clock::Clock;
8use reifydb_type::{
9 fragment::Fragment,
10 params::Params,
11 value::{Value, datetime::DateTime, r#type::Type},
12};
13
14use super::set::extract_millis;
15use crate::routine::{Routine, RoutineInfo, context::ProcedureContext, error::RoutineError};
16
17static INFO: LazyLock<RoutineInfo> = LazyLock::new(|| RoutineInfo::new("clock::advance"));
18
19pub struct ClockAdvanceProcedure;
20
21impl Default for ClockAdvanceProcedure {
22 fn default() -> Self {
23 Self::new()
24 }
25}
26
27impl ClockAdvanceProcedure {
28 pub fn new() -> Self {
29 Self
30 }
31}
32
33impl<'a, 'tx> Routine<ProcedureContext<'a, 'tx>> for ClockAdvanceProcedure {
34 fn info(&self) -> &RoutineInfo {
35 &INFO
36 }
37
38 fn return_type(&self, _input_types: &[Type]) -> Type {
39 Type::DateTime
40 }
41
42 fn execute(&self, ctx: &mut ProcedureContext<'a, 'tx>, _args: &Columns) -> Result<Columns, RoutineError> {
43 let arg = match ctx.params {
44 Params::Positional(args) if args.len() == 1 => &args[0],
45 Params::Positional(args) => {
46 return Err(RoutineError::ProcedureArityMismatch {
47 procedure: Fragment::internal("clock::advance"),
48 expected: 1,
49 actual: args.len(),
50 });
51 }
52 _ => {
53 return Err(RoutineError::ProcedureArityMismatch {
54 procedure: Fragment::internal("clock::advance"),
55 expected: 1,
56 actual: 0,
57 });
58 }
59 };
60
61 match &ctx.runtime_context.clock {
62 Clock::Mock(mock) => {
63 match arg {
64 Value::Duration(dur) => {
65 if dur.get_months() == 0 && dur.get_days() == 0 {
66 let nanos = dur.get_nanos();
67 if nanos >= 0 {
68 mock.advance_nanos(nanos as u64);
69 } else {
70 let current = mock.now_nanos();
71 let abs_nanos = nanos.unsigned_abs();
72 if abs_nanos > current {
73 return Err(RoutineError::ProcedureExecutionFailed {
74 procedure: Fragment::internal("clock::advance"),
75 reason: "clock cannot be set before Unix epoch".to_string(),
76 });
77 }
78 mock.set_nanos(current - abs_nanos);
79 }
80 } else {
81 let current_nanos = mock.now_nanos();
82 let current_dt = DateTime::from_nanos(current_nanos);
83 let new_dt = current_dt.add_duration(dur)?;
84 mock.set_nanos(new_dt.to_nanos());
85 }
86 }
87 other => {
88 let millis = extract_millis(other).ok_or_else(|| {
89 RoutineError::ProcedureInvalidArgumentType {
90 procedure: Fragment::internal("clock::advance"),
91 argument_index: 0,
92 expected: EXPECTED_ADVANCE_TYPES.to_vec(),
93 actual: other.get_type(),
94 }
95 })?;
96 mock.advance_millis(millis);
97 }
98 }
99 let current_nanos = mock.now_nanos();
100 let dt = DateTime::from_nanos(current_nanos);
101 Ok(Columns::single_row([("clock", Value::DateTime(dt))]))
102 }
103 Clock::Real => Err(RoutineError::ProcedureExecutionFailed {
104 procedure: Fragment::internal("clock::advance"),
105 reason: "clock::advance can only be used with a mock clock".to_string(),
106 }),
107 }
108 }
109}
110
111const EXPECTED_ADVANCE_TYPES: &[Type] = &[
112 Type::Duration,
113 Type::Int1,
114 Type::Int2,
115 Type::Int4,
116 Type::Int8,
117 Type::Int16,
118 Type::Uint1,
119 Type::Uint2,
120 Type::Uint4,
121 Type::Uint8,
122 Type::Uint16,
123];