Don't require use of iterators

The common case is to simply want a vec of rows to work with, so make
that the main API, wrapping the old version returning an iterator.
This commit is contained in:
Steven Fackler
2019-03-10 17:27:28 -07:00
parent 1cbf9a1b8d
commit 2a80118e87
7 changed files with 91 additions and 125 deletions

View File

@@ -3,13 +3,13 @@ use futures::{Async, Future, Poll, Stream};
use std::io::{self, Read};
use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
use tokio_postgres::types::{ToSql, Type};
use tokio_postgres::Error;
#[cfg(feature = "runtime")]
use tokio_postgres::Socket;
use tokio_postgres::{Error, Row, SimpleQueryMessage};
#[cfg(feature = "runtime")]
use crate::Config;
use crate::{CopyOutReader, Query, SimpleQuery, Statement, ToStatement, Transaction};
use crate::{CopyOutReader, QueryIter, SimpleQueryIter, Statement, ToStatement, Transaction};
pub struct Client(tokio_postgres::Client);
@@ -46,12 +46,23 @@ impl Client {
self.0.execute(&statement.0, params).wait()
}
pub fn query<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<Query<'_>, Error>
pub fn query<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<Vec<Row>, Error>
where
T: ?Sized + ToStatement,
{
self.query_iter(query, params)?.collect()
}
pub fn query_iter<T>(
&mut self,
query: &T,
params: &[&dyn ToSql],
) -> Result<QueryIter<'_>, Error>
where
T: ?Sized + ToStatement,
{
let statement = query.__statement(self)?;
Ok(Query::new(self.0.query(&statement.0, params)))
Ok(QueryIter::new(self.0.query(&statement.0, params)))
}
pub fn copy_in<T, R>(
@@ -83,12 +94,16 @@ impl Client {
CopyOutReader::new(stream)
}
pub fn simple_query(&mut self, query: &str) -> Result<SimpleQuery<'_>, Error> {
Ok(SimpleQuery::new(self.0.simple_query(query)))
pub fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
self.simple_query_iter(query)?.collect()
}
pub fn simple_query_iter(&mut self, query: &str) -> Result<SimpleQueryIter<'_>, Error> {
Ok(SimpleQueryIter::new(self.0.simple_query(query)))
}
pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
self.simple_query("BEGIN")?.count()?;
self.simple_query("BEGIN")?;
Ok(Transaction::new(self))
}

View File

@@ -8,9 +8,9 @@ mod client;
mod config;
mod copy_out_reader;
mod portal;
mod query;
mod query_portal;
mod simple_query;
mod query_iter;
mod query_portal_iter;
mod simple_query_iter;
mod statement;
mod to_statement;
mod transaction;
@@ -24,9 +24,9 @@ pub use crate::client::*;
pub use crate::config::*;
pub use crate::copy_out_reader::*;
pub use crate::portal::*;
pub use crate::query::*;
pub use crate::query_portal::*;
pub use crate::simple_query::*;
pub use crate::query_iter::*;
pub use crate::query_portal_iter::*;
pub use crate::simple_query_iter::*;
pub use crate::statement::*;
pub use crate::to_statement::*;
pub use crate::transaction::*;

View File

@@ -4,31 +4,26 @@ use std::marker::PhantomData;
use tokio_postgres::impls;
use tokio_postgres::{Error, Row};
pub struct Query<'a> {
pub struct QueryIter<'a> {
it: stream::Wait<impls::Query>,
_p: PhantomData<&'a mut ()>,
}
// no-op impl to extend the borrow until drop
impl<'a> Drop for Query<'a> {
impl<'a> Drop for QueryIter<'a> {
fn drop(&mut self) {}
}
impl<'a> Query<'a> {
pub(crate) fn new(stream: impls::Query) -> Query<'a> {
Query {
impl<'a> QueryIter<'a> {
pub(crate) fn new(stream: impls::Query) -> QueryIter<'a> {
QueryIter {
it: stream.wait(),
_p: PhantomData,
}
}
/// A convenience API which collects the resulting rows into a `Vec` and returns them.
pub fn into_vec(self) -> Result<Vec<Row>, Error> {
self.collect()
}
}
impl<'a> FallibleIterator for Query<'a> {
impl<'a> FallibleIterator for QueryIter<'a> {
type Item = Row;
type Error = Error;

View File

@@ -4,31 +4,26 @@ use std::marker::PhantomData;
use tokio_postgres::impls;
use tokio_postgres::{Error, Row};
pub struct QueryPortal<'a> {
pub struct QueryPortalIter<'a> {
it: stream::Wait<impls::QueryPortal>,
_p: PhantomData<&'a mut ()>,
}
// no-op impl to extend the borrow until drop
impl<'a> Drop for QueryPortal<'a> {
impl<'a> Drop for QueryPortalIter<'a> {
fn drop(&mut self) {}
}
impl<'a> QueryPortal<'a> {
pub(crate) fn new(stream: impls::QueryPortal) -> QueryPortal<'a> {
QueryPortal {
impl<'a> QueryPortalIter<'a> {
pub(crate) fn new(stream: impls::QueryPortal) -> QueryPortalIter<'a> {
QueryPortalIter {
it: stream.wait(),
_p: PhantomData,
}
}
/// A convenience API which collects the resulting rows into a `Vec` and returns them.
pub fn into_vec(self) -> Result<Vec<Row>, Error> {
self.collect()
}
}
impl<'a> FallibleIterator for QueryPortal<'a> {
impl<'a> FallibleIterator for QueryPortalIter<'a> {
type Item = Row;
type Error = Error;

View File

@@ -4,31 +4,26 @@ use std::marker::PhantomData;
use tokio_postgres::impls;
use tokio_postgres::{Error, SimpleQueryMessage};
pub struct SimpleQuery<'a> {
pub struct SimpleQueryIter<'a> {
it: stream::Wait<impls::SimpleQuery>,
_p: PhantomData<&'a mut ()>,
}
// no-op impl to extend borrow until drop
impl<'a> Drop for SimpleQuery<'a> {
impl<'a> Drop for SimpleQueryIter<'a> {
fn drop(&mut self) {}
}
impl<'a> SimpleQuery<'a> {
pub(crate) fn new(stream: impls::SimpleQuery) -> SimpleQuery<'a> {
SimpleQuery {
impl<'a> SimpleQueryIter<'a> {
pub(crate) fn new(stream: impls::SimpleQuery) -> SimpleQueryIter<'a> {
SimpleQueryIter {
it: stream.wait(),
_p: PhantomData,
}
}
/// A convenience API which collects the resulting messages into a `Vec` and returns them.
pub fn into_vec(self) -> Result<Vec<SimpleQueryMessage>, Error> {
self.collect()
}
}
impl<'a> FallibleIterator for SimpleQuery<'a> {
impl<'a> FallibleIterator for SimpleQueryIter<'a> {
type Item = SimpleQueryMessage;
type Error = Error;

View File

@@ -1,4 +1,3 @@
use fallible_iterator::FallibleIterator;
use std::io::Read;
use tokio_postgres::types::Type;
use tokio_postgres::NoTls;
@@ -21,11 +20,7 @@ fn query_prepared() {
let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
let stmt = client.prepare("SELECT $1::TEXT").unwrap();
let rows = client
.query(&stmt, &[&"hello"])
.unwrap()
.into_vec()
.unwrap();
let rows = client.query(&stmt, &[&"hello"]).unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].get::<_, &str>(0), "hello");
}
@@ -34,11 +29,7 @@ fn query_prepared() {
fn query_unprepared() {
let mut client = Client::connect("host=localhost port=5433 user=postgres", NoTls).unwrap();
let rows = client
.query("SELECT $1::TEXT", &[&"hello"])
.unwrap()
.into_vec()
.unwrap();
let rows = client.query("SELECT $1::TEXT", &[&"hello"]).unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].get::<_, &str>(0), "hello");
}
@@ -49,8 +40,6 @@ fn transaction_commit() {
client
.simple_query("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
.unwrap()
.count()
.unwrap();
let mut transaction = client.transaction().unwrap();
@@ -61,11 +50,7 @@ fn transaction_commit() {
transaction.commit().unwrap();
let rows = client
.query("SELECT * FROM foo", &[])
.unwrap()
.into_vec()
.unwrap();
let rows = client.query("SELECT * FROM foo", &[]).unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].get::<_, i32>(0), 1);
}
@@ -76,8 +61,6 @@ fn transaction_rollback() {
client
.simple_query("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
.unwrap()
.count()
.unwrap();
let mut transaction = client.transaction().unwrap();
@@ -88,11 +71,7 @@ fn transaction_rollback() {
transaction.rollback().unwrap();
let rows = client
.query("SELECT * FROM foo", &[])
.unwrap()
.into_vec()
.unwrap();
let rows = client.query("SELECT * FROM foo", &[]).unwrap();
assert_eq!(rows.len(), 0);
}
@@ -102,8 +81,6 @@ fn transaction_drop() {
client
.simple_query("CREATE TEMPORARY TABLE foo (id SERIAL PRIMARY KEY)")
.unwrap()
.count()
.unwrap();
let mut transaction = client.transaction().unwrap();
@@ -114,11 +91,7 @@ fn transaction_drop() {
drop(transaction);
let rows = client
.query("SELECT * FROM foo", &[])
.unwrap()
.into_vec()
.unwrap();
let rows = client.query("SELECT * FROM foo", &[]).unwrap();
assert_eq!(rows.len(), 0);
}
@@ -128,8 +101,6 @@ fn nested_transactions() {
client
.simple_query("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY)")
.unwrap()
.count()
.unwrap();
let mut transaction = client.transaction().unwrap();
@@ -148,8 +119,6 @@ fn nested_transactions() {
let rows = transaction
.query("SELECT id FROM foo ORDER BY id", &[])
.unwrap()
.into_vec()
.unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].get::<_, i32>(0), 1);
@@ -170,11 +139,7 @@ fn nested_transactions() {
transaction3.commit().unwrap();
transaction.commit().unwrap();
let rows = client
.query("SELECT id FROM foo ORDER BY id", &[])
.unwrap()
.into_vec()
.unwrap();
let rows = client.query("SELECT id FROM foo ORDER BY id", &[]).unwrap();
assert_eq!(rows.len(), 3);
assert_eq!(rows[0].get::<_, i32>(0), 1);
assert_eq!(rows[1].get::<_, i32>(0), 3);
@@ -187,8 +152,6 @@ fn copy_in() {
client
.simple_query("CREATE TEMPORARY TABLE foo (id INT, name TEXT)")
.unwrap()
.count()
.unwrap();
client
@@ -201,8 +164,6 @@ fn copy_in() {
let rows = client
.query("SELECT id, name FROM foo ORDER BY id", &[])
.unwrap()
.into_vec()
.unwrap();
assert_eq!(rows.len(), 2);
@@ -221,8 +182,6 @@ fn copy_out() {
"CREATE TEMPORARY TABLE foo (id INT, name TEXT);
INSERT INTO foo (id, name) VALUES (1, 'steven'), (2, 'timothy');",
)
.unwrap()
.count()
.unwrap();
let mut reader = client
@@ -234,7 +193,7 @@ fn copy_out() {
assert_eq!(s, "1\tsteven\n2\ttimothy\n");
client.simple_query("SELECT 1").unwrap().count().unwrap();
client.simple_query("SELECT 1").unwrap();
}
#[test]
@@ -246,8 +205,6 @@ fn portal() {
"CREATE TEMPORARY TABLE foo (id INT);
INSERT INTO foo (id) VALUES (1), (2), (3);",
)
.unwrap()
.count()
.unwrap();
let mut transaction = client.transaction().unwrap();
@@ -256,20 +213,12 @@ fn portal() {
.bind("SELECT * FROM foo ORDER BY id", &[])
.unwrap();
let rows = transaction
.query_portal(&portal, 2)
.unwrap()
.into_vec()
.unwrap();
let rows = transaction.query_portal(&portal, 2).unwrap();
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].get::<_, i32>(0), 1);
assert_eq!(rows[1].get::<_, i32>(0), 2);
let rows = transaction
.query_portal(&portal, 2)
.unwrap()
.into_vec()
.unwrap();
let rows = transaction.query_portal(&portal, 2).unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].get::<_, i32>(0), 3);
}

View File

@@ -2,10 +2,11 @@ use fallible_iterator::FallibleIterator;
use futures::Future;
use std::io::Read;
use tokio_postgres::types::{ToSql, Type};
use tokio_postgres::Error;
use tokio_postgres::{Error, Row, SimpleQueryMessage};
use crate::{
Client, CopyOutReader, Portal, Query, QueryPortal, SimpleQuery, Statement, ToStatement,
Client, CopyOutReader, Portal, QueryIter, QueryPortalIter, SimpleQueryIter, Statement,
ToStatement,
};
pub struct Transaction<'a> {
@@ -33,13 +34,12 @@ impl<'a> Transaction<'a> {
pub fn commit(mut self) -> Result<(), Error> {
self.done = true;
let it = if self.depth == 0 {
self.client.simple_query("COMMIT")?
if self.depth == 0 {
self.client.simple_query("COMMIT")?;
} else {
self.client
.simple_query(&format!("RELEASE sp{}", self.depth))?
};
it.count()?;
.simple_query(&format!("RELEASE sp{}", self.depth))?;
}
Ok(())
}
@@ -49,13 +49,12 @@ impl<'a> Transaction<'a> {
}
fn rollback_inner(&mut self) -> Result<(), Error> {
let it = if self.depth == 0 {
self.client.simple_query("ROLLBACK")?
if self.depth == 0 {
self.client.simple_query("ROLLBACK")?;
} else {
self.client
.simple_query(&format!("ROLLBACK TO sp{}", self.depth))?
};
it.count()?;
.simple_query(&format!("ROLLBACK TO sp{}", self.depth))?;
}
Ok(())
}
@@ -74,13 +73,24 @@ impl<'a> Transaction<'a> {
self.client.execute(query, params)
}
pub fn query<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<Query<'_>, Error>
pub fn query<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<Vec<Row>, Error>
where
T: ?Sized + ToStatement,
{
self.client.query(query, params)
}
pub fn query_iter<T>(
&mut self,
query: &T,
params: &[&dyn ToSql],
) -> Result<QueryIter<'_>, Error>
where
T: ?Sized + ToStatement,
{
self.client.query_iter(query, params)
}
pub fn bind<T>(&mut self, query: &T, params: &[&dyn ToSql]) -> Result<Portal, Error>
where
T: ?Sized + ToStatement,
@@ -93,12 +103,16 @@ impl<'a> Transaction<'a> {
.map(Portal)
}
pub fn query_portal(
pub fn query_portal(&mut self, portal: &Portal, max_rows: i32) -> Result<Vec<Row>, Error> {
self.query_portal_iter(portal, max_rows)?.collect()
}
pub fn query_portal_iter(
&mut self,
portal: &Portal,
max_rows: i32,
) -> Result<QueryPortal<'_>, Error> {
Ok(QueryPortal::new(
) -> Result<QueryPortalIter<'_>, Error> {
Ok(QueryPortalIter::new(
self.client.get_mut().query_portal(&portal.0, max_rows),
))
}
@@ -127,15 +141,18 @@ impl<'a> Transaction<'a> {
self.client.copy_out(query, params)
}
pub fn simple_query(&mut self, query: &str) -> Result<SimpleQuery<'_>, Error> {
pub fn simple_query(&mut self, query: &str) -> Result<Vec<SimpleQueryMessage>, Error> {
self.client.simple_query(query)
}
pub fn simple_query_iter(&mut self, query: &str) -> Result<SimpleQueryIter<'_>, Error> {
self.client.simple_query_iter(query)
}
pub fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
let depth = self.depth + 1;
self.client
.simple_query(&format!("SAVEPOINT sp{}", depth))?
.count()?;
.simple_query(&format!("SAVEPOINT sp{}", depth))?;
Ok(Transaction {
client: self.client,
depth,