adapter: Automatically cast parameters of EXECUTE

Author: ggevay

src/adapter/src/coord/sql.rs

@@ -68,12 +68,12 @@ impl Coordinator {
         now: EpochMillis,
     ) -> Result<(), AdapterError> {
         let param_types = params
-            .types
+            .actual_types
             .iter()
             .map(|ty| Some(ty.clone()))
             .collect::<Vec<_>>();
         let desc = describe(catalog, stmt.clone(), &param_types, session)?;
-        let params = params.datums.into_iter().zip(params.types).collect();
+        let params = params.datums.into_iter().zip(params.actual_types).collect();
         let result_formats = vec![mz_pgwire_common::Format::Text; desc.arity()];
         let logging = session.mint_logging(sql, Some(&stmt), now);
         session.set_portal(

src/adapter/src/coord/statement_logging.rs

@@ -727,7 +727,7 @@ impl Coordinator {
             now,
         );
 
-        let params = std::iter::zip(params.types.iter(), params.datums.iter())
+        let params = std::iter::zip(params.actual_types.iter(), params.datums.iter())
             .map(|(r#type, datum)| {
                 mz_pgrepr::Value::from_datum(datum, r#type).map(|val| {
                     let mut buf = BytesMut::new();

src/adapter/src/session.rs

@@ -710,6 +710,7 @@ impl<T: TimestampManipulation> Session<T> {
         if !portal_name.is_empty() && self.portals.contains_key(&portal_name) {
             return Err(AdapterError::DuplicateCursor(portal_name));
         }
+        let param_types = desc.param_types.clone();
         self.portals.insert(
             portal_name,
             Portal {
@@ -718,7 +719,8 @@ impl<T: TimestampManipulation> Session<T> {
                 catalog_revision,
                 parameters: Params {
                     datums: Row::pack(params.iter().map(|(d, _t)| d)),
-                    types: params.into_iter().map(|(_d, t)| t).collect(),
+                    actual_types: params.into_iter().map(|(_d, t)| t).collect(),
+                    expected_types: param_types,
                 },
                 result_formats,
                 state: PortalState::NotStarted,

src/sql/src/plan.rs

@@ -2042,15 +2042,17 @@ impl TryFrom<ClusterAlterOptionExtracted> for AlterClusterPlanStrategy {
 #[derive(Debug, Clone)]
 pub struct Params {
     pub datums: Row,
-    pub types: Vec<ScalarType>,
+    pub actual_types: Vec<ScalarType>,
+    pub expected_types: Vec<ScalarType>,
 }
 
 impl Params {
     /// Returns a `Params` with no parameters.
     pub fn empty() -> Params {
         Params {
             datums: Row::pack_slice(&[]),
-            types: vec![],
+            actual_types: vec![],
+            expected_types: vec![],
         }
     }
 }

src/sql/src/plan/error.rs

@@ -101,6 +101,7 @@ pub enum PlanError {
     },
     UnknownParameter(usize),
     ParameterNotAllowed(String),
+    WrongParameterType(usize, String, String),
     RecursionLimit(RecursionLimitError),
     StrconvParse(strconv::ParseError),
     Catalog(CatalogError),
@@ -479,6 +480,9 @@ impl PlanError {
             Self::NetworkPolicyInUse => {
                 Some("Use ALTER SYSTEM SET 'network_policy' to change the default network policy.".into())
             }
+            Self::WrongParameterType(_, _, _) => {
+                Some("EXECUTE automatically inserts only such casts that are allowed in an assignment cast context. Try adding an explicit cast.".into())
+            }
             _ => None,
         }
     }
@@ -577,6 +581,7 @@ impl fmt::Display for PlanError {
             }
             Self::UnknownParameter(n) => write!(f, "there is no parameter ${}", n),
             Self::ParameterNotAllowed(object_type) => write!(f, "{} cannot have parameters", object_type),
+            Self::WrongParameterType(i, expected_ty, actual_ty) => write!(f, "unable to cast given parameter ${}: expected {}, got {}", i, expected_ty, actual_ty),
             Self::RecursionLimit(e) => write!(f, "{}", e),
             Self::StrconvParse(e) => write!(f, "{}", e),
             Self::Catalog(e) => write!(f, "{}", e),

src/sql/src/plan/hir.rs

@@ -36,10 +36,10 @@ use mz_repr::adt::numeric::NumericMaxScale;
 use mz_repr::*;
 use serde::{Deserialize, Serialize};
 
-use crate::plan::Params;
 use crate::plan::error::PlanError;
-use crate::plan::query::ExprContext;
-use crate::plan::typeconv::{self, CastContext};
+use crate::plan::query::{EXECUTE_CAST_CONTEXT, ExprContext, execute_expr_context};
+use crate::plan::typeconv::{self, CastContext, plan_cast};
+use crate::plan::{Params, QueryContext, QueryLifetime, StatementContext};
 
 use super::plan_utils::GroupSizeHints;
 
@@ -167,24 +167,18 @@ pub enum HirRelationExpr {
         /// Column indices used to order rows within groups.
         order_key: Vec<ColumnOrder>,
         /// Number of records to retain.
-        /// It is of ScalarType::Int64. (It's not entirely clear why not UInt64, see below at
-        /// `offset`.)
+        /// It is of ScalarType::Int64.
+        /// (UInt64 would make sense in theory: Then we wouldn't need to manually check
+        /// non-negativity, but would just get this for free when casting to UInt64. However, Int64
+        /// is better for Postgres compat. This is because if there is a $1 here, then when external
+        /// tools `describe` the prepared statement, they discover this type. If what they find
+        /// were UInt64, then they might have trouble calling the prepared statement, because the
+        /// unsigned types are non-standard, and also don't exist even in Postgres.)
         limit: Option<HirScalarExpr>,
         /// Number of records to skip.
         /// It is of ScalarType::Int64.
         /// This can contain parameters at first, but by the time we reach lowering, this should
         /// already be simply a Literal.
-        ///
-        /// TODO: It's not clear why this is Int64 instead of UInt64. If it were UInt64, we wouldn't
-        /// need to manually check non-negativity, but would just get this for free when casting to
-        /// UInt64.
-        /// There are some arguments for Postgres-compatibility, because Postgres expects an
-        /// Int64. But it's not clear whether there is an actual situation where having a UInt64
-        /// here would introduce a Postgres-incompatibility.
-        /// Another thing is that our prepared statements currently have the limitation that the
-        /// argument passed to EXECUTE must have exactly the same type as the parameter, i.e., it's
-        /// not enough if it's castable. Since UInt64 is not a standard type, various tools might
-        /// have trouble passing it.
         offset: HirScalarExpr,
         /// User-supplied hint: how many rows will have the same group key.
         expected_group_size: Option<u64>,
@@ -2112,10 +2106,15 @@ impl HirRelationExpr {
 
     /// Replaces any parameter references in the expression with the
     /// corresponding datum from `params`.
-    pub fn bind_parameters(&mut self, params: &Params) -> Result<(), PlanError> {
+    pub fn bind_parameters(
+        &mut self,
+        scx: &StatementContext,
+        lifetime: QueryLifetime,
+        params: &Params,
+    ) -> Result<(), PlanError> {
         #[allow(deprecated)]
         self.visit_scalar_expressions_mut(0, &mut |e: &mut HirScalarExpr, _: usize| {
-            e.bind_parameters(params)
+            e.bind_parameters(scx, lifetime, params)
         })
     }
 
@@ -2984,15 +2983,20 @@ impl HirScalarExpr {
 
     /// Replaces any parameter references in the expression with the
     /// corresponding datum in `params`.
-    pub fn bind_parameters(&mut self, params: &Params) -> Result<(), PlanError> {
+    pub fn bind_parameters(
+        &mut self,
+        scx: &StatementContext,
+        lifetime: QueryLifetime,
+        params: &Params,
+    ) -> Result<(), PlanError> {
         #[allow(deprecated)]
         self.visit_recursively_mut(0, &mut |_: usize, e: &mut HirScalarExpr| {
             if let HirScalarExpr::Parameter(n, name) = e {
                 let datum = match params.datums.iter().nth(*n - 1) {
                     None => return Err(PlanError::UnknownParameter(*n)),
                     Some(datum) => datum,
                 };
-                let scalar_type = &params.types[*n - 1];
+                let scalar_type = &params.actual_types[*n - 1];
                 let row = Row::pack([datum]);
                 let column_type = scalar_type.clone().nullable(datum.is_null());
 
@@ -3002,7 +3006,16 @@ impl HirScalarExpr {
                     Some(Arc::from(format!("${n}")))
                 };
 
-                *e = HirScalarExpr::Literal(row, column_type, TreatAsEqual(name));
+                let qcx = QueryContext::root(scx, lifetime);
+                let ecx = execute_expr_context(&qcx);
+
+                *e = plan_cast(
+                    &ecx,
+                    *EXECUTE_CAST_CONTEXT,
+                    HirScalarExpr::Literal(row, column_type, TreatAsEqual(name)),
+                    &params.expected_types[*n - 1],
+                )
+                .expect("checked in plan_params");
             }
             Ok(())
         })

src/sql/src/plan/query.rs

@@ -42,7 +42,7 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::convert::{TryFrom, TryInto};
 use std::num::NonZeroU64;
 use std::rc::Rc;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 use std::{iter, mem};
 
 use itertools::Itertools;
@@ -97,7 +97,7 @@ use crate::plan::hir::{
 use crate::plan::plan_utils::{self, GroupSizeHints, JoinSide};
 use crate::plan::scope::{Scope, ScopeItem, ScopeUngroupedColumn};
 use crate::plan::statement::{StatementContext, StatementDesc, show};
-use crate::plan::typeconv::{self, CastContext};
+use crate::plan::typeconv::{self, CastContext, plan_hypothetical_cast};
 use crate::plan::{
     Params, PlanContext, QueryWhen, ShowCreatePlan, WebhookValidation, WebhookValidationSecret,
     literal, transform_ast,
@@ -1312,43 +1312,60 @@ pub fn plan_params<'a>(
     }
 
     let qcx = QueryContext::root(scx, QueryLifetime::OneShot);
-    let scope = Scope::empty();
-    let rel_type = RelationType::empty();
 
     let mut datums = Row::default();
     let mut packer = datums.packer();
     let mut actual_types = Vec::new();
     let temp_storage = &RowArena::new();
-    for (mut expr, expected_ty) in params.into_iter().zip(&desc.param_types) {
+    for (i, (mut expr, expected_ty)) in params.into_iter().zip(&desc.param_types).enumerate() {
         transform_ast::transform(scx, &mut expr)?;
 
-        let ecx = &ExprContext {
-            qcx: &qcx,
-            name: "EXECUTE",
-            scope: &scope,
-            relation_type: &rel_type,
-            allow_aggregates: false,
-            allow_subqueries: false,
-            allow_parameters: false,
-            allow_windows: false,
-        };
-        let ex = plan_expr(ecx, &expr)?.type_as_any(ecx)?;
+        let ecx = execute_expr_context(&qcx);
+        let ex = plan_expr(&ecx, &expr)?.type_as_any(&ecx)?;
         let actual_ty = ecx.scalar_type(&ex);
-        if actual_ty != *expected_ty {
-            sql_bail!(
-                "mismatched parameter type: expected {}, got {}",
+        if plan_hypothetical_cast(&ecx, *EXECUTE_CAST_CONTEXT, &actual_ty, expected_ty).is_none() {
+            return Err(PlanError::WrongParameterType(
+                i + 1,
                 ecx.humanize_scalar_type(expected_ty, false),
                 ecx.humanize_scalar_type(&actual_ty, false),
-            );
+            ));
         }
         let ex = ex.lower_uncorrelated()?;
         let evaled = ex.eval(&[], temp_storage)?;
         packer.push(evaled);
         actual_types.push(actual_ty);
     }
-    Ok(Params { datums, types: actual_types })
+    Ok(Params {
+        datums,
+        actual_types,
+        expected_types: desc.param_types.clone(),
+    })
+}
+
+static EXECUTE_CONTEXT_SCOPE: LazyLock<Scope> = LazyLock::new(Scope::empty);
+static EXECUTE_CONTEXT_REL_TYPE: LazyLock<RelationType> = LazyLock::new(RelationType::empty);
+
+/// Returns an `ExprContext` for the expressions in the parameters of an EXECUTE statement.
+pub(crate) fn execute_expr_context<'a>(qcx: &'a QueryContext<'a>) -> ExprContext<'a> {
+    ExprContext {
+        qcx,
+        name: "EXECUTE",
+        scope: &EXECUTE_CONTEXT_SCOPE,
+        relation_type: &EXECUTE_CONTEXT_REL_TYPE,
+        allow_aggregates: false,
+        allow_subqueries: false,
+        allow_parameters: false,
+        allow_windows: false,
+    }
 }
 
+/// The CastContext used when matching up the types of parameters passed to EXECUTE.
+///
+/// This is an assignment cast also in Postgres, see
+/// <https://github.com/MaterializeInc/database-issues/issues/9266>
+pub(crate) static EXECUTE_CAST_CONTEXT: LazyLock<CastContext> =
+    LazyLock::new(|| CastContext::Assignment);
+
 pub fn plan_index_exprs<'a>(
     scx: &'a StatementContext,
     on_desc: &RelationDesc,

src/sql/src/plan/side_effecting_func.rs

@@ -102,7 +102,7 @@ pub fn plan_select_if_side_effecting(
     let temp_storage = RowArena::new();
     let mut args = vec![];
     for mut arg in sef_call.args {
-        arg.bind_parameters(params)?;
+        arg.bind_parameters(scx, QueryLifetime::OneShot, params)?;
         let arg = arg.lower_uncorrelated()?;
         args.push(arg);
     }

src/sql/src/plan/statement.rs

@@ -278,7 +278,10 @@ pub fn plan(
     resolved_ids: &ResolvedIds,
 ) -> Result<Plan, PlanError> {
     let param_types = params
-        .types
+        // We need the `expected_types` here, not the `actual_types`! This is because
+        // `expected_types` is how the parameter expression (e.g. `$1`) looks "from the outside":
+        // `bind_parameters` will insert a cast from the actual type to the expected type.
+        .expected_types
         .iter()
         .enumerate()
         .map(|(i, ty)| (i + 1, ty.clone()))