Fix: Preserve decimal/float values in INSERT and UPDATE statements (#630)

This pull request improves support for decimal/float values in SQL
parsing and sharding logic, especially for `INSERT` and `UPDATE`
statements. It introduces a new `Float` variant to the `Value` and
`AssignmentValue` enums, ensures decimals are preserved (not coerced to
integers or null), and adds handling and tests for these cases.
Additionally, it prevents the use of float/decimal columns as sharding
keys, returning clear errors or routing to all shards when necessary.

**Decimal/Float Value Handling Improvements:**

* Added a `Float` variant to the `Value` and `AssignmentValue` enums,
ensuring decimal values are preserved as strings throughout parsing and
assignment (`pgdog/src/frontend/router/parser/value.rs`,
`pgdog/src/frontend/router/parser/rewrite/shard_key.rs`).
[[1]](diffhunk://#diff-bd5edaf0e41ddd3af4577a1e7e53af5953d62db4ef5e6531db38dccaf0c87b91R15)
[[2]](diffhunk://#diff-7f25c6defa22d112850e0f942520463df6cd1d80af656ec435abcdeea3cd4f99R11)
* Updated parsing logic to map float/decimal literals to the new `Float`
variant instead of attempting to parse as integers or returning null
(`pgdog/src/frontend/router/parser/value.rs`).
* Modified assignment and value handling to correctly propagate and
stringify float values in query rewriting and sharding logic
(`pgdog/src/frontend/client/query_engine/shard_key_rewrite.rs`,
`pgdog/src/frontend/router/parser/insert.rs`).
[[1]](diffhunk://#diff-fabbbd1dac63eca483df4ee327365faebbcf441cb474b442c97e7db89d4bf854R424)
[[2]](diffhunk://#diff-55ec6e26408c3763f93de5033481bdd1c9a202cfd351f48338a53d554c60a9d5R377)
[[3]](diffhunk://#diff-314df04f8984799d59642fa3ea4027c9536ee9783f7a68ff7b01ddae3c5b80b3R323)

**Sharding Logic and Error Handling:**

* Updated sharding logic to reject float/decimal columns as sharding
keys for split inserts, returning a clear error message, and to route
updates with float sharding keys to all shards
(`pgdog/src/frontend/router/parser/insert.rs`,
`pgdog/src/frontend/router/parser/query/update.rs`).
[[1]](diffhunk://#diff-55ec6e26408c3763f93de5033481bdd1c9a202cfd351f48338a53d554c60a9d5R247-R256)
[[2]](diffhunk://#diff-314df04f8984799d59642fa3ea4027c9536ee9783f7a68ff7b01ddae3c5b80b3R266-R270)

**Testing and Validation:**

* Added comprehensive tests to verify that decimal values are preserved
(not quoted unless originally quoted), that quoted decimals are treated
as strings, and that float sharding keys are rejected or handled safely
in both insert and update scenarios
(`pgdog/src/frontend/router/parser/insert.rs`,
`pgdog/src/frontend/router/parser/query/update.rs`).
[[1]](diffhunk://#diff-55ec6e26408c3763f93de5033481bdd1c9a202cfd351f48338a53d554c60a9d5R892-R1074)
[[2]](diffhunk://#diff-314df04f8984799d59642fa3ea4027c9536ee9783f7a68ff7b01ddae3c5b80b3R385-R470)

Author: phillipvanheerden

pgdog/src/frontend/client/query_engine/shard_key_rewrite.rs

@@ -421,6 +421,7 @@ fn apply_assignments(
 
         let new_value = match assignment.value() {
             AssignmentValue::Integer(value) => Some(value.to_string()),
+            AssignmentValue::Float(value) => Some(value.clone()),
             AssignmentValue::String(value) => Some(value.clone()),
             AssignmentValue::Boolean(value) => Some(value.to_string()),
             AssignmentValue::Null => None,

pgdog/src/frontend/router/parser/insert.rs

@@ -244,6 +244,16 @@ impl<'a> Insert<'a> {
                 .shards(schema.shards)
                 .build()?
                 .apply()?,
+            Value::Float(_) => {
+                return Err(Error::SplitInsertNotSupported {
+                    table: table.name.to_owned(),
+                    reason: format!(
+                        "row {} uses a float/decimal value for sharding column \"{}\", which is not supported as a sharding key",
+                        tuple_index + 1,
+                        key.table.column
+                    ),
+                })
+            }
             Value::Null => {
                 return Err(Error::SplitInsertNotSupported {
                     table: table.name.to_owned(),
@@ -364,6 +374,7 @@ impl<'a> Insert<'a> {
     ) -> Result<StdString, Error> {
         match value {
             Value::Integer(int) => Ok(int.to_string()),
+            Value::Float(float) => Ok((*float).to_string()),
             Value::String(string) => Ok(format_literal(string)),
             Value::Boolean(boolean) => Ok(if *boolean {
                 StdString::from("TRUE")
@@ -878,4 +889,187 @@ mod test {
             _ => panic!("not an insert"),
         }
     }
+
+    #[test]
+    fn split_insert_preserves_decimal_values() {
+        let query = parse(
+            "INSERT INTO transactions (txn_id, user_id, amount, status) VALUES \
+             (1001, 1, 50.00, 'completed'), \
+             (1002, 2, 20.00, 'failed'), \
+             (1003, 3, 25.75, 'pending')",
+        )
+        .unwrap();
+        let select = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap();
+        let schema = ShardingSchema {
+            shards: 2,
+            tables: ShardedTables::new(
+                vec![ShardedTable {
+                    name: Some("transactions".into()),
+                    column: "user_id".into(),
+                    ..Default::default()
+                }],
+                vec![],
+            ),
+            ..Default::default()
+        };
+
+        match &select.node {
+            Some(NodeEnum::InsertStmt(stmt)) => {
+                let insert = Insert::new(stmt);
+                let routing = insert
+                    .shard(&schema, None, true, RewriteMode::Rewrite)
+                    .unwrap();
+
+                match routing {
+                    InsertRouting::Split(plan) => {
+                        let rows = plan.rows();
+                        // Check that decimal values are preserved without quotes
+                        assert_eq!(rows[0].values()[2], "50.00");
+                        assert_eq!(rows[1].values()[2], "20.00");
+                        assert_eq!(rows[2].values()[2], "25.75");
+
+                        // Verify strings are quoted
+                        assert_eq!(rows[0].values()[3], "'completed'");
+                        assert_eq!(rows[1].values()[3], "'failed'");
+                        assert_eq!(rows[2].values()[3], "'pending'");
+                    }
+                    InsertRouting::Routed(_) => panic!("expected split plan"),
+                }
+            }
+            _ => panic!("not an insert"),
+        }
+    }
+
+    #[test]
+    fn split_insert_with_quoted_decimal_values() {
+        let query = parse(
+            "INSERT INTO transactions (txn_id, user_id, amount, status) VALUES \
+             (1001, 1, '50.00', 'completed'), \
+             (1002, 2, '20.00', 'failed'), \
+             (1003, 3, '25.75', 'pending')",
+        )
+        .unwrap();
+        let select = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap();
+        let schema = ShardingSchema {
+            shards: 2,
+            tables: ShardedTables::new(
+                vec![ShardedTable {
+                    name: Some("transactions".into()),
+                    column: "user_id".into(),
+                    ..Default::default()
+                }],
+                vec![],
+            ),
+            ..Default::default()
+        };
+
+        match &select.node {
+            Some(NodeEnum::InsertStmt(stmt)) => {
+                let insert = Insert::new(stmt);
+                let routing = insert
+                    .shard(&schema, None, true, RewriteMode::Rewrite)
+                    .unwrap();
+
+                match routing {
+                    InsertRouting::Split(plan) => {
+                        let rows = plan.rows();
+                        // Quoted decimals should be preserved as strings
+                        assert_eq!(rows[0].values()[2], "'50.00'");
+                        assert_eq!(rows[1].values()[2], "'20.00'");
+                        assert_eq!(rows[2].values()[2], "'25.75'");
+
+                        // Verify strings are quoted
+                        assert_eq!(rows[0].values()[3], "'completed'");
+                        assert_eq!(rows[1].values()[3], "'failed'");
+                        assert_eq!(rows[2].values()[3], "'pending'");
+                    }
+                    InsertRouting::Routed(_) => panic!("expected split plan"),
+                }
+            }
+            _ => panic!("not an insert"),
+        }
+    }
+
+    #[test]
+    fn debug_decimal_parsing() {
+        let query = parse(
+            "INSERT INTO transactions (txn_id, user_id, amount, status) VALUES \
+             (1001, 101, 50.00, 'completed')",
+        )
+        .unwrap();
+        let select = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap();
+
+        match &select.node {
+            Some(NodeEnum::InsertStmt(stmt)) => {
+                let insert = Insert::new(stmt);
+                let tuples = insert.tuples();
+                println!("Tuples: {:?}", tuples);
+                assert_eq!(tuples.len(), 1);
+                println!("Values: {:?}", tuples[0].values);
+            }
+            _ => panic!("not an insert"),
+        }
+    }
+
+    #[test]
+    fn reproduce_decimal_null_issue() {
+        let query = parse(
+            "INSERT INTO transactions (txn_id, user_id, amount, status) VALUES \
+             (1001, 101, 50.00, 'completed')",
+        )
+        .unwrap();
+        let select = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap();
+
+        match &select.node {
+            Some(NodeEnum::InsertStmt(stmt)) => {
+                let insert = Insert::new(stmt);
+                let tuples = insert.tuples();
+                println!("Values: {:?}", tuples[0].values);
+
+                // After fix, this should be Float("50.00"), not Null
+                assert_eq!(tuples[0].values[2], Value::Float("50.00"));
+            }
+            _ => panic!("not an insert"),
+        }
+    }
+
+    #[test]
+    fn split_insert_rejects_float_sharding_key() {
+        let query = parse(
+            "INSERT INTO transactions (txn_id, amount, status) VALUES \
+             (1001, 50.00, 'completed'), \
+             (1002, 20.00, 'failed')",
+        )
+        .unwrap();
+        let select = query.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap();
+        let schema = ShardingSchema {
+            shards: 2,
+            tables: ShardedTables::new(
+                vec![ShardedTable {
+                    name: Some("transactions".into()),
+                    column: "amount".into(),
+                    ..Default::default()
+                }],
+                vec![],
+            ),
+            ..Default::default()
+        };
+
+        match &select.node {
+            Some(NodeEnum::InsertStmt(stmt)) => {
+                let insert = Insert::new(stmt);
+                let result = insert.shard(&schema, None, true, RewriteMode::Rewrite);
+
+                match result {
+                    Err(Error::SplitInsertNotSupported { table, reason }) => {
+                        assert_eq!(table, "transactions");
+                        assert!(reason.contains("float/decimal"));
+                        assert!(reason.contains("not supported as a sharding key"));
+                    }
+                    other => panic!("expected error for float sharding key, got {:?}", other),
+                }
+            }
+            _ => panic!("not an insert"),
+        }
+    }
 }

pgdog/src/frontend/router/parser/query/update.rs

@@ -263,6 +263,11 @@ impl QueryParser {
                     .build()?;
                 Ok(context_builder.apply()?)
             }
+            AssignmentValue::Float(_) => {
+                // Floats are not supported as sharding keys
+                // Return Shard::All to route to all shards (safe but not optimal)
+                Ok(Shard::All)
+            }
             AssignmentValue::String(text) => {
                 let context_builder = ContextBuilder::new(sharded_table)
                     .data(text.as_str())
@@ -315,6 +320,7 @@ impl QueryParser {
             if let Ok(value) = Value::try_from(&val.node) {
                 return match value {
                     Value::Integer(i) => Ok(AssignmentValue::Integer(i)),
+                    Value::Float(f) => Ok(AssignmentValue::Float(f.to_owned())),
                     Value::String(s) => Ok(AssignmentValue::String(s.to_owned())),
                     Value::Boolean(b) => Ok(AssignmentValue::Boolean(b)),
                     Value::Null => Ok(AssignmentValue::Null),
@@ -376,6 +382,92 @@ mod tests {
         let column = QueryParser::res_target_column(target).expect("column");
         assert_eq!(column, "id");
     }
+
+    #[test]
+    fn update_preserves_decimal_values() {
+        let parsed = pgdog_plugin::pg_query::parse(
+            "UPDATE transactions SET amount = 50.00, status = 'completed' WHERE id = 1",
+        )
+        .expect("parse");
+
+        let stmt = parsed
+            .protobuf
+            .stmts
+            .first()
+            .and_then(|node| node.stmt.as_ref())
+            .and_then(|node| node.node.as_ref())
+            .expect("statement node");
+
+        let update = match stmt {
+            NodeEnum::UpdateStmt(update) => update,
+            _ => panic!("expected update stmt"),
+        };
+
+        // Check that we can extract assignment values including decimals
+        let mut found_decimal = false;
+        let mut found_string = false;
+
+        for target in &update.target_list {
+            if let Some(NodeEnum::ResTarget(res)) = &target.node {
+                if let Some(val) = &res.val {
+                    if let Ok(value) = Value::try_from(&val.node) {
+                        match value {
+                            Value::Float(f) => {
+                                assert_eq!(f, "50.00");
+                                found_decimal = true;
+                            }
+                            Value::String(s) => {
+                                assert_eq!(s, "completed");
+                                found_string = true;
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+            }
+        }
+        assert!(found_decimal, "Should have found decimal value");
+        assert!(found_string, "Should have found string value");
+    }
+
+    #[test]
+    fn update_with_quoted_decimal() {
+        let parsed =
+            pgdog_plugin::pg_query::parse("UPDATE transactions SET amount = '50.00' WHERE id = 1")
+                .expect("parse");
+
+        let stmt = parsed
+            .protobuf
+            .stmts
+            .first()
+            .and_then(|node| node.stmt.as_ref())
+            .and_then(|node| node.node.as_ref())
+            .expect("statement node");
+
+        let update = match stmt {
+            NodeEnum::UpdateStmt(update) => update,
+            _ => panic!("expected update stmt"),
+        };
+
+        // Quoted decimals should be treated as strings
+        let mut found_string = false;
+        for target in &update.target_list {
+            if let Some(NodeEnum::ResTarget(res)) = &target.node {
+                if let Some(val) = &res.val {
+                    if let Ok(value) = Value::try_from(&val.node) {
+                        match value {
+                            Value::String(s) => {
+                                assert_eq!(s, "50.00");
+                                found_string = true;
+                            }
+                            _ => {}
+                        }
+                    }
+                }
+            }
+        }
+        assert!(found_string, "Should have found string value");
+    }
 }
 
 impl QueryParser {

pgdog/src/frontend/router/parser/rewrite/shard_key.rs

@@ -8,6 +8,7 @@ use super::super::table::OwnedTable;
 pub enum AssignmentValue {
     Parameter(i32),
     Integer(i64),
+    Float(String),
     String(String),
     Boolean(bool),
     Null,

pgdog/src/frontend/router/parser/value.rs

@@ -12,6 +12,7 @@ use crate::net::{messages::Vector, vector::str_to_vector};
 pub enum Value<'a> {
     String(&'a str),
     Integer(i64),
+    Float(&'a str),
     Boolean(bool),
     Null,
     Placeholder(i32),
@@ -53,13 +54,7 @@ impl<'a> From<&'a AConst> for Value<'a> {
             }
             Some(Val::Boolval(b)) => Value::Boolean(b.boolval),
             Some(Val::Ival(i)) => Value::Integer(i.ival as i64),
-            Some(Val::Fval(Float { fval })) => {
-                if let Ok(val) = fval.parse() {
-                    Value::Integer(val)
-                } else {
-                    Value::Null
-                }
-            }
+            Some(Val::Fval(Float { fval })) => Value::Float(fval.as_str()),
             _ => Value::Null,
         }
     }