atuin_scripts/
store.rs

1use eyre::{Result, eyre};
2
3use atuin_client::record::sqlite_store::SqliteStore;
4use atuin_client::record::{encryption::PASETO_V4, store::Store};
5use atuin_common::record::{Host, HostId, Record, RecordId, RecordIdx};
6use record::ScriptRecord;
7use script::{SCRIPT_TAG, SCRIPT_VERSION, Script};
8
9use crate::database::Database;
10
11pub mod record;
12pub mod script;
13
14#[derive(Debug, Clone)]
15pub struct ScriptStore {
16    pub store: SqliteStore,
17    pub host_id: HostId,
18    pub encryption_key: [u8; 32],
19}
20
21impl ScriptStore {
22    pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> Self {
23        ScriptStore {
24            store,
25            host_id,
26            encryption_key,
27        }
28    }
29
30    async fn push_record(&self, record: ScriptRecord) -> Result<(RecordId, RecordIdx)> {
31        let bytes = record.serialize()?;
32        let idx = self
33            .store
34            .last(self.host_id, SCRIPT_TAG)
35            .await?
36            .map_or(0, |p| p.idx + 1);
37
38        let record = Record::builder()
39            .host(Host::new(self.host_id))
40            .version(SCRIPT_VERSION.to_string())
41            .tag(SCRIPT_TAG.to_string())
42            .idx(idx)
43            .data(bytes)
44            .build();
45
46        let id = record.id;
47
48        self.store
49            .push(&record.encrypt::<PASETO_V4>(&self.encryption_key))
50            .await?;
51
52        Ok((id, idx))
53    }
54
55    pub async fn create(&self, script: Script) -> Result<()> {
56        let record = ScriptRecord::Create(script);
57        self.push_record(record).await?;
58        Ok(())
59    }
60
61    pub async fn update(&self, script: Script) -> Result<()> {
62        let record = ScriptRecord::Update(script);
63        self.push_record(record).await?;
64        Ok(())
65    }
66
67    pub async fn delete(&self, script_id: uuid::Uuid) -> Result<()> {
68        let record = ScriptRecord::Delete(script_id);
69        self.push_record(record).await?;
70        Ok(())
71    }
72
73    pub async fn scripts(&self) -> Result<Vec<ScriptRecord>> {
74        let records = self.store.all_tagged(SCRIPT_TAG).await?;
75        let mut ret = Vec::with_capacity(records.len());
76        let mut skipped = 0;
77
78        for record in records.into_iter() {
79            // Skip records we can't decrypt or decode, rather than failing the entire build.
80            let script = match record.version.as_str() {
81                SCRIPT_VERSION => {
82                    record
83                        .decrypt::<PASETO_V4>(&self.encryption_key)
84                        .and_then(|decrypted| {
85                            ScriptRecord::deserialize(&decrypted.data, SCRIPT_VERSION)
86                        })
87                }
88                version => Err(eyre!("unknown script version {version:?}")),
89            };
90
91            match script {
92                Ok(script) => ret.push(script),
93                Err(e) => {
94                    tracing::warn!("failed to decode script record, skipping: {e}");
95                    skipped += 1;
96                }
97            }
98        }
99
100        if skipped > 0 {
101            // library code that may run under the TUI or shell hooks, so no stderr here
102            tracing::warn!(
103                "skipped {skipped} script records that could not be decrypted or decoded"
104            );
105        }
106
107        Ok(ret)
108    }
109
110    pub async fn build(&self, database: Database) -> Result<()> {
111        // Clear existing data before replaying all records from the store.
112        // Without this, stale rows can cause unique constraint violations
113        // when records are replayed (eg name conflicts from renamed scripts).
114        database.clear().await?;
115
116        // Get all the scripts from the store - they are already sorted by timestamp
117        let scripts = self.scripts().await?;
118
119        for script in scripts {
120            match script {
121                ScriptRecord::Create(script) => {
122                    database.save(&script).await?;
123                }
124                ScriptRecord::Update(script) => database.update(&script).await?,
125                ScriptRecord::Delete(id) => database.delete(&id.to_string()).await?,
126            }
127        }
128
129        Ok(())
130    }
131}
atuin_scripts/store.rs

atuin_scripts/
store.rs
