use fallible_iterator::FallibleIterator;
use postgres_protocol::message::frontend;
use postgres_shared::rows::RowData;
use std::ascii::AsciiExt;
use std::collections::VecDeque;
use std::fmt;
use std::io;
use std::ops::Deref;
use std::slice;
use {Result, RowsNew, LazyRowsNew, StatementInternals};
use transaction::Transaction;
use types::{FromSql, SessionInfo, WrongType};
use stmt::{Statement, Column};
use error::Error;
enum MaybeOwned<'a, T: 'a> {
Borrowed(&'a T),
Owned(T),
}
impl<'a, T> Deref for MaybeOwned<'a, T> {
type Target = T;
fn deref(&self) -> &T {
match *self {
MaybeOwned::Borrowed(s) => s,
MaybeOwned::Owned(ref s) => s,
}
}
}
pub struct Rows<'stmt> {
stmt: MaybeOwned<'stmt, Statement<'stmt>>,
data: Vec<RowData>,
}
impl<'a> RowsNew<'a> for Rows<'a> {
fn new(stmt: &'a Statement<'a>, data: Vec<RowData>) -> Rows<'a> {
Rows {
stmt: MaybeOwned::Borrowed(stmt),
data: data,
}
}
fn new_owned(stmt: Statement<'a>, data: Vec<RowData>) -> Rows<'a> {
Rows {
stmt: MaybeOwned::Owned(stmt),
data: data,
}
}
}
impl<'a> fmt::Debug for Rows<'a> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("Rows")
.field("columns", &self.columns())
.field("rows", &self.data.len())
.finish()
}
}
impl<'stmt> Rows<'stmt> {
pub fn columns(&self) -> &[Column] {
self.stmt.columns()
}
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn get<'a>(&'a self, idx: usize) -> Row<'a> {
Row {
stmt: &*self.stmt,
data: MaybeOwned::Borrowed(&self.data[idx]),
}
}
pub fn iter<'a>(&'a self) -> Iter<'a> {
Iter {
stmt: &*self.stmt,
iter: self.data.iter(),
}
}
}
impl<'a> IntoIterator for &'a Rows<'a> {
type Item = Row<'a>;
type IntoIter = Iter<'a>;
fn into_iter(self) -> Iter<'a> {
self.iter()
}
}
pub struct Iter<'a> {
stmt: &'a Statement<'a>,
iter: slice::Iter<'a, RowData>,
}
impl<'a> Iterator for Iter<'a> {
type Item = Row<'a>;
fn next(&mut self) -> Option<Row<'a>> {
self.iter.next().map(|row| {
Row {
stmt: &*self.stmt,
data: MaybeOwned::Borrowed(row),
}
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
}
impl<'a> DoubleEndedIterator for Iter<'a> {
fn next_back(&mut self) -> Option<Row<'a>> {
self.iter.next_back().map(|row| {
Row {
stmt: &*self.stmt,
data: MaybeOwned::Borrowed(row),
}
})
}
}
impl<'a> ExactSizeIterator for Iter<'a> {}
pub struct Row<'a> {
stmt: &'a Statement<'a>,
data: MaybeOwned<'a, RowData>,
}
impl<'a> fmt::Debug for Row<'a> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("Row")
.field("statement", self.stmt)
.finish()
}
}
impl<'a> Row<'a> {
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn columns(&self) -> &[Column] {
self.stmt.columns()
}
pub fn get<I, T>(&self, idx: I) -> T
where I: RowIndex + fmt::Debug,
T: FromSql
{
match self.get_inner(&idx) {
Some(Ok(ok)) => ok,
Some(Err(err)) => panic!("error retrieving column {:?}: {:?}", idx, err),
None => panic!("no such column {:?}", idx),
}
}
pub fn get_opt<I, T>(&self, idx: I) -> Option<Result<T>>
where I: RowIndex,
T: FromSql
{
self.get_inner(&idx)
}
fn get_inner<I, T>(&self, idx: &I) -> Option<Result<T>>
where I: RowIndex,
T: FromSql
{
let idx = match idx.idx(self.stmt) {
Some(idx) => idx,
None => return None,
};
let ty = self.stmt.columns()[idx].type_();
if !<T as FromSql>::accepts(ty) {
return Some(Err(Error::Conversion(Box::new(WrongType::new(ty.clone())))));
}
let conn = self.stmt.conn().0.borrow();
let value = FromSql::from_sql_nullable(ty,
self.data.get(idx),
&SessionInfo::new(&conn.parameters));
Some(value.map_err(Error::Conversion))
}
pub fn get_bytes<I>(&self, idx: I) -> Option<&[u8]>
where I: RowIndex + fmt::Debug
{
match idx.idx(self.stmt) {
Some(idx) => self.data.get(idx),
None => panic!("invalid index {:?}", idx),
}
}
}
pub trait RowIndex {
fn idx(&self, stmt: &Statement) -> Option<usize>;
}
impl RowIndex for usize {
#[inline]
fn idx(&self, stmt: &Statement) -> Option<usize> {
if *self >= stmt.columns().len() {
None
} else {
Some(*self)
}
}
}
impl<'a> RowIndex for &'a str {
#[inline]
fn idx(&self, stmt: &Statement) -> Option<usize> {
if let Some(idx) = stmt.columns().iter().position(|d| d.name() == *self) {
return Some(idx);
};
stmt.columns().iter().position(|d| d.name().eq_ignore_ascii_case(*self))
}
}
pub struct LazyRows<'trans, 'stmt> {
stmt: &'stmt Statement<'stmt>,
data: VecDeque<RowData>,
name: String,
row_limit: i32,
more_rows: bool,
finished: bool,
_trans: &'trans Transaction<'trans>,
}
impl<'trans, 'stmt> LazyRowsNew<'trans, 'stmt> for LazyRows<'trans, 'stmt> {
fn new(stmt: &'stmt Statement<'stmt>,
data: VecDeque<RowData>,
name: String,
row_limit: i32,
more_rows: bool,
finished: bool,
trans: &'trans Transaction<'trans>)
-> LazyRows<'trans, 'stmt> {
LazyRows {
stmt: stmt,
data: data,
name: name,
row_limit: row_limit,
more_rows: more_rows,
finished: finished,
_trans: trans,
}
}
}
impl<'a, 'b> Drop for LazyRows<'a, 'b> {
fn drop(&mut self) {
if !self.finished {
let _ = self.finish_inner();
}
}
}
impl<'a, 'b> fmt::Debug for LazyRows<'a, 'b> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("LazyRows")
.field("name", &self.name)
.field("row_limit", &self.row_limit)
.field("remaining_rows", &self.data.len())
.field("more_rows", &self.more_rows)
.finish()
}
}
impl<'trans, 'stmt> LazyRows<'trans, 'stmt> {
fn finish_inner(&mut self) -> Result<()> {
let mut conn = self.stmt.conn().0.borrow_mut();
check_desync!(conn);
conn.close_statement(&self.name, b'P')
}
fn execute(&mut self) -> Result<()> {
let mut conn = self.stmt.conn().0.borrow_mut();
try!(conn.stream.write_message(|buf| frontend::execute(&self.name, self.row_limit, buf)));
try!(conn.stream.write_message(|buf| Ok::<(), io::Error>(frontend::sync(buf))));
try!(conn.stream.flush());
conn.read_rows(|row| self.data.push_back(row)).map(|more_rows| self.more_rows = more_rows)
}
pub fn columns(&self) -> &[Column] {
self.stmt.columns()
}
pub fn finish(mut self) -> Result<()> {
self.finish_inner()
}
}
impl<'trans, 'stmt> FallibleIterator for LazyRows<'trans, 'stmt> {
type Item = Row<'stmt>;
type Error = Error;
fn next(&mut self) -> Result<Option<Row<'stmt>>> {
if self.data.is_empty() && self.more_rows {
try!(self.execute());
}
let row = self.data
.pop_front()
.map(|r| {
Row {
stmt: self.stmt,
data: MaybeOwned::Owned(r),
}
});
Ok(row)
}
fn size_hint(&self) -> (usize, Option<usize>) {
let lower = self.data.len();
let upper = if self.more_rows { None } else { Some(lower) };
(lower, upper)
}
}