feat: SELECT non-deterministic transaction-time rewrites

integration/rust/Cargo.toml

@@ -27,3 +27,4 @@ bytes.workspace = true
 rust_decimal = { version = "1.42.0", features = ["macros"] }
 chrono-tz = "0.10.4"
 pgdog-stats = { path = "../../pgdog-stats" }
+itertools = "0.15"

integration/rust/tests/integration/non_deterministic_funcs.rs

@@ -31,7 +31,7 @@ use sqlx::{Executor, Row};
 /// This tests a case where we're performing two INSERTs which resolve to different Shards. This uses 2 connections.
 /// Previously, the now() values would be different in each.
 #[tokio::test]
-async fn two_conns_transaction_time_reuse() {
+async fn two_conns_transaction_time_reuse_insert() {
     let single_sharded_list_pool = PgPoolOptions::new()
         .max_connections(1)
         .connect("postgres://pgdog:pgdog@127.0.0.1:6432/single_sharded_list?application_name=sqlx")
@@ -65,6 +65,109 @@ async fn two_conns_transaction_time_reuse() {
     transaction.rollback().await.unwrap();
 }
 
+/// Ensure that SELECT now() and CURRENT_TIMESTAMP return the correct type, are consistent
+/// within a transaction across multiple queries, and return within the correct column name
+/// despite being re-written.
+#[tokio::test]
+async fn transaction_time_select_equality() {
+    run_and_assert_eq::<DateTime<Utc>>("SELECT now();", "now").await;
+    run_and_assert_eq::<DateTime<Utc>>("SELECT CURRENT_TIMESTAMP;", "current_timestamp").await;
+}
+
+/// Helper method to run a query twice, in one transaction, in both simple and extended protocol,
+/// and assert that the response is equal (and of the correct type)
+async fn run_and_assert_eq<T>(query: &str, col: &str)
+where
+    T: PartialEq
+        + std::fmt::Debug
+        + sqlx::Type<sqlx::Postgres>
+        + for<'a> sqlx::Decode<'a, sqlx::Postgres>,
+{
+    let conn = connections_sqlx().await;
+    let conn = conn.get(1).unwrap();
+    let mut transaction = conn.begin().await.unwrap();
+
+    for simple_protocol in [false, true] {
+        let mut rows = vec![];
+        for _ in 0..2 {
+            rows.push(if simple_protocol {
+                sqlx::raw_sql(query)
+                    .fetch_one(&mut *transaction)
+                    .await
+                    .unwrap()
+            } else {
+                sqlx::query(query)
+                    .fetch_one(&mut *transaction)
+                    .await
+                    .unwrap()
+            });
+        }
+
+        assert_eq!(rows[0].get::<T, &str>(col), rows[1].get::<T, &str>(col));
+    }
+}
+
+/// Ensure that now() functionality isn't broken for use within a query
+/// Tests things like comparisons and addition with an interval in a SELECT query
+#[tokio::test]
+async fn transaction_time_usage_in_select_query() {
+    let conn = connections_sqlx().await;
+    let conn = conn.get(1).unwrap();
+
+    let mut transaction = conn.begin().await.unwrap();
+
+    // 'created_at' col has default now()
+    sqlx::query("INSERT INTO sharded (id) VALUES (1)")
+        .execute(&mut *transaction)
+        .await
+        .unwrap();
+
+    for use_simple_protocol in [false, true] {
+        // Test comparisons
+        let query_tries = [
+            ("SELECT * FROM sharded WHERE created_at = now()", 1),
+            ("SELECT * FROM sharded WHERE created_at < now()", 0),
+            ("SELECT * FROM sharded WHERE created_at > now()", 0),
+        ];
+        for (query, expected_rows) in query_tries {
+            let rows_fetched = if use_simple_protocol {
+                sqlx::raw_sql(query)
+                    .fetch_all(&mut *transaction)
+                    .await
+                    .unwrap()
+            } else {
+                sqlx::query(query)
+                    .fetch_all(&mut *transaction)
+                    .await
+                    .unwrap()
+            };
+            assert_eq!(rows_fetched.len(), expected_rows);
+        }
+
+        // Test addition
+        let query = "SELECT now() + INTERVAL '1 day' AS this_time_tomorrow;";
+
+        let time_plus_1_day: DateTime<Utc> = if use_simple_protocol {
+            let row: PgRow = sqlx::raw_sql(query)
+                .fetch_one(&mut *transaction)
+                .await
+                .unwrap();
+            row.get::<_, &str>("this_time_tomorrow")
+        } else {
+            sqlx::query_scalar(query)
+                .fetch_one(&mut *transaction)
+                .await
+                .unwrap()
+        };
+
+        let now: DateTime<Utc> = Utc::now();
+        assert!(
+            time_plus_1_day > now + Duration::hours(23)
+                && time_plus_1_day < now + Duration::hours(25)
+        );
+    }
+}
+
 /// Test that an INSERT into an omnisharded table which uses UUID functions is
 /// re-written to a constant to be consistent across all shards.
 ///

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

@@ -193,11 +193,15 @@ impl<'a> StatementRewrite<'a> {
 
         if nd_function_rewrite {
             match stmt.stmt_mut() {
-                NodeMut::InsertStmt(_) => {
+                // TODO: we could also support UPDATE / etc
+                NodeMut::InsertStmt(_) | NodeMut::SelectStmt(_) => {
                     self.rewrite_nd_functions(stmt.stmt_mut(), mem, &mut plan.bind_params)?;
                 }
                 NodeMut::PrepareStmt(mut prepare) => {
-                    if matches!(prepare.query_mut(), NodeMut::InsertStmt(_)) {
+                    if matches!(
+                        prepare.query_mut(),
+                        NodeMut::InsertStmt(_) | NodeMut::SelectStmt(_)
+                    ) {
                         self.rewrite_nd_functions(prepare.query_mut(), mem, &mut plan.bind_params)?;
                     }
                 }