feat: re-write uuid functions (calls + DEFAULT) in omni insert

Cargo.toml

@@ -31,7 +31,7 @@ indexmap = { version = "2.14", features = ["serde"] }
 serde_with = { version = "3.22", features = ["macros", "schemars_1"] }
 reqwest = { version = "0.13", features = ["json"] }
 bytes = "1"
-uuid = { version = "1", features = ["v4", "serde"] }
+uuid = { version = "1", features = ["v4", "v7", "serde"] }
 parking_lot = "0.12"
 futures-util = "0.3.32"
 

integration/rust/Cargo.toml

@@ -10,7 +10,7 @@ test = true
 tokio-postgres = {version = "0.7.13", features = ["with-uuid-1"]}
 postgres-native-tls = "0.5"
 native-tls = "0.2"
-sqlx = { version = "0.8.6", features = ["postgres", "runtime-tokio", "tls-native-tls", "bigdecimal", "chrono", "json", "rust_decimal"]}
+sqlx = { version = "0.8.6", features = ["postgres", "runtime-tokio", "tls-native-tls", "bigdecimal", "chrono", "json", "rust_decimal", "uuid"]}
 tokio = { version = "1", features = ["full"]}
 futures-util.workspace = true
 uuid.workspace = true

integration/rust/tests/integration/mod.rs

@@ -24,7 +24,7 @@ pub mod max;
 pub mod multi_set;
 pub mod notify;
 pub mod offset;
-pub mod omni_timestamps;
+pub mod omni_non_deterministic_funcs;
 pub mod partial_req;
 pub mod per_stmt_routing;
 pub mod prepared;

integration/rust/tests/integration/omni_timestamps.rs → integration/rust/tests/integration/omni_non_deterministic_funcs.rs

@@ -24,6 +24,96 @@ use sqlx::{Executor, Row};
 // TODO: Test to make sure this doesn't affect harded tables (it doesn't; but doesn't hurt to assert that)
 // TODO: Assert what happens if we don't explicitly set timezone
 
+/// Test that an INSERT into an omnisharded table which uses UUID functions is
+/// re-written to a constant to be consistent across all shards.
+///
+/// - Verifies DEFAULT works (in schema, in VALUES)
+/// - Verifies functions in VALUES work.
+/// - Tests uuidv4(), uuidv7(), gen_random_uuid()
+///
+/// It also asserts that uuidv7(interval) does **NOT** work right now, as we don't have an easy way to parse intervals.
+#[tokio::test]
+async fn omni_uuid_rewrite() {
+    let sharded_conn = connections_sqlx().await;
+    let sharded_conn = sharded_conn.get(1).unwrap();
+
+    // TODO: Check that UUID v7 shift interval is parsed correctly (and works without errors)
+
+    // Create a test table through PgDog, and then reload, so that the Schema is loaded.
+    {
+        sharded_conn
+            .execute("DROP TABLE IF EXISTS test_omni_uuid")
+            .await
+            .unwrap();
+        sharded_conn
+            .execute(
+                "CREATE TABLE IF NOT EXISTS test_omni_uuid(
+            id BIGSERIAL PRIMARY KEY,
+            uuid4 uuid,
+            uuid7 uuid,
+            uuid4_default uuid DEFAULT gen_random_uuid(),
+            uuid7_default uuid DEFAULT uuidv7(),
+            uuid7_default_explicit uuid DEFAULT uuidv7())",
+            )
+            .await
+            .unwrap();
+
+        admin_sqlx().await.execute("RELOAD").await.unwrap();
+    }
+
+    // Other functions below test that general rewrites work in all protocols.
+    // It would be redundant to test that here, as they all share the same re-usable structure.
+    {
+        let mut transaction = sharded_conn.begin().await.unwrap();
+
+        transaction
+            .execute(
+                "INSERT INTO test_omni_uuid(id, uuid4, uuid7, uuid7_default_explicit)
+            VALUES(1, uuidv4(), uuidv7(), DEFAULT)",
+            )
+            .await
+            .unwrap();
+
+        let (shard_0_row, shard_1_row) = (
+            transaction
+                .fetch_one("/* pgdog_shard: 0 */ SELECT * FROM test_omni_uuid")
+                .await
+                .unwrap(),
+            transaction
+                .fetch_one("/* pgdog_shard: 1 */ SELECT * FROM test_omni_uuid")
+                .await
+                .unwrap(),
+        );
+
+        assert_eq!(shard_0_row.columns().len(), shard_1_row.columns().len());
+
+        // Iterate through all columns ensuring the UUIDs are consistent across shards for each case.
+        for col_num in 1..shard_0_row.columns().len() {
+            let (shard_0_uuid, shard_1_uuid) = (
+                shard_0_row.get::<sqlx::types::Uuid, usize>(col_num),
+                shard_1_row.get::<sqlx::types::Uuid, usize>(col_num),
+            );
+
+            assert_eq!(shard_0_uuid, shard_1_uuid);
+        }
+
+        // Specify an INTERVAL as an argument within uuidv7. This should fail with an Error.
+        // It would require us to parse Postgres intervals (possible, but not supported yet)
+        let err = transaction
+            .execute(
+                "INSERT INTO test_omni_uuid(id, uuid4, uuid7, uuid7_default_explicit)
+            VALUES(2, uuidv4(), uuidv7(INTERVAL '-2 weeks'), DEFAULT)",
+            )
+            .await
+            .err()
+            .unwrap();
+
+        assert!(err.to_string().contains("parser: rewrite: could not determine how to parse the argument passed in uuidv7; it is likely not supported yet"));
+
+        transaction.rollback().await.unwrap();
+    }
+}
+
 /// LOCAL_TIME testing
 /// - Case 1: `test_time_text` has no DEFAULT w/ precision arg & text col.
 /// - Case 2: `test_time_regular` has DEFAULT w/ no precision arg & time col.

pgdog/src/frontend/router/parser/rewrite/statement/error.rs

@@ -54,4 +54,9 @@ pub(crate) enum Error {
 
     #[error("could not determine the session TimeZone for time functions on omnisharded tables")]
     UnknownTimeZone,
+
+    #[error(
+        "could not determine how to parse the argument passed in {0}; it is likely not supported yet"
+    )]
+    UnsupportedArgument(String),
 }

pgdog/src/frontend/router/parser/rewrite/statement/mod.rs

@@ -14,10 +14,10 @@ pub(crate) mod auto_id;
 pub(crate) mod error;
 pub(crate) mod insert;
 pub(crate) mod nextval;
+pub(crate) mod non_deterministic_funcs;
 pub(crate) mod offset;
 pub(crate) mod plan;
 pub(crate) mod simple_prepared;
-pub(crate) mod timestamp;
 pub(crate) mod unique_id;
 pub(crate) mod update;