more datetime functions

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"

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 <https://www.rfc-editor.org/rfc/rfc3339>`_.
+
+    If a second string argument is provided, it is interpreted as a `timezone <https://en.wikipedia.org/wiki/Tz_database>`_ and used to format the timestamp.

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,
     })
 }

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(_))
+            | (Bool(_), Bool(_))
-        | (Num(_), Num(_))
+            | (Num(_), Num(_))
-        | (Str(_), Str(_))
+            | (Str(_), Str(_))
-        | (Bytes(_), Bytes(_))
+            | (Bytes(_), Bytes(_))
-        | (Regex(_), Regex(_))
+            | (Regex(_), Regex(_))
-        | (List(_), List(_))
+            | (List(_), List(_))
-        | (Set(_), Set(_))
+            | (Set(_), Set(_))
-        | (Rev(_), Rev(_))
+            | (Rev(_), Rev(_))
-        | (Guard, Guard)
+            | (Guard, Guard)
-        | (Bot, Bot)) {
+            | (Bot, Bot)
-        bail!("comparison can only be done between the same datatypes, got {:?} and {:?}", a, b)
+    ) {
+        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<DataValue> {
 define_op!(OP_AND, 0, true);
 pub(crate) fn op_and(args: &[DataValue]) -> Result<DataValue> {
     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<DataValue> {
 define_op!(OP_OR, 0, true);
 pub(crate) fn op_or(args: &[DataValue]) -> Result<DataValue> {
     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<DataValue> {
     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::sin((lat1 - lat2) / 2.).powi(2)
-            + f64::cos(lat1) * f64::cos(lat2) * f64::sin((lon1 - lon2) / 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<DataValue> {
     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::sin((lat1 - lat2) / 2.).powi(2)
-            + f64::cos(lat1) * f64::cos(lat2) * f64::sin((lon1 - lon2) / 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<DataValue> {
             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<DataValue> {
+    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<DataValue> {
+    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<DataValue> {
 define_op!(OP_UUID_TIMESTAMP, 1, false);
 pub(crate) fn op_uuid_timestamp(args: &[DataValue]) -> Result<DataValue> {
     Ok(match &args[0] {
-        DataValue::Uuid(UuidWrapper(id)) => {
+        DataValue::Uuid(UuidWrapper(id)) => match id.get_timestamp() {
-            match id.get_timestamp() {
+            None => DataValue::Null,
-                None => DataValue::Null,
+            Some(t) => (t.to_unix_nanos() as f64 / 1_000_000_000.).into(),
-                Some(t) => (t.to_unix_nanos() as f64 / 1_000_000_000.).into()
+        },
-            }
+        _ => bail!("not an UUID"),
-        }
-        _ => bail!("not an UUID")
     })
-}
+}

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();
 }