use core::marker::PhantomData;
use std::io::BufRead;
use super::CommonOptions;
use super::CopyFormat;
use super::CopyTarget;
use crate::AppearsOnTable;
use crate::Connection;
use crate::Expression;
use crate::QueryResult;
use crate::Selectable;
use crate::deserialize::FromSqlRow;
use crate::pg::Pg;
#[cfg(feature = "postgres")]
use crate::pg::value::TypeOidLookup;
use crate::query_builder::QueryFragment;
use crate::query_builder::QueryId;
use crate::row::Row;
#[cfg(feature = "postgres")]
use crate::row::{self, Field, PartialRow, RowIndex, RowSealed};
#[derive(Default, Debug)]
pub struct CopyToOptions {
common: CommonOptions,
header: Option<bool>,
}
impl CopyToOptions {
fn any_set(&self) -> bool {
self.common.any_set() || self.header.is_some()
}
}
impl QueryFragment<Pg> for CopyToOptions {
fn walk_ast<'b>(
&'b self,
mut pass: crate::query_builder::AstPass<'_, 'b, Pg>,
) -> crate::QueryResult<()> {
if self.any_set() {
let mut comma = "";
pass.push_sql(" WITH (");
self.common.walk_ast(pass.reborrow(), &mut comma);
if let Some(header_is_set) = self.header {
pass.push_sql(comma);
pass.push_sql("HEADER ");
pass.push_sql(if header_is_set { "1" } else { "0" });
}
pass.push_sql(")");
}
Ok(())
}
}
#[derive(Debug)]
pub struct CopyToCommand<S> {
options: CopyToOptions,
p: PhantomData<S>,
}
impl<S> QueryId for CopyToCommand<S>
where
S: CopyTarget,
{
type QueryId = ();
const HAS_STATIC_QUERY_ID: bool = false;
}
impl<S> QueryFragment<Pg> for CopyToCommand<S>
where
S: CopyTarget,
{
fn walk_ast<'b>(
&'b self,
mut pass: crate::query_builder::AstPass<'_, 'b, Pg>,
) -> crate::QueryResult<()> {
pass.unsafe_to_cache_prepared();
pass.push_sql("COPY ");
S::walk_target(pass.reborrow())?;
pass.push_sql(" TO STDOUT");
self.options.walk_ast(pass.reborrow())?;
Ok(())
}
}
#[derive(Debug, Clone, Copy)]
pub struct NotSet;
pub trait CopyToMarker: Sized {
fn setup_options<T>(q: CopyToQuery<T, Self>) -> CopyToQuery<T, CopyToOptions>;
}
impl CopyToMarker for NotSet {
fn setup_options<T>(q: CopyToQuery<T, Self>) -> CopyToQuery<T, CopyToOptions> {
CopyToQuery {
target: q.target,
options: CopyToOptions::default(),
}
}
}
impl CopyToMarker for CopyToOptions {
fn setup_options<T>(q: CopyToQuery<T, Self>) -> CopyToQuery<T, CopyToOptions> {
q
}
}
#[derive(Debug)]
#[must_use = "`COPY TO` statements are only executed when calling `.load()` or `load_raw()`."]
#[cfg(feature = "postgres_backend")]
pub struct CopyToQuery<T, O> {
target: T,
options: O,
}
#[cfg(feature = "postgres")]
struct CopyRow<'a> {
buffers: Vec<Option<&'a [u8]>>,
result: &'a crate::pg::connection::PgResult,
}
#[cfg(feature = "postgres")]
struct CopyField<'a> {
field: &'a Option<&'a [u8]>,
result: &'a crate::pg::connection::PgResult,
col_idx: usize,
}
#[cfg(feature = "postgres")]
impl<'f> Field<'f, Pg> for CopyField<'f> {
fn field_name(&self) -> Option<&str> {
None
}
fn value(&self) -> Option<<Pg as crate::backend::Backend>::RawValue<'_>> {
let value = self.field.as_deref()?;
Some(crate::pg::PgValue::new_internal(value, self))
}
}
#[cfg(feature = "postgres")]
impl TypeOidLookup for CopyField<'_> {
fn lookup(&self) -> core::num::NonZeroU32 {
self.result.column_type(self.col_idx)
}
}
#[cfg(feature = "postgres")]
impl RowSealed for CopyRow<'_> {}
#[cfg(feature = "postgres")]
impl RowIndex<usize> for CopyRow<'_> {
fn idx(&self, idx: usize) -> Option<usize> {
if idx < self.field_count() {
Some(idx)
} else {
None
}
}
}
#[cfg(feature = "postgres")]
impl<'a> RowIndex<&'a str> for CopyRow<'_> {
fn idx(&self, _idx: &'a str) -> Option<usize> {
None
}
}
#[cfg(feature = "postgres")]
impl<'a> Row<'a, Pg> for CopyRow<'_> {
type Field<'f>
= CopyField<'f>
where
'a: 'f,
Self: 'f;
type InnerPartialRow = Self;
fn field_count(&self) -> usize {
self.buffers.len()
}
fn get<'b, I>(&'b self, idx: I) -> Option<Self::Field<'b>>
where
'a: 'b,
Self: RowIndex<I>,
{
let idx = self.idx(idx)?;
let buffer = self.buffers.get(idx)?;
Some(CopyField {
field: buffer,
result: self.result,
col_idx: idx,
})
}
fn partial_row(
&self,
range: core::ops::Range<usize>,
) -> row::PartialRow<'_, Self::InnerPartialRow> {
PartialRow::new(self, range)
}
}
pub trait ExecuteCopyToConnection: Connection<Backend = Pg> {
type CopyToBuffer<'a>: BufRead;
fn make_row<'a, 'b>(
out: &'a Self::CopyToBuffer<'_>,
buffers: Vec<Option<&'a [u8]>>,
) -> impl Row<'b, Pg> + 'a;
fn get_buffer<'a>(out: &'a Self::CopyToBuffer<'_>) -> &'a [u8];
fn execute<T>(&mut self, command: CopyToCommand<T>) -> QueryResult<Self::CopyToBuffer<'_>>
where
T: CopyTarget;
}
#[cfg(feature = "postgres")]
impl ExecuteCopyToConnection for crate::PgConnection {
type CopyToBuffer<'a> = crate::pg::connection::copy::CopyToBuffer<'a>;
fn make_row<'a, 'b>(
out: &'a Self::CopyToBuffer<'_>,
buffers: Vec<Option<&'a [u8]>>,
) -> impl Row<'b, Pg> + 'a {
CopyRow {
buffers,
result: out.get_result(),
}
}
fn get_buffer<'a>(out: &'a Self::CopyToBuffer<'_>) -> &'a [u8] {
out.data_slice()
}
fn execute<T>(&mut self, command: CopyToCommand<T>) -> QueryResult<Self::CopyToBuffer<'_>>
where
T: CopyTarget,
{
self.copy_to(command)
}
}
#[cfg(feature = "r2d2")]
impl<C> ExecuteCopyToConnection for crate::r2d2::PooledConnection<crate::r2d2::ConnectionManager<C>>
where
C: ExecuteCopyToConnection + crate::r2d2::R2D2Connection + 'static,
{
type CopyToBuffer<'a> = C::CopyToBuffer<'a>;
fn make_row<'a, 'b>(
out: &'a Self::CopyToBuffer<'_>,
buffers: Vec<Option<&'a [u8]>>,
) -> impl Row<'b, Pg> + 'a {
C::make_row(out, buffers)
}
fn get_buffer<'a>(out: &'a Self::CopyToBuffer<'_>) -> &'a [u8] {
C::get_buffer(out)
}
fn execute<T>(&mut self, command: CopyToCommand<T>) -> QueryResult<Self::CopyToBuffer<'_>>
where
T: CopyTarget,
{
C::execute(&mut **self, command)
}
}
impl<T> CopyToQuery<T, NotSet>
where
T: CopyTarget,
{
pub fn load<U, C>(self, conn: &mut C) -> QueryResult<impl Iterator<Item = QueryResult<U>> + '_>
where
U: FromSqlRow<<U::SelectExpression as Expression>::SqlType, Pg> + Selectable<Pg>,
U::SelectExpression: AppearsOnTable<T::Table> + CopyTarget<Table = T::Table>,
C: ExecuteCopyToConnection,
{
let command = CopyToCommand {
p: PhantomData::<U::SelectExpression>,
options: CopyToOptions {
header: None,
common: CommonOptions {
format: Some(CopyFormat::Binary),
..Default::default()
},
},
};
let out = ExecuteCopyToConnection::execute(conn, command)?;
decode_data::<U, C>(out)
}
}
fn decode_data<'a, U, C>(
mut out: C::CopyToBuffer<'a>,
) -> QueryResult<impl Iterator<Item = QueryResult<U>> + 'a>
where
U: FromSqlRow<<U::SelectExpression as Expression>::SqlType, Pg> + Selectable<Pg>,
C: ExecuteCopyToConnection + 'a,
{
let io_result_mapper = |e| crate::result::Error::DeserializationError(Box::new(e));
let buffer = out.fill_buf().map_err(io_result_mapper)?;
expect_buffer_len(buffer, super::COPY_MAGIC_HEADER.len() + 8)?;
if buffer[..super::COPY_MAGIC_HEADER.len()] != super::COPY_MAGIC_HEADER {
return Err(crate::result::Error::DeserializationError(
"Unexpected protocol header".into(),
));
}
let flags_backward_incompatible = i16::from_be_bytes(
(&buffer[super::COPY_MAGIC_HEADER.len() + 2..super::COPY_MAGIC_HEADER.len() + 4])
.try_into()
.expect("Exactly 2 byte"),
);
if flags_backward_incompatible != 0 {
return Err(crate::result::Error::DeserializationError(
format!("Unexpected flag value: {flags_backward_incompatible:x}").into(),
));
}
let header_size = usize::try_from(i32::from_be_bytes(
(&buffer[super::COPY_MAGIC_HEADER.len() + 4..super::COPY_MAGIC_HEADER.len() + 8])
.try_into()
.expect("Exactly 4 byte"),
))
.map_err(|e| crate::result::Error::DeserializationError(Box::new(e)))?;
out.consume(super::COPY_MAGIC_HEADER.len() + 8 + header_size);
let mut len = None;
let mut finished = false;
Ok(core::iter::from_fn(move || {
if finished {
return None;
}
if let Some(len) = len {
out.consume(len);
if let Err(e) = out.fill_buf().map_err(io_result_mapper) {
finished = true;
return Some(Err(e));
}
}
let buffer = C::get_buffer(&out);
len = Some(buffer.len());
if let Err(e) = expect_buffer_len(buffer, 2) {
finished = true;
return Some(Err(e));
}
let tuple_count = i16::from_be_bytes((&buffer[..2]).try_into().expect("Exactly 2 bytes"));
if tuple_count > 0 {
let tuple_count = match usize::try_from(tuple_count) {
Ok(o) => o,
Err(e) => {
return Some(Err(crate::result::Error::DeserializationError(Box::new(e))));
}
};
let mut buffers = Vec::with_capacity(tuple_count);
let mut offset = 2;
for _t in 0..tuple_count {
if let Err(e) = expect_buffer_len(buffer, offset + 4) {
finished = true;
return Some(Err(e));
}
let data_size = i32::from_be_bytes(
(&buffer[offset..offset + 4])
.try_into()
.expect("Exactly 4 bytes"),
);
if data_size < 0 {
offset += 4;
buffers.push(None);
} else {
match usize::try_from(data_size) {
Ok(data_size) => {
if let Err(e) = expect_buffer_len(buffer, offset + 4 + data_size) {
finished = true;
return Some(Err(e));
}
buffers.push(Some(&buffer[offset + 4..offset + 4 + data_size]));
offset += data_size + 4;
}
Err(e) => {
return Some(Err(crate::result::Error::DeserializationError(
Box::new(e),
)));
}
}
}
}
let row = C::make_row(&out, buffers);
Some(U::build_from_row(&row).map_err(crate::result::Error::DeserializationError))
} else {
out.consume(len.expect("It's always set here"));
finished = true;
if let Err(e) = out.fill_buf().map_err(io_result_mapper) {
return Some(Err(e));
}
None
}
}))
}
fn expect_buffer_len(buffer: &[u8], expected: usize) -> QueryResult<()> {
if buffer.len() < expected {
Err(crate::result::Error::DeserializationError(
format!(
"Expected at least {expected} bytes but got {} bytes",
buffer.len()
)
.into(),
))
} else {
Ok(())
}
}
impl<T, O> CopyToQuery<T, O>
where
O: CopyToMarker,
T: CopyTarget,
{
pub fn load_raw<C>(self, conn: &mut C) -> QueryResult<impl BufRead + '_>
where
C: ExecuteCopyToConnection,
{
let q = O::setup_options(self);
let command = CopyToCommand {
p: PhantomData::<T>,
options: q.options,
};
ExecuteCopyToConnection::execute(conn, command)
}
pub fn with_format(self, format: CopyFormat) -> CopyToQuery<T, CopyToOptions> {
let mut out = O::setup_options(self);
out.options.common.format = Some(format);
out
}
pub fn with_freeze(self, freeze: bool) -> CopyToQuery<T, CopyToOptions> {
let mut out = O::setup_options(self);
out.options.common.freeze = Some(freeze);
out
}
pub fn with_delimiter(self, delimiter: char) -> CopyToQuery<T, CopyToOptions> {
let mut out = O::setup_options(self);
out.options.common.delimiter = Some(delimiter);
out
}
pub fn with_null(self, null: impl Into<String>) -> CopyToQuery<T, CopyToOptions> {
let mut out = O::setup_options(self);
out.options.common.null = Some(null.into());
out
}
pub fn with_quote(self, quote: char) -> CopyToQuery<T, CopyToOptions> {
let mut out = O::setup_options(self);
out.options.common.quote = Some(quote);
out
}
pub fn with_escape(self, escape: char) -> CopyToQuery<T, CopyToOptions> {
let mut out = O::setup_options(self);
out.options.common.escape = Some(escape);
out
}
pub fn with_header(self, set: bool) -> CopyToQuery<T, CopyToOptions> {
let mut out = O::setup_options(self);
out.options.header = Some(set);
out
}
}
#[cfg(feature = "postgres_backend")]
pub fn copy_to<T>(target: T) -> CopyToQuery<T, NotSet>
where
T: CopyTarget,
{
CopyToQuery {
target,
options: NotSet,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::connection::Instrumentation;
use crate::connection::{AnsiTransactionManager, TransactionManager};
use crate::pg::query_builder::copy::COPY_MAGIC_HEADER;
use std::collections::VecDeque;
use std::io::Read;
struct Wrapper;
struct Messages {
messages: VecDeque<Vec<u8>>,
offset: usize,
}
impl Messages {
fn new(messages: Vec<Vec<u8>>) -> Self {
Self {
messages: messages.into(),
offset: 0,
}
}
fn data_slice(&self) -> &[u8] {
self.messages
.front()
.map(|m| &m[self.offset..])
.unwrap_or_default()
}
}
impl Read for Messages {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let data = self.fill_buf()?;
let len = usize::min(buf.len(), data.len());
buf[..len].copy_from_slice(&data[..len]);
self.consume(len);
Ok(len)
}
}
impl BufRead for Messages {
fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
if self.data_slice().is_empty() && self.messages.pop_front().is_some() {
self.offset = 0;
}
Ok(self.data_slice())
}
fn consume(&mut self, amt: usize) {
let len = self.messages.front().map_or(0, Vec::len);
self.offset = usize::min(len, self.offset + amt);
}
}
impl crate::connection::SimpleConnection for Wrapper {
fn batch_execute(&mut self, _query: &str) -> QueryResult<()> {
unreachable!()
}
}
impl crate::connection::ConnectionSealed for Wrapper {}
impl crate::connection::Connection for Wrapper {
type Backend = crate::pg::Pg;
type TransactionManager = crate::connection::AnsiTransactionManager;
fn establish(_database_url: &str) -> crate::prelude::ConnectionResult<Self> {
unreachable!()
}
fn execute_returning_count<T>(&mut self, _source: &T) -> QueryResult<usize>
where
T: QueryFragment<Self::Backend> + QueryId,
{
unreachable!()
}
fn transaction_state(
&mut self,
) -> &mut <AnsiTransactionManager as TransactionManager<Self>>::TransactionStateData
{
unreachable!()
}
fn instrumentation(&mut self) -> &mut dyn Instrumentation {
unreachable!()
}
fn set_instrumentation(
&mut self,
_instrumentation: impl crate::connection::Instrumentation,
) {
unreachable!()
}
fn set_prepared_statement_cache_size(&mut self, _size: crate::connection::CacheSize) {
unreachable!()
}
}
impl ExecuteCopyToConnection for Wrapper {
type CopyToBuffer<'a> = Messages;
fn make_row<'a, 'b>(
_out: &'a Self::CopyToBuffer<'_>,
_buffers: Vec<Option<&'a [u8]>>,
) -> impl Row<'b, Pg> + 'a {
unreachable!() as CopyRow<'a>
}
fn get_buffer<'a>(out: &'a Self::CopyToBuffer<'_>) -> &'a [u8] {
out.data_slice()
}
fn execute<T>(&mut self, _command: CopyToCommand<T>) -> QueryResult<Self::CopyToBuffer<'_>>
where
T: CopyTarget,
{
unreachable!()
}
}
crate::table! {
users {
id -> Integer,
name -> Text,
hair_color -> Nullable<Text>,
age -> Nullable<Integer>,
}
}
#[derive(diesel::Queryable, diesel::Selectable, Debug)]
#[diesel(table_name = users)]
#[allow(dead_code, reason = "Only used as structure for the test")]
struct Dummy {
id: i32,
name: String,
hair_color: Option<String>,
age: Option<i32>,
}
fn decode(messages: Vec<Vec<u8>>) -> QueryResult<impl Iterator<Item = QueryResult<Dummy>>> {
decode_data::<Dummy, Wrapper>(Messages::new(messages))
}
fn tuple_count(n: i16) -> [u8; 2] {
n.to_be_bytes()
}
fn field_len(n: i32) -> [u8; 4] {
n.to_be_bytes()
}
#[track_caller]
fn first_row_error(first_message: &[&[u8]]) -> String {
let first_message = [
COPY_MAGIC_HEADER.as_slice(),
[0; 8].as_slice(),
&first_message.concat(),
]
.concat();
let mut rows = decode(vec![first_message]).unwrap();
let error = match rows.next() {
Some(Err(e)) => e.to_string(),
Some(Ok(row)) => panic!("expected an error, got {row:?}"),
None => panic!("expected an error, got the end of the data"),
};
assert!(
rows.next().is_none(),
"the iterator must end after an error"
);
error
}
#[track_caller]
fn check_decode_error(input: Vec<Vec<u8>>, error: &str) {
let r = decode(input);
assert!(r.is_err());
let Err(e) = r else { unreachable!() };
assert_eq!(e.to_string(), error);
}
#[test]
fn check_empty_buffer_returns_error() {
check_decode_error(
vec![Vec::new()],
"Expected at least 19 bytes but got 0 bytes",
);
}
#[test]
fn check_header_without_flags_returns_error() {
check_decode_error(
vec![COPY_MAGIC_HEADER.to_vec()],
"Expected at least 19 bytes but got 11 bytes",
);
}
#[test]
fn check_header_without_extension_length_returns_error() {
check_decode_error(
vec![[COPY_MAGIC_HEADER.as_slice(), &[0; 7]].concat().to_vec()],
"Expected at least 19 bytes but got 18 bytes",
);
}
#[test]
fn check_missing_trailer_returns_error() {
assert_eq!(
first_row_error(&[]),
"Expected at least 2 bytes but got 0 bytes",
);
}
#[test]
fn check_truncated_tuple_count_returns_error() {
assert_eq!(
first_row_error(&[&[0]]),
"Expected at least 2 bytes but got 1 bytes"
);
}
#[test]
fn check_missing_field_length_returns_error() {
assert_eq!(
first_row_error(&[&tuple_count(4)]),
"Expected at least 6 bytes but got 2 bytes"
);
}
#[test]
fn check_truncated_field_length_returns_error() {
assert_eq!(
first_row_error(&[&tuple_count(4), &[0, 0]]),
"Expected at least 6 bytes but got 4 bytes"
);
}
#[test]
fn check_missing_field_length_after_value_returns_error() {
assert_eq!(
first_row_error(&[&tuple_count(4), &field_len(4), &1_i32.to_be_bytes()]),
"Expected at least 14 bytes but got 10 bytes"
);
}
#[test]
fn check_missing_field_length_after_null_returns_error() {
assert_eq!(
first_row_error(&[&tuple_count(4), &field_len(-1)]),
"Expected at least 10 bytes but got 6 bytes"
);
}
#[test]
fn check_missing_field_data_returns_error() {
assert_eq!(
first_row_error(&[&tuple_count(4), &field_len(4)]),
"Expected at least 10 bytes but got 6 bytes"
);
}
#[test]
fn check_truncated_field_data_returns_error() {
assert_eq!(
first_row_error(&[&tuple_count(4), &field_len(100), b"abc"]),
"Expected at least 106 bytes but got 9 bytes"
);
}
#[test]
fn check_truncated_later_field_data_returns_error() {
assert_eq!(
first_row_error(&[
&tuple_count(4),
&field_len(4),
&1_i32.to_be_bytes(),
&field_len(5),
b"Sean",
]),
"Expected at least 19 bytes but got 18 bytes"
);
}
#[test]
fn check_field_length_i32_max_returns_error() {
assert_eq!(
first_row_error(&[&tuple_count(4), &field_len(i32::MAX), b"abc"]),
format!(
"Expected at least {} bytes but got 9 bytes",
6 + i32::MAX as usize
)
);
}
}