diff --git a/Cargo.toml b/Cargo.toml index 0d7a3717..a8a3704a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,7 @@ rmp-serde = "1.1.0" rmpv = "1.0.0" base64 = "0.13.0" chrono = "0.4.19" +chrono-tz = "0.6.3" priority-queue = "1.2.3" ordered-float = "3.0.0" num-traits = "0.2.15" diff --git a/docs/source/functions.rst b/docs/source/functions.rst index 334872e7..18618d73 100644 --- a/docs/source/functions.rst +++ b/docs/source/functions.rst @@ -624,3 +624,18 @@ Empty matches:: \z only the end of the text \b a Unicode word boundary (\w on one side and \W, \A, or \z on the other) \B not a Unicode word boundary + + +-------------------- +Misc functions +-------------------- + +.. function:: now() + + Returns the current timestamp as seconds since the UNIX epoch. + +.. function:: format_timestamp(ts, tz?) + + Interpret ``ts`` as seconds since the epoch and format as a string according to `RFC3339 `_. + + If a second string argument is provided, it is interpreted as a `timezone `_ and used to format the timestamp. \ No newline at end of file diff --git a/src/data/expr.rs b/src/data/expr.rs index 4f48fca4..ef31a870 100644 --- a/src/data/expr.rs +++ b/src/data/expr.rs @@ -651,6 +651,8 @@ pub(crate) fn get_op(name: &str) -> Option<&'static Op> { "rand_uuid_v1" => &OP_RAND_UUID_V1, "rand_uuid_v4" => &OP_RAND_UUID_V4, "uuid_timestamp" => &OP_UUID_TIMESTAMP, + "now" => &OP_NOW, + "format_timestamp" => &OP_FORMAT_TIMESTAMP, _ => return None, }) } diff --git a/src/data/functions.rs b/src/data/functions.rs index 8c82b07c..01c21a01 100644 --- a/src/data/functions.rs +++ b/src/data/functions.rs @@ -3,6 +3,7 @@ use std::ops::{Div, Rem}; use std::str::FromStr; use std::time::{SystemTime, UNIX_EPOCH}; +use chrono::{TimeZone, Utc}; use itertools::Itertools; use miette::{bail, ensure, miette, Result}; use num_traits::FloatConst; @@ -26,22 +27,27 @@ macro_rules! define_op { }; } - fn ensure_same_value_type(a: &DataValue, b: &DataValue) -> Result<()> { use DataValue::*; - if !matches!((a, b), + if !matches!( + (a, b), (Null, Null) - | (Bool(_), Bool(_)) - | (Num(_), Num(_)) - | (Str(_), Str(_)) - | (Bytes(_), Bytes(_)) - | (Regex(_), Regex(_)) - | (List(_), List(_)) - | (Set(_), Set(_)) - | (Rev(_), Rev(_)) - | (Guard, Guard) - | (Bot, Bot)) { - bail!("comparison can only be done between the same datatypes, got {:?} and {:?}", a, b) + | (Bool(_), Bool(_)) + | (Num(_), Num(_)) + | (Str(_), Str(_)) + | (Bytes(_), Bytes(_)) + | (Regex(_), Regex(_)) + | (List(_), List(_)) + | (Set(_), Set(_)) + | (Rev(_), Rev(_)) + | (Guard, Guard) + | (Bot, Bot) + ) { + bail!( + "comparison can only be done between the same datatypes, got {:?} and {:?}", + a, + b + ) } Ok(()) } @@ -504,7 +510,10 @@ pub(crate) fn op_mod(args: &[DataValue]) -> Result { define_op!(OP_AND, 0, true); pub(crate) fn op_and(args: &[DataValue]) -> Result { for arg in args { - if !arg.get_bool().ok_or_else(|| miette!("'and' requires booleans"))? { + if !arg + .get_bool() + .ok_or_else(|| miette!("'and' requires booleans"))? + { return Ok(DataValue::Bool(false)); } } @@ -514,7 +523,10 @@ pub(crate) fn op_and(args: &[DataValue]) -> Result { define_op!(OP_OR, 0, true); pub(crate) fn op_or(args: &[DataValue]) -> Result { for arg in args { - if arg.get_bool().ok_or_else(|| miette!("'or' requires booleans"))? { + if arg + .get_bool() + .ok_or_else(|| miette!("'or' requires booleans"))? + { return Ok(DataValue::Bool(true)); } } @@ -990,9 +1002,9 @@ pub(crate) fn op_haversine(args: &[DataValue]) -> Result { let lon2 = args[3].get_float().ok_or_else(miette)?; let ret = 2. * f64::asin(f64::sqrt( - f64::sin((lat1 - lat2) / 2.).powi(2) - + f64::cos(lat1) * f64::cos(lat2) * f64::sin((lon1 - lon2) / 2.).powi(2), - )); + f64::sin((lat1 - lat2) / 2.).powi(2) + + f64::cos(lat1) * f64::cos(lat2) * f64::sin((lon1 - lon2) / 2.).powi(2), + )); Ok(DataValue::from(ret)) } @@ -1005,9 +1017,9 @@ pub(crate) fn op_haversine_deg_input(args: &[DataValue]) -> Result { let lon2 = args[3].get_float().ok_or_else(miette)? * f64::PI() / 180.; let ret = 2. * f64::asin(f64::sqrt( - f64::sin((lat1 - lat2) / 2.).powi(2) - + f64::cos(lat1) * f64::cos(lat2) * f64::sin((lon1 - lon2) / 2.).powi(2), - )); + f64::sin((lat1 - lat2) / 2.).powi(2) + + f64::cos(lat1) * f64::cos(lat2) * f64::sin((lon1 - lon2) / 2.).powi(2), + )); Ok(DataValue::from(ret)) } @@ -1381,7 +1393,39 @@ pub(crate) fn op_to_uuid(args: &[DataValue]) -> Result { let id = uuid::Uuid::try_parse(s).map_err(|_| miette!("invalid UUID"))?; Ok(DataValue::uuid(id)) } - _ => bail!("'to_uuid' requires a string") + _ => bail!("'to_uuid' requires a string"), + } +} + +define_op!(OP_NOW, 0, false); +pub(crate) fn op_now(_args: &[DataValue]) -> Result { + let now = SystemTime::now(); + Ok(DataValue::from( + now.duration_since(UNIX_EPOCH).unwrap().as_secs_f64(), + )) +} + +define_op!(OP_FORMAT_TIMESTAMP, 1, true); +pub(crate) fn op_format_timestamp(args: &[DataValue]) -> Result { + let f = args[0] + .get_float() + .ok_or_else(|| miette!("'format_timestamp' expects a number"))?; + let millis = (f * 1000.) as i64; + let dt = Utc.timestamp_millis(millis); + match args.get(1) { + Some(tz_v) => { + let tz_s = tz_v.get_string().ok_or_else(|| { + miette!("'format_timestamp' timezone specification requires a string") + })?; + let tz = chrono_tz::Tz::from_str(tz_s).map_err(|_| miette!("bad timezone specification"))?; + let dt_tz = dt.with_timezone(&tz); + let s = SmartString::from(dt_tz.to_rfc3339()); + Ok(DataValue::Str(s)) + } + None => { + let s = SmartString::from(dt.to_rfc3339()); + Ok(DataValue::Str(s)) + } } } @@ -1407,12 +1451,10 @@ pub(crate) fn op_rand_uuid_v4(_args: &[DataValue]) -> Result { define_op!(OP_UUID_TIMESTAMP, 1, false); pub(crate) fn op_uuid_timestamp(args: &[DataValue]) -> Result { Ok(match &args[0] { - DataValue::Uuid(UuidWrapper(id)) => { - match id.get_timestamp() { - None => DataValue::Null, - Some(t) => (t.to_unix_nanos() as f64 / 1_000_000_000.).into() - } - } - _ => bail!("not an UUID") + DataValue::Uuid(UuidWrapper(id)) => match id.get_timestamp() { + None => DataValue::Null, + Some(t) => (t.to_unix_nanos() as f64 / 1_000_000_000.).into(), + }, + _ => bail!("not an UUID"), }) -} \ No newline at end of file +} diff --git a/src/data/tests/functions.rs b/src/data/tests/functions.rs index 33ed3627..59224fb6 100644 --- a/src/data/tests/functions.rs +++ b/src/data/tests/functions.rs @@ -1322,4 +1322,11 @@ fn test_uuid() { assert!(op_uuid_timestamp(&[v1]).unwrap().get_float().is_some()); assert!(op_to_uuid(&[DataValue::Str(SmartString::from(""))]).is_err()); assert!(op_to_uuid(&[DataValue::Str(SmartString::from("f3b4958c-52a1-11e7-802a-010203040506"))]).is_ok()); +} + +#[test] +fn test_now() { + let now = op_now(&[]).unwrap(); + assert!(matches!(now, DataValue::Num(_))); + op_format_timestamp(&[now]).unwrap(); } \ No newline at end of file