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

reifydb_routine/procedure/clock/
set.rs
