From 5169820d6a82f2f20bb0aa91981ea00b34ad7e28 Mon Sep 17 00:00:00 2001 From: Steven Fackler Date: Fri, 28 Dec 2018 20:39:32 -0800 Subject: [PATCH] Return iterators from query in sync API --- postgres/Cargo.toml | 1 + postgres/src/client.rs | 8 +++--- postgres/src/lib.rs | 4 +++ postgres/src/query.rs | 36 ++++++++++++++++++++++++ postgres/src/query_portal.rs | 36 ++++++++++++++++++++++++ postgres/src/test.rs | 53 ++++++++++++++++++++++++++++++------ postgres/src/transaction.rs | 22 ++++++++------- 7 files changed, 138 insertions(+), 22 deletions(-) create mode 100644 postgres/src/query.rs create mode 100644 postgres/src/query_portal.rs diff --git a/postgres/Cargo.toml b/postgres/Cargo.toml index a59c0794..2a203e3a 100644 --- a/postgres/Cargo.toml +++ b/postgres/Cargo.toml @@ -11,6 +11,7 @@ runtime = ["tokio-postgres/runtime", "tokio", "lazy_static", "log"] [dependencies] bytes = "0.4" +fallible-iterator = "0.1" futures = "0.1" tokio-postgres = { version = "0.3", path = "../tokio-postgres", default-features = false } diff --git a/postgres/src/client.rs b/postgres/src/client.rs index c2f691d2..4ccb9b5a 100644 --- a/postgres/src/client.rs +++ b/postgres/src/client.rs @@ -4,13 +4,13 @@ use futures::{Async, Future, Poll, Stream}; use std::io::{self, BufRead, Cursor, Read}; use std::marker::PhantomData; use tokio_postgres::types::{ToSql, Type}; -use tokio_postgres::{Error, Row}; +use tokio_postgres::Error; #[cfg(feature = "runtime")] use tokio_postgres::{MakeTlsMode, Socket, TlsMode}; #[cfg(feature = "runtime")] use crate::Builder; -use crate::{Statement, ToStatement, Transaction}; +use crate::{Query, Statement, ToStatement, Transaction}; pub struct Client(tokio_postgres::Client); @@ -48,12 +48,12 @@ impl Client { self.0.execute(&statement.0, params).wait() } - pub fn query(&mut self, query: &T, params: &[&dyn ToSql]) -> Result, Error> + pub fn query(&mut self, query: &T, params: &[&dyn ToSql]) -> Result, Error> where T: ?Sized + ToStatement, { let statement = query.__statement(self)?; - self.0.query(&statement.0, params).collect().wait() + Ok(Query::new(self.0.query(&statement.0, params))) } pub fn copy_in( diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs index 197be698..638ef26c 100644 --- a/postgres/src/lib.rs +++ b/postgres/src/lib.rs @@ -7,6 +7,8 @@ use tokio::runtime::{self, Runtime}; mod builder; mod client; mod portal; +mod query; +mod query_portal; mod statement; mod to_statement; mod transaction; @@ -19,6 +21,8 @@ mod test; pub use crate::builder::*; pub use crate::client::*; pub use crate::portal::*; +pub use crate::query::*; +pub use crate::query_portal::*; pub use crate::statement::*; pub use crate::to_statement::*; pub use crate::transaction::*; diff --git a/postgres/src/query.rs b/postgres/src/query.rs new file mode 100644 index 00000000..e7eb7948 --- /dev/null +++ b/postgres/src/query.rs @@ -0,0 +1,36 @@ +use fallible_iterator::FallibleIterator; +use futures::stream::{self, Stream}; +use std::marker::PhantomData; +use tokio_postgres::{Error, Row}; + +pub struct Query<'a> { + it: stream::Wait, + _p: PhantomData<&'a mut ()>, +} + +// no-op impl to extend the borrow until drop +impl<'a> Drop for Query<'a> { + fn drop(&mut self) {} +} + +impl<'a> Query<'a> { + pub(crate) fn new(stream: tokio_postgres::Query) -> Query<'a> { + Query { + it: stream.wait(), + _p: PhantomData, + } + } +} + +impl<'a> FallibleIterator for Query<'a> { + type Item = Row; + type Error = Error; + + fn next(&mut self) -> Result, Error> { + match self.it.next() { + Some(Ok(row)) => Ok(Some(row)), + Some(Err(e)) => Err(e), + None => Ok(None), + } + } +} diff --git a/postgres/src/query_portal.rs b/postgres/src/query_portal.rs new file mode 100644 index 00000000..0ed8250a --- /dev/null +++ b/postgres/src/query_portal.rs @@ -0,0 +1,36 @@ +use fallible_iterator::FallibleIterator; +use futures::stream::{self, Stream}; +use std::marker::PhantomData; +use tokio_postgres::{Error, Row}; + +pub struct QueryPortal<'a> { + it: stream::Wait, + _p: PhantomData<&'a mut ()>, +} + +// no-op impl to extend the borrow until drop +impl<'a> Drop for QueryPortal<'a> { + fn drop(&mut self) {} +} + +impl<'a> QueryPortal<'a> { + pub(crate) fn new(stream: tokio_postgres::QueryPortal) -> QueryPortal<'a> { + QueryPortal { + it: stream.wait(), + _p: PhantomData, + } + } +} + +impl<'a> FallibleIterator for QueryPortal<'a> { + type Item = Row; + type Error = Error; + + fn next(&mut self) -> Result, Error> { + match self.it.next() { + Some(Ok(row)) => Ok(Some(row)), + Some(Err(e)) => Err(e), + None => Ok(None), + } + } +} diff --git a/postgres/src/test.rs b/postgres/src/test.rs index e86f243d..ed6242a1 100644 --- a/postgres/src/test.rs +++ b/postgres/src/test.rs @@ -1,3 +1,4 @@ +use fallible_iterator::FallibleIterator; use std::io::Read; use tokio_postgres::types::Type; use tokio_postgres::NoTls; @@ -20,7 +21,11 @@ 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(); + let rows = client + .query(&stmt, &[&"hello"]) + .unwrap() + .collect::>() + .unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].get::<_, &str>(0), "hello"); } @@ -29,7 +34,11 @@ 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(); + let rows = client + .query("SELECT $1::TEXT", &[&"hello"]) + .unwrap() + .collect::>() + .unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].get::<_, &str>(0), "hello"); } @@ -50,7 +59,11 @@ fn transaction_commit() { transaction.commit().unwrap(); - let rows = client.query("SELECT * FROM foo", &[]).unwrap(); + let rows = client + .query("SELECT * FROM foo", &[]) + .unwrap() + .collect::>() + .unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].get::<_, i32>(0), 1); } @@ -71,7 +84,11 @@ fn transaction_rollback() { transaction.rollback().unwrap(); - let rows = client.query("SELECT * FROM foo", &[]).unwrap(); + let rows = client + .query("SELECT * FROM foo", &[]) + .unwrap() + .collect::>() + .unwrap(); assert_eq!(rows.len(), 0); } @@ -91,7 +108,11 @@ fn transaction_drop() { drop(transaction); - let rows = client.query("SELECT * FROM foo", &[]).unwrap(); + let rows = client + .query("SELECT * FROM foo", &[]) + .unwrap() + .collect::>() + .unwrap(); assert_eq!(rows.len(), 0); } @@ -119,6 +140,8 @@ fn nested_transactions() { let rows = transaction .query("SELECT id FROM foo ORDER BY id", &[]) + .unwrap() + .collect::>() .unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].get::<_, i32>(0), 1); @@ -139,7 +162,11 @@ fn nested_transactions() { transaction3.commit().unwrap(); transaction.commit().unwrap(); - let rows = client.query("SELECT id FROM foo ORDER BY id", &[]).unwrap(); + let rows = client + .query("SELECT id FROM foo ORDER BY id", &[]) + .unwrap() + .collect::>() + .unwrap(); assert_eq!(rows.len(), 3); assert_eq!(rows[0].get::<_, i32>(0), 1); assert_eq!(rows[1].get::<_, i32>(0), 3); @@ -164,6 +191,8 @@ fn copy_in() { let rows = client .query("SELECT id, name FROM foo ORDER BY id", &[]) + .unwrap() + .collect::>() .unwrap(); assert_eq!(rows.len(), 2); @@ -219,12 +248,20 @@ fn portal() { .bind("SELECT * FROM foo ORDER BY id", &[]) .unwrap(); - let rows = transaction.query_portal(&portal, 2).unwrap(); + let rows = transaction + .query_portal(&portal, 2) + .unwrap() + .collect::>() + .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(); + let rows = transaction + .query_portal(&portal, 2) + .unwrap() + .collect::>() + .unwrap(); assert_eq!(rows.len(), 1); assert_eq!(rows[0].get::<_, i32>(0), 3); } diff --git a/postgres/src/transaction.rs b/postgres/src/transaction.rs index 2c2e180c..4453f88c 100644 --- a/postgres/src/transaction.rs +++ b/postgres/src/transaction.rs @@ -1,9 +1,9 @@ -use futures::{Future, Stream}; +use futures::Future; use std::io::Read; use tokio_postgres::types::{ToSql, Type}; -use tokio_postgres::{Error, Row}; +use tokio_postgres::Error; -use crate::{Client, CopyOutReader, Portal, Statement, ToStatement}; +use crate::{Client, CopyOutReader, Portal, Query, QueryPortal, Statement, ToStatement}; pub struct Transaction<'a> { client: &'a mut Client, @@ -67,7 +67,7 @@ impl<'a> Transaction<'a> { self.client.execute(query, params) } - pub fn query(&mut self, query: &T, params: &[&dyn ToSql]) -> Result, Error> + pub fn query(&mut self, query: &T, params: &[&dyn ToSql]) -> Result, Error> where T: ?Sized + ToStatement, { @@ -86,12 +86,14 @@ impl<'a> Transaction<'a> { .map(Portal) } - pub fn query_portal(&mut self, portal: &Portal, max_rows: i32) -> Result, Error> { - self.client - .get_mut() - .query_portal(&portal.0, max_rows) - .collect() - .wait() + pub fn query_portal( + &mut self, + portal: &Portal, + max_rows: i32, + ) -> Result, Error> { + Ok(QueryPortal::new( + self.client.get_mut().query_portal(&portal.0, max_rows), + )) } pub fn copy_in(