diff --git a/Cargo.toml b/Cargo.toml index 4752836a..4d8dbe78 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ "postgres-protocol", "postgres-types", "tokio-postgres", + "tokio-postgres-binary-copy", ] [profile.release] diff --git a/tokio-postgres-binary-copy/Cargo.toml b/tokio-postgres-binary-copy/Cargo.toml new file mode 100644 index 00000000..e6a70baf --- /dev/null +++ b/tokio-postgres-binary-copy/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "tokio-postgres-binary-copy" +version = "0.1.0" +authors = ["Steven Fackler "] +edition = "2018" + +[dependencies] +bytes = "0.4" +futures-preview = "=0.3.0-alpha.19" +parking_lot = "0.9" +pin-project-lite = "0.1" +tokio-postgres = { version = "=0.5.0-alpha.1", default-features = false, path = "../tokio-postgres" } + +[dev-dependencies] +tokio = "=0.2.0-alpha.6" +tokio-postgres = { version = "=0.5.0-alpha.1", path = "../tokio-postgres" } + diff --git a/tokio-postgres-binary-copy/src/lib.rs b/tokio-postgres-binary-copy/src/lib.rs new file mode 100644 index 00000000..9e0677d2 --- /dev/null +++ b/tokio-postgres-binary-copy/src/lib.rs @@ -0,0 +1,123 @@ +use bytes::{BigEndian, BufMut, ByteOrder, Bytes, BytesMut}; +use futures::{future, Stream}; +use parking_lot::Mutex; +use pin_project_lite::pin_project; +use std::convert::TryFrom; +use std::error::Error; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use tokio_postgres::types::{IsNull, ToSql, Type}; + +#[cfg(test)] +mod test; + +const BLOCK_SIZE: usize = 4096; + +pin_project! { + pub struct BinaryCopyStream { + #[pin] + future: F, + buf: Arc>, + done: bool, + } +} + +impl BinaryCopyStream +where + F: Future>>, +{ + pub fn new(types: &[Type], write_values: M) -> BinaryCopyStream + where + M: FnOnce(BinaryCopyWriter) -> F, + { + let mut buf = BytesMut::new(); + buf.reserve(11 + 4 + 4); + buf.put_slice(b"PGCOPY\n\xff\r\n\0"); // magic + buf.put_i32_be(0); // flags + buf.put_i32_be(0); // header extension + + let buf = Arc::new(Mutex::new(buf)); + let writer = BinaryCopyWriter { + buf: buf.clone(), + types: types.to_vec(), + idx: 0, + }; + + BinaryCopyStream { + future: write_values(writer), + buf, + done: false, + } + } +} + +impl Stream for BinaryCopyStream +where + F: Future>>, +{ + type Item = Result>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.project(); + + if *this.done { + return Poll::Ready(None); + } + + *this.done = this.future.poll(cx)?.is_ready(); + + let mut buf = this.buf.lock(); + if *this.done { + buf.reserve(2); + buf.put_i16_be(-1); + Poll::Ready(Some(Ok(buf.take().freeze()))) + } else if buf.len() > BLOCK_SIZE { + Poll::Ready(Some(Ok(buf.take().freeze()))) + } else { + Poll::Pending + } + } +} + +// FIXME this should really just take a reference to the buffer, but that requires HKT :( +pub struct BinaryCopyWriter { + buf: Arc>, + types: Vec, + idx: usize, +} + +impl BinaryCopyWriter { + pub async fn write( + &mut self, + value: &(dyn ToSql + Send), + ) -> Result<(), Box> { + future::poll_fn(|_| { + if self.buf.lock().len() > BLOCK_SIZE { + Poll::Pending + } else { + Poll::Ready(()) + } + }) + .await; + + let mut buf = self.buf.lock(); + if self.idx == 0 { + buf.reserve(2); + buf.put_i16_be(self.types.len() as i16); + } + let idx = buf.len(); + buf.reserve(4); + buf.put_i32_be(0); + let len = match value.to_sql_checked(&self.types[self.idx], &mut buf)? { + IsNull::Yes => -1, + IsNull::No => i32::try_from(buf.len() - idx - 4)?, + }; + BigEndian::write_i32(&mut buf[idx..], len); + + self.idx = (self.idx + 1) % self.types.len(); + + Ok(()) + } +} diff --git a/tokio-postgres-binary-copy/src/test.rs b/tokio-postgres-binary-copy/src/test.rs new file mode 100644 index 00000000..dcf91b69 --- /dev/null +++ b/tokio-postgres-binary-copy/src/test.rs @@ -0,0 +1,101 @@ +use crate::BinaryCopyStream; +use tokio_postgres::types::Type; +use tokio_postgres::{Client, NoTls}; + +async fn connect() -> Client { + let (client, connection) = + tokio_postgres::connect("host=localhost port=5433 user=postgres", NoTls) + .await + .unwrap(); + tokio::spawn(async { + connection.await.unwrap(); + }); + client +} + +#[tokio::test] +async fn write_basic() { + let client = connect().await; + + client + .batch_execute("CREATE TEMPORARY TABLE foo (id INT, bar TEXT)") + .await + .unwrap(); + + let stream = BinaryCopyStream::new(&[Type::INT4, Type::TEXT], |mut w| { + async move { + w.write(&1i32).await?; + w.write(&"foobar").await?; + + w.write(&2i32).await?; + w.write(&None::<&str>).await?; + + Ok(()) + } + }); + + client + .copy_in("COPY foo (id, bar) FROM STDIN BINARY", &[], stream) + .await + .unwrap(); + + let rows = client.query("SELECT id, bar FROM foo ORDER BY id", &[]).await.unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].get::<_, i32>(0), 1); + assert_eq!(rows[0].get::<_, Option<&str>>(1), Some("foobar")); + assert_eq!(rows[1].get::<_, i32>(0), 2); + assert_eq!(rows[1].get::<_, Option<&str>>(1), None); +} + +#[tokio::test] +async fn write_many_rows() { + let client = connect().await; + + client + .batch_execute("CREATE TEMPORARY TABLE foo (id INT, bar TEXT)") + .await + .unwrap(); + + let stream = BinaryCopyStream::new(&[Type::INT4, Type::TEXT], |mut w| async move { + for i in 0..10_000i32 { + w.write(&i).await?; + w.write(&format!("the value for {}", i)).await?; + } + + Ok(()) + }); + + client.copy_in("COPY foo (id, bar) FROM STDIN BINARY", &[], stream).await.unwrap(); + + let rows = client.query("SELECT id, bar FROM foo ORDER BY id", &[]).await.unwrap(); + for (i, row) in rows.iter().enumerate() { + assert_eq!(row.get::<_, i32>(0), i as i32); + assert_eq!(row.get::<_, &str>(1), format!("the value for {}", i)); + } +} + +#[tokio::test] +async fn write_big_rows() { + let client = connect().await; + + client.batch_execute("CREATE TEMPORARY TABLE foo (id INT, bar BYTEA)").await.unwrap(); + + let stream = BinaryCopyStream::new(&[Type::INT4, Type::BYTEA], |mut w| { + async move { + for i in 0..2i32 { + w.write(&i).await.unwrap(); + w.write(&vec![i as u8; 128 * 1024]).await.unwrap(); + } + + Ok(()) + } + }); + + client.copy_in("COPY foo (id, bar) FROM STDIN BINARY", &[], stream).await.unwrap(); + + let rows = client.query("SELECT id, bar FROM foo ORDER BY id", &[]).await.unwrap(); + for (i, row) in rows.iter().enumerate() { + assert_eq!(row.get::<_, i32>(0), i as i32); + assert_eq!(row.get::<_, &[u8]>(1), &*vec![i as u8; 128 * 1024]); + } +}