Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit dcde61c

Browse files
committed
Add bind/portal
1 parent a4bdcb1 commit dcde61c

File tree

5 files changed

+174
-7
lines changed

5 files changed

+174
-7
lines changed

tokio-postgres/src/lib.rs

+29-3
Original file line numberDiff line numberDiff line change
@@ -45,10 +45,14 @@ pub mod error;
4545
mod proto;
4646
pub mod tls;
4747

48-
static NEXT_STATEMENT_ID: AtomicUsize = AtomicUsize::new(0);
49-
5048
fn next_statement() -> String {
51-
format!("s{}", NEXT_STATEMENT_ID.fetch_add(1, Ordering::SeqCst))
49+
static ID: AtomicUsize = AtomicUsize::new(0);
50+
format!("s{}", ID.fetch_add(1, Ordering::SeqCst))
51+
}
52+
53+
fn next_portal() -> String {
54+
static ID: AtomicUsize = AtomicUsize::new(0);
55+
format!("p{}", ID.fetch_add(1, Ordering::SeqCst))
5256
}
5357

5458
pub enum TlsMode {
@@ -84,6 +88,10 @@ impl Client {
8488
Query(self.0.query(&statement.0, params))
8589
}
8690

91+
pub fn bind(&mut self, statement: &Statement, params: &[&ToSql]) -> Bind {
92+
Bind(self.0.bind(&statement.0, next_portal(), params))
93+
}
94+
8795
pub fn copy_in<S>(&mut self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyIn<S>
8896
where
8997
S: Stream,
@@ -226,6 +234,24 @@ impl Stream for Query {
226234
}
227235
}
228236

237+
#[must_use = "futures do nothing unless polled"]
238+
pub struct Bind(proto::BindFuture);
239+
240+
impl Future for Bind {
241+
type Item = Portal;
242+
type Error = Error;
243+
244+
fn poll(&mut self) -> Poll<Portal, Error> {
245+
match self.0.poll() {
246+
Ok(Async::Ready(portal)) => Ok(Async::Ready(Portal(portal))),
247+
Ok(Async::NotReady) => Ok(Async::NotReady),
248+
Err(e) => Err(e),
249+
}
250+
}
251+
}
252+
253+
pub struct Portal(proto::Portal);
254+
229255
#[must_use = "futures do nothing unless polled"]
230256
pub struct CopyIn<S>(proto::CopyInFuture<S>)
231257
where

tokio-postgres/src/proto/bind.rs

+72
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
use futures::sync::mpsc;
2+
use futures::{Poll, Stream};
3+
use postgres_protocol::message::backend::Message;
4+
use proto::client::{Client, PendingRequest};
5+
use proto::portal::Portal;
6+
use proto::statement::Statement;
7+
use state_machine_future::RentToOwn;
8+
use Error;
9+
10+
#[derive(StateMachineFuture)]
11+
pub enum Bind {
12+
#[state_machine_future(start, transitions(ReadBindComplete))]
13+
Start {
14+
client: Client,
15+
request: PendingRequest,
16+
name: String,
17+
statement: Statement,
18+
},
19+
#[state_machine_future(transitions(Finished))]
20+
ReadBindComplete {
21+
receiver: mpsc::Receiver<Message>,
22+
client: Client,
23+
name: String,
24+
statement: Statement,
25+
},
26+
#[state_machine_future(ready)]
27+
Finished(Portal),
28+
#[state_machine_future(error)]
29+
Failed(Error),
30+
}
31+
32+
impl PollBind for Bind {
33+
fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll<AfterStart, Error> {
34+
let state = state.take();
35+
let receiver = state.client.send(state.request)?;
36+
37+
transition!(ReadBindComplete {
38+
receiver,
39+
client: state.client,
40+
name: state.name,
41+
statement: state.statement,
42+
})
43+
}
44+
45+
fn poll_read_bind_complete<'a>(
46+
state: &'a mut RentToOwn<'a, ReadBindComplete>,
47+
) -> Poll<AfterReadBindComplete, Error> {
48+
let message = try_ready_receive!(state.receiver.poll());
49+
let state = state.take();
50+
51+
match message {
52+
Some(Message::BindComplete) => transition!(Finished(Portal::new(
53+
state.client.downgrade(),
54+
state.name,
55+
state.statement,
56+
))),
57+
Some(_) => Err(Error::unexpected_message()),
58+
None => Err(Error::closed()),
59+
}
60+
}
61+
}
62+
63+
impl BindFuture {
64+
pub fn new(
65+
client: Client,
66+
request: PendingRequest,
67+
name: String,
68+
statement: Statement,
69+
) -> BindFuture {
70+
Bind::start(client, request, name, statement)
71+
}
72+
}

tokio-postgres/src/proto/client.rs

+31-4
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use std::collections::HashMap;
88
use std::error::Error as StdError;
99
use std::sync::{Arc, Weak};
1010

11+
use proto::bind::BindFuture;
1112
use proto::connection::{Request, RequestMessages};
1213
use proto::copy_in::{CopyInFuture, CopyInReceiver, CopyMessage};
1314
use proto::copy_out::CopyOutStream;
@@ -140,6 +141,15 @@ impl Client {
140141
QueryStream::new(self.clone(), pending, statement.clone())
141142
}
142143

144+
pub fn bind(&self, statement: &Statement, name: String, params: &[&ToSql]) -> BindFuture {
145+
let mut buf = self.bind_message(statement, &name, params);
146+
if let Ok(ref mut buf) = buf {
147+
frontend::sync(buf);
148+
}
149+
let pending = PendingRequest(buf.map(RequestMessages::Single));
150+
BindFuture::new(self.clone(), pending, name, statement.clone())
151+
}
152+
143153
pub fn copy_in<S>(&self, statement: &Statement, params: &[&ToSql], stream: S) -> CopyInFuture<S>
144154
where
145155
S: Stream,
@@ -169,8 +179,16 @@ impl Client {
169179
}
170180

171181
pub fn close_statement(&self, name: &str) {
182+
self.close(b'S', name)
183+
}
184+
185+
pub fn close_portal(&self, name: &str) {
186+
self.close(b'P', name)
187+
}
188+
189+
fn close(&self, ty: u8, name: &str) {
172190
let mut buf = vec![];
173-
frontend::close(b'S', name, &mut buf).expect("statement name not valid");
191+
frontend::close(ty, name, &mut buf).expect("statement name not valid");
174192
frontend::sync(&mut buf);
175193
let (sender, _) = mpsc::channel(0);
176194
let _ = self.0.sender.unbounded_send(Request {
@@ -179,10 +197,15 @@ impl Client {
179197
});
180198
}
181199

182-
fn excecute_message(&self, statement: &Statement, params: &[&ToSql]) -> Result<Vec<u8>, Error> {
200+
fn bind_message(
201+
&self,
202+
statement: &Statement,
203+
name: &str,
204+
params: &[&ToSql],
205+
) -> Result<Vec<u8>, Error> {
183206
let mut buf = vec![];
184207
let r = frontend::bind(
185-
"",
208+
name,
186209
statement.name(),
187210
Some(1),
188211
params.iter().zip(statement.params()),
@@ -195,10 +218,14 @@ impl Client {
195218
&mut buf,
196219
);
197220
match r {
198-
Ok(()) => {}
221+
Ok(()) => Ok(buf),
199222
Err(frontend::BindError::Conversion(e)) => return Err(Error::to_sql(e)),
200223
Err(frontend::BindError::Serialization(e)) => return Err(Error::encode(e)),
201224
}
225+
}
226+
227+
fn excecute_message(&self, statement: &Statement, params: &[&ToSql]) -> Result<Vec<u8>, Error> {
228+
let mut buf = self.bind_message(statement, "", params)?;
202229
frontend::execute("", 0, &mut buf).map_err(Error::parse)?;
203230
frontend::sync(&mut buf);
204231
Ok(buf)

tokio-postgres/src/proto/mod.rs

+4
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ macro_rules! try_ready_closed {
1818
};
1919
}
2020

21+
mod bind;
2122
mod cancel;
2223
mod client;
2324
mod codec;
@@ -27,6 +28,7 @@ mod copy_in;
2728
mod copy_out;
2829
mod execute;
2930
mod handshake;
31+
mod portal;
3032
mod prepare;
3133
mod query;
3234
mod row;
@@ -38,6 +40,7 @@ mod typeinfo;
3840
mod typeinfo_composite;
3941
mod typeinfo_enum;
4042

43+
pub use proto::bind::BindFuture;
4144
pub use proto::cancel::CancelFuture;
4245
pub use proto::client::Client;
4346
pub use proto::codec::PostgresCodec;
@@ -46,6 +49,7 @@ pub use proto::copy_in::CopyInFuture;
4649
pub use proto::copy_out::CopyOutStream;
4750
pub use proto::execute::ExecuteFuture;
4851
pub use proto::handshake::HandshakeFuture;
52+
pub use proto::portal::Portal;
4953
pub use proto::prepare::PrepareFuture;
5054
pub use proto::query::QueryStream;
5155
pub use proto::row::Row;

tokio-postgres/src/proto/portal.rs

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
use std::sync::Arc;
2+
3+
use proto::client::WeakClient;
4+
use proto::statement::Statement;
5+
6+
struct Inner {
7+
client: WeakClient,
8+
name: String,
9+
statement: Statement,
10+
}
11+
12+
pub struct Portal(Arc<Inner>);
13+
14+
impl Drop for Portal {
15+
fn drop(&mut self) {
16+
if let Some(client) = self.0.client.upgrade() {
17+
client.close_portal(&self.0.name);
18+
}
19+
}
20+
}
21+
22+
impl Portal {
23+
pub fn new(client: WeakClient, name: String, statement: Statement) -> Portal {
24+
Portal(Arc::new(Inner {
25+
client,
26+
name,
27+
statement,
28+
}))
29+
}
30+
31+
pub fn name(&self) -> &str {
32+
&self.0.name
33+
}
34+
35+
pub fn statement(&self) -> &Statement {
36+
&self.0.statement
37+
}
38+
}

0 commit comments

Comments
 (0)