atuin_client/import/
xonsh_sqlite.rs

1use std::env;
2use std::path::PathBuf;
3
4use async_trait::async_trait;
5use directories::BaseDirs;
6use eyre::{eyre, Result};
7use futures::TryStreamExt;
8use sqlx::{sqlite::SqlitePool, FromRow, Row};
9use time::OffsetDateTime;
10use uuid::timestamp::{context::NoContext, Timestamp};
11use uuid::Uuid;
12
13use super::{get_histpath, Importer, Loader};
14use crate::history::History;
15use crate::utils::get_host_user;
16
17#[derive(Debug, FromRow)]
18struct HistDbEntry {
19    inp: String,
20    rtn: Option<i64>,
21    tsb: f64,
22    tse: f64,
23    cwd: String,
24    session_start: f64,
25}
26
27impl HistDbEntry {
28    fn into_hist_with_hostname(self, hostname: String) -> History {
29        let ts_nanos = (self.tsb * 1_000_000_000_f64) as i128;
30        let timestamp = OffsetDateTime::from_unix_timestamp_nanos(ts_nanos).unwrap();
31
32        let session_ts_seconds = self.session_start.trunc() as u64;
33        let session_ts_nanos = (self.session_start.fract() * 1_000_000_000_f64) as u32;
34        let session_ts = Timestamp::from_unix(NoContext, session_ts_seconds, session_ts_nanos);
35        let session_id = Uuid::new_v7(session_ts).to_string();
36        let duration = (self.tse - self.tsb) * 1_000_000_000_f64;
37
38        if let Some(exit) = self.rtn {
39            let imported = History::import()
40                .timestamp(timestamp)
41                .duration(duration.trunc() as i64)
42                .exit(exit)
43                .command(self.inp)
44                .cwd(self.cwd)
45                .session(session_id)
46                .hostname(hostname);
47            imported.build().into()
48        } else {
49            let imported = History::import()
50                .timestamp(timestamp)
51                .duration(duration.trunc() as i64)
52                .command(self.inp)
53                .cwd(self.cwd)
54                .session(session_id)
55                .hostname(hostname);
56            imported.build().into()
57        }
58    }
59}
60
61fn xonsh_db_path(xonsh_data_dir: Option<String>) -> Result<PathBuf> {
62    // if running within xonsh, this will be available
63    if let Some(d) = xonsh_data_dir {
64        let mut path = PathBuf::from(d);
65        path.push("xonsh-history.sqlite");
66        return Ok(path);
67    }
68
69    // otherwise, fall back to default
70    let base = BaseDirs::new().ok_or_else(|| eyre!("Could not determine home directory"))?;
71
72    let hist_file = base.data_dir().join("xonsh/xonsh-history.sqlite");
73    if hist_file.exists() || cfg!(test) {
74        Ok(hist_file)
75    } else {
76        Err(eyre!(
77            "Could not find xonsh history db at: {}",
78            hist_file.to_string_lossy()
79        ))
80    }
81}
82
83#[derive(Debug)]
84pub struct XonshSqlite {
85    pool: SqlitePool,
86    hostname: String,
87}
88
89#[async_trait]
90impl Importer for XonshSqlite {
91    const NAME: &'static str = "xonsh_sqlite";
92
93    async fn new() -> Result<Self> {
94        // wrap xonsh-specific path resolver in general one so that it respects $HISTPATH
95        let xonsh_data_dir = env::var("XONSH_DATA_DIR").ok();
96        let db_path = get_histpath(|| xonsh_db_path(xonsh_data_dir))?;
97        let connection_str = db_path.to_str().ok_or_else(|| {
98            eyre!(
99                "Invalid path for SQLite database: {}",
100                db_path.to_string_lossy()
101            )
102        })?;
103
104        let pool = SqlitePool::connect(connection_str).await?;
105        let hostname = get_host_user();
106        Ok(XonshSqlite { pool, hostname })
107    }
108
109    async fn entries(&mut self) -> Result<usize> {
110        let query = "SELECT COUNT(*) FROM xonsh_history";
111        let row = sqlx::query(query).fetch_one(&self.pool).await?;
112        let count: u32 = row.get(0);
113        Ok(count as usize)
114    }
115
116    async fn load(self, loader: &mut impl Loader) -> Result<()> {
117        let query = r#"
118            SELECT inp, rtn, tsb, tse, cwd,
119            MIN(tsb) OVER (PARTITION BY sessionid) AS session_start
120            FROM xonsh_history
121            ORDER BY rowid
122        "#;
123
124        let mut entries = sqlx::query_as::<_, HistDbEntry>(query).fetch(&self.pool);
125
126        let mut count = 0;
127        while let Some(entry) = entries.try_next().await? {
128            let hist = entry.into_hist_with_hostname(self.hostname.clone());
129            loader.push(hist).await?;
130            count += 1;
131        }
132
133        println!("Loaded: {count}");
134        Ok(())
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use time::macros::datetime;
141
142    use super::*;
143
144    use crate::history::History;
145    use crate::import::tests::TestLoader;
146
147    #[test]
148    fn test_db_path_xonsh() {
149        let db_path = xonsh_db_path(Some("/home/user/xonsh_data".to_string())).unwrap();
150        assert_eq!(
151            db_path,
152            PathBuf::from("/home/user/xonsh_data/xonsh-history.sqlite")
153        );
154    }
155
156    #[tokio::test]
157    async fn test_import() {
158        let connection_str = "tests/data/xonsh-history.sqlite";
159        let xonsh_sqlite = XonshSqlite {
160            pool: SqlitePool::connect(connection_str).await.unwrap(),
161            hostname: "box:user".to_string(),
162        };
163
164        let mut loader = TestLoader::default();
165        xonsh_sqlite.load(&mut loader).await.unwrap();
166
167        for (actual, expected) in loader.buf.iter().zip(expected_hist_entries().iter()) {
168            assert_eq!(actual.timestamp, expected.timestamp);
169            assert_eq!(actual.command, expected.command);
170            assert_eq!(actual.cwd, expected.cwd);
171            assert_eq!(actual.exit, expected.exit);
172            assert_eq!(actual.duration, expected.duration);
173            assert_eq!(actual.hostname, expected.hostname);
174        }
175    }
176
177    fn expected_hist_entries() -> [History; 4] {
178        [
179            History::import()
180                .timestamp(datetime!(2024-02-6 17:56:21.130956288 +00:00:00))
181                .command("echo hello world!".to_string())
182                .cwd("/home/user/Documents/code/atuin".to_string())
183                .exit(0)
184                .duration(2628564)
185                .hostname("box:user".to_string())
186                .build()
187                .into(),
188            History::import()
189                .timestamp(datetime!(2024-02-06 17:56:28.190406144 +00:00:00))
190                .command("ls -l".to_string())
191                .cwd("/home/user/Documents/code/atuin".to_string())
192                .exit(0)
193                .duration(9371519)
194                .hostname("box:user".to_string())
195                .build()
196                .into(),
197            History::import()
198                .timestamp(datetime!(2024-02-06 17:56:46.989020928 +00:00:00))
199                .command("false".to_string())
200                .cwd("/home/user/Documents/code/atuin".to_string())
201                .exit(1)
202                .duration(17337560)
203                .hostname("box:user".to_string())
204                .build()
205                .into(),
206            History::import()
207                .timestamp(datetime!(2024-02-06 17:56:48.218384128 +00:00:00))
208                .command("exit".to_string())
209                .cwd("/home/user/Documents/code/atuin".to_string())
210                .exit(0)
211                .duration(4599094)
212                .hostname("box:user".to_string())
213                .build()
214                .into(),
215        ]
216    }
217}
atuin_client/import/xonsh_sqlite.rs

atuin_client/import/
xonsh_sqlite.rs
