1use anyhow::Result;
2use rusqlite::{params, Connection, OptionalExtension};
3
4use std::{
5 ffi::OsStr,
6 fmt::Debug,
7 os::unix::prelude::OsStrExt,
8 path::{Path, PathBuf},
9 sync::Arc,
10};
11
12use crate::pane::{PaneGroupId, PaneId, SerializedPane, SerializedPaneGroup};
13
14use super::Db;
15
16pub(crate) const WORKSPACE_M_1: &str = "
17CREATE TABLE workspaces(
18 workspace_id INTEGER PRIMARY KEY AUTOINCREMENT,
19 timestamp TEXT DEFAULT CURRENT_TIMESTAMP
20) STRICT;
21
22CREATE TABLE worktree_roots(
23 worktree_root BLOB NOT NULL,
24 workspace_id INTEGER NOT NULL,
25 FOREIGN KEY(workspace_id) REFERENCES workspaces(workspace_id) ON DELETE CASCADE
26 PRIMARY KEY(worktree_root, workspace_id)
27) STRICT;
28";
29
30#[derive(Debug, PartialEq, Eq, Copy, Clone, Default)]
31pub struct WorkspaceId(i64);
32
33struct WorkspaceRow {
34 pub center_group_id: PaneGroupId,
35 pub dock_pane_id: PaneId,
36}
37
38#[derive(Default, Debug)]
39pub struct SerializedWorkspace {
40 pub workspace_id: WorkspaceId,
41 // pub center_group: SerializedPaneGroup,
42 // pub dock_pane: Option<SerializedPane>,
43}
44
45impl Db {
46 /// Finds or creates a workspace id for the given set of worktree roots. If the passed worktree roots is empty, return the
47 /// the last workspace id
48 pub fn workspace_for_worktree_roots(
49 &self,
50 worktree_roots: &[Arc<Path>],
51 ) -> SerializedWorkspace {
52 // Find the workspace id which is uniquely identified by this set of paths
53 // return it if found
54 let mut workspace_id = self.workspace_id(worktree_roots);
55 if workspace_id.is_none() && worktree_roots.len() == 0 {
56 workspace_id = self.last_workspace_id();
57 }
58
59 if let Some(workspace_id) = workspace_id {
60 // TODO
61 // let workspace_row = self.get_workspace_row(workspace_id);
62 // let center_group = self.get_pane_group(workspace_row.center_group_id);
63 // let dock_pane = self.get_pane(workspace_row.dock_pane_id);
64
65 SerializedWorkspace {
66 workspace_id,
67 // center_group,
68 // dock_pane: Some(dock_pane),
69 }
70 } else {
71 self.make_new_workspace()
72 }
73 }
74
75 fn make_new_workspace(&self) -> SerializedWorkspace {
76 self.real()
77 .map(|db| {
78 let lock = db.connection.lock();
79 // No need to waste the memory caching this, should happen rarely.
80 match lock.execute("INSERT INTO workspaces DEFAULT VALUES", []) {
81 Ok(_) => SerializedWorkspace {
82 workspace_id: WorkspaceId(lock.last_insert_rowid()),
83 },
84 Err(_) => Default::default(),
85 }
86 })
87 .unwrap_or_default()
88 }
89
90 fn workspace_id<P>(&self, worktree_roots: &[P]) -> Option<WorkspaceId>
91 where
92 P: AsRef<Path> + Debug,
93 {
94 self.real()
95 .map(|db| {
96 let lock = db.connection.lock();
97
98 get_workspace_id(worktree_roots, &lock)
99 })
100 .unwrap_or(None)
101 }
102
103 // fn get_workspace_row(&self, workspace_id: WorkspaceId) -> WorkspaceRow {
104 // unimplemented!()
105 // }
106
107 /// Updates the open paths for the given workspace id. Will garbage collect items from
108 /// any workspace ids which are no replaced by the new workspace id. Updates the timestamps
109 /// in the workspace id table
110 pub fn update_worktree_roots<P>(&self, workspace_id: &WorkspaceId, worktree_roots: &[P])
111 where
112 P: AsRef<Path> + Debug,
113 {
114 fn logic<P>(
115 connection: &mut Connection,
116 worktree_roots: &[P],
117 workspace_id: &WorkspaceId,
118 ) -> Result<()>
119 where
120 P: AsRef<Path> + Debug,
121 {
122 let tx = connection.transaction()?;
123 {
124 // Lookup any old WorkspaceIds which have the same set of roots, and delete them.
125 let preexisting_id = get_workspace_id(worktree_roots, &tx);
126 if let Some(preexisting_id) = preexisting_id {
127 if preexisting_id != *workspace_id {
128 // Should also delete fields in other tables with cascading updates
129 tx.execute(
130 "DELETE FROM workspaces WHERE workspace_id = ?",
131 [preexisting_id.0],
132 )?;
133 }
134 }
135
136 tx.execute(
137 "DELETE FROM worktree_roots WHERE workspace_id = ?",
138 [workspace_id.0],
139 )?;
140
141 for root in worktree_roots {
142 let path = root.as_ref().as_os_str().as_bytes();
143
144 tx.execute(
145 "INSERT INTO worktree_roots(workspace_id, worktree_root) VALUES (?, ?)",
146 params![workspace_id.0, path],
147 )?;
148 }
149
150 tx.execute(
151 "UPDATE workspaces SET timestamp = CURRENT_TIMESTAMP WHERE workspace_id = ?",
152 [workspace_id.0],
153 )?;
154 }
155 tx.commit()?;
156 Ok(())
157 }
158
159 self.real().map(|db| {
160 let mut lock = db.connection.lock();
161
162 match logic(&mut lock, worktree_roots, workspace_id) {
163 Ok(_) => {}
164 Err(err) => {
165 log::error!(
166 "Failed to update the worktree roots for {:?}, roots: {:?}, error: {}",
167 workspace_id,
168 worktree_roots,
169 err
170 );
171 }
172 }
173 });
174 }
175
176 pub fn last_workspace_id(&self) -> Option<WorkspaceId> {
177 fn logic(connection: &mut Connection) -> Result<Option<WorkspaceId>> {
178 let mut stmt = connection
179 .prepare("SELECT workspace_id FROM workspaces ORDER BY timestamp DESC LIMIT 1")?;
180
181 Ok(stmt
182 .query_row([], |row| Ok(WorkspaceId(row.get(0)?)))
183 .optional()?)
184 }
185
186 self.real()
187 .map(|db| {
188 let mut lock = db.connection.lock();
189
190 match logic(&mut lock) {
191 Ok(result) => result,
192 Err(err) => {
193 log::error!("Failed to get last workspace id, err: {}", err);
194 None
195 }
196 }
197 })
198 .unwrap_or(None)
199 }
200
201 /// Returns the previous workspace ids sorted by last modified along with their opened worktree roots
202 pub fn recent_workspaces(&self, limit: usize) -> Vec<(WorkspaceId, Vec<Arc<Path>>)> {
203 fn logic(
204 connection: &mut Connection,
205 limit: usize,
206 ) -> Result<Vec<(WorkspaceId, Vec<Arc<Path>>)>, anyhow::Error> {
207 let tx = connection.transaction()?;
208 let result = {
209 let mut stmt = tx.prepare(
210 "SELECT workspace_id FROM workspaces ORDER BY timestamp DESC LIMIT ?",
211 )?;
212
213 let workspace_ids = stmt
214 .query_map([limit], |row| Ok(WorkspaceId(row.get(0)?)))?
215 .collect::<Result<Vec<_>, rusqlite::Error>>()?;
216
217 let mut result = Vec::new();
218 let mut stmt =
219 tx.prepare("SELECT worktree_root FROM worktree_roots WHERE workspace_id = ?")?;
220 for workspace_id in workspace_ids {
221 let roots = stmt
222 .query_map([workspace_id.0], |row| {
223 let row = row.get::<_, Vec<u8>>(0)?;
224 Ok(PathBuf::from(OsStr::from_bytes(&row)).into())
225 })?
226 .collect::<Result<Vec<_>, rusqlite::Error>>()?;
227 result.push((workspace_id, roots))
228 }
229
230 result
231 };
232 tx.commit()?;
233 return Ok(result);
234 }
235
236 self.real()
237 .map(|db| {
238 let mut lock = db.connection.lock();
239
240 match logic(&mut lock, limit) {
241 Ok(result) => result,
242 Err(err) => {
243 log::error!("Failed to get recent workspaces, err: {}", err);
244 Vec::new()
245 }
246 }
247 })
248 .unwrap_or_else(|| Vec::new())
249 }
250}
251
252fn get_workspace_id<P>(worktree_roots: &[P], connection: &Connection) -> Option<WorkspaceId>
253where
254 P: AsRef<Path> + Debug,
255{
256 fn logic<P>(
257 worktree_roots: &[P],
258 connection: &Connection,
259 ) -> Result<Option<WorkspaceId>, anyhow::Error>
260 where
261 P: AsRef<Path> + Debug,
262 {
263 // Prepare the array binding string. SQL doesn't have syntax for this, so
264 // we have to do it ourselves.
265 let mut array_binding_stmt = "(".to_string();
266 for i in 0..worktree_roots.len() {
267 array_binding_stmt.push_str(&format!("?{}", (i + 1))); //sqlite is 1-based
268 if i < worktree_roots.len() - 1 {
269 array_binding_stmt.push(',');
270 array_binding_stmt.push(' ');
271 }
272 }
273 array_binding_stmt.push(')');
274 // Any workspace can have multiple independent paths, and these paths
275 // can overlap in the database. Take this test data for example:
276 //
277 // [/tmp, /tmp2] -> 1
278 // [/tmp] -> 2
279 // [/tmp2, /tmp3] -> 3
280 //
281 // This would be stred in the database like so:
282 //
283 // ID PATH
284 // 1 /tmp
285 // 1 /tmp2
286 // 2 /tmp
287 // 3 /tmp2
288 // 3 /tmp3
289 //
290 // Note how both /tmp and /tmp2 are associated with multiple workspace IDs.
291 // So, given an array of worktree roots, how can we find the exactly matching ID?
292 // Let's analyze what happens when querying for [/tmp, /tmp2], from the inside out:
293 // - We start with a join of this table on itself, generating every possible
294 // pair of ((path, ID), (path, ID)), and filtering the join down to just the
295 // *overlapping* workspace IDs. For this small data set, this would look like:
296 //
297 // wt1.ID wt1.PATH | wt2.ID wt2.PATH
298 // 3 /tmp3 3 /tmp2
299 //
300 // - Moving one SELECT out, we use the first pair's ID column to invert the selection,
301 // meaning we now have a list of all the entries for our array and *subsets*
302 // of our array:
303 //
304 // ID PATH
305 // 1 /tmp
306 // 2 /tmp
307 // 2 /tmp2
308 //
309 // - To trim out the subsets, we need to exploit the fact that there can be no duplicate
310 // entries in this table. We can just use GROUP BY, COUNT, and a WHERE clause that checks
311 // for the length of our array:
312 //
313 // ID num_matching
314 // 1 2
315 //
316 // And we're done! We've found the matching ID correctly :D
317 // However, due to limitations in sqlite's query binding, we still have to do some string
318 // substitution to generate the correct query
319 // 47,116,109,112,50
320 // 2F746D7032
321
322 let query = format!(
323 r#"
324 SELECT workspace_id
325 FROM (SELECT count(workspace_id) as num_matching, workspace_id FROM worktree_roots
326 WHERE worktree_root in {array_bind} AND workspace_id NOT IN
327 (SELECT wt1.workspace_id FROM worktree_roots as wt1
328 JOIN worktree_roots as wt2
329 ON wt1.workspace_id = wt2.workspace_id
330 WHERE wt1.worktree_root NOT in {array_bind} AND wt2.worktree_root in {array_bind})
331 GROUP BY workspace_id)
332 WHERE num_matching = ?
333 "#,
334 array_bind = array_binding_stmt
335 );
336
337 // This will only be called on start up and when root workspaces change, no need to waste memory
338 // caching it.
339 let mut stmt = connection.prepare(&query)?;
340 // Make sure we bound the parameters correctly
341 debug_assert!(worktree_roots.len() + 1 == stmt.parameter_count());
342
343 for i in 0..worktree_roots.len() {
344 let path = &worktree_roots[i].as_ref().as_os_str().as_bytes();
345 stmt.raw_bind_parameter(i + 1, path)?
346 }
347 // No -1, because SQLite is 1 based
348 stmt.raw_bind_parameter(worktree_roots.len() + 1, worktree_roots.len())?;
349
350 let mut rows = stmt.raw_query();
351 let row = rows.next();
352 let result = if let Ok(Some(row)) = row {
353 Ok(Some(WorkspaceId(row.get(0)?)))
354 } else {
355 Ok(None)
356 };
357
358 // Ensure that this query only returns one row. The PRIMARY KEY constraint should catch this case
359 // but this is here to catch if someone refactors that constraint out.
360 debug_assert!(matches!(rows.next(), Ok(None)));
361
362 result
363 }
364
365 match logic(worktree_roots, connection) {
366 Ok(result) => result,
367 Err(err) => {
368 log::error!(
369 "Failed to get the workspace ID for paths {:?}, err: {}",
370 worktree_roots,
371 err
372 );
373 None
374 }
375 }
376}
377
378#[cfg(test)]
379mod tests {
380
381 use std::{
382 path::{Path, PathBuf},
383 sync::Arc,
384 thread::sleep,
385 time::Duration,
386 };
387
388 use crate::Db;
389
390 use super::WorkspaceId;
391
392 #[test]
393 fn test_empty_worktrees() {
394 // TODO determine update_worktree_roots(), workspace_id(), recent_workspaces()
395 // semantics for this case
396 }
397
398 #[test]
399 fn test_more_workspace_ids() {
400 let data = &[
401 (WorkspaceId(1), vec!["/tmp1"]),
402 (WorkspaceId(2), vec!["/tmp1", "/tmp2"]),
403 (WorkspaceId(3), vec!["/tmp1", "/tmp2", "/tmp3"]),
404 (WorkspaceId(4), vec!["/tmp2", "/tmp3"]),
405 (WorkspaceId(5), vec!["/tmp2", "/tmp3", "/tmp4"]),
406 (WorkspaceId(6), vec!["/tmp2", "/tmp4"]),
407 (WorkspaceId(7), vec!["/tmp2"]),
408 ];
409
410 let db = Db::open_in_memory();
411
412 for (workspace_id, entries) in data {
413 db.make_new_workspace();
414 db.update_worktree_roots(workspace_id, entries);
415 }
416
417 assert_eq!(Some(WorkspaceId(1)), db.workspace_id(&["/tmp1"]));
418 assert_eq!(db.workspace_id(&["/tmp1", "/tmp2"]), Some(WorkspaceId(2)));
419 assert_eq!(
420 db.workspace_id(&["/tmp1", "/tmp2", "/tmp3"]),
421 Some(WorkspaceId(3))
422 );
423 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), Some(WorkspaceId(4)));
424 assert_eq!(
425 db.workspace_id(&["/tmp2", "/tmp3", "/tmp4"]),
426 Some(WorkspaceId(5))
427 );
428 assert_eq!(db.workspace_id(&["/tmp2", "/tmp4"]), Some(WorkspaceId(6)));
429 assert_eq!(db.workspace_id(&["/tmp2"]), Some(WorkspaceId(7)));
430
431 assert_eq!(db.workspace_id(&["/tmp1", "/tmp5"]), None);
432 assert_eq!(db.workspace_id(&["/tmp5"]), None);
433 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3", "/tmp4", "/tmp5"]), None);
434 }
435
436 #[test]
437 fn test_detect_workspace_id() {
438 let data = &[
439 (WorkspaceId(1), vec!["/tmp"]),
440 (WorkspaceId(2), vec!["/tmp", "/tmp2"]),
441 (WorkspaceId(3), vec!["/tmp", "/tmp2", "/tmp3"]),
442 ];
443
444 let db = Db::open_in_memory();
445
446 for (workspace_id, entries) in data {
447 db.make_new_workspace();
448 db.update_worktree_roots(workspace_id, entries);
449 }
450
451 assert_eq!(db.workspace_id(&["/tmp2"]), None);
452 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), None);
453 assert_eq!(db.workspace_id(&["/tmp"]), Some(WorkspaceId(1)));
454 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), Some(WorkspaceId(2)));
455 assert_eq!(
456 db.workspace_id(&["/tmp", "/tmp2", "/tmp3"]),
457 Some(WorkspaceId(3))
458 );
459 }
460
461 fn arc_path(path: &'static str) -> Arc<Path> {
462 PathBuf::from(path).into()
463 }
464
465 #[test]
466 fn test_tricky_overlapping_updates() {
467 // DB state:
468 // (/tree) -> ID: 1
469 // (/tree, /tree2) -> ID: 2
470 // (/tree2, /tree3) -> ID: 3
471
472 // -> User updates 2 to: (/tree2, /tree3)
473
474 // DB state:
475 // (/tree) -> ID: 1
476 // (/tree2, /tree3) -> ID: 2
477 // Get rid of 3 for garbage collection
478
479 let data = &[
480 (WorkspaceId(1), vec!["/tmp"]),
481 (WorkspaceId(2), vec!["/tmp", "/tmp2"]),
482 (WorkspaceId(3), vec!["/tmp2", "/tmp3"]),
483 ];
484
485 let db = Db::open_in_memory();
486
487 // Load in the test data
488 for (workspace_id, entries) in data {
489 db.make_new_workspace();
490 db.update_worktree_roots(workspace_id, entries);
491 }
492
493 // Make sure the timestamp updates
494 sleep(Duration::from_secs(1));
495
496 // Execute the update
497 db.update_worktree_roots(&WorkspaceId(2), &["/tmp2", "/tmp3"]);
498
499 // Make sure that workspace 3 doesn't exist
500 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), Some(WorkspaceId(2)));
501
502 // And that workspace 1 was untouched
503 assert_eq!(db.workspace_id(&["/tmp"]), Some(WorkspaceId(1)));
504
505 // And that workspace 2 is no longer registered under these roots
506 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), None);
507
508 assert_eq!(Some(WorkspaceId(2)), db.last_workspace_id());
509
510 let recent_workspaces = db.recent_workspaces(10);
511 assert_eq!(
512 recent_workspaces.get(0).unwrap(),
513 &(WorkspaceId(2), vec![arc_path("/tmp2"), arc_path("/tmp3")])
514 );
515 assert_eq!(
516 recent_workspaces.get(1).unwrap(),
517 &(WorkspaceId(1), vec![arc_path("/tmp")])
518 );
519 }
520}