diff --git a/src/lib.rs b/src/lib.rs index 7f6d02c..2768973 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,6 +17,7 @@ pub(crate) struct EtcdFdw { client: Client, rt: Runtime, fetch_results: Vec, + tgt_cols: Vec, fetch_key: bool, fetch_value: bool, } @@ -231,6 +232,7 @@ impl ForeignDataWrapper for EtcdFdw { client, rt, fetch_results, + tgt_cols: Vec::new(), fetch_key: false, fetch_value: false, }) @@ -434,6 +436,7 @@ impl ForeignDataWrapper for EtcdFdw { }; let result_vec = result_unwrapped.take_kvs(); self.fetch_results = result_vec; + self.tgt_cols = columns.to_vec(); Ok(()) } @@ -448,11 +451,13 @@ impl ForeignDataWrapper for EtcdFdw { let value = x .value_str() .expect("Expected a value, but the value was empty"); - if self.fetch_key { - row.push("key", Some(Cell::String(key.to_string()))); - } - if self.fetch_value { - row.push("value", Some(Cell::String(value.to_string()))); + for tgt_col in &self.tgt_cols { + if tgt_col.name == "key" { + row.push(&tgt_col.name, Some(Cell::String(key.to_string()))); + } + if tgt_col.name == "value" { + row.push(&tgt_col.name, Some(Cell::String(value.to_string()))); + } } })) } @@ -766,4 +771,62 @@ mod tests { assert_eq!(Err(spi::SpiError::InvalidPosition), query_result); } + + #[pg_test] + fn test_select_value_only_with_key_filter() { + // Test for issue #26: selecting only value column with WHERE clause on key + let (_container, url) = create_container(); + + create_fdt(url); + + // Insert test data + Spi::run("INSERT INTO test (key, value) VALUES ('key1','value1'),('key2','value2'),('key3','value3')") + .expect("INSERT should work"); + + // Test 1: SELECT value WHERE key = 'key2' should return the value + let query_result = Spi::get_one::("SELECT value FROM test WHERE key = 'key2'") + .expect("SELECT value with key filter should work"); + + assert_eq!(Some(format!("value2")), query_result); + + // Test 2: SELECT key WHERE key = 'key1' should also work + let query_result = Spi::get_one::("SELECT key FROM test WHERE key = 'key1'") + .expect("SELECT key with key filter should work"); + + assert_eq!(Some(format!("key1")), query_result); + + // Test 3: SELECT * WHERE key = 'key3' should work (baseline) + let query_result = Spi::get_two::("SELECT * FROM test WHERE key = 'key3'") + .expect("SELECT * with key filter should work"); + + assert_eq!((Some(format!("key3")), Some(format!("value3"))), query_result); + } + + #[pg_test] + fn test_update_value_only_with_key_filter() { + // Test for issue #26: UPDATE with only value column when key is in WHERE clause + let (_container, url) = create_container(); + + create_fdt(url); + + // Insert test data + Spi::run("INSERT INTO test (key, value) VALUES ('gather/78','original_value'),('gather/84','other_value')") + .expect("INSERT should work"); + + // Update without including key column in SET clause + Spi::run("UPDATE test SET value = 'updated_value' WHERE key = 'gather/84'") + .expect("UPDATE with key filter should work"); + + // Verify the update worked + let query_result = Spi::get_one::("SELECT value FROM test WHERE key = 'gather/84'") + .expect("SELECT after UPDATE should work"); + + assert_eq!(Some(format!("updated_value")), query_result); + + // Verify other row was not affected + let query_result = Spi::get_one::("SELECT value FROM test WHERE key = 'gather/78'") + .expect("SELECT other row should work"); + + assert_eq!(Some(format!("original_value")), query_result); + } }