gremlin-client 0.2.0

A Rust client for Apache TinkerPop™
Documentation
use crate::{GremlinError, GremlinResult};
use std::net::TcpStream;
use websocket::{sync::Client, ClientBuilder, OwnedMessage};

enum ConnectionStream {
    TCP(Client<TcpStream>),
    #[allow(dead_code)]
    SSL,
}

impl std::fmt::Debug for ConnectionStream {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "Connection")
    }
}

impl ConnectionStream {
    fn connect(options: ConnectionOptions) -> GremlinResult<Self> {
        let address = format!("ws://{}:{}/gremlin", options.host, options.port);

        let client = ClientBuilder::new(&address)
            .map_err(|e| GremlinError::Generic(e.to_string()))?
            .connect_insecure()?;

        Ok(ConnectionStream::TCP(client))
    }

    fn send(&mut self, payload: Vec<u8>) -> GremlinResult<()> {
        match self {
            ConnectionStream::TCP(s) => s
                .send_message(&OwnedMessage::Binary(payload))
                .map_err(GremlinError::from),
            ConnectionStream::SSL => unimplemented!(),
        }
    }

    fn recv(&mut self) -> GremlinResult<Vec<u8>> {
        let message = match self {
            ConnectionStream::TCP(s) => s.recv_message()?,
            ConnectionStream::SSL => unimplemented!(),
        };

        match message {
            OwnedMessage::Binary(binary) => Ok(binary),
            _ => unimplemented!(),
        }
    }
}

#[derive(Debug)]
pub(crate) struct Connection {
    stream: ConnectionStream,
}

impl Into<ConnectionOptions> for (&str, u16) {
    fn into(self) -> ConnectionOptions {
        ConnectionOptions {
            host: String::from(self.0),
            port: self.1,
            ..Default::default()
        }
    }
}

impl Into<ConnectionOptions> for &str {
    fn into(self) -> ConnectionOptions {
        ConnectionOptions {
            host: String::from(self),
            ..Default::default()
        }
    }
}

pub struct ConnectionOptionsBuilder(ConnectionOptions);

impl ConnectionOptionsBuilder {
    pub fn host<T>(mut self, host: T) -> Self
    where
        T: Into<String>,
    {
        self.0.host = host.into();
        self
    }

    pub fn port(mut self, port: u16) -> Self {
        self.0.port = port;
        self
    }

    pub fn pool_size(mut self, pool_size: u32) -> Self {
        self.0.pool_size = pool_size;
        self
    }

    pub fn build(self) -> ConnectionOptions {
        self.0
    }
}

#[derive(Clone, Debug)]
pub struct ConnectionOptions {
    pub(crate) host: String,
    pub(crate) port: u16,
    pub(crate) pool_size: u32,
}

impl Default for ConnectionOptions {
    fn default() -> ConnectionOptions {
        ConnectionOptions {
            host: String::from("localhost"),
            port: 8182,
            pool_size: 10,
        }
    }
}

impl ConnectionOptions {
    pub fn builder() -> ConnectionOptionsBuilder {
        ConnectionOptionsBuilder(ConnectionOptions::default())
    }
}

impl Connection {
    pub fn connect<T>(options: T) -> GremlinResult<Connection>
    where
        T: Into<ConnectionOptions>,
    {
        Ok(Connection {
            stream: ConnectionStream::connect(options.into())?,
        })
    }

    pub fn send(&mut self, payload: Vec<u8>) -> GremlinResult<()> {
        self.stream.send(payload)
    }

    pub fn recv(&mut self) -> GremlinResult<Vec<u8>> {
        self.stream.recv()
    }
}

#[cfg(test)]
mod tests {

    use super::Connection;

    #[test]
    fn it_should_connect() {
        Connection::connect(("localhost", 8182)).unwrap();
    }
}