reifydb_routine/procedure/clock/
set.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 crate::routine::{Routine, RoutineInfo, context::ProcedureContext, error::RoutineError};
15
16static INFO: LazyLock<RoutineInfo> = LazyLock::new(|| RoutineInfo::new("clock::set"));
17
18/// Native procedure that sets the mock clock to a specific time.
19///
20/// Accepts 1 positional argument: a DateTime, Duration (since epoch), or integer milliseconds.
21pub struct ClockSetProcedure;
22
23impl Default for ClockSetProcedure {
24	fn default() -> Self {
25		Self::new()
26	}
27}
28
29impl ClockSetProcedure {
30	pub fn new() -> Self {
31		Self
32	}
33}
34
35impl<'a, 'tx> Routine<ProcedureContext<'a, 'tx>> for ClockSetProcedure {
36	fn info(&self) -> &RoutineInfo {
37		&INFO
38	}
39
40	fn return_type(&self, _input_types: &[Type]) -> Type {
41		Type::DateTime
42	}
43
44	fn execute(&self, ctx: &mut ProcedureContext<'a, 'tx>, _args: &Columns) -> Result<Columns, RoutineError> {
45		let arg = match ctx.params {
46			Params::Positional(args) if args.len() == 1 => &args[0],
47			Params::Positional(args) => {
48				return Err(RoutineError::ProcedureArityMismatch {
49					procedure: Fragment::internal("clock::set"),
50					expected: 1,
51					actual: args.len(),
52				});
53			}
54			_ => {
55				return Err(RoutineError::ProcedureArityMismatch {
56					procedure: Fragment::internal("clock::set"),
57					expected: 1,
58					actual: 0,
59				});
60			}
61		};
62
63		match &ctx.runtime_context.clock {
64			Clock::Mock(mock) => {
65				match arg {
66					Value::DateTime(dt) => {
67						mock.set_nanos(dt.to_nanos());
68					}
69					Value::Duration(dur) => {
70						let epoch = DateTime::default(); // 1970-01-01T00:00:00Z
71						let target = epoch.add_duration(dur)?;
72						mock.set_nanos(target.to_nanos());
73					}
74					other => {
75						let millis = extract_millis(other).ok_or_else(|| {
76							RoutineError::ProcedureInvalidArgumentType {
77								procedure: Fragment::internal("clock::set"),
78								argument_index: 0,
79								expected: EXPECTED_SET_TYPES.to_vec(),
80								actual: other.get_type(),
81							}
82						})?;
83						mock.set_millis(millis);
84					}
85				}
86				let current_nanos = mock.now_nanos();
87				let dt = DateTime::from_nanos(current_nanos);
88				Ok(Columns::single_row([("clock", Value::DateTime(dt))]))
89			}
90			Clock::Real => Err(RoutineError::ProcedureExecutionFailed {
91				procedure: Fragment::internal("clock::set"),
92				reason: "clock::set can only be used with a mock clock".to_string(),
93			}),
94		}
95	}
96}
97
98const EXPECTED_SET_TYPES: &[Type] = &[
99	Type::DateTime,
100	Type::Duration,
101	Type::Int1,
102	Type::Int2,
103	Type::Int4,
104	Type::Int8,
105	Type::Int16,
106	Type::Uint1,
107	Type::Uint2,
108	Type::Uint4,
109	Type::Uint8,
110	Type::Uint16,
111];
112
113pub fn extract_millis(value: &Value) -> Option<u64> {
114	match value {
115		Value::Int1(v) => Some(*v as u64),
116		Value::Int2(v) => Some(*v as u64),
117		Value::Int4(v) => Some(*v as u64),
118		Value::Int8(v) => Some(*v as u64),
119		Value::Int16(v) => Some(*v as u64),
120		Value::Uint1(v) => Some(*v as u64),
121		Value::Uint2(v) => Some(*v as u64),
122		Value::Uint4(v) => Some(*v as u64),
123		Value::Uint8(v) => Some(*v),
124		Value::Uint16(v) => Some(*v as u64),
125		_ => None,
126	}
127}
reifydb_routine/procedure/clock/set.rs

reifydb_routine/procedure/clock/
set.rs
