atuin_client/import/
zsh_histdb.rs1use 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
73async 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 let user_dirs = UserDirs::new().unwrap(); 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 const NAME: &'static str = "zsh_histdb";
125
126 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, };
148 let cwd = match std::str::from_utf8(&entry.dir) {
149 Ok(s) => s.trim_end(),
150 Err(_) => continue, };
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 assert_eq!(env::var(key).unwrap(), test_env_db.to_string());
188
189 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 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 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}