Rust PG server: fix-one batch snapshot of pg-batch (1ada2a53, 2 commits)

changelog.d/rp-876-quote-function-arg-names.md

@@ -0,0 +1,13 @@
+### A routine's parameter and result column names print quoted
+
+`pg_get_function_arguments`, `pg_get_function_identity_arguments` and
+`pg_get_function_result` now pass each parameter's and `RETURNS TABLE`
+column's name through `quote_identifier`, as PostgreSQL's
+`print_function_arguments` and `print_function_rettype` do: a parameter
+named `"A b"` or `"select"` prints quoted, so the text reads back as valid SQL.
+
+#### Fixed
+
+- Parameter and TABLE column names in the function-signature helpers are
+  quoted when needed.
+  Fixes RP-876 (tests/backlog/test_rp_found_rp515.py::test_argument_and_column_names_are_quoted)

changelog.d/rp-905-dml-portal.md

@@ -0,0 +1,17 @@
+### A PL/pgSQL FOR loop or OPEN walks a DML statement's RETURNING rows
+
+`for x in update t set a = a + 1 returning a loop` and `open c for insert
+into t values (5) returning a` used to be 0A000 on the Rust PostgreSQL
+server. As on PostgreSQL, where `SPI_cursor_open` takes any statement that
+returns rows, the statement now runs and its RETURNING rows become the
+portal. Without RETURNING it is 42P11 `cannot open UPDATE query as cursor`,
+and a procedure's COMMIT inside a loop over such a portal is 55000, since
+only a read-only portal can be held across it.
+
+#### Fixed
+
+- A PL/pgSQL FOR loop or OPEN over an INSERT, UPDATE or DELETE with
+  RETURNING walks its rows; without RETURNING it is 42P11.
+- A COMMIT or ROLLBACK inside a loop over such a portal is 55000 `cannot
+  perform transaction commands inside a cursor loop that is not read-only`.
+- Fixes RP-905 (tests/backlog/test_rp_found_rp828.py::test_a_dml_statement_opens_as_a_portal)

changelog.d/rp-906-dynamic-portal-using.md

@@ -0,0 +1,13 @@
+### A FOR loop or an OPEN over `EXECUTE ... USING` binds its values
+
+`for x in execute 'select $1' using 5 loop` and `open c for execute '...'
+using 7` used to run the dynamic query with no parameters, so the first
+`$1` was 42P02 `there is no parameter $1`. The Rust PostgreSQL server now
+evaluates the USING values as PostgreSQL does (`exec_eval_using_params`)
+and binds them to the portal, as it already did for `EXECUTE ... USING`.
+
+#### Fixed
+
+- A PL/pgSQL `FOR ... IN EXECUTE ... USING` loop and an `OPEN ... FOR
+  EXECUTE ... USING` bind their USING values.
+- Fixes RP-906 (tests/backlog/test_rp_found_rp828.py::test_a_dynamic_portal_binds_its_using_values)

crates/secantus-pg/src/catalog_fill.rs

@@ -865,7 +865,10 @@ impl PgHandler {
                         let bare = if name.is_empty() {
                             format!("{prefix}{ty}")
                         } else {
-                            format!("{prefix}{name} {ty}")
+                            format!(
+                                "{prefix}{} {ty}",
+                                secantus_pgplan::scalar::quote_identifier(name)
+                            )
                         };
                         let default = defaults
                             .get(i)
@@ -896,7 +899,9 @@ impl PgHandler {
                                     let c = c.as_document()?;
                                     Some(format!(
                                         "{} {}",
-                                        c.get_str("name").ok()?,
+                                        secantus_pgplan::scalar::quote_identifier(
+                                            c.get_str("name").ok()?
+                                        ),
                                         self.display_type_name(c.get_str("type_tag").ok()?)
                                     ))
                                 })

crates/secantus-pg/src/lib.rs

@@ -2532,7 +2532,7 @@ impl plpgsql_fn::Host for PlHost<'_> {
         }
         // `_SPI_commit` / `_SPI_rollback`: a FOR loop's pinned cursor is
         // held first, so the loop reads on in the next transaction.
-        self.h.hold_pinned_cursors();
+        self.h.hold_pinned_cursors().map_err(|e| pl_error(&e))?;
         self.h
             .restart_transaction(commit, chain)
             .map_err(|e| pl_error(&e))
@@ -4833,6 +4833,10 @@ struct CursorState {
     /// across every later transaction end until the loop closes it, as
     /// `AtCommit_Portals` and `AtAbort_Portals` pass an auto-held portal by.
     auto_held: bool,
+    /// A PL/pgSQL portal over a DML statement's RETURNING rows
+    /// (`PORTAL_ONE_RETURNING`, RP-905): not read-only, so a COMMIT or
+    /// ROLLBACK while it is pinned is 55000 (`HoldPinnedPortals`).
+    returning: bool,
 }
 
 /// What `WHERE CURRENT OF` needs of a cursor (RP-072): the one table it
@@ -10278,6 +10282,31 @@ impl PgHandler {
         &self,
         stmt: Statement,
     ) -> PgWireResult<(ColumnTypes, Vec<Vec<Bson>>)> {
+        let (schema, rows) = self.dml_returning_rows(stmt)?;
+        let columns = schema
+            .iter()
+            .map(|f| {
+                let ty = secantus_pgplan::pgtypes::name_of_oid(i64::from(f.datatype().oid()))
+                    .unwrap_or("text")
+                    .to_string();
+                (f.name().to_string(), ty)
+            })
+            .collect();
+        let rows = rows
+            .into_iter()
+            .map(|r| r.into_iter().map(|v| v.unwrap_or(Bson::Null)).collect())
+            .collect();
+        Ok((columns, rows))
+    }
+
+    /// Run a DML statement and answer its RETURNING rows with their wire
+    /// description (none when it has no RETURNING): what
+    /// [`Self::dml_returning_values`] reads, and the rows of a PL/pgSQL
+    /// portal over a DML statement (RP-905).
+    pub(crate) fn dml_returning_rows(
+        &self,
+        stmt: Statement,
+    ) -> PgWireResult<(Vec<FieldInfo>, CapturedRows)> {
         let previous = self
             .cursor_capture
             .lock()
@@ -10300,31 +10329,16 @@ impl PgHandler {
                 return Err(e);
             }
         };
-        let mut columns = Vec::new();
+        let mut schema = Vec::new();
         for r in responses {
             if let Response::Query(q) = r {
-                columns = q
-                    .row_schema
-                    .iter()
-                    .map(|f| {
-                        let ty =
-                            secantus_pgplan::pgtypes::name_of_oid(i64::from(f.datatype().oid()))
-                                .unwrap_or("text")
-                                .to_string();
-                        (f.name().to_string(), ty)
-                    })
-                    .collect();
+                schema = q.row_schema.iter().cloned().collect();
                 // The rows are produced as the stream is read; reading it is
                 // what fills the capture.
                 futures::executor::block_on(q.data_rows.try_collect::<Vec<_>>())?;
             }
         }
-        let captured = restore(self).unwrap_or_default();
-        let rows = captured
-            .into_iter()
-            .map(|r| r.into_iter().map(|v| v.unwrap_or(Bson::Null)).collect())
-            .collect();
-        Ok((columns, rows))
+        Ok((schema, restore(self).unwrap_or_default()))
     }
 
     /// A wire error as a planner error, its SQLSTATE kept.
@@ -33894,11 +33908,22 @@ impl PgHandler {
     /// in the next transaction. Its rows are already materialised (a
     /// PL/pgSQL portal is, `register_portal`), which is what `HoldPortal`
     /// makes of the rest of a portal's rows.
-    pub(crate) fn hold_pinned_cursors(&self) {
+    ///
+    /// Only a read-only portal can be held: one over a DML statement's
+    /// RETURNING rows is 55000 (RP-905).
+    pub(crate) fn hold_pinned_cursors(&self) -> PgWireResult<()> {
         let mut cursors = self.cursors.lock().unwrap_or_else(|e| e.into_inner());
         for c in cursors.values_mut().filter(|c| c.pinned) {
+            if c.returning {
+                return Err(Self::user_error(
+                    "55000",
+                    "cannot perform transaction commands inside a cursor loop that is not read-only"
+                        .into(),
+                ));
+            }
             c.auto_held = true;
         }
+        Ok(())
     }
 
     /// Pin or unpin the cursor `name` (`PinPortal` / `UnpinPortal`).
@@ -35276,6 +35301,7 @@ impl PgHandler {
                                 held: false,
                                 pinned: false,
                                 auto_held: false,
+                                returning: false,
                             },
                         );
                     if holdable {
@@ -35381,6 +35407,7 @@ impl PgHandler {
                             held: false,
                             pinned: false,
                             auto_held: false,
+                            returning: false,
                         },
                     );
                 if holdable {
@@ -35438,6 +35465,7 @@ impl PgHandler {
                         held: holdable && !in_block,
                         pinned: false,
                         auto_held: false,
+                        returning: false,
                     },
                 );
             if holdable {

crates/secantus-pg/src/plpgsql_fn.rs

@@ -3380,14 +3380,7 @@ impl Interp<'_> {
             "PLpgSQL_stmt_dynexecute" => {
                 let sql = self.dynamic_sql(body.get("query").unwrap_or(&Value::Null))?;
                 self.dyn_text = Some(sql.clone());
-                let using: Vec<Bson> = body
-                    .get("params")
-                    .and_then(Value::as_array)
-                    .cloned()
-                    .unwrap_or_default()
-                    .iter()
-                    .map(|p| self.eval(p))
-                    .collect::<Result<_, _>>()?;
+                let using = self.using_values(body)?;
                 let types: Vec<String> = Vec::new();
                 if body.get("into").and_then(Value::as_bool).unwrap_or(false) {
                     let out = self.host.query(&sql, &using, &types)?;

crates/secantus-pg/src/plpgsql_fn/open_cursor.rs

@@ -70,6 +70,16 @@ impl Interp<'_> {
         bound
     }
 
+    /// A dynamic statement's `USING` values (`exec_eval_using_params`),
+    /// each evaluated once, in order, when the statement runs: `$1` is the
+    /// first. Their types are the values' own.
+    pub(super) fn using_values(&self, body: &Value) -> Result<Vec<Bson>, PlError> {
+        body.get("params")
+            .and_then(Value::as_array)
+            .map(|ps| ps.iter().map(|p| self.eval(p)).collect())
+            .unwrap_or_else(|| Ok(Vec::new()))
+    }
+
     pub(super) fn open_cursor(&mut self, body: &Value) -> Result<Flow, PlError> {
         let var = cursor_var(body);
         // A bound cursor declares its options (`c NO SCROLL CURSOR FOR ...`).
@@ -84,7 +94,8 @@ impl Interp<'_> {
                 (q.0.to_string(), self.bind(q.0)?)
             } else if let Some(d) = body.get("dynquery") {
                 let text = self.dynamic_sql(d)?;
-                (text.clone(), (text, Vec::new(), Vec::new()))
+                let using = self.using_values(body)?;
+                (text.clone(), (text, using, Vec::new()))
             } else if let Some(q) = self.cursor_exprs.get(&var).cloned() {
                 let bound = self.bind_bound_cursor(var, body, &q)?;
                 (q, bound)
@@ -181,7 +192,8 @@ impl Interp<'_> {
         let query = body.get("query").unwrap_or(&Value::Null);
         let (statement, sql, params, types) = if dynamic {
             let sql = self.dynamic_sql(query)?;
-            (sql.clone(), sql, Vec::new(), Vec::new())
+            let using = self.using_values(body)?;
+            (sql.clone(), sql, using, Vec::new())
         } else {
             let (text, _) =
                 expr_query(query).ok_or_else(|| PlError::new("XX000", "a FOR without a query"))?;

crates/secantus-pg/src/plpgsql_portals.rs

@@ -4,6 +4,24 @@
 
 use super::*;
 
+/// The command of a DML statement (`INSERT`, `UPDATE`, `DELETE`, `MERGE`)
+/// and whether it returns rows (a RETURNING list), as SPI reads the
+/// statement's raw parse; `None` for any other statement.
+fn dml_command(sql: &str) -> Option<(&'static str, bool)> {
+    use pg_query::protobuf::node::Node as N;
+    let tree = pg_query::parse(sql).ok()?;
+    let [stmt] = tree.protobuf.stmts.as_slice() else {
+        return None;
+    };
+    Some(match stmt.stmt.as_ref()?.node.as_ref()? {
+        N::InsertStmt(i) => ("INSERT", !i.returning_list.is_empty()),
+        N::UpdateStmt(u) => ("UPDATE", !u.returning_list.is_empty()),
+        N::DeleteStmt(d) => ("DELETE", !d.returning_list.is_empty()),
+        N::MergeStmt(_) => ("MERGE", false),
+        _ => return None,
+    })
+}
+
 impl PlHost<'_> {
     pub(crate) fn open_portal(
         &self,
@@ -14,10 +32,50 @@ impl PlHost<'_> {
         types: &[String],
         scroll: bool,
     ) -> Result<String, plpgsql_fn::PlError> {
+        // `SPI_cursor_open` takes any statement that returns rows: a DML
+        // statement's RETURNING rows are a `PORTAL_ONE_RETURNING` portal,
+        // and one with no RETURNING has none to open (RP-905).
+        let dml = dml_command(sql);
+        if let Some((command, false)) = dml {
+            return Err(pl_error(&PgHandler::user_error(
+                "42P11",
+                format!("cannot open {command} query as cursor"),
+            )));
+        }
+        if dml.is_some() {
+            self.check_volatility(sql, params, types)?;
+        }
         // A statement of its own, as `query` runs one: it sees what the
         // body wrote before it (a FOR loop's query reads here, RP-828).
         let level = statement_views::StatementLevel::enter(self.h);
         let stmt = self.plan(sql, params, types)?;
+        if dml.is_some() {
+            let read = self.joined(|| {
+                self.atomic(|| {
+                    self.h
+                        .capture_for_savepoints(&stmt)
+                        .and_then(|()| self.h.dml_returning_rows(stmt))
+                        .map_err(|e| pl_error(&e))
+                })
+            });
+            drop(level);
+            let (schema, rows) = read?;
+            let domains = vec![None; schema.len()];
+            let name = self
+                .h
+                .register_portal(name, statement, (schema, domains), rows, scroll, None)
+                .map_err(|e| pl_error(&e))?;
+            if let Some(c) = self
+                .h
+                .cursors
+                .lock()
+                .unwrap_or_else(|e| e.into_inner())
+                .get_mut(&name)
+            {
+                c.returning = true;
+            }
+            return Ok(name);
+        }
         // A simply updatable scan keeps each row's id, for `WHERE CURRENT
         // OF` (`current_of.rs`).
         let updatable = match &stmt {
@@ -162,6 +220,7 @@ impl PgHandler {
                 held: false,
                 pinned: false,
                 auto_held: false,
+                returning: false,
             },
         );
         Ok(name)

crates/secantus-pg/src/portal_stream.rs

@@ -1911,6 +1911,7 @@ impl PgHandler {
             failed: false,
             pinned: false,
             auto_held: false,
+            returning: false,
         }))
     }