Skip to content

Commit 13fcea7

Browse files
sfacklersfackler
committed
Working statement preparation
1 parent 0d0435f commit 13fcea7

10 files changed

Lines changed: 336 additions & 32 deletions

File tree

‎Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,5 @@ tokio-uds = { git = "https://github.com/sfackler/tokio" }
1414
tokio-io = { git = "https://github.com/sfackler/tokio" }
1515
tokio-timer = { git = "https://github.com/sfackler/tokio" }
1616
tokio-codec = { git = "https://github.com/sfackler/tokio" }
17+
tokio-reactor = { git = "https://github.com/sfackler/tokio" }
18+
tokio-executor = { git = "https://github.com/sfackler/tokio" }

‎postgres-shared/src/params/mod.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@
22
use std::error::Error;
33
use std::mem;
44
use std::path::PathBuf;
5+
use std::str::FromStr;
56
use std::time::Duration;
67

8+
use error;
79
use params::url::Url;
810

911
mod url;
@@ -96,6 +98,14 @@ impl ConnectParams {
9698
}
9799
}
98100

101+
impl FromStr for ConnectParams {
102+
type Err = error::Error;
103+
104+
fn from_str(s: &str) -> Result<ConnectParams, error::Error> {
105+
s.into_connect_params().map_err(error::connect)
106+
}
107+
}
108+
99109
/// A builder for `ConnectParams`.
100110
pub struct Builder {
101111
port: u16,

‎tokio-postgres/Cargo.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ fallible-iterator = "0.1.3"
3737
futures = "0.1.7"
3838
futures-cpupool = "0.1"
3939
lazy_static = "1.0"
40+
log = "0.4"
4041
postgres-protocol = { version = "0.3.0", path = "../postgres-protocol" }
4142
postgres-shared = { version = "0.4.0", path = "../postgres-shared" }
4243
state_machine_future = "0.1.7"
@@ -48,3 +49,7 @@ want = "0.0.5"
4849

4950
[target.'cfg(unix)'.dependencies]
5051
tokio-uds = "0.2"
52+
53+
[dev-dependencies]
54+
tokio = "0.1.7"
55+
env_logger = "0.5"

‎tokio-postgres/src/lib.rs‎

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,16 @@ extern crate futures;
1414
#[macro_use]
1515
extern crate lazy_static;
1616
#[macro_use]
17+
extern crate log;
18+
#[macro_use]
1719
extern crate state_machine_future;
1820

1921
#[cfg(unix)]
2022
extern crate tokio_uds;
2123

2224
use futures::{Async, Future, Poll};
2325
use std::io;
26+
use std::sync::atomic::{AtomicUsize, Ordering};
2427

2528
#[doc(inline)]
2629
pub use postgres_shared::stmt::Column;
@@ -31,9 +34,12 @@ pub use postgres_shared::{CancelData, Notification};
3134

3235
use error::Error;
3336
use params::ConnectParams;
37+
use types::Type;
3438

3539
mod proto;
3640

41+
static NEXT_STATEMENT_ID: AtomicUsize = AtomicUsize::new(0);
42+
3743
fn bad_response() -> Error {
3844
Error::from(io::Error::new(
3945
io::ErrorKind::InvalidInput,
@@ -48,13 +54,13 @@ fn disconnected() -> Error {
4854
))
4955
}
5056

57+
pub fn connect(params: ConnectParams) -> Handshake {
58+
Handshake(proto::HandshakeFuture::new(params))
59+
}
60+
5161
pub struct Client(proto::Client);
5262

5363
impl Client {
54-
pub fn connect(params: ConnectParams) -> Handshake {
55-
Handshake(proto::HandshakeFuture::new(params))
56-
}
57-
5864
/// Polls to to determine whether the connection is ready to send new requests to the backend.
5965
///
6066
/// Requests are unboundedly buffered to enable pipelining, but this risks unbounded memory consumption if requests
@@ -64,6 +70,15 @@ impl Client {
6470
pub fn poll_ready(&mut self) -> Poll<(), Error> {
6571
self.0.poll_ready()
6672
}
73+
74+
pub fn prepare(&mut self, query: &str) -> Prepare {
75+
self.prepare_typed(query, &[])
76+
}
77+
78+
pub fn prepare_typed(&mut self, query: &str, param_types: &[Type]) -> Prepare {
79+
let name = format!("s{}", NEXT_STATEMENT_ID.fetch_add(1, Ordering::SeqCst));
80+
Prepare(self.0.prepare(name, query, param_types))
81+
}
6782
}
6883

6984
pub struct Connection(proto::Connection);
@@ -99,3 +114,28 @@ impl Future for Handshake {
99114
Ok(Async::Ready((Client(client), Connection(connection))))
100115
}
101116
}
117+
118+
pub struct Prepare(proto::PrepareFuture);
119+
120+
impl Future for Prepare {
121+
type Item = Statement;
122+
type Error = Error;
123+
124+
fn poll(&mut self) -> Poll<Statement, Error> {
125+
let statement = try_ready!(self.0.poll());
126+
127+
Ok(Async::Ready(Statement(statement)))
128+
}
129+
}
130+
131+
pub struct Statement(proto::Statement);
132+
133+
impl Statement {
134+
pub fn params(&self) -> &[Type] {
135+
self.0.params()
136+
}
137+
138+
pub fn columns(&self) -> &[Column] {
139+
self.0.columns()
140+
}
141+
}

‎tokio-postgres/src/proto/client.rs‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
use futures::sync::mpsc;
22
use futures::Poll;
33
use postgres_protocol::message::backend::Message;
4+
use postgres_protocol::message::frontend;
45
use want::Giver;
56

67
use disconnected;
78
use error::Error;
89
use proto::connection::Request;
10+
use proto::prepare::PrepareFuture;
11+
use types::Type;
912

1013
pub struct Client {
1114
sender: mpsc::UnboundedSender<Request>,
@@ -21,7 +24,18 @@ impl Client {
2124
self.giver.poll_want().map_err(|_| disconnected())
2225
}
2326

24-
pub fn send(&mut self, messages: Vec<u8>) -> Result<mpsc::Receiver<Message>, Error> {
27+
pub fn prepare(&mut self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {
28+
let mut buf = vec![];
29+
let receiver = frontend::parse(&name, query, param_types.iter().map(|t| t.oid()), &mut buf)
30+
.and_then(|()| frontend::describe(b'S', &name, &mut buf))
31+
.and_then(|()| Ok(frontend::sync(&mut buf)))
32+
.map_err(Into::into)
33+
.and_then(|()| self.send(buf));
34+
35+
PrepareFuture::new(self.sender.clone(), receiver, name)
36+
}
37+
38+
fn send(&mut self, messages: Vec<u8>) -> Result<mpsc::Receiver<Message>, Error> {
2539
let (sender, receiver) = mpsc::channel(0);
2640
self.giver.give();
2741
self.sender

‎tokio-postgres/src/proto/connection.rs‎

Lines changed: 58 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use std::io;
77
use tokio_codec::Framed;
88
use want::Taker;
99

10+
use disconnected;
1011
use error::{self, Error};
1112
use proto::codec::PostgresCodec;
1213
use proto::socket::Socket;
@@ -17,7 +18,7 @@ pub struct Request {
1718
pub sender: mpsc::Sender<Message>,
1819
}
1920

20-
#[derive(PartialEq)]
21+
#[derive(PartialEq, Debug)]
2122
enum State {
2223
Active,
2324
Terminating,
@@ -67,22 +68,28 @@ impl Connection {
6768

6869
fn poll_response(&mut self) -> Poll<Option<Message>, io::Error> {
6970
if let Some(message) = self.pending_response.take() {
71+
trace!("retrying pending response");
7072
return Ok(Async::Ready(Some(message)));
7173
}
7274

7375
self.stream.poll()
7476
}
7577

76-
fn poll_read(&mut self) -> Result<(), Error> {
78+
fn poll_read(&mut self) -> Poll<(), Error> {
79+
if self.state != State::Active {
80+
trace!("poll_read: done");
81+
return Ok(Async::Ready(()));
82+
}
83+
7784
loop {
7885
let message = match self.poll_response()? {
7986
Async::Ready(Some(message)) => message,
8087
Async::Ready(None) => {
81-
return Err(Error::from(io::Error::from(io::ErrorKind::UnexpectedEof)));
88+
return Err(disconnected());
8289
}
8390
Async::NotReady => {
84-
self.taker.want();
85-
return Ok(());
91+
trace!("poll_read: waiting on response");
92+
return Ok(Async::NotReady);
8693
}
8794
};
8895

@@ -123,65 +130,95 @@ impl Connection {
123130
Ok(AsyncSink::NotReady(message)) => {
124131
self.responses.push_front(sender);
125132
self.pending_response = Some(message);
126-
return Ok(());
133+
trace!("poll_read: waiting on socket");
134+
return Ok(Async::NotReady);
127135
}
128136
}
129137
}
130138
}
131139

132140
fn poll_request(&mut self) -> Poll<Option<Vec<u8>>, Error> {
133141
if let Some(message) = self.pending_request.take() {
142+
trace!("retrying pending request");
134143
return Ok(Async::Ready(Some(message)));
135144
}
136145

137-
match self.receiver.poll() {
138-
Ok(Async::Ready(Some(request))) => {
146+
match try_receive!(self.receiver.poll()) {
147+
Some(request) => {
148+
trace!("polled new request");
139149
self.responses.push_back(request.sender);
140150
Ok(Async::Ready(Some(request.messages)))
141151
}
142-
Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
143-
Ok(Async::NotReady) => Ok(Async::NotReady),
144-
Err(()) => unreachable!("mpsc::Receiver doesn't return errors"),
152+
None => Ok(Async::Ready(None)),
145153
}
146154
}
147155

148-
fn poll_write(&mut self) -> Result<(), Error> {
156+
fn poll_write(&mut self) -> Poll<(), Error> {
149157
loop {
158+
if self.state == State::Closing {
159+
trace!("poll_write: done");
160+
return Ok(Async::Ready(()));
161+
}
162+
150163
let request = match self.poll_request()? {
151164
Async::Ready(Some(request)) => request,
152165
Async::Ready(None) if self.responses.is_empty() && self.state == State::Active => {
166+
trace!("poll_write: at eof, terminating");
153167
self.state = State::Terminating;
154168
let mut request = vec![];
155169
frontend::terminate(&mut request);
156170
request
157171
}
158-
Async::Ready(None) => return Ok(()),
159-
Async::NotReady => return Ok(()),
172+
Async::Ready(None) => {
173+
trace!(
174+
"poll_write: at eof, pending responses {}",
175+
self.responses.len(),
176+
);
177+
return Ok(Async::Ready(()));
178+
}
179+
Async::NotReady => {
180+
trace!("poll_write: waiting on request");
181+
self.taker.want();
182+
return Ok(Async::NotReady);
183+
}
160184
};
161185

162186
match self.stream.start_send(request)? {
163187
AsyncSink::Ready => {
164188
if self.state == State::Terminating {
189+
trace!("poll_write: sent eof, closing");
165190
self.state = State::Closing;
166191
}
167192
}
168193
AsyncSink::NotReady(request) => {
194+
trace!("poll_write: waiting on socket");
169195
self.pending_request = Some(request);
170-
return Ok(());
196+
return Ok(Async::NotReady);
171197
}
172198
}
173199
}
174200
}
175201

176-
fn poll_flush(&mut self) -> Result<(), Error> {
177-
self.stream.poll_complete()?;
178-
Ok(())
202+
fn poll_flush(&mut self) -> Poll<(), Error> {
203+
trace!("flushing");
204+
self.stream.poll_complete().map_err(Into::into)
179205
}
180206

181207
fn poll_shutdown(&mut self) -> Poll<(), Error> {
182-
match self.state {
183-
State::Active | State::Terminating => Ok(Async::NotReady),
184-
State::Closing => self.stream.close().map_err(Into::into),
208+
if self.state != State::Closing {
209+
return Ok(Async::NotReady);
210+
}
211+
212+
match self.stream.close() {
213+
Ok(Async::Ready(())) => {
214+
trace!("poll_shutdown: complete");
215+
Ok(Async::Ready(()))
216+
}
217+
Ok(Async::NotReady) => {
218+
trace!("poll_shutdown: waiting on socket");
219+
Ok(Async::NotReady)
220+
}
221+
Err(e) => Err(Error::from(e)),
185222
}
186223
}
187224
}

‎tokio-postgres/src/proto/mod.rs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,13 @@
1+
macro_rules! try_receive {
2+
($e:expr) => {
3+
match $e {
4+
Ok(::futures::Async::Ready(v)) => v,
5+
Ok(::futures::Async::NotReady) => return Ok(::futures::Async::NotReady),
6+
Err(()) => unreachable!("mpsc::Receiver doesn't return errors"),
7+
}
8+
};
9+
}
10+
111
mod client;
212
mod codec;
313
mod connection;
@@ -12,4 +22,4 @@ pub use proto::connection::Connection;
1222
pub use proto::handshake::HandshakeFuture;
1323
pub use proto::prepare::PrepareFuture;
1424
pub use proto::socket::Socket;
15-
pub use statement::Statement;
25+
pub use proto::statement::Statement;

0 commit comments

Comments
 (0)