use std::fmt::{self, Write};
use std::ops::{Deref, DerefMut};
use crate::arguments::Arguments;
use crate::encode::{Encode, IsNull};
use crate::error::Error;
use crate::ext::ustr::UStr;
use crate::postgres::{PgConnection, PgTypeInfo, Postgres};
use crate::types::Type;
#[derive(Default)]
pub struct PgArgumentBuffer {
buffer: Vec<u8>,
count: usize,
patches: Vec<(
usize, // offset
usize, // argument index
Box<dyn Fn(&mut [u8], &PgTypeInfo) + 'static + Send + Sync>,
)>,
type_holes: Vec<(usize, UStr)>, }
#[derive(Default)]
pub struct PgArguments {
pub(crate) types: Vec<PgTypeInfo>,
pub(crate) buffer: PgArgumentBuffer,
}
impl PgArguments {
pub(crate) fn add<'q, T>(&mut self, value: T)
where
T: Encode<'q, Postgres> + Type<Postgres>,
{
self.types
.push(value.produces().unwrap_or_else(T::type_info));
self.buffer.encode(value);
self.buffer.count += 1;
}
pub(crate) async fn apply_patches(
&mut self,
conn: &mut PgConnection,
parameters: &[PgTypeInfo],
) -> Result<(), Error> {
let PgArgumentBuffer {
ref patches,
ref type_holes,
ref mut buffer,
..
} = self.buffer;
for (offset, ty, callback) in patches {
let buf = &mut buffer[*offset..];
let ty = ¶meters[*ty];
callback(buf, ty);
}
for (offset, name) in type_holes {
let oid = conn.fetch_type_id_by_name(&*name).await?;
buffer[*offset..(*offset + 4)].copy_from_slice(&oid.0.to_be_bytes());
}
Ok(())
}
}
impl<'q> Arguments<'q> for PgArguments {
type Database = Postgres;
fn reserve(&mut self, additional: usize, size: usize) {
self.types.reserve(additional);
self.buffer.reserve(size);
}
fn add<T>(&mut self, value: T)
where
T: Encode<'q, Self::Database> + Type<Self::Database>,
{
self.add(value)
}
fn format_placeholder<W: Write>(&self, writer: &mut W) -> fmt::Result {
write!(writer, "${}", self.buffer.count)
}
}
impl PgArgumentBuffer {
pub(crate) fn encode<'q, T>(&mut self, value: T)
where
T: Encode<'q, Postgres>,
{
let offset = self.len();
self.extend(&[0; 4]);
let len = if let IsNull::No = value.encode(self) {
(self.len() - offset - 4) as i32
} else {
debug_assert_eq!(self.len(), offset + 4);
-1_i32
};
self[offset..(offset + 4)].copy_from_slice(&len.to_be_bytes());
}
#[allow(dead_code)]
pub(crate) fn patch<F>(&mut self, callback: F)
where
F: Fn(&mut [u8], &PgTypeInfo) + 'static + Send + Sync,
{
let offset = self.len();
let index = self.count;
self.patches.push((offset, index, Box::new(callback)));
}
pub(crate) fn patch_type_by_name(&mut self, type_name: &UStr) {
let offset = self.len();
self.extend_from_slice(&0_u32.to_be_bytes());
self.type_holes.push((offset, type_name.clone()));
}
}
impl Deref for PgArgumentBuffer {
type Target = Vec<u8>;
#[inline]
fn deref(&self) -> &Self::Target {
&self.buffer
}
}
impl DerefMut for PgArgumentBuffer {
#[inline]
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.buffer
}
}