Skip to content

Commit ae0335b

Browse files
YohDeadfallYohDeadfall
authored
Made SPI query arguments type safe (pgcentralfoundation#1858)
1 parent 039c24f commit ae0335b

30 files changed

+267
-294
lines changed

‎pgrx-examples/bgworker/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ pub extern "C" fn background_worker_main(arg: pg_sys::Datum) {
6969
let tuple_table = client.select(
7070
"SELECT 'Hi', id, ''||a FROM (SELECT id, 42 from generate_series(1,10) id) a ",
7171
None,
72-
None,
72+
&[],
7373
)?;
7474
for tuple in tuple_table {
7575
let a = tuple.get_datum_by_ordinal(1)?.value::<String>()?;

‎pgrx-examples/custom_sql/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ mod tests {
8989
let buf = Spi::connect(|client| {
9090
Ok::<_, spi::Error>(
9191
client
92-
.select("SELECT * FROM extension_sql", None, None)?
92+
.select("SELECT * FROM extension_sql", None, &[])?
9393
.flat_map(|tup| {
9494
tup.get_datum_by_ordinal(1)
9595
.ok()

‎pgrx-examples/schemas/src/lib.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,10 +103,10 @@ mod tests {
103103
fn test_my_some_schema_type() -> Result<(), spi::Error> {
104104
Spi::connect(|mut c| {
105105
// "MySomeSchemaType" is in 'some_schema', so it needs to be discoverable
106-
c.update("SET search_path TO some_schema,public", None, None)?;
106+
c.update("SET search_path TO some_schema,public", None, &[])?;
107107
assert_eq!(
108108
String::from("test"),
109-
c.select("SELECT '\"test\"'::MySomeSchemaType", None, None)?
109+
c.select("SELECT '\"test\"'::MySomeSchemaType", None, &[])?
110110
.first()
111111
.get_one::<MySomeSchemaType>()
112112
.expect("get_one::<MySomeSchemaType>() failed")

‎pgrx-examples/spi/src/lib.rs‎

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ fn spi_return_query() -> Result<
4848

4949
Spi::connect(|client| {
5050
client
51-
.select(query, None, None)?
51+
.select(query, None, &[])?
5252
.map(|row| Ok((row["oid"].value()?, row[2].value()?)))
5353
.collect::<Result<Vec<_>, _>>()
5454
})
@@ -62,21 +62,14 @@ fn spi_query_random_id() -> Result<Option<i64>, pgrx::spi::Error> {
6262

6363
#[pg_extern]
6464
fn spi_query_title(title: &str) -> Result<Option<i64>, pgrx::spi::Error> {
65-
Spi::get_one_with_args(
66-
"SELECT id FROM spi.spi_example WHERE title = $1;",
67-
vec![(PgBuiltInOids::TEXTOID.oid(), title.into_datum())],
68-
)
65+
Spi::get_one_with_args("SELECT id FROM spi.spi_example WHERE title = $1;", &[title.into()])
6966
}
7067

7168
#[pg_extern]
7269
fn spi_query_by_id(id: i64) -> Result<Option<String>, spi::Error> {
7370
let (returned_id, title) = Spi::connect(|client| {
7471
let tuptable = client
75-
.select(
76-
"SELECT id, title FROM spi.spi_example WHERE id = $1",
77-
None,
78-
Some(vec![(PgBuiltInOids::INT8OID.oid(), id.into_datum())]),
79-
)?
72+
.select("SELECT id, title FROM spi.spi_example WHERE id = $1", None, &[id.into()])?
8073
.first();
8174

8275
tuptable.get_two::<i64, String>()
@@ -90,7 +83,7 @@ fn spi_query_by_id(id: i64) -> Result<Option<String>, spi::Error> {
9083
fn spi_insert_title(title: &str) -> Result<Option<i64>, spi::Error> {
9184
Spi::get_one_with_args(
9285
"INSERT INTO spi.spi_example(title) VALUES ($1) RETURNING id",
93-
vec![(PgBuiltInOids::TEXTOID.oid(), title.into_datum())],
86+
&[title.into()],
9487
)
9588
}
9689

@@ -100,7 +93,7 @@ fn spi_insert_title2(
10093
) -> TableIterator<(name!(id, Option<i64>), name!(title, Option<String>))> {
10194
let tuple = Spi::get_two_with_args(
10295
"INSERT INTO spi.spi_example(title) VALUES ($1) RETURNING id, title",
103-
vec![(PgBuiltInOids::TEXTOID.oid(), title.into_datum())],
96+
&[title.into()],
10497
)
10598
.unwrap();
10699

@@ -110,7 +103,7 @@ fn spi_insert_title2(
110103
#[pg_extern]
111104
fn issue1209_fixed() -> Result<Option<String>, Box<dyn std::error::Error>> {
112105
let res = Spi::connect(|c| {
113-
let mut cursor = c.open_cursor("SELECT 'hello' FROM generate_series(1, 10000)", None);
106+
let mut cursor = c.try_open_cursor("SELECT 'hello' FROM generate_series(1, 10000)", &[])?;
114107
let table = cursor.fetch(10000)?;
115108
table.into_iter().map(|row| row.get::<&str>(1)).collect::<Result<Vec<_>, _>>()
116109
})?;

‎pgrx-examples/spi_srf/src/lib.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ fn calculate_human_years() -> Result<
5555

5656
Spi::connect(|client| {
5757
let mut results = Vec::new();
58-
let tup_table = client.select(query, None, None)?;
58+
let tup_table = client.select(query, None, &[])?;
5959

6060
for row in tup_table {
6161
let dog_name = row["dog_name"].value::<String>();
@@ -89,10 +89,10 @@ fn filter_by_breed(
8989
*/
9090

9191
let query = "SELECT * FROM spi_srf.dog_daycare WHERE dog_breed = $1;";
92-
let args = vec![(PgBuiltInOids::TEXTOID.oid(), breed.into_datum())];
92+
let args = vec![breed.into()];
9393

9494
Spi::connect(|client| {
95-
let tup_table = client.select(query, None, Some(args))?;
95+
let tup_table = client.select(query, None, &args)?;
9696

9797
let filtered = tup_table
9898
.map(|row| {

‎pgrx-tests/src/tests/aggregate_tests.rs‎

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -263,15 +263,9 @@ mod tests {
263263
fn aggregate_first_json() -> Result<(), pgrx::spi::Error> {
264264
let retval = Spi::get_one_with_args::<pgrx::Json>(
265265
"SELECT FirstJson(value) FROM UNNEST(ARRAY [$1, $2]) as value;",
266-
vec![
267-
(
268-
PgBuiltInOids::JSONOID.oid(),
269-
pgrx::Json(serde_json::json!({ "foo": "one" })).into_datum(),
270-
),
271-
(
272-
PgBuiltInOids::JSONOID.oid(),
273-
pgrx::Json(serde_json::json!({ "foo": "two" })).into_datum(),
274-
),
266+
&[
267+
pgrx::Json(serde_json::json!({ "foo": "one" })).into(),
268+
pgrx::Json(serde_json::json!({ "foo": "two" })).into(),
275269
],
276270
)?
277271
.map(|json| json.0);
@@ -285,15 +279,9 @@ mod tests {
285279
fn aggregate_first_jsonb() -> Result<(), pgrx::spi::Error> {
286280
let retval = Spi::get_one_with_args::<pgrx::JsonB>(
287281
"SELECT FirstJsonB(value) FROM UNNEST(ARRAY [$1, $2]) as value;",
288-
vec![
289-
(
290-
PgBuiltInOids::JSONBOID.oid(),
291-
pgrx::JsonB(serde_json::json!({ "foo": "one" })).into_datum(),
292-
),
293-
(
294-
PgBuiltInOids::JSONBOID.oid(),
295-
pgrx::JsonB(serde_json::json!({ "foo": "two" })).into_datum(),
296-
),
282+
&[
283+
pgrx::JsonB(serde_json::json!({ "foo": "one" })).into(),
284+
pgrx::JsonB(serde_json::json!({ "foo": "two" })).into(),
297285
],
298286
)?
299287
.map(|json| json.0);

‎pgrx-tests/src/tests/anyelement_tests.rs‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,13 @@ mod tests {
1111
#[allow(unused_imports)]
1212
use crate as pgrx_tests;
1313

14-
use pgrx::{prelude::*, AnyElement};
14+
use pgrx::{datum::DatumWithOid, prelude::*, AnyElement};
1515

1616
#[pg_test]
1717
fn test_anyelement_arg() -> Result<(), pgrx::spi::Error> {
18-
let element = Spi::get_one_with_args::<AnyElement>(
19-
"SELECT anyelement_arg($1);",
20-
vec![(PgBuiltInOids::ANYELEMENTOID.oid(), 123.into_datum())],
21-
)?
18+
let element = Spi::get_one_with_args::<AnyElement>("SELECT anyelement_arg($1);", unsafe {
19+
&[DatumWithOid::new(123, AnyElement::type_oid())]
20+
})?
2221
.map(|e| e.datum());
2322

2423
assert_eq!(element, 123.into_datum());

‎pgrx-tests/src/tests/anynumeric_tests.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,9 @@ mod tests {
1515

1616
#[pg_test]
1717
fn test_anynumeric_arg() -> Result<(), pgrx::spi::Error> {
18-
let numeric = Spi::get_one_with_args::<AnyNumeric>(
19-
"SELECT anynumeric_arg($1);",
20-
vec![(PgBuiltInOids::INT4OID.oid(), 123.into_datum())],
21-
)?
22-
.map(|n| n.normalize().to_string());
18+
let numeric =
19+
Spi::get_one_with_args::<AnyNumeric>("SELECT anynumeric_arg($1);", &[123.into()])?
20+
.map(|n| n.normalize().to_string());
2321

2422
assert_eq!(numeric, Some("123".to_string()));
2523

‎pgrx-tests/src/tests/array_tests.rs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ mod tests {
191191

192192
use super::ArrayTestEnum;
193193
use pgrx::prelude::*;
194-
use pgrx::{IntoDatum, Json};
194+
use pgrx::Json;
195195
use serde_json::json;
196196

197197
#[pg_test]
@@ -300,10 +300,7 @@ mod tests {
300300
.select(
301301
"SELECT serde_serialize_array_i32($1)",
302302
None,
303-
Some(vec![(
304-
PgBuiltInOids::INT4ARRAYOID.oid(),
305-
owned_vec.as_slice().into_datum(),
306-
)]),
303+
&[owned_vec.as_slice().into()],
307304
)?
308305
.first()
309306
.get_one::<Json>()

‎pgrx-tests/src/tests/bgworker_tests.rs‎

Lines changed: 4 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,7 @@ pub extern "C" fn bgworker(arg: pg_sys::Datum) {
2727
Spi::run("CREATE TABLE tests.bgworker_test (v INTEGER);")?;
2828
Spi::connect(|mut client| {
2929
client
30-
.update(
31-
"INSERT INTO tests.bgworker_test VALUES ($1);",
32-
None,
33-
Some(vec![(PgOid::BuiltIn(PgBuiltInOids::INT4OID), arg.into_datum())]),
34-
)
30+
.update("INSERT INTO tests.bgworker_test VALUES ($1);", None, &[arg.into()])
3531
.map(|_| ())
3632
})
3733
})
@@ -61,10 +57,7 @@ pub extern "C" fn bgworker_return_value(arg: pg_sys::Datum) {
6157
let val = if arg > 0 {
6258
BackgroundWorker::transaction(|| {
6359
Spi::run("CREATE TABLE tests.bgworker_test_return (v INTEGER);")?;
64-
Spi::get_one_with_args::<i32>(
65-
"SELECT $1",
66-
vec![(PgOid::BuiltIn(PgBuiltInOids::INT4OID), arg.into_datum())],
67-
)
60+
Spi::get_one_with_args::<i32>("SELECT $1", &[arg.into()])
6861
})
6962
.expect("bgworker transaction failed")
7063
.unwrap()
@@ -74,12 +67,8 @@ pub extern "C" fn bgworker_return_value(arg: pg_sys::Datum) {
7467
while BackgroundWorker::wait_latch(Some(Duration::from_millis(100))) {}
7568
BackgroundWorker::transaction(|| {
7669
Spi::connect(|mut c| {
77-
c.update(
78-
"INSERT INTO tests.bgworker_test_return VALUES ($1)",
79-
None,
80-
Some(vec![(PgOid::BuiltIn(PgBuiltInOids::INT4OID), val.into_datum())]),
81-
)
82-
.map(|_| ())
70+
c.update("INSERT INTO tests.bgworker_test_return VALUES ($1)", None, &[val.into()])
71+
.map(|_| ())
8372
})
8473
})
8574
.expect("bgworker transaction failed");

0 commit comments

Comments
 (0)