2015-08-16 06:21:39 +00:00
|
|
|
//! Prepared statements
|
|
|
|
|
2016-12-19 00:11:38 +00:00
|
|
|
use fallible_iterator::FallibleIterator;
|
2015-08-16 06:21:39 +00:00
|
|
|
use std::cell::{Cell, RefMut};
|
2015-08-16 04:39:46 +00:00
|
|
|
use std::collections::VecDeque;
|
|
|
|
use std::fmt;
|
2015-09-16 06:11:14 +00:00
|
|
|
use std::io::{self, Read, Write};
|
2016-01-03 05:02:54 +00:00
|
|
|
use std::sync::Arc;
|
2016-09-14 02:30:28 +00:00
|
|
|
use postgres_protocol::message::{backend, frontend};
|
2016-12-21 03:50:44 +00:00
|
|
|
use postgres_shared::RowData;
|
2015-08-16 04:39:46 +00:00
|
|
|
|
2016-12-21 03:50:44 +00:00
|
|
|
use error::Error;
|
2016-02-16 06:25:52 +00:00
|
|
|
use types::{SessionInfo, Type, ToSql};
|
2015-08-16 04:39:46 +00:00
|
|
|
use rows::{Rows, LazyRows};
|
2016-03-04 05:48:25 +00:00
|
|
|
use transaction::Transaction;
|
2016-12-21 03:50:44 +00:00
|
|
|
use {bad_response, err, Connection, StatementInternals, Result, RowsNew, InnerConnection,
|
2016-12-21 00:07:45 +00:00
|
|
|
SessionInfoNew, LazyRowsNew, ColumnNew, StatementInfo, TransactionInternals};
|
2015-08-16 04:39:46 +00:00
|
|
|
|
|
|
|
/// A prepared statement.
|
|
|
|
pub struct Statement<'conn> {
|
|
|
|
conn: &'conn Connection,
|
2016-01-03 05:02:54 +00:00
|
|
|
info: Arc<StatementInfo>,
|
2015-08-16 04:39:46 +00:00
|
|
|
next_portal_id: Cell<u32>,
|
|
|
|
finished: bool,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'a> fmt::Debug for Statement<'a> {
|
|
|
|
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
|
2015-09-19 05:55:32 +00:00
|
|
|
fmt.debug_struct("Statement")
|
2016-09-13 04:48:49 +00:00
|
|
|
.field("name", &self.info.name)
|
|
|
|
.field("parameter_types", &self.info.param_types)
|
|
|
|
.field("columns", &self.info.columns)
|
|
|
|
.finish()
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'conn> Drop for Statement<'conn> {
|
|
|
|
fn drop(&mut self) {
|
|
|
|
let _ = self.finish_inner();
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'conn> StatementInternals<'conn> for Statement<'conn> {
|
|
|
|
fn new(conn: &'conn Connection,
|
2016-01-03 05:02:54 +00:00
|
|
|
info: Arc<StatementInfo>,
|
2015-08-16 04:39:46 +00:00
|
|
|
next_portal_id: Cell<u32>,
|
2015-11-16 03:51:04 +00:00
|
|
|
finished: bool)
|
|
|
|
-> Statement<'conn> {
|
2015-08-16 04:39:46 +00:00
|
|
|
Statement {
|
|
|
|
conn: conn,
|
2016-01-03 05:02:54 +00:00
|
|
|
info: info,
|
2015-08-16 04:39:46 +00:00
|
|
|
next_portal_id: next_portal_id,
|
|
|
|
finished: finished,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
fn conn(&self) -> &'conn Connection {
|
|
|
|
self.conn
|
|
|
|
}
|
2015-11-14 03:34:11 +00:00
|
|
|
|
|
|
|
fn into_query(self, params: &[&ToSql]) -> Result<Rows<'conn>> {
|
|
|
|
check_desync!(self.conn);
|
2016-12-19 03:41:23 +00:00
|
|
|
let mut rows = vec![];
|
|
|
|
try!(self.inner_query("", 0, params, |row| rows.push(row)));
|
|
|
|
Ok(Rows::new_owned(self, rows))
|
2015-11-14 03:34:11 +00:00
|
|
|
}
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl<'conn> Statement<'conn> {
|
|
|
|
fn finish_inner(&mut self) -> Result<()> {
|
2016-04-06 03:52:34 +00:00
|
|
|
if self.finished {
|
|
|
|
Ok(())
|
|
|
|
} else {
|
2015-08-16 04:39:46 +00:00
|
|
|
self.finished = true;
|
2016-09-25 02:36:41 +00:00
|
|
|
let mut conn = self.conn.0.borrow_mut();
|
2015-08-16 04:39:46 +00:00
|
|
|
check_desync!(conn);
|
2016-01-03 05:02:54 +00:00
|
|
|
conn.close_statement(&self.info.name, b'S')
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2016-02-20 23:26:33 +00:00
|
|
|
#[allow(type_complexity)]
|
2016-12-19 03:41:23 +00:00
|
|
|
fn inner_query<F>(&self,
|
|
|
|
portal_name: &str,
|
|
|
|
row_limit: i32,
|
|
|
|
params: &[&ToSql],
|
|
|
|
acceptor: F)
|
|
|
|
-> Result<bool>
|
|
|
|
where F: FnMut(RowData)
|
|
|
|
{
|
2016-09-25 02:36:41 +00:00
|
|
|
let mut conn = self.conn.0.borrow_mut();
|
2016-02-16 06:25:52 +00:00
|
|
|
|
2016-02-22 04:02:34 +00:00
|
|
|
try!(conn.raw_execute(&self.info.name,
|
|
|
|
portal_name,
|
|
|
|
row_limit,
|
|
|
|
self.param_types(),
|
|
|
|
params));
|
2015-08-16 04:39:46 +00:00
|
|
|
|
2016-12-19 03:41:23 +00:00
|
|
|
conn.read_rows(acceptor)
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns a slice containing the expected parameter types.
|
|
|
|
pub fn param_types(&self) -> &[Type] {
|
2016-01-03 05:02:54 +00:00
|
|
|
&self.info.param_types
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns a slice describing the columns of the result of the query.
|
|
|
|
pub fn columns(&self) -> &[Column] {
|
2016-01-03 05:02:54 +00:00
|
|
|
&self.info.columns
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Executes the prepared statement, returning the number of rows modified.
|
|
|
|
///
|
|
|
|
/// If the statement does not modify any rows (e.g. SELECT), 0 is returned.
|
|
|
|
///
|
2015-11-29 05:05:13 +00:00
|
|
|
/// # Panics
|
2015-08-16 04:39:46 +00:00
|
|
|
///
|
|
|
|
/// Panics if the number of parameters provided does not match the number
|
|
|
|
/// expected.
|
|
|
|
///
|
2015-11-29 05:05:13 +00:00
|
|
|
/// # Example
|
2015-08-16 04:39:46 +00:00
|
|
|
///
|
|
|
|
/// ```rust,no_run
|
2016-07-01 19:21:12 +00:00
|
|
|
/// # use postgres::{Connection, TlsMode};
|
|
|
|
/// # let conn = Connection::connect("", TlsMode::None).unwrap();
|
2015-08-16 04:39:46 +00:00
|
|
|
/// # let bar = 1i32;
|
|
|
|
/// # let baz = true;
|
|
|
|
/// let stmt = conn.prepare("UPDATE foo SET bar = $1 WHERE baz = $2").unwrap();
|
2015-11-30 04:00:25 +00:00
|
|
|
/// let rows_updated = stmt.execute(&[&bar, &baz]).unwrap();
|
|
|
|
/// println!("{} rows updated", rows_updated);
|
2015-08-16 04:39:46 +00:00
|
|
|
/// ```
|
|
|
|
pub fn execute(&self, params: &[&ToSql]) -> Result<u64> {
|
2016-09-25 02:36:41 +00:00
|
|
|
let mut conn = self.conn.0.borrow_mut();
|
2016-02-16 06:25:52 +00:00
|
|
|
check_desync!(conn);
|
|
|
|
try!(conn.raw_execute(&self.info.name, "", 0, self.param_types(), params));
|
|
|
|
|
2015-08-16 04:39:46 +00:00
|
|
|
let num;
|
|
|
|
loop {
|
|
|
|
match try!(conn.read_message()) {
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::DataRow(_) => {}
|
|
|
|
backend::Message::ErrorResponse(body) => {
|
2015-08-16 04:39:46 +00:00
|
|
|
try!(conn.wait_for_ready());
|
2016-12-21 03:50:44 +00:00
|
|
|
return Err(err(&mut body.fields()));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CommandComplete(body) => {
|
|
|
|
num = parse_update_count(try!(body.tag()));
|
2015-08-16 04:39:46 +00:00
|
|
|
break;
|
|
|
|
}
|
2016-09-14 02:30:28 +00:00
|
|
|
backend::Message::EmptyQueryResponse => {
|
2015-08-16 04:39:46 +00:00
|
|
|
num = 0;
|
|
|
|
break;
|
|
|
|
}
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CopyInResponse(_) => {
|
2016-09-20 20:58:04 +00:00
|
|
|
try!(conn.stream.write_message(|buf| {
|
2016-09-20 13:40:59 +00:00
|
|
|
frontend::copy_fail("COPY queries cannot be directly executed", buf)
|
2016-09-11 23:48:19 +00:00
|
|
|
}));
|
2016-09-20 20:33:26 +00:00
|
|
|
try!(conn.stream
|
2016-09-20 20:58:04 +00:00
|
|
|
.write_message(|buf| Ok::<(), io::Error>(frontend::sync(buf))));
|
2016-09-11 23:48:19 +00:00
|
|
|
try!(conn.stream.flush());
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CopyOutResponse(_) => {
|
2015-08-16 04:39:46 +00:00
|
|
|
loop {
|
|
|
|
match try!(conn.read_message()) {
|
2016-09-14 02:30:28 +00:00
|
|
|
backend::Message::CopyDone => break,
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::ErrorResponse(body) => {
|
2015-08-16 04:39:46 +00:00
|
|
|
try!(conn.wait_for_ready());
|
2016-12-21 03:50:44 +00:00
|
|
|
return Err(err(&mut body.fields()));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
_ => {}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
num = 0;
|
|
|
|
break;
|
|
|
|
}
|
|
|
|
_ => {
|
|
|
|
conn.desynchronized = true;
|
2015-12-26 03:14:09 +00:00
|
|
|
return Err(Error::Io(bad_response()));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
try!(conn.wait_for_ready());
|
|
|
|
|
|
|
|
Ok(num)
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Executes the prepared statement, returning the resulting rows.
|
|
|
|
///
|
2015-11-29 05:05:13 +00:00
|
|
|
/// # Panics
|
2015-08-16 04:39:46 +00:00
|
|
|
///
|
|
|
|
/// Panics if the number of parameters provided does not match the number
|
|
|
|
/// expected.
|
|
|
|
///
|
2015-11-29 05:05:13 +00:00
|
|
|
/// # Example
|
2015-08-16 04:39:46 +00:00
|
|
|
///
|
|
|
|
/// ```rust,no_run
|
2016-07-01 19:21:12 +00:00
|
|
|
/// # use postgres::{Connection, TlsMode};
|
|
|
|
/// # let conn = Connection::connect("", TlsMode::None).unwrap();
|
2015-08-16 04:39:46 +00:00
|
|
|
/// let stmt = conn.prepare("SELECT foo FROM bar WHERE baz = $1").unwrap();
|
|
|
|
/// # let baz = true;
|
2015-11-13 06:56:05 +00:00
|
|
|
/// for row in &stmt.query(&[&baz]).unwrap() {
|
2015-08-16 04:39:46 +00:00
|
|
|
/// let foo: i32 = row.get("foo");
|
|
|
|
/// println!("foo: {}", foo);
|
|
|
|
/// }
|
|
|
|
/// ```
|
|
|
|
pub fn query<'a>(&'a self, params: &[&ToSql]) -> Result<Rows<'a>> {
|
|
|
|
check_desync!(self.conn);
|
2016-12-19 03:41:23 +00:00
|
|
|
let mut rows = vec![];
|
|
|
|
try!(self.inner_query("", 0, params, |row| rows.push(row)));
|
|
|
|
Ok(Rows::new(self, rows))
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Executes the prepared statement, returning a lazily loaded iterator
|
|
|
|
/// over the resulting rows.
|
|
|
|
///
|
|
|
|
/// No more than `row_limit` rows will be stored in memory at a time. Rows
|
|
|
|
/// will be pulled from the database in batches of `row_limit` as needed.
|
|
|
|
/// If `row_limit` is less than or equal to 0, `lazy_query` is equivalent
|
|
|
|
/// to `query`.
|
|
|
|
///
|
|
|
|
/// This can only be called inside of a transaction, and the `Transaction`
|
|
|
|
/// object representing the active transaction must be passed to
|
|
|
|
/// `lazy_query`.
|
|
|
|
///
|
2015-11-29 05:05:13 +00:00
|
|
|
/// # Panics
|
2015-08-16 04:39:46 +00:00
|
|
|
///
|
|
|
|
/// Panics if the provided `Transaction` is not associated with the same
|
|
|
|
/// `Connection` as this `Statement`, if the `Transaction` is not
|
|
|
|
/// active, or if the number of parameters provided does not match the
|
|
|
|
/// number of parameters expected.
|
|
|
|
pub fn lazy_query<'trans, 'stmt>(&'stmt self,
|
|
|
|
trans: &'trans Transaction,
|
|
|
|
params: &[&ToSql],
|
|
|
|
row_limit: i32)
|
|
|
|
-> Result<LazyRows<'trans, 'stmt>> {
|
2016-02-24 07:43:28 +00:00
|
|
|
assert!(self.conn as *const _ == trans.conn() as *const _,
|
2015-08-16 04:39:46 +00:00
|
|
|
"the `Transaction` passed to `lazy_query` must be associated with the same \
|
|
|
|
`Connection` as the `Statement`");
|
2016-09-25 02:36:41 +00:00
|
|
|
let conn = self.conn.0.borrow();
|
2015-08-16 04:39:46 +00:00
|
|
|
check_desync!(conn);
|
2016-02-24 07:43:28 +00:00
|
|
|
assert!(conn.trans_depth == trans.depth(),
|
2015-08-16 04:39:46 +00:00
|
|
|
"`lazy_query` must be passed the active transaction");
|
|
|
|
drop(conn);
|
|
|
|
|
|
|
|
let id = self.next_portal_id.get();
|
|
|
|
self.next_portal_id.set(id + 1);
|
2016-01-03 05:02:54 +00:00
|
|
|
let portal_name = format!("{}p{}", self.info.name, id);
|
2015-08-16 04:39:46 +00:00
|
|
|
|
2016-12-19 03:41:23 +00:00
|
|
|
let mut rows = VecDeque::new();
|
|
|
|
let more_rows = try!(self.inner_query(&portal_name,
|
|
|
|
row_limit,
|
|
|
|
params,
|
|
|
|
|row| rows.push_back(row)));
|
|
|
|
Ok(LazyRows::new(self, rows, portal_name, row_limit, more_rows, false, trans))
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Executes a `COPY FROM STDIN` statement, returning the number of rows
|
|
|
|
/// added.
|
|
|
|
///
|
|
|
|
/// The contents of the provided reader are passed to the Postgres server
|
|
|
|
/// verbatim; it is the caller's responsibility to ensure it uses the
|
|
|
|
/// proper format. See the
|
|
|
|
/// [Postgres documentation](http://www.postgresql.org/docs/9.4/static/sql-copy.html)
|
|
|
|
/// for details.
|
|
|
|
///
|
|
|
|
/// If the statement is not a `COPY FROM STDIN` statement it will still be
|
|
|
|
/// executed and this method will return an error.
|
|
|
|
///
|
|
|
|
/// # Examples
|
|
|
|
///
|
|
|
|
/// ```rust,no_run
|
2016-07-01 19:21:12 +00:00
|
|
|
/// # use postgres::{Connection, TlsMode};
|
|
|
|
/// # let conn = Connection::connect("", TlsMode::None).unwrap();
|
2015-08-16 04:39:46 +00:00
|
|
|
/// conn.batch_execute("CREATE TABLE people (id INT PRIMARY KEY, name VARCHAR)").unwrap();
|
|
|
|
/// let stmt = conn.prepare("COPY people FROM STDIN").unwrap();
|
|
|
|
/// stmt.copy_in(&[], &mut "1\tjohn\n2\tjane\n".as_bytes()).unwrap();
|
|
|
|
/// ```
|
|
|
|
pub fn copy_in<R: ReadWithInfo>(&self, params: &[&ToSql], r: &mut R) -> Result<u64> {
|
2016-09-25 02:36:41 +00:00
|
|
|
let mut conn = self.conn.0.borrow_mut();
|
2016-02-16 06:25:52 +00:00
|
|
|
try!(conn.raw_execute(&self.info.name, "", 0, self.param_types(), params));
|
2015-08-16 04:39:46 +00:00
|
|
|
|
2015-08-22 20:50:20 +00:00
|
|
|
let (format, column_formats) = match try!(conn.read_message()) {
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CopyInResponse(body) => {
|
|
|
|
let format = body.format();
|
|
|
|
let column_formats = try!(body.column_formats()
|
|
|
|
.map(|f| Format::from_u16(f))
|
|
|
|
.collect());
|
|
|
|
(format, column_formats)
|
|
|
|
}
|
|
|
|
backend::Message::ErrorResponse(body) => {
|
2015-10-24 05:55:39 +00:00
|
|
|
try!(conn.wait_for_ready());
|
2016-12-21 03:50:44 +00:00
|
|
|
return Err(err(&mut body.fields()));
|
2015-10-24 05:55:39 +00:00
|
|
|
}
|
2015-08-16 04:39:46 +00:00
|
|
|
_ => {
|
|
|
|
loop {
|
2016-12-19 00:11:38 +00:00
|
|
|
if let backend::Message::ReadyForQuery(_) = try!(conn.read_message()) {
|
2016-02-20 23:05:48 +00:00
|
|
|
return Err(Error::Io(io::Error::new(io::ErrorKind::InvalidInput,
|
|
|
|
"called `copy_in` on a \
|
|
|
|
non-`COPY FROM STDIN` \
|
|
|
|
statement")));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2015-08-22 20:50:20 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
let mut info = CopyInfo {
|
|
|
|
conn: conn,
|
|
|
|
format: Format::from_u16(format as u16),
|
2016-12-19 00:11:38 +00:00
|
|
|
column_formats: column_formats,
|
2015-08-22 20:50:20 +00:00
|
|
|
};
|
2015-08-16 04:39:46 +00:00
|
|
|
|
|
|
|
let mut buf = [0; 16 * 1024];
|
|
|
|
loop {
|
2015-08-22 20:50:20 +00:00
|
|
|
match fill_copy_buf(&mut buf, r, &info) {
|
2015-08-16 04:39:46 +00:00
|
|
|
Ok(0) => break,
|
|
|
|
Ok(len) => {
|
2016-09-20 20:58:04 +00:00
|
|
|
try!(info
|
|
|
|
.conn
|
|
|
|
.stream.write_message(|out| frontend::copy_data(&buf[..len], out)));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
Err(err) => {
|
2016-09-20 20:58:04 +00:00
|
|
|
try!(info.conn.stream.write_message(|buf| frontend::copy_fail("", buf)));
|
2016-09-20 20:33:26 +00:00
|
|
|
try!(info.conn
|
|
|
|
.stream
|
2016-09-20 20:58:04 +00:00
|
|
|
.write_message(|buf| Ok::<(), io::Error>(frontend::copy_done(buf))));
|
2016-09-20 20:33:26 +00:00
|
|
|
try!(info.conn
|
|
|
|
.stream
|
2016-09-20 20:58:04 +00:00
|
|
|
.write_message(|buf| Ok::<(), io::Error>(frontend::sync(buf))));
|
2016-09-11 23:48:19 +00:00
|
|
|
try!(info.conn.stream.flush());
|
2015-08-22 20:50:20 +00:00
|
|
|
match try!(info.conn.read_message()) {
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::ErrorResponse(_) => {
|
2015-11-16 03:51:04 +00:00
|
|
|
// expected from the CopyFail
|
|
|
|
}
|
2015-08-16 04:39:46 +00:00
|
|
|
_ => {
|
2015-08-22 20:50:20 +00:00
|
|
|
info.conn.desynchronized = true;
|
2015-12-26 03:14:09 +00:00
|
|
|
return Err(Error::Io(bad_response()));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
}
|
2015-08-22 20:50:20 +00:00
|
|
|
try!(info.conn.wait_for_ready());
|
2015-12-26 03:14:09 +00:00
|
|
|
return Err(Error::Io(err));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2016-09-20 20:58:04 +00:00
|
|
|
try!(info.conn.stream.write_message(|buf| Ok::<(), io::Error>(frontend::copy_done(buf))));
|
|
|
|
try!(info.conn.stream.write_message(|buf| Ok::<(), io::Error>(frontend::sync(buf))));
|
2016-09-11 23:48:19 +00:00
|
|
|
try!(info.conn.stream.flush());
|
2015-08-16 04:39:46 +00:00
|
|
|
|
2015-08-22 20:50:20 +00:00
|
|
|
let num = match try!(info.conn.read_message()) {
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CommandComplete(body) => parse_update_count(try!(body.tag())),
|
|
|
|
backend::Message::ErrorResponse(body) => {
|
2015-08-22 20:50:20 +00:00
|
|
|
try!(info.conn.wait_for_ready());
|
2016-12-21 03:50:44 +00:00
|
|
|
return Err(err(&mut body.fields()));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
_ => {
|
2015-08-22 20:50:20 +00:00
|
|
|
info.conn.desynchronized = true;
|
2015-12-26 03:14:09 +00:00
|
|
|
return Err(Error::Io(bad_response()));
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
};
|
|
|
|
|
2015-08-22 20:50:20 +00:00
|
|
|
try!(info.conn.wait_for_ready());
|
2015-08-16 04:39:46 +00:00
|
|
|
Ok(num)
|
|
|
|
}
|
|
|
|
|
2015-09-16 06:11:14 +00:00
|
|
|
/// Executes a `COPY TO STDOUT` statement, passing the resulting data to
|
|
|
|
/// the provided writer and returning the number of rows received.
|
2015-08-16 06:21:39 +00:00
|
|
|
///
|
|
|
|
/// See the [Postgres documentation](http://www.postgresql.org/docs/9.4/static/sql-copy.html)
|
|
|
|
/// for details on the data format.
|
|
|
|
///
|
|
|
|
/// If the statement is not a `COPY TO STDOUT` statement it will still be
|
|
|
|
/// executed and this method will return an error.
|
|
|
|
///
|
|
|
|
/// # Examples
|
|
|
|
///
|
|
|
|
/// ```rust,no_run
|
2016-07-01 19:21:12 +00:00
|
|
|
/// # use postgres::{Connection, TlsMode};
|
|
|
|
/// # let conn = Connection::connect("", TlsMode::None).unwrap();
|
2015-08-16 06:21:39 +00:00
|
|
|
/// conn.batch_execute("
|
|
|
|
/// CREATE TABLE people (id INT PRIMARY KEY, name VARCHAR);
|
|
|
|
/// INSERT INTO people (id, name) VALUES (1, 'john'), (2, 'jane');").unwrap();
|
|
|
|
/// let stmt = conn.prepare("COPY people TO STDOUT").unwrap();
|
|
|
|
/// let mut buf = vec![];
|
2015-11-30 05:43:47 +00:00
|
|
|
/// stmt.copy_out(&[], &mut buf).unwrap();
|
2015-08-16 06:21:39 +00:00
|
|
|
/// assert_eq!(buf, b"1\tjohn\n2\tjane\n");
|
|
|
|
/// ```
|
2015-09-16 06:11:14 +00:00
|
|
|
pub fn copy_out<'a, W: WriteWithInfo>(&'a self, params: &[&ToSql], w: &mut W) -> Result<u64> {
|
2016-09-25 02:36:41 +00:00
|
|
|
let mut conn = self.conn.0.borrow_mut();
|
2016-02-16 06:25:52 +00:00
|
|
|
try!(conn.raw_execute(&self.info.name, "", 0, self.param_types(), params));
|
2015-08-16 06:21:39 +00:00
|
|
|
|
|
|
|
let (format, column_formats) = match try!(conn.read_message()) {
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CopyOutResponse(body) => {
|
|
|
|
let format = body.format();
|
|
|
|
let column_formats = try!(body.column_formats()
|
|
|
|
.map(|f| Format::from_u16(f))
|
|
|
|
.collect());
|
2016-09-16 03:51:26 +00:00
|
|
|
(format, column_formats)
|
|
|
|
}
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CopyInResponse(_) => {
|
2016-09-20 20:58:04 +00:00
|
|
|
try!(conn.stream.write_message(|buf| frontend::copy_fail("", buf)));
|
2016-09-20 20:33:26 +00:00
|
|
|
try!(conn.stream
|
2016-09-20 20:58:04 +00:00
|
|
|
.write_message(|buf| Ok::<(), io::Error>(frontend::copy_done(buf))));
|
|
|
|
try!(conn.stream.write_message(|buf| Ok::<(), io::Error>(frontend::sync(buf))));
|
2016-09-11 23:48:19 +00:00
|
|
|
try!(conn.stream.flush());
|
2015-08-16 06:21:39 +00:00
|
|
|
match try!(conn.read_message()) {
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::ErrorResponse(_) => {
|
2015-11-16 03:51:04 +00:00
|
|
|
// expected from the CopyFail
|
|
|
|
}
|
2015-08-16 06:21:39 +00:00
|
|
|
_ => {
|
|
|
|
conn.desynchronized = true;
|
2015-12-26 03:14:09 +00:00
|
|
|
return Err(Error::Io(bad_response()));
|
2015-08-16 06:21:39 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
try!(conn.wait_for_ready());
|
2015-12-26 03:14:09 +00:00
|
|
|
return Err(Error::Io(io::Error::new(io::ErrorKind::InvalidInput,
|
2015-12-26 03:20:28 +00:00
|
|
|
"called `copy_out` on a non-`COPY TO \
|
|
|
|
STDOUT` statement")));
|
2015-08-16 06:21:39 +00:00
|
|
|
}
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::ErrorResponse(body) => {
|
2015-10-24 05:55:39 +00:00
|
|
|
try!(conn.wait_for_ready());
|
2016-12-21 03:50:44 +00:00
|
|
|
return Err(err(&mut body.fields()));
|
2015-10-24 05:55:39 +00:00
|
|
|
}
|
2015-08-16 06:21:39 +00:00
|
|
|
_ => {
|
|
|
|
loop {
|
2016-12-19 00:11:38 +00:00
|
|
|
if let backend::Message::ReadyForQuery(_) = try!(conn.read_message()) {
|
2016-02-20 23:05:48 +00:00
|
|
|
return Err(Error::Io(io::Error::new(io::ErrorKind::InvalidInput,
|
|
|
|
"called `copy_out` on a \
|
|
|
|
non-`COPY TO STDOUT` statement")));
|
2015-08-16 06:21:39 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
};
|
|
|
|
|
2015-09-16 06:11:14 +00:00
|
|
|
let mut info = CopyInfo {
|
|
|
|
conn: conn,
|
|
|
|
format: Format::from_u16(format as u16),
|
2016-12-19 00:11:38 +00:00
|
|
|
column_formats: column_formats,
|
2015-09-16 06:11:14 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
let count;
|
|
|
|
loop {
|
|
|
|
match try!(info.conn.read_message()) {
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CopyData(body) => {
|
|
|
|
let mut data = body.data();
|
2015-09-16 06:11:14 +00:00
|
|
|
while !data.is_empty() {
|
|
|
|
match w.write_with_info(data, &info) {
|
|
|
|
Ok(n) => data = &data[n..],
|
|
|
|
Err(e) => {
|
|
|
|
loop {
|
2016-12-19 00:11:38 +00:00
|
|
|
if let backend::Message::ReadyForQuery(_) =
|
2016-09-20 20:33:26 +00:00
|
|
|
try!(info.conn.read_message()) {
|
2016-02-20 23:05:48 +00:00
|
|
|
return Err(Error::Io(e));
|
2015-09-16 06:11:14 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2016-09-14 02:30:28 +00:00
|
|
|
backend::Message::CopyDone => {}
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::CommandComplete(body) => {
|
|
|
|
count = parse_update_count(try!(body.tag()));
|
2015-09-16 06:11:14 +00:00
|
|
|
break;
|
|
|
|
}
|
2016-12-19 00:11:38 +00:00
|
|
|
backend::Message::ErrorResponse(body) => {
|
2015-09-16 06:11:14 +00:00
|
|
|
loop {
|
2016-12-19 00:11:38 +00:00
|
|
|
if let backend::Message::ReadyForQuery(_) =
|
2016-09-20 20:33:26 +00:00
|
|
|
try!(info.conn.read_message()) {
|
2016-12-21 03:50:44 +00:00
|
|
|
return Err(err(&mut body.fields()));
|
2015-09-16 06:11:14 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
_ => {
|
|
|
|
loop {
|
2016-12-19 00:11:38 +00:00
|
|
|
if let backend::Message::ReadyForQuery(_) =
|
2016-09-20 20:33:26 +00:00
|
|
|
try!(info.conn.read_message()) {
|
2016-02-20 23:05:48 +00:00
|
|
|
return Err(Error::Io(bad_response()));
|
2015-09-16 06:11:14 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
try!(info.conn.wait_for_ready());
|
|
|
|
Ok(count)
|
2015-08-16 06:21:39 +00:00
|
|
|
}
|
|
|
|
|
2015-08-16 04:39:46 +00:00
|
|
|
/// Consumes the statement, clearing it from the Postgres session.
|
|
|
|
///
|
|
|
|
/// If this statement was created via the `prepare_cached` method, `finish`
|
|
|
|
/// does nothing.
|
|
|
|
///
|
|
|
|
/// Functionally identical to the `Drop` implementation of the
|
|
|
|
/// `Statement` except that it returns any error to the caller.
|
|
|
|
pub fn finish(mut self) -> Result<()> {
|
|
|
|
self.finish_inner()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2015-11-16 03:51:04 +00:00
|
|
|
fn fill_copy_buf<R: ReadWithInfo>(buf: &mut [u8], r: &mut R, info: &CopyInfo) -> io::Result<usize> {
|
2015-08-16 04:39:46 +00:00
|
|
|
let mut nread = 0;
|
|
|
|
while nread < buf.len() {
|
|
|
|
match r.read_with_info(&mut buf[nread..], info) {
|
|
|
|
Ok(0) => break,
|
|
|
|
Ok(n) => nread += n,
|
|
|
|
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
|
|
|
|
Err(e) => return Err(e),
|
|
|
|
}
|
|
|
|
}
|
|
|
|
Ok(nread)
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Information about a column of the result of a query.
|
|
|
|
#[derive(PartialEq, Eq, Clone, Debug)]
|
|
|
|
pub struct Column {
|
|
|
|
name: String,
|
2015-11-16 03:51:04 +00:00
|
|
|
type_: Type,
|
2015-08-16 04:39:46 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
impl ColumnNew for Column {
|
|
|
|
fn new(name: String, type_: Type) -> Column {
|
|
|
|
Column {
|
|
|
|
name: name,
|
|
|
|
type_: type_,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2015-12-16 06:12:43 +00:00
|
|
|
impl Column {
|
|
|
|
/// The name of the column.
|
|
|
|
pub fn name(&self) -> &str {
|
|
|
|
&self.name
|
|
|
|
}
|
|
|
|
|
|
|
|
/// The type of the data in the column.
|
|
|
|
pub fn type_(&self) -> &Type {
|
|
|
|
&self.type_
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2015-08-22 20:50:20 +00:00
|
|
|
/// A struct containing information relevant for a `COPY` operation.
|
|
|
|
pub struct CopyInfo<'a> {
|
|
|
|
conn: RefMut<'a, InnerConnection>,
|
|
|
|
format: Format,
|
|
|
|
column_formats: Vec<Format>,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<'a> CopyInfo<'a> {
|
|
|
|
/// Returns the format of the overall data.
|
|
|
|
pub fn format(&self) -> Format {
|
|
|
|
self.format
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the format of the individual columns.
|
|
|
|
pub fn column_formats(&self) -> &[Format] {
|
|
|
|
&self.column_formats
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns session info for the associated connection.
|
|
|
|
pub fn session_info<'b>(&'b self) -> SessionInfo<'b> {
|
2016-09-20 03:46:25 +00:00
|
|
|
SessionInfo::new(&self.conn.parameters)
|
2015-08-22 20:50:20 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Like `Read` except that a `CopyInfo` object is provided as well.
|
|
|
|
///
|
|
|
|
/// All types that implement `Read` also implement this trait.
|
|
|
|
pub trait ReadWithInfo {
|
|
|
|
/// Like `Read::read`.
|
|
|
|
fn read_with_info(&mut self, buf: &mut [u8], info: &CopyInfo) -> io::Result<usize>;
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<R: Read> ReadWithInfo for R {
|
|
|
|
fn read_with_info(&mut self, buf: &mut [u8], _: &CopyInfo) -> io::Result<usize> {
|
|
|
|
self.read(buf)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2015-09-16 06:11:14 +00:00
|
|
|
/// Like `Write` except that a `CopyInfo` object is provided as well.
|
|
|
|
///
|
|
|
|
/// All types that implement `Write` also implement this trait.
|
|
|
|
pub trait WriteWithInfo {
|
|
|
|
/// Like `Write::write`.
|
|
|
|
fn write_with_info(&mut self, buf: &[u8], info: &CopyInfo) -> io::Result<usize>;
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<W: Write> WriteWithInfo for W {
|
|
|
|
fn write_with_info(&mut self, buf: &[u8], _: &CopyInfo) -> io::Result<usize> {
|
|
|
|
self.write(buf)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2015-08-16 06:21:39 +00:00
|
|
|
/// The format of a portion of COPY query data.
|
|
|
|
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
|
|
|
|
pub enum Format {
|
|
|
|
/// A text based format.
|
|
|
|
Text,
|
|
|
|
/// A binary format.
|
|
|
|
Binary,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl Format {
|
|
|
|
fn from_u16(value: u16) -> Format {
|
|
|
|
match value {
|
|
|
|
0 => Format::Text,
|
|
|
|
_ => Format::Binary,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
2016-03-04 05:21:00 +00:00
|
|
|
|
2016-12-19 00:11:38 +00:00
|
|
|
fn parse_update_count(tag: &str) -> u64 {
|
2016-03-04 05:21:00 +00:00
|
|
|
tag.split(' ').last().unwrap().parse().unwrap_or(0)
|
|
|
|
}
|