reifydb_routine/procedure/clock/
set.rs1use 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
15pub 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(); 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}