oppi.li
personal fork of atuin without AI stuff, sync stuff, script management stuff
1use std::{env, path::PathBuf};
2
3use clap::Parser;
4use eyre::{eyre, Result};
5use indicatif::ProgressBar;
6
7use atuin_client::{
8 database::Database,
9 history::History,
10 import::{bash::Bash, fish::Fish, resh::Resh, zsh::Zsh, Importer},
11};
12
13#[derive(Parser)]
14#[clap(infer_subcommands = true)]
15pub enum Cmd {
16 /// Import history for the current shell
17 Auto,
18
19 /// Import history from the zsh history file
20 Zsh,
21
22 /// Import history from the bash history file
23 Bash,
24
25 /// Import history from the resh history file
26 Resh,
27
28 /// Import history from the fish history file
29 Fish,
30}
31
32const BATCH_SIZE: usize = 100;
33
34impl Cmd {
35 pub async fn run(&self, db: &mut (impl Database + Send + Sync)) -> Result<()> {
36 println!(" Atuin ");
37 println!("======================");
38 println!(" \u{1f30d} ");
39 println!(" \u{1f418}\u{1f418}\u{1f418}\u{1f418} ");
40 println!(" \u{1f422} ");
41 println!("======================");
42 println!("Importing history...");
43
44 match self {
45 Self::Auto => {
46 let shell = env::var("SHELL").unwrap_or_else(|_| String::from("NO_SHELL"));
47
48 if shell.ends_with("/zsh") {
49 println!("Detected ZSH");
50 import::<Zsh<_>, _>(db, BATCH_SIZE).await
51 } else if shell.ends_with("/fish") {
52 println!("Detected Fish");
53 import::<Fish<_>, _>(db, BATCH_SIZE).await
54 } else if shell.ends_with("/bash") {
55 println!("Detected Bash");
56 import::<Bash<_>, _>(db, BATCH_SIZE).await
57 } else {
58 println!("cannot import {} history", shell);
59 Ok(())
60 }
61 }
62
63 Self::Zsh => import::<Zsh<_>, _>(db, BATCH_SIZE).await,
64 Self::Bash => import::<Bash<_>, _>(db, BATCH_SIZE).await,
65 Self::Resh => import::<Resh, _>(db, BATCH_SIZE).await,
66 Self::Fish => import::<Fish<_>, _>(db, BATCH_SIZE).await,
67 }
68 }
69}
70
71async fn import<I: Importer + Send, DB: Database + Send + Sync>(
72 db: &mut DB,
73 buf_size: usize,
74) -> Result<()>
75where
76 I::IntoIter: Send,
77{
78 println!("Importing history from {}", I::NAME);
79
80 let histpath = get_histpath::<I>()?;
81 let contents = I::parse(histpath)?;
82
83 let iter = contents.into_iter();
84 let progress = if let (_, Some(upper_bound)) = iter.size_hint() {
85 ProgressBar::new(upper_bound as u64)
86 } else {
87 ProgressBar::new_spinner()
88 };
89
90 let mut buf = Vec::<History>::with_capacity(buf_size);
91 let mut iter = progress.wrap_iter(iter);
92 loop {
93 // fill until either no more entries
94 // or until the buffer is full
95 let done = fill_buf(&mut buf, &mut iter);
96
97 // flush
98 db.save_bulk(&buf).await?;
99
100 if done {
101 break;
102 }
103 }
104
105 println!("Import complete!");
106
107 Ok(())
108}
109
110fn get_histpath<I: Importer>() -> Result<PathBuf> {
111 if let Ok(p) = env::var("HISTFILE") {
112 is_file(PathBuf::from(p))
113 } else {
114 is_file(I::histpath()?)
115 }
116}
117
118fn is_file(p: PathBuf) -> Result<PathBuf> {
119 if p.is_file() {
120 Ok(p)
121 } else {
122 Err(eyre!(
123 "Could not find history file {:?}. Try setting $HISTFILE",
124 p
125 ))
126 }
127}
128
129fn fill_buf<T, E>(buf: &mut Vec<T>, iter: &mut impl Iterator<Item = Result<T, E>>) -> bool {
130 buf.clear();
131 loop {
132 match iter.next() {
133 Some(Ok(t)) => buf.push(t),
134 Some(Err(_)) => (),
135 None => break true,
136 }
137
138 if buf.len() == buf.capacity() {
139 break false;
140 }
141 }
142}
143
144#[cfg(test)]
145mod tests {
146 use super::fill_buf;
147
148 #[test]
149 fn test_fill_buf() {
150 let mut buf = Vec::with_capacity(4);
151 let mut iter = vec![
152 Ok(1),
153 Err(2),
154 Ok(3),
155 Ok(4),
156 Err(5),
157 Ok(6),
158 Ok(7),
159 Err(8),
160 Ok(9),
161 ]
162 .into_iter();
163
164 assert!(!fill_buf(&mut buf, &mut iter));
165 assert_eq!(buf, vec![1, 3, 4, 6]);
166
167 assert!(fill_buf(&mut buf, &mut iter));
168 assert_eq!(buf, vec![7, 9]);
169 }
170}