atuin_client/import/
zsh_histdb.rs

1// import old shell history from zsh-histdb!
2// automatically hoover up all that we can find
3
4// As far as i can tell there are no version numbers in the histdb sqlite DB, so we're going based
5// on the schema from 2022-05-01
6//
7// I have run into some histories that will not import b/c of non UTF-8 characters.
8//
9
10//
11// An Example sqlite query for hsitdb data:
12//
13//id|session|command_id|place_id|exit_status|start_time|duration|id|argv|id|host|dir
14//
15//
16//  select
17//    history.id,
18//    history.start_time,
19//    places.host,
20//    places.dir,
21//    commands.argv
22//  from history
23//    left join commands on history.command_id = commands.id
24//    left join places on history.place_id = places.id ;
25//
26// CREATE TABLE history  (id integer primary key autoincrement,
27//                       session int,
28//                       command_id int references commands (id),
29//                       place_id int references places (id),
30//                       exit_status int,
31//                       start_time int,
32//                       duration int);
33//
34
35use std::collections::HashMap;
36use std::path::{Path, PathBuf};
37
38use async_trait::async_trait;
39use atuin_common::utils::uuid_v7;
40use directories::UserDirs;
41use eyre::{eyre, Result};
42use sqlx::{sqlite::SqlitePool, Pool};
43use time::PrimitiveDateTime;
44
45use super::Importer;
46use crate::history::History;
47use crate::import::Loader;
48use crate::utils::{get_hostname, get_username};
49
50#[derive(sqlx::FromRow, Debug)]
51pub struct HistDbEntryCount {
52    pub count: usize,
53}
54
55#[derive(sqlx::FromRow, Debug)]
56pub struct HistDbEntry {
57    pub id: i64,
58    pub start_time: PrimitiveDateTime,
59    pub host: Vec<u8>,
60    pub dir: Vec<u8>,
61    pub argv: Vec<u8>,
62    pub duration: i64,
63    pub exit_status: i64,
64    pub session: i64,
65}
66
67#[derive(Debug)]
68pub struct ZshHistDb {
69    histdb: Vec<HistDbEntry>,
70    username: String,
71}
72
73/// Read db at given file, return vector of entries.
74async fn hist_from_db(dbpath: PathBuf) -> Result<Vec<HistDbEntry>> {
75    let pool = SqlitePool::connect(dbpath.to_str().unwrap()).await?;
76    hist_from_db_conn(pool).await
77}
78
79async fn hist_from_db_conn(pool: Pool<sqlx::Sqlite>) -> Result<Vec<HistDbEntry>> {
80    let query = r#"
81        SELECT
82            history.id, history.start_time, history.duration, places.host, places.dir,
83            commands.argv, history.exit_status, history.session
84        FROM history
85        LEFT JOIN commands ON history.command_id = commands.id
86        LEFT JOIN places ON history.place_id = places.id
87        ORDER BY history.start_time
88    "#;
89    let histdb_vec: Vec<HistDbEntry> = sqlx::query_as::<_, HistDbEntry>(query)
90        .fetch_all(&pool)
91        .await?;
92    Ok(histdb_vec)
93}
94
95impl ZshHistDb {
96    pub fn histpath_candidate() -> PathBuf {
97        // By default histdb database is `${HOME}/.histdb/zsh-history.db`
98        // This can be modified by ${HISTDB_FILE}
99        //
100        //  if [[ -z ${HISTDB_FILE} ]]; then
101        //      typeset -g HISTDB_FILE="${HOME}/.histdb/zsh-history.db"
102        let user_dirs = UserDirs::new().unwrap(); // should catch error here?
103        let home_dir = user_dirs.home_dir();
104        std::env::var("HISTDB_FILE")
105            .as_ref()
106            .map(|x| Path::new(x).to_path_buf())
107            .unwrap_or_else(|_err| home_dir.join(".histdb/zsh-history.db"))
108    }
109    pub fn histpath() -> Result<PathBuf> {
110        let histdb_path = ZshHistDb::histpath_candidate();
111        if histdb_path.exists() {
112            Ok(histdb_path)
113        } else {
114            Err(eyre!(
115                "Could not find history file. Try setting $HISTDB_FILE"
116            ))
117        }
118    }
119}
120
121#[async_trait]
122impl Importer for ZshHistDb {
123    // Not sure how this is used
124    const NAME: &'static str = "zsh_histdb";
125
126    /// Creates a new ZshHistDb and populates the history based on the pre-populated data
127    /// structure.
128    async fn new() -> Result<Self> {
129        let dbpath = ZshHistDb::histpath()?;
130        let histdb_entry_vec = hist_from_db(dbpath).await?;
131        Ok(Self {
132            histdb: histdb_entry_vec,
133            username: get_username(),
134        })
135    }
136
137    async fn entries(&mut self) -> Result<usize> {
138        Ok(self.histdb.len())
139    }
140
141    async fn load(self, h: &mut impl Loader) -> Result<()> {
142        let mut session_map = HashMap::new();
143        for entry in self.histdb {
144            let command = match std::str::from_utf8(&entry.argv) {
145                Ok(s) => s.trim_end(),
146                Err(_) => continue, // we can skip past things like invalid utf8
147            };
148            let cwd = match std::str::from_utf8(&entry.dir) {
149                Ok(s) => s.trim_end(),
150                Err(_) => continue, // we can skip past things like invalid utf8
151            };
152            let hostname = format!(
153                "{}:{}",
154                String::from_utf8(entry.host).unwrap_or_else(|_e| get_hostname()),
155                self.username
156            );
157            let session = session_map.entry(entry.session).or_insert_with(uuid_v7);
158
159            let imported = History::import()
160                .timestamp(entry.start_time.assume_utc())
161                .command(command)
162                .cwd(cwd)
163                .duration(entry.duration * 1_000_000_000)
164                .exit(entry.exit_status)
165                .session(session.as_simple().to_string())
166                .hostname(hostname)
167                .build();
168            h.push(imported.into()).await?;
169        }
170        Ok(())
171    }
172}
173
174#[cfg(test)]
175mod test {
176
177    use super::*;
178    use sqlx::sqlite::SqlitePoolOptions;
179    use std::env;
180    #[tokio::test(flavor = "multi_thread")]
181    async fn test_env_vars() {
182        let test_env_db = "nonstd-zsh-history.db";
183        let key = "HISTDB_FILE";
184        env::set_var(key, test_env_db);
185
186        // test the env got set
187        assert_eq!(env::var(key).unwrap(), test_env_db.to_string());
188
189        // test histdb returns the proper db from previous step
190        let histdb_path = ZshHistDb::histpath_candidate();
191        assert_eq!(histdb_path.to_str().unwrap(), test_env_db);
192    }
193
194    #[tokio::test(flavor = "multi_thread")]
195    async fn test_import() {
196        let pool: SqlitePool = SqlitePoolOptions::new()
197            .min_connections(2)
198            .connect(":memory:")
199            .await
200            .unwrap();
201
202        // sql dump directly from a test database.
203        let db_sql = r#"
204        PRAGMA foreign_keys=OFF;
205        BEGIN TRANSACTION;
206        CREATE TABLE commands (id integer primary key autoincrement, argv text, unique(argv) on conflict ignore);
207        INSERT INTO commands VALUES(1,'pwd');
208        INSERT INTO commands VALUES(2,'curl google.com');
209        INSERT INTO commands VALUES(3,'bash');
210        CREATE TABLE places   (id integer primary key autoincrement, host text, dir text, unique(host, dir) on conflict ignore);
211        INSERT INTO places VALUES(1,'mbp16.local','/home/noyez');
212        CREATE TABLE history  (id integer primary key autoincrement,
213                               session int,
214                               command_id int references commands (id),
215                               place_id int references places (id),
216                               exit_status int,
217                               start_time int,
218                               duration int);
219        INSERT INTO history VALUES(1,0,1,1,0,1651497918,1);
220        INSERT INTO history VALUES(2,0,2,1,0,1651497923,1);
221        INSERT INTO history VALUES(3,0,3,1,NULL,1651497930,NULL);
222        DELETE FROM sqlite_sequence;
223        INSERT INTO sqlite_sequence VALUES('commands',3);
224        INSERT INTO sqlite_sequence VALUES('places',3);
225        INSERT INTO sqlite_sequence VALUES('history',3);
226        CREATE INDEX hist_time on history(start_time);
227        CREATE INDEX place_dir on places(dir);
228        CREATE INDEX place_host on places(host);
229        CREATE INDEX history_command_place on history(command_id, place_id);
230        COMMIT; "#;
231
232        sqlx::query(db_sql).execute(&pool).await.unwrap();
233
234        // test histdb iterator
235        let histdb_vec = hist_from_db_conn(pool).await.unwrap();
236        let histdb = ZshHistDb {
237            histdb: histdb_vec,
238            username: get_username(),
239        };
240
241        println!("h: {:#?}", histdb.histdb);
242        println!("counter: {:?}", histdb.histdb.len());
243        for i in histdb.histdb {
244            println!("{i:?}");
245        }
246    }
247}