diff --git a/src/db.rs b/src/db.rs index 4f6573d1e3c..3f95a3e4c76 100644 --- a/src/db.rs +++ b/src/db.rs @@ -3,7 +3,7 @@ use std::mem; use std::sync::Arc; use pg; -use pg::{PostgresConnection, PostgresStatement, PostgresResult}; +use pg::{PostgresConnection, PostgresTransaction, PostgresStatement, PostgresResult}; use pg::types::ToSql; use r2d2::{mod, LoggingErrorHandler}; use r2d2_postgres::PostgresPoolManager; @@ -154,6 +154,7 @@ impl<'a> RequestTransaction<'a> for &'a Request + 'a { pub trait Connection { fn prepare<'a>(&'a self, query: &str) -> PostgresResult>; fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult; + fn transaction<'a>(&'a self) -> PostgresResult>; } impl Connection for pg::PostgresConnection { @@ -163,6 +164,9 @@ impl Connection for pg::PostgresConnection { fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult { self.execute(query, params) } + fn transaction<'a>(&'a self) -> PostgresResult> { + self.transaction() + } } // // impl Connection for pg::pool::PooledPostgresConnection { @@ -181,4 +185,7 @@ impl<'a> Connection for pg::PostgresTransaction<'a> { fn execute(&self, query: &str, params: &[&ToSql]) -> PostgresResult { self.execute(query, params) } + fn transaction<'a>(&'a self) -> PostgresResult> { + self.transaction() + } } diff --git a/src/krate.rs b/src/krate.rs index 74c233258d2..4b700c78b50 100644 --- a/src/krate.rs +++ b/src/krate.rs @@ -1,3 +1,4 @@ +use std::io::MemWriter; use std::collections::hashmap::{HashMap, Occupied, Vacant}; use std::sync::Arc; use serialize::json; @@ -133,20 +134,47 @@ impl Crate { pub fn encode_many(conn: &Connection, crates: Vec) -> CargoResult> { - // TODO: can rust-postgres do this escaping? - let crateids: Vec = crates.iter().map(|p| p.id).collect(); - let mut map = HashMap::new(); - let query = format!("'{{{:#}}}'::int[]", crateids.as_slice()); - let stmt = try!(conn.prepare(format!("SELECT id, crate_id FROM versions \ - WHERE crate_id = ANY({})", - query).as_slice())); - for row in try!(stmt.query(&[])) { - match map.entry(row.get("crate_id")) { - Occupied(e) => e.into_mut(), - Vacant(e) => e.set(Vec::new()), - }.push(row.get("id")); + if crates.is_empty() { + return Ok(vec![]); } + let trans = try!(conn.transaction()); + + try!(trans.execute("CREATE TEMPORARY TABLE crateids ( + id INT PRIMARY KEY + ) ON COMMIT DROP", [])); + + let mut map = { + let mut query = MemWriter::new(); + let _ = write!(query, "INSERT INTO crateids (id) VALUES "); + let mut crateids: Vec<&ToSql> = vec![]; + let mut first = true; + for (i, krate) in crates.iter().enumerate() { + if !first { + let _ = write!(query, ", "); + } + first = false; + let _ = write!(query, "(${})", i+1); + crateids.push(&krate.id); + } + let query = String::from_utf8(query.unwrap()).unwrap(); + try!(trans.execute(query.as_slice(), crateids.as_slice())); + + let stmt = try!(trans.prepare("SELECT v.id, v.crate_id FROM versions v + INNER JOIN crateids c ON v.id = c.id")); + + let mut map = HashMap::new(); + for row in try!(stmt.query(&[])) { + match map.entry(row.get("crate_id")) { + Occupied(e) => e.into_mut(), + Vacant(e) => e.set(Vec::new()), + }.push(row.get("id")); + } + map + }; + + try!(trans.finish()); + Ok(crates.into_iter().map(|p| { let id = p.id; p.encodable(map.pop(&id).unwrap())