feat(pushid): add timestamp extraction & parsing

src/ksuid.rs

@@ -41,7 +41,7 @@ fn idkit_ksuid_extract_timestamptz(val: String) -> pgrx::TimestampWithTimeZone {
 #[pg_schema]
 mod tests {
     use chrono::{DateTime, Utc};
-    use pgrx::*;
+    use pgrx::pg_test;
 
     use crate::ksuid::idkit_ksuid_extract_timestamptz;
     use crate::ksuid::idkit_ksuid_generate;

src/lib.rs

@@ -10,6 +10,7 @@ mod timeflake;
 mod ulid;
 mod uuid_v6;
 mod uuid_v7;
+mod vendor;
 mod xid;
 
 use pgrx::pg_module_magic;

src/pushid.rs

@@ -1,6 +1,10 @@
+use chrono::NaiveDateTime;
 use pgrx::*;
-use pushid::PushId;
-use pushid::PushIdGen;
+
+use crate::common::naive_datetime_to_pg_timestamptz;
+use crate::common::OrPgrxError;
+use crate::vendor::pushid::PushId;
+use crate::vendor::pushid::PushIdGen;
 
 /// Generate a random PushID UUID
 #[pg_extern]
@@ -14,19 +18,59 @@ fn idkit_pushid_generate_text() -> String {
     idkit_pushid_generate()
 }
 
+/// Retrieve a `timestamptz` (with millisecond precision) from a given textual KSUID
+///
+/// # Panics
+///
+/// This function panics (with a [`pgrx::error`]) when the timezone can't be created
+#[pg_extern]
+fn idkit_pushid_extract_timestamptz(val: String) -> pgrx::TimestampWithTimeZone {
+    let pushid =
+        PushId::from_str(val.as_ref()).or_pgrx_error(format!("[{val}] is an invalid PushID"));
+    naive_datetime_to_pg_timestamptz(
+        NaiveDateTime::from_timestamp_opt(pushid.timestamp_seconds(), 0)
+            .or_pgrx_error("failed to create timestamp from PushID [{val}]")
+            .and_utc(),
+        format!("failed to convert timestamp for PUSHID [{val}]"),
+    )
+}
+
 //////////
 // Test //
 //////////
 
 #[cfg(any(test, feature = "pg_test"))]
 #[pg_schema]
 mod tests {
-    use pgrx::*;
+    use chrono::{DateTime, Utc};
+    use pgrx::pg_test;
+
+    use crate::pushid::idkit_pushid_extract_timestamptz;
+    use crate::pushid::idkit_pushid_generate;
 
     /// Basic length test
     #[pg_test]
     fn test_pushid_len() {
         let generated = crate::pushid::idkit_pushid_generate();
         assert_eq!(generated.len(), 20);
     }
+
+    /// Ensure timestamps extracted from CUIDs are valid
+    #[pg_test]
+    fn test_pushid_extract_timestamptz() {
+        let timestamp = idkit_pushid_extract_timestamptz(idkit_pushid_generate());
+        let parsed: DateTime<Utc> = DateTime::parse_from_rfc3339(&timestamp.to_iso_string())
+            .expect("extracted timestamp as ISO string parsed to UTC DateTime")
+            .into();
+        assert!(
+            Utc::now().signed_duration_since(parsed).num_seconds() < 3,
+            "extracted, printed & re-parsed pushid timestamp is from recent past (within 3s)"
+        );
+    }
+
+    /// Ensure an existing, hardcoded timestamp works for extraction
+    #[pg_test]
+    fn test_pushid_extract_timestamptz_existing() {
+        idkit_pushid_extract_timestamptz("1srOrx2ZWZBpBUvZwXKQmoEYga2".into());
+    }
 }

src/vendor/mod.rs

@@ -0,0 +1 @@
+pub(crate) mod pushid;

src/vendor/pushid.rs

@@ -0,0 +1,129 @@
+extern crate rand;
+
+use rand::Rng;
+
+use std::time::{SystemTime, UNIX_EPOCH};
+
+const PUSH_CHARS: &str = "-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz";
+
+pub trait PushIdGen {
+    fn get_id(&mut self) -> String;
+}
+
+pub struct PushId {
+    /// Seconds since the UNIX epoch
+    last_time: u64,
+    previous_indices: [usize; 12],
+}
+
+impl PushId {
+    pub fn new() -> Self {
+        let random_indices = PushId::generate_random_indices();
+        PushId {
+            last_time: 0,
+            previous_indices: random_indices,
+        }
+    }
+
+    fn gen_random_indices(&self, is_duplicate_time: bool) -> [usize; 12] {
+        if is_duplicate_time {
+            // If the timestamp hasn't changed since last push, use the same random number, except incremented by 1.
+            let mut indices_copy = self.previous_indices.clone();
+
+            for x in (0..12).rev() {
+                if indices_copy[x] == 63 {
+                    indices_copy[x] = 0;
+                } else {
+                    indices_copy[x] = indices_copy[x] + 1;
+                    break;
+                }
+            }
+            indices_copy
+        } else {
+            PushId::generate_random_indices()
+        }
+    }
+
+    fn generate_random_indices() -> [usize; 12] {
+        let mut rng = rand::thread_rng();
+        let mut random_indices = [0; 12];
+        for i in 0..12 {
+            let n = rng.gen::<f64>() * 64 as f64;
+            random_indices[i] = n as usize;
+        }
+        random_indices
+    }
+
+    fn gen_time_based_prefix(now: u64, mut acc: [usize; 8], i: u8) -> [usize; 8] {
+        let index = (now % 64) as usize;
+        acc[i as usize] = index;
+
+        match now / 64 {
+            new_now if new_now > 0 => PushId::gen_time_based_prefix(new_now, acc, i - 1),
+            _ => acc, // We've reached the end of "time". Return the indices
+        }
+    }
+
+    fn indices_to_characters(indices: Vec<&usize>) -> String {
+        indices.iter().fold(String::from(""), |acc, &&x| {
+            acc + &PUSH_CHARS
+                .chars()
+                .nth(x)
+                .expect("Index out of range")
+                .to_string()
+        })
+    }
+
+    /// Retrieve the number of milliseconds since
+    fn get_now() -> u64 {
+        let start = SystemTime::now();
+        let since_the_epoch = start
+            .duration_since(UNIX_EPOCH)
+            .expect("Unexpected time seed, EPOCH is not in the past");
+        since_the_epoch.as_secs() * 1000 + since_the_epoch.subsec_nanos() as u64 / 1_000_000
+    }
+
+    /// Get the milliseconds since UNIX epoch for the PushId
+    pub fn last_time_millis(&self) -> u64 {
+        self.last_time.into()
+    }
+}
+
+impl PushIdGen for PushId {
+    fn get_id(&mut self) -> String {
+        let now = PushId::get_now();
+        let is_duplicate_time = now == self.last_time;
+        let prefix = PushId::gen_time_based_prefix(now, [0; 8], 7);
+        let suffix = PushId::gen_random_indices(self, is_duplicate_time);
+        self.previous_indices = suffix;
+        self.last_time = PushId::get_now();
+        let all = prefix.iter().chain(suffix.iter()).collect::<Vec<&usize>>();
+        PushId::indices_to_characters(all)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    use crate::{PushId, PushIdGen};
+
+    /// Ensure that timestamps work properly
+    #[test]
+    fn test_timestamp() {
+        let mut pushid = PushId::new();
+        let id = pushid.get_id();
+        assert!(!id.is_empty(), "generated pushid");
+
+        let now = SystemTime::now();
+        let millis_since = now
+            .duration_since(UNIX_EPOCH)
+            .expect("invalid epoch")
+            .as_millis();
+        let millis_since_pushid = pushid.last_time_millis() as u128;
+        assert!(
+            millis_since - millis_since_pushid < 10,
+            "retrieved pushid generation time was within 10ms from now()"
+        );
+    }
+}