Skip to main content

reifydb_routine/procedure/clock/
advance.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use 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];