Skip to main content

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 ClockSetProcedure {
21	pub fn new() -> Self {
22		Self
23	}
24}
25
26impl Procedure for ClockSetProcedure {
27	fn call(&self, ctx: &ProcedureContext, _tx: &mut Transaction<'_>) -> Result<Columns, ProcedureError> {
28		let arg = match ctx.params {
29			Params::Positional(args) if args.len() == 1 => &args[0],
30			Params::Positional(args) => {
31				return Err(ProcedureError::ArityMismatch {
32					procedure: Fragment::internal("clock::set"),
33					expected: 1,
34					actual: args.len(),
35				});
36			}
37			_ => {
38				return Err(ProcedureError::ArityMismatch {
39					procedure: Fragment::internal("clock::set"),
40					expected: 1,
41					actual: 0,
42				});
43			}
44		};
45
46		match &ctx.runtime_context.clock {
47			Clock::Mock(mock) => {
48				match arg {
49					Value::DateTime(dt) => {
50						mock.set_nanos(dt.to_nanos_since_epoch_u128());
51					}
52					Value::Duration(dur) => {
53						let epoch = DateTime::default(); // 1970-01-01T00:00:00Z
54						let target = epoch.add_duration(dur)?;
55						mock.set_nanos(target.to_nanos_since_epoch_u128());
56					}
57					other => {
58						let millis = extract_millis(other).ok_or_else(|| {
59							ProcedureError::InvalidArgumentType {
60								procedure: Fragment::internal("clock::set"),
61								argument_index: 0,
62								expected: EXPECTED_SET_TYPES.to_vec(),
63								actual: other.get_type(),
64							}
65						})?;
66						mock.set_millis(millis);
67					}
68				}
69				let current_nanos = mock.now_nanos();
70				let dt = DateTime::from_timestamp_nanos(current_nanos)?;
71				Ok(Columns::single_row([("clock", Value::DateTime(dt))]))
72			}
73			Clock::Real => Err(ProcedureError::ExecutionFailed {
74				procedure: Fragment::internal("clock::set"),
75				reason: "clock::set can only be used with a mock clock".to_string(),
76			}),
77		}
78	}
79}
80
81const EXPECTED_SET_TYPES: &[Type] = &[
82	Type::DateTime,
83	Type::Duration,
84	Type::Int1,
85	Type::Int2,
86	Type::Int4,
87	Type::Int8,
88	Type::Int16,
89	Type::Uint1,
90	Type::Uint2,
91	Type::Uint4,
92	Type::Uint8,
93	Type::Uint16,
94];
95
96pub fn extract_millis(value: &Value) -> Option<u64> {
97	match value {
98		Value::Int1(v) => Some(*v as u64),
99		Value::Int2(v) => Some(*v as u64),
100		Value::Int4(v) => Some(*v as u64),
101		Value::Int8(v) => Some(*v as u64),
102		Value::Int16(v) => Some(*v as u64),
103		Value::Uint1(v) => Some(*v as u64),
104		Value::Uint2(v) => Some(*v as u64),
105		Value::Uint4(v) => Some(*v as u64),
106		Value::Uint8(v) => Some(*v),
107		Value::Uint16(v) => Some(*v as u64),
108		_ => None,
109	}
110}