reifydb_routine/procedure/clock/
set.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::value::column::columns::Columns;
5use reifydb_runtime::context::clock::Clock;
6use reifydb_transaction::transaction::Transaction;
7use reifydb_type::{
8	fragment::Fragment,
9	params::Params,
10	value::{Value, datetime::DateTime, r#type::Type},
11};
12
13use crate::procedure::{Procedure, context::ProcedureContext, error::ProcedureError};
14
15/// Native procedure that sets the mock clock to a specific time.
16///
17/// Accepts 1 positional argument: a DateTime, Duration (since epoch), or integer milliseconds.
18pub struct ClockSetProcedure;
19
20impl Default for ClockSetProcedure {
21	fn default() -> Self {
22		Self::new()
23	}
24}
25
26impl ClockSetProcedure {
27	pub fn new() -> Self {
28		Self
29	}
30}
31
32impl Procedure for ClockSetProcedure {
33	fn call(&self, ctx: &ProcedureContext, _tx: &mut Transaction<'_>) -> Result<Columns, ProcedureError> {
34		let arg = match ctx.params {
35			Params::Positional(args) if args.len() == 1 => &args[0],
36			Params::Positional(args) => {
37				return Err(ProcedureError::ArityMismatch {
38					procedure: Fragment::internal("clock::set"),
39					expected: 1,
40					actual: args.len(),
41				});
42			}
43			_ => {
44				return Err(ProcedureError::ArityMismatch {
45					procedure: Fragment::internal("clock::set"),
46					expected: 1,
47					actual: 0,
48				});
49			}
50		};
51
52		match &ctx.runtime_context.clock {
53			Clock::Mock(mock) => {
54				match arg {
55					Value::DateTime(dt) => {
56						mock.set_nanos(dt.to_nanos());
57					}
58					Value::Duration(dur) => {
59						let epoch = DateTime::default(); // 1970-01-01T00:00:00Z
60						let target = epoch.add_duration(dur)?;
61						mock.set_nanos(target.to_nanos());
62					}
63					other => {
64						let millis = extract_millis(other).ok_or_else(|| {
65							ProcedureError::InvalidArgumentType {
66								procedure: Fragment::internal("clock::set"),
67								argument_index: 0,
68								expected: EXPECTED_SET_TYPES.to_vec(),
69								actual: other.get_type(),
70							}
71						})?;
72						mock.set_millis(millis);
73					}
74				}
75				let current_nanos = mock.now_nanos();
76				let dt = DateTime::from_nanos(current_nanos);
77				Ok(Columns::single_row([("clock", Value::DateTime(dt))]))
78			}
79			Clock::Real => Err(ProcedureError::ExecutionFailed {
80				procedure: Fragment::internal("clock::set"),
81				reason: "clock::set can only be used with a mock clock".to_string(),
82			}),
83		}
84	}
85}
86
87const EXPECTED_SET_TYPES: &[Type] = &[
88	Type::DateTime,
89	Type::Duration,
90	Type::Int1,
91	Type::Int2,
92	Type::Int4,
93	Type::Int8,
94	Type::Int16,
95	Type::Uint1,
96	Type::Uint2,
97	Type::Uint4,
98	Type::Uint8,
99	Type::Uint16,
100];
101
102pub fn extract_millis(value: &Value) -> Option<u64> {
103	match value {
104		Value::Int1(v) => Some(*v as u64),
105		Value::Int2(v) => Some(*v as u64),
106		Value::Int4(v) => Some(*v as u64),
107		Value::Int8(v) => Some(*v as u64),
108		Value::Int16(v) => Some(*v as u64),
109		Value::Uint1(v) => Some(*v as u64),
110		Value::Uint2(v) => Some(*v as u64),
111		Value::Uint4(v) => Some(*v as u64),
112		Value::Uint8(v) => Some(*v),
113		Value::Uint16(v) => Some(*v as u64),
114		_ => None,
115	}
116}
reifydb_routine/procedure/clock/set.rs

reifydb_routine/procedure/clock/
set.rs
