atuin_daemon/
client.rs

1use eyre::{Context, Result};
2#[cfg(windows)]
3use tokio::net::TcpStream;
4use tonic::transport::{Channel, Endpoint, Uri};
5use tower::service_fn;
6
7use hyper_util::rt::TokioIo;
8
9#[cfg(unix)]
10use tokio::net::UnixStream;
11
12use atuin_client::history::History;
13
14use crate::history::{
15    EndHistoryRequest, StartHistoryRequest, history_client::HistoryClient as HistoryServiceClient,
16};
17
18pub struct HistoryClient {
19    client: HistoryServiceClient<Channel>,
20}
21
22// Wrap the grpc client
23impl HistoryClient {
24    #[cfg(unix)]
25    pub async fn new(path: String) -> Result<Self> {
26        let log_path = path.clone();
27        let channel = Endpoint::try_from("http://atuin_local_daemon:0")?
28            .connect_with_connector(service_fn(move |_: Uri| {
29                let path = path.clone();
30
31                async move {
32                    Ok::<_, std::io::Error>(TokioIo::new(UnixStream::connect(path.clone()).await?))
33                }
34            }))
35            .await
36            .wrap_err_with(|| {
37                format!(
38                    "failed to connect to local atuin daemon at {}. Is it running?",
39                    &log_path
40                )
41            })?;
42
43        let client = HistoryServiceClient::new(channel);
44
45        Ok(HistoryClient { client })
46    }
47
48    #[cfg(not(unix))]
49    pub async fn new(port: u64) -> Result<Self> {
50        let channel = Endpoint::try_from("http://atuin_local_daemon:0")?
51            .connect_with_connector(service_fn(move |_: Uri| {
52                let url = format!("127.0.0.1:{port}");
53
54                async move {
55                    Ok::<_, std::io::Error>(TokioIo::new(TcpStream::connect(url.clone()).await?))
56                }
57            }))
58            .await
59            .wrap_err_with(|| {
60                format!(
61                    "failed to connect to local atuin daemon at 127.0.0.1:{port}. Is it running?"
62                )
63            })?;
64
65        let client = HistoryServiceClient::new(channel);
66
67        Ok(HistoryClient { client })
68    }
69
70    pub async fn start_history(&mut self, h: History) -> Result<String> {
71        let req = StartHistoryRequest {
72            command: h.command,
73            cwd: h.cwd,
74            hostname: h.hostname,
75            session: h.session,
76            timestamp: h.timestamp.unix_timestamp_nanos() as u64,
77        };
78
79        let resp = self.client.start_history(req).await?;
80
81        Ok(resp.into_inner().id)
82    }
83
84    pub async fn end_history(
85        &mut self,
86        id: String,
87        duration: u64,
88        exit: i64,
89    ) -> Result<(String, u64)> {
90        let req = EndHistoryRequest { id, duration, exit };
91
92        let resp = self.client.end_history(req).await?;
93        let resp = resp.into_inner();
94
95        Ok((resp.id, resp.idx))
96    }
97}