rust-postgres/src/lib.rs

167 lines
4.5 KiB
Rust
Raw Normal View History

2013-08-22 05:52:15 +00:00
extern mod extra;
2013-08-04 02:17:32 +00:00
2013-08-22 05:52:15 +00:00
use std::cell::Cell;
use std::hashmap::HashMap;
use std::rt::io::io_error;
use std::rt::io::net::ip::SocketAddr;
use std::rt::io::net::tcp::TcpStream;
use extra::url::Url;
2013-08-04 02:17:32 +00:00
2013-08-22 05:52:15 +00:00
use message::*;
2013-07-25 07:10:18 +00:00
2013-08-22 05:52:15 +00:00
mod message;
2013-07-25 07:10:18 +00:00
2013-08-18 03:30:31 +00:00
pub struct PostgresConnection {
2013-08-22 05:52:15 +00:00
priv stream: Cell<TcpStream>,
priv next_stmt_id: Cell<int>
2013-08-04 02:17:32 +00:00
}
2013-08-18 03:30:31 +00:00
impl Drop for PostgresConnection {
2013-08-04 02:17:32 +00:00
fn drop(&self) {
2013-08-22 07:12:35 +00:00
self.write_message(&Terminate);
2013-08-18 03:42:40 +00:00
}
2013-08-04 02:17:32 +00:00
}
2013-07-25 07:10:18 +00:00
2013-08-18 03:30:31 +00:00
impl PostgresConnection {
2013-08-23 02:47:06 +00:00
pub fn connect(url: &str) -> PostgresConnection {
2013-08-22 05:52:15 +00:00
let parsed_url: Url = FromStr::from_str(url).unwrap();
let socket_url = fmt!("%s:%s", parsed_url.host,
2013-08-23 02:47:06 +00:00
parsed_url.port.get_ref().as_slice());
2013-08-22 05:52:15 +00:00
let addr: SocketAddr = FromStr::from_str(socket_url).unwrap();
let conn = PostgresConnection {
stream: Cell::new(TcpStream::connect(addr).unwrap()),
2013-08-18 03:42:40 +00:00
next_stmt_id: Cell::new(0)
};
2013-08-22 06:41:26 +00:00
let mut args = HashMap::new();
2013-08-23 02:47:06 +00:00
args.insert(&"user", parsed_url.user.get_ref().user.as_slice());
2013-08-22 06:41:26 +00:00
conn.write_message(&StartupMessage(args));
2013-08-04 05:21:16 +00:00
2013-08-22 06:41:26 +00:00
match conn.read_message() {
2013-08-22 05:52:15 +00:00
AuthenticationOk => (),
2013-08-22 06:41:26 +00:00
resp => fail!("Bad response: %?", resp)
2013-08-04 02:17:32 +00:00
}
2013-08-22 05:52:15 +00:00
conn.finish_connect();
2013-08-04 02:17:32 +00:00
2013-08-22 05:52:15 +00:00
conn
2013-07-25 07:10:18 +00:00
}
2013-08-22 05:52:15 +00:00
fn finish_connect(&self) {
loop {
2013-08-22 06:41:26 +00:00
match self.read_message() {
2013-08-22 05:52:15 +00:00
ParameterStatus(param, value) =>
printfln!("Param %s = %s", param, value),
2013-08-22 06:41:26 +00:00
BackendKeyData(*) => (),
ReadyForQuery(*) => break,
2013-08-22 05:52:15 +00:00
resp => fail!("Bad response: %?", resp)
}
2013-08-05 00:48:48 +00:00
}
}
2013-08-22 06:41:26 +00:00
fn write_message(&self, message: &FrontendMessage) {
2013-08-22 05:52:15 +00:00
do self.stream.with_mut_ref |s| {
2013-08-22 06:41:26 +00:00
s.write_message(message);
2013-07-25 07:10:18 +00:00
}
2013-08-22 06:41:26 +00:00
}
2013-08-04 02:17:32 +00:00
2013-08-22 06:41:26 +00:00
fn read_message(&self) -> BackendMessage {
2013-08-22 05:52:15 +00:00
do self.stream.with_mut_ref |s| {
2013-08-22 06:41:26 +00:00
s.read_message()
2013-08-04 05:21:16 +00:00
}
2013-08-22 06:41:26 +00:00
}
pub fn prepare<'a>(&'a self, query: &str) -> PostgresStatement<'a> {
let id = self.next_stmt_id.take();
let stmt_name = ifmt!("statement_{}", id);
self.next_stmt_id.put_back(id + 1);
2013-08-05 00:48:48 +00:00
2013-08-22 06:41:26 +00:00
let types = [];
self.write_message(&Parse(stmt_name, query, types));
self.write_message(&Sync);
match self.read_message() {
2013-08-22 05:52:15 +00:00
ParseComplete => (),
ErrorResponse(ref data) => fail!("Error: %?", data),
resp => fail!("Bad response: %?", resp)
2013-08-05 00:48:48 +00:00
}
2013-08-22 06:41:26 +00:00
self.wait_for_ready();
self.write_message(&Describe('S' as u8, stmt_name));
self.write_message(&Sync);
let num_params = match self.read_message() {
ParameterDescription(ref types) => types.len(),
resp => fail!("Bad response: %?", resp)
};
match self.read_message() {
RowDescription(*) | NoData => (),
2013-08-22 05:52:15 +00:00
resp => fail!("Bad response: %?", resp)
2013-08-05 00:48:48 +00:00
}
2013-08-17 22:09:26 +00:00
2013-08-22 06:41:26 +00:00
self.wait_for_ready();
PostgresStatement {
conn: self,
name: stmt_name,
2013-08-22 07:12:35 +00:00
num_params: num_params,
next_portal_id: Cell::new(0)
2013-08-22 06:41:26 +00:00
}
}
fn wait_for_ready(&self) {
match self.read_message() {
ReadyForQuery(*) => (),
resp => fail!("Bad response: %?", resp)
}
2013-08-17 22:09:26 +00:00
}
}
2013-08-22 05:52:15 +00:00
pub struct PostgresStatement<'self> {
priv conn: &'self PostgresConnection,
2013-08-22 06:41:26 +00:00
priv name: ~str,
2013-08-22 07:12:35 +00:00
priv num_params: uint,
priv next_portal_id: Cell<uint>
}
#[unsafe_destructor]
impl<'self> Drop for PostgresStatement<'self> {
fn drop(&self) {
self.conn.write_message(&Close('S' as u8, self.name.as_slice()));
self.conn.write_message(&Sync);
self.conn.read_message(); // CloseComplete or ErrorResponse
self.conn.wait_for_ready();
}
}
2013-08-22 05:52:15 +00:00
impl<'self> PostgresStatement<'self> {
2013-08-22 07:12:35 +00:00
pub fn num_params(&self) -> uint {
self.num_params
}
pub fn query(&self) {
let id = self.next_portal_id.take();
let portal_name = ifmt!("{:s}_portal_{}", self.name.as_slice(), id);
self.next_portal_id.put_back(id + 1);
2013-08-22 07:12:35 +00:00
let formats = [];
let values = [];
let result_formats = [];
self.conn.write_message(&Bind(portal_name, self.name.as_slice(),
formats, values, result_formats));
self.conn.write_message(&Sync);
match self.conn.read_message() {
BindComplete => (),
ErrorResponse(ref data) => fail!("Error: %?", data),
resp => fail!("Bad response: %?", resp)
}
self.conn.wait_for_ready();
}
}