session: store a pool of connections
use anyhow::Result;
-use tokio::net::ToSocketAddrs;
-use std::collections::HashMap;
+use tokio::net::{lookup_host, ToSocketAddrs};
+use std::net::SocketAddr;
+use std::collections::{BTreeMap, HashMap};
use crate::frame::response::Response;
use crate::frame::response::result;
@@ -9,12 +10,20 @@ use crate::prepared_statement::PreparedStatement;
use crate::transport::connection::Connection;
use crate::transport::Compression;
+#[derive(PartialEq, Eq, PartialOrd, Ord, Hash)]
+struct Node {
+ // TODO: a potentially node may have multiple addresses, remember them?
+ // but we need an Ord instance on Node
+ addr: SocketAddr,
+}
+
pub struct Session {
- connection: Connection,
+ // invariant: nonempty
+ pool: HashMap<Node, Connection>,
}
impl Session {
- pub async fn connect(addr: impl ToSocketAddrs, compression: Option<Compression>) -> Result<Self> {
+ pub async fn connect(addr: impl ToSocketAddrs + Clone, compression: Option<Compression>) -> Result<Self> {
let mut options = HashMap::new();
if let Some(compression) = &compression {
let val = match compression {
@@ -23,11 +32,14 @@ impl Session {
options.insert("COMPRESSION".to_string(), val.to_string());
}
+ let resolved = lookup_host(addr.clone()).await?.next().map_or(
+ Err(anyhow!("no addresses found")), |a| Ok(a))?;
let connection = Connection::new(addr, compression).await?;
-
connection.startup(options).await?;
- Ok(Session { connection })
+ let pool = vec![(Node{ addr: resolved }, connection)].into_iter().collect();
+
+ Ok(Session { pool })
}
// TODO: Should return an iterator over results
@@ -35,7 +47,7 @@ impl Session {
// But maybe "INSERT" and "SELECT" should go through different methods,
// so we expect "SELECT" to always return Vec<result::Row>?
pub async fn query(&self, query: impl Into<Query>) -> Result<Option<Vec<result::Row>>> {
- let result = self.connection.query(&query.into()).await?;
+ let result = self.any_connection().query(&query.into()).await?;
match result {
Response::Error(err) => {
Err(err.into())
@@ -51,7 +63,7 @@ impl Session {
}
pub async fn prepare(&self, query: String) -> Result<PreparedStatement> {
- let result = self.connection.prepare(query.clone()).await?;
+ let result = self.any_connection().prepare(query.clone()).await?;
match result {
Response::Error(err) => {
Err(err.into())
@@ -63,4 +75,8 @@ impl Session {
_ => return Err(anyhow!("Unexpected frame received")),
}
}
+
+ fn any_connection(&self) -> &Connection {
+ self.pool.values().next().unwrap()
+ }
}
