feat: Placeholders in prepared statements

crates/pgrepr/src/error.rs

@@ -12,6 +12,15 @@ pub enum PgReprError {
     #[error(transparent)]
     Io(#[from] std::io::Error),
 
+    #[error(transparent)]
+    Utf8Error(#[from] std::str::Utf8Error),
+
+    #[error("Binary read unimplemented.")]
+    BinaryReadUnimplemented,
+
+    #[error("Failed to parse: {0}")]
+    ParseError(Box<dyn std::error::Error + Sync + Send>),
+
     #[error("arrow type '{0}' not supported")]
     UnsupportedArrowType(datafusion::arrow::datatypes::DataType),
 

crates/pgrepr/src/lib.rs

@@ -3,4 +3,5 @@ pub mod format;
 pub mod oid;
 pub mod types;
 
+mod reader;
 mod writer;

crates/pgrepr/src/reader.rs

@@ -0,0 +1,98 @@
+use crate::error::{PgReprError, Result};
+use std::str::FromStr;
+
+/// Reader defines the interface for the different kinds of values that can be
+/// decoded as a postgres type.
+pub(crate) trait Reader {
+    fn read_bool(buf: &[u8]) -> Result<bool>;
+
+    fn read_int2(buf: &[u8]) -> Result<i16>;
+    fn read_int4(buf: &[u8]) -> Result<i32>;
+    fn read_int8(buf: &[u8]) -> Result<i64>;
+    fn read_float4(buf: &[u8]) -> Result<f32>;
+    fn read_float8(buf: &[u8]) -> Result<f64>;
+
+    fn read_text(buf: &[u8]) -> Result<String>;
+}
+
+#[derive(Debug)]
+pub(crate) struct TextReader;
+
+impl TextReader {
+    fn parse<E: std::error::Error + Sync + Send + 'static, F: FromStr<Err = E>>(
+        buf: &[u8],
+    ) -> Result<F> {
+        std::str::from_utf8(buf)?
+            .parse::<F>()
+            .map_err(|e| PgReprError::ParseError(Box::new(e)))
+    }
+}
+
+impl Reader for TextReader {
+    fn read_bool(buf: &[u8]) -> Result<bool> {
+        TextReader::parse::<_, SqlBool>(buf).map(|b| b.0)
+    }
+
+    fn read_int2(buf: &[u8]) -> Result<i16> {
+        TextReader::parse(buf)
+    }
+
+    fn read_int4(buf: &[u8]) -> Result<i32> {
+        TextReader::parse(buf)
+    }
+
+    fn read_int8(buf: &[u8]) -> Result<i64> {
+        TextReader::parse(buf)
+    }
+
+    fn read_float4(buf: &[u8]) -> Result<f32> {
+        TextReader::parse(buf)
+    }
+
+    fn read_float8(buf: &[u8]) -> Result<f64> {
+        TextReader::parse(buf)
+    }
+
+    fn read_text(buf: &[u8]) -> Result<String> {
+        TextReader::parse(buf)
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+#[error("String was not 't', 'true', 'f', or 'false'")]
+struct ParseSqlBoolError;
+
+struct SqlBool(bool);
+
+impl FromStr for SqlBool {
+    type Err = ParseSqlBoolError;
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "true" | "t" => Ok(SqlBool(true)),
+            "false" | "f" => Ok(SqlBool(false)),
+            _ => Err(ParseSqlBoolError),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn read_sql_bool() {
+        let v = TextReader::read_bool("t".as_bytes()).unwrap();
+        assert!(v);
+
+        let v = TextReader::read_bool("true".as_bytes()).unwrap();
+        assert!(v);
+
+        let v = TextReader::read_bool("f".as_bytes()).unwrap();
+        assert!(!v);
+
+        let v = TextReader::read_bool("false".as_bytes()).unwrap();
+        assert!(!v);
+
+        let _ = TextReader::read_bool("none".as_bytes()).unwrap_err();
+    }
+}

crates/pgrepr/src/types.rs

@@ -8,6 +8,7 @@ use tokio_postgres::types::Type as PgType;
 
 use crate::error::{PgReprError, Result};
 use crate::format::Format;
+use crate::reader::{Reader, TextReader};
 use crate::writer::{BinaryWriter, TextWriter, Writer};
 
 /// Returns a compatible postgres type for the arrow datatype.
@@ -64,6 +65,36 @@ pub fn encode_array_value(
     Ok(())
 }
 
+/// Decodes a scalar value using the provided format and arrow type.
+pub fn decode_scalar_value(
+    buf: Option<&[u8]>,
+    format: Format,
+    arrow_type: &ArrowType,
+) -> Result<ScalarValue> {
+    match buf {
+        Some(buf) => match format {
+            Format::Text => decode_not_null_value::<TextReader>(buf, arrow_type),
+            Format::Binary => Err(PgReprError::BinaryReadUnimplemented),
+        },
+        None => Ok(ScalarValue::Null),
+    }
+}
+
+fn decode_not_null_value<R: Reader>(buf: &[u8], arrow_type: &ArrowType) -> Result<ScalarValue> {
+    Ok(match arrow_type {
+        &ArrowType::Boolean => R::read_bool(buf)?.into(),
+        &ArrowType::Int8 => R::read_int2(buf)?.into(),
+        &ArrowType::Int16 => R::read_int2(buf)?.into(),
+        &ArrowType::Int32 => R::read_int4(buf)?.into(),
+        &ArrowType::Int64 => R::read_int8(buf)?.into(),
+        &ArrowType::Float16 => R::read_float4(buf)?.into(),
+        &ArrowType::Float32 => R::read_float4(buf)?.into(),
+        &ArrowType::Float64 => R::read_float8(buf)?.into(),
+        &ArrowType::Utf8 => ScalarValue::Utf8(Some(R::read_text(buf)?)),
+        other => return Err(PgReprError::UnsupportedArrowType(other.clone())),
+    })
+}
+
 /// Per writer implementation for encoding non-null array values.
 fn encode_array_not_null_value<W: Writer>(
     buf: &mut BytesMut,

crates/pgsrv/src/handler.rs

@@ -6,9 +6,12 @@ use crate::messages::{
 };
 use crate::proxy::{ProxyKey, GLAREDB_DATABASE_ID_KEY, GLAREDB_USER_ID_KEY};
 use crate::ssl::{Connection, SslConfig};
+use datafusion::arrow::datatypes::DataType;
 use datafusion::physical_plan::SendableRecordBatchStream;
+use datafusion::scalar::ScalarValue;
 use futures::StreamExt;
 use pgrepr::format::Format;
+use pgrepr::types::decode_scalar_value;
 use sqlexec::context::{OutputFields, Portal, PreparedStatement};
 use sqlexec::{
     engine::Engine,
@@ -369,13 +372,9 @@ where
             );
 
             // Bind...
-            if let Err(e) = session.bind_statement(
-                UNNAMED,
-                &UNNAMED,
-                Vec::new(),
-                Vec::new(),
-                all_text_formats(num_fields),
-            ) {
+            if let Err(e) =
+                session.bind_statement(UNNAMED, &UNNAMED, Vec::new(), all_text_formats(num_fields))
+            {
                 self.send_error(e.into()).await?;
                 return self.ready_for_query().await;
             }
@@ -462,7 +461,14 @@ where
             Err(e) => return self.send_error(e.into()).await,
         };
 
-        // TODO: Check and parse param formats and values.
+        // Read scalars for query parameters.
+        let scalars = match stmt.input_paramaters() {
+            Some(types) => match decode_param_scalars(param_formats, param_values, types) {
+                Ok(scalars) => scalars,
+                Err(e) => return self.send_error(e).await,
+            },
+            None => Vec::new(), // Would only happen with an empty query.
+        };
 
         // Extend out the result formats.
         let result_formats = match extend_formats(
@@ -473,13 +479,10 @@ where
             Err(e) => return self.send_error(e).await,
         };
 
-        match self.session.bind_statement(
-            portal,
-            &statement,
-            param_formats,
-            param_values,
-            result_formats,
-        ) {
+        match self
+            .session
+            .bind_statement(portal, &statement, scalars, result_formats)
+        {
             Ok(_) => self.conn.send(BackendMessage::BindComplete).await,
             Err(e) => self.send_error(e.into()).await,
         }
@@ -648,6 +651,55 @@ where
     }
 }
 
+/// Decodes inputs for a prepared query into the appropriate scalar values.
+fn decode_param_scalars(
+    param_formats: Vec<Format>,
+    param_values: Vec<Option<Vec<u8>>>,
+    types: &HashMap<String, Option<DataType>>,
+) -> Result<Vec<ScalarValue>, ErrorResponse> {
+    let param_formats = extend_formats(param_formats, param_values.len())?;
+
+    if param_values.len() != types.len() {
+        return Err(ErrorResponse::error_internal(format!(
+            "Invalid number of values provided. Expected: {}, got: {}",
+            types.len(),
+            param_values.len(),
+        )));
+    }
+
+    let mut scalars = Vec::with_capacity(param_values.len());
+    for (idx, (val, format)) in param_values
+        .into_iter()
+        .zip(param_formats.into_iter())
+        .enumerate()
+    {
+        // Parameter types keyed by '$n'.
+        let str_id = format!("${}", idx + 1);
+
+        let typ = types.get(&str_id).ok_or_else(|| {
+            ErrorResponse::error_internal(format!(
+                "Missing type for param value at index {}, input types: {:?}",
+                idx, types
+            ))
+        })?;
+
+        match typ {
+            Some(typ) => {
+                let scalar = decode_scalar_value(val.as_deref(), format, typ)?;
+                scalars.push(scalar);
+            }
+            None => {
+                return Err(ErrorResponse::error_internal(format!(
+                    "Unknown type at index {}, input types: {:?}",
+                    idx, types
+                )))
+            }
+        }
+    }
+
+    Ok(scalars)
+}
+
 /// Parse a sql string, returning an error response if failed to parse.
 fn parse_sql(sql: &str) -> Result<VecDeque<StatementWithExtensions>, ErrorResponse> {
     parser::parse_sql(sql).map_err(|e| ErrorResponse::error(SqlState::SyntaxError, e.to_string()))
@@ -684,3 +736,88 @@ fn get_encoding_state(portal: &Portal) -> Vec<(PgType, Format)> {
             .collect(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn decode_params_success() {
+        // Success test cases for decoding params.
+
+        struct TestCase {
+            values: Vec<Option<Vec<u8>>>,
+            types: Vec<(&'static str, Option<DataType>)>,
+            expected: Vec<ScalarValue>,
+        }
+
+        let test_cases = vec![
+            // No params.
+            TestCase {
+                values: Vec::new(),
+                types: Vec::new(),
+                expected: Vec::new(),
+            },
+            // One param of type int64.
+            TestCase {
+                values: vec![Some(vec![49])],
+                types: vec![("$1", Some(DataType::Int64))],
+                expected: vec![ScalarValue::Int64(Some(1))],
+            },
+            // Two params param of type string.
+            TestCase {
+                values: vec![Some(vec![49, 48]), Some(vec![50, 48])],
+                types: vec![("$1", Some(DataType::Utf8)), ("$2", Some(DataType::Utf8))],
+                expected: vec![
+                    ScalarValue::Utf8(Some("10".to_string())),
+                    ScalarValue::Utf8(Some("20".to_string())),
+                ],
+            },
+        ];
+
+        for test_case in test_cases {
+            let types: HashMap<_, _> = test_case
+                .types
+                .into_iter()
+                .map(|(k, v)| (k.to_string(), v))
+                .collect();
+
+            let scalars = decode_param_scalars(Vec::new(), test_case.values, &types).unwrap();
+            assert_eq!(test_case.expected, scalars);
+        }
+    }
+
+    #[test]
+    fn decode_params_fail() {
+        // Failure test cases for decoding params (all cases should result in an
+        // error).
+
+        struct TestCase {
+            values: Vec<Option<Vec<u8>>>,
+            types: Vec<(&'static str, Option<DataType>)>,
+        }
+
+        let test_cases = vec![
+            // Params provided, none expected.
+            TestCase {
+                values: vec![Some(vec![49])],
+                types: Vec::new(),
+            },
+            // No params provided, one expected.
+            TestCase {
+                values: Vec::new(),
+                types: vec![("$1", Some(DataType::Int64))],
+            },
+        ];
+
+        for test_case in test_cases {
+            let types: HashMap<_, _> = test_case
+                .types
+                .into_iter()
+                .map(|(k, v)| (k.to_string(), v))
+                .collect();
+
+            decode_param_scalars(Vec::new(), test_case.values, &types).unwrap_err();
+        }
+    }
+}

crates/pgsrv/src/messages.rs

@@ -1,4 +1,5 @@
 use datafusion::arrow::record_batch::RecordBatch;
+use pgrepr::error::PgReprError;
 use pgrepr::format::Format;
 use sqlexec::errors::ExecError;
 use std::collections::HashMap;
@@ -256,6 +257,13 @@ impl From<&PgSrvError> for ErrorResponse {
     }
 }
 
+impl From<PgReprError> for ErrorResponse {
+    fn from(e: PgReprError) -> Self {
+        // TODO: Actually set appropriate codes.
+        ErrorResponse::error_internal(e.to_string())
+    }
+}
+
 #[derive(Debug)]
 pub enum NoticeSeverity {
     Warning,