Mutable SQL commands should be used with Spi::update(). (#1114)

higuoxing · beeender · web-flow · commit a46825ae4391 · 2023-10-23T08:53:37.000-07:00
Co-authored-by: Chen Mulong &lt;chenmulong@gmail.com&gt;
diff --git a/pgml-extension/src/orm/dataset.rs b/pgml-extension/src/orm/dataset.rs
@@ -95,7 +95,7 @@ impl Display for TextDataset {
 fn drop_table_if_exists(table_name: &str) {
     // Avoid the existence for DROP TABLE IF EXISTS warning by checking the schema for the table first
     let table_count = Spi::get_one_with_args::<i64>("SELECT COUNT(*) FROM information_schema.tables WHERE table_name = $1 AND table_schema = 'pgml'", vec![
-        (PgBuiltInOids::TEXTOID.oid(), table_name.clone().into_datum())
+        (PgBuiltInOids::TEXTOID.oid(), table_name.into_datum())
     ]).unwrap().unwrap();
     if table_count == 1 {
         Spi::run(&format!(r#"DROP TABLE pgml.{table_name} CASCADE"#)).unwrap();
diff --git a/pgml-extension/src/orm/model.rs b/pgml-extension/src/orm/model.rs
@@ -96,8 +96,8 @@ impl Model {
         let dataset = snapshot.tabular_dataset();
         let status = Status::in_progress;
         // Create the model record.
-        Spi::connect(|client| {
-            let result = client.select("
+        Spi::connect(|mut client| {
+            let result = client.update("
           INSERT INTO pgml.models (project_id, snapshot_id, algorithm, runtime, hyperparams, status, search, search_params, search_args, num_features)
           VALUES ($1, $2, cast($3 AS pgml.algorithm), cast($4 AS pgml.runtime), $5, cast($6 as pgml.status), $7, $8, $9, $10)
           RETURNING id, project_id, snapshot_id, algorithm, runtime, hyperparams, status, metrics, search, search_params, search_args, created_at, updated_at;",
@@ -168,8 +168,8 @@ impl Model {
         let dataset = snapshot.text_dataset();
 
         // Create the model record.
-        Spi::connect(|client| {
-            let result = client.select("
+        Spi::connect(|mut client| {
+            let result = client.update("
           INSERT INTO pgml.models (project_id, snapshot_id, algorithm, runtime, hyperparams, status, search, search_params, search_args, num_features)
           VALUES ($1, $2, cast($3 AS pgml.algorithm), cast($4 AS pgml.runtime), $5, cast($6 as pgml.status), $7, $8, $9, $10)
           RETURNING id, project_id, snapshot_id, algorithm, runtime::TEXT, hyperparams, status, metrics, search, search_params, search_args, created_at, updated_at;",
diff --git a/pgml-extension/src/orm/snapshot.rs b/pgml-extension/src/orm/snapshot.rs
@@ -508,7 +508,7 @@ impl Snapshot {
         let preprocessors: HashMap<String, Preprocessor> =
             serde_json::from_value(preprocess.0).expect("is valid");
 
-        Spi::connect(|client| {
+        Spi::connect(|mut client| {
             let mut columns: Vec<Column> = Vec::new();
             client.select("SELECT column_name::TEXT, udt_name::TEXT, is_nullable::BOOLEAN, ordinal_position::INTEGER FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position ASC",
                 None,
@@ -587,7 +587,7 @@ impl Snapshot {
                 }
             }
 
-            let result = client.select("INSERT INTO pgml.snapshots (relation_name, y_column_name, test_size, test_sampling, status, columns, materialized) VALUES ($1, $2, $3, $4::pgml.sampling, $5::pgml.status, $6, $7) RETURNING id, relation_name, y_column_name, test_size, test_sampling::TEXT, status::TEXT, columns, analysis, created_at, updated_at;",
+            let result = client.update("INSERT INTO pgml.snapshots (relation_name, y_column_name, test_size, test_sampling, status, columns, materialized) VALUES ($1, $2, $3, $4::pgml.sampling, $5::pgml.status, $6, $7) RETURNING id, relation_name, y_column_name, test_size, test_sampling::TEXT, status::TEXT, columns, analysis, created_at, updated_at;",
                 Some(1),
                 Some(vec![
                     (PgBuiltInOids::TEXTOID.oid(), relation_name.into_datum()),
@@ -623,7 +623,7 @@ impl Snapshot {
                 if s.test_sampling == Sampling::random {
                     sql += " ORDER BY random()";
                 }
-                client.select(&sql, None, None).unwrap();
+                client.update(&sql, None, None).unwrap();
             }
             snapshot = Some(s);
         });

Original file line number	Diff line number	Diff line change
`@@ -508,7 +508,7 @@ impl Snapshot {`
`508`	`508`	`let preprocessors: HashMap<String, Preprocessor> =`
`509`	`509`	`serde_json::from_value(preprocess.0).expect("is valid");`
`510`	`510`
`511`		`- Spi::connect(\|client\| {`
	`511`	`+ Spi::connect(\|mut client\| {`
`512`	`512`	`let mut columns: Vec<Column> = Vec::new();`
`513`	`513`	`client.select("SELECT column_name::TEXT, udt_name::TEXT, is_nullable::BOOLEAN, ordinal_position::INTEGER FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2 ORDER BY ordinal_position ASC",`
`514`	`514`	`None,`
`@@ -587,7 +587,7 @@ impl Snapshot {`
`587`	`587`	`}`
`588`	`588`	`}`
`589`	`589`
`590`		`- let result = client.select("INSERT INTO pgml.snapshots (relation_name, y_column_name, test_size, test_sampling, status, columns, materialized) VALUES ($1, $2, $3, $4::pgml.sampling, $5::pgml.status, $6, $7) RETURNING id, relation_name, y_column_name, test_size, test_sampling::TEXT, status::TEXT, columns, analysis, created_at, updated_at;",`
	`590`	`+ let result = client.update("INSERT INTO pgml.snapshots (relation_name, y_column_name, test_size, test_sampling, status, columns, materialized) VALUES ($1, $2, $3, $4::pgml.sampling, $5::pgml.status, $6, $7) RETURNING id, relation_name, y_column_name, test_size, test_sampling::TEXT, status::TEXT, columns, analysis, created_at, updated_at;",`
`591`	`591`	`Some(1),`
`592`	`592`	`Some(vec![`
`593`	`593`	`(PgBuiltInOids::TEXTOID.oid(), relation_name.into_datum()),`
`@@ -623,7 +623,7 @@ impl Snapshot {`
`623`	`623`	`if s.test_sampling == Sampling::random {`
`624`	`624`	`sql += " ORDER BY random()";`
`625`	`625`	`}`
`626`		`- client.select(&sql, None, None).unwrap();`
	`626`	`+ client.update(&sql, None, None).unwrap();`
`627`	`627`	`}`
`628`	`628`	`snapshot = Some(s);`
`629`	`629`	`});`