1use anyhow::Result;
2
3use std::{
4 ffi::OsStr,
5 fmt::Debug,
6 os::unix::prelude::OsStrExt,
7 path::{Path, PathBuf},
8 sync::Arc,
9 time::{SystemTime, UNIX_EPOCH},
10};
11
12use anyhow::Result;
13use indoc::indoc;
14use sqlez::{connection::Connection, migrations::Migration};
15
16use crate::pane::SerializedDockPane;
17
18use super::Db;
19
20// If you need to debug the worktree root code, change 'BLOB' here to 'TEXT' for easier debugging
21// you might want to update some of the parsing code as well, I've left the variations in but commented
22// out. This will panic if run on an existing db that has already been migrated
23const WORKSPACES_MIGRATION: Migration = Migration::new(
24 "migrations",
25 &[indoc! {"
26 CREATE TABLE workspaces(
27 workspace_id INTEGER PRIMARY KEY,
28 timestamp TEXT DEFAULT CURRENT_TIMESTAMP NOT NULL
29 ) STRICT;
30
31 CREATE TABLE worktree_roots(
32 worktree_root BLOB NOT NULL,
33 workspace_id INTEGER NOT NULL,
34 FOREIGN KEY(workspace_id) REFERENCES workspaces(workspace_id) ON DELETE CASCADE
35 PRIMARY KEY(worktree_root, workspace_id)
36 ) STRICT;"}],
37);
38
39#[derive(Debug, PartialEq, Eq, Copy, Clone, Default)]
40pub struct WorkspaceId(i64);
41
42impl WorkspaceId {
43 pub fn raw_id(&self) -> i64 {
44 self.0
45 }
46}
47
48#[derive(Default, Debug)]
49pub struct SerializedWorkspace {
50 pub workspace_id: WorkspaceId,
51 // pub center_group: SerializedPaneGroup,
52 pub dock_pane: Option<SerializedDockPane>,
53}
54
55impl Db {
56 /// Finds or creates a workspace id for the given set of worktree roots. If the passed worktree roots is empty, return the
57 /// the last workspace id
58 pub fn workspace_for_roots<P>(&self, worktree_roots: &[P]) -> SerializedWorkspace
59 where
60 P: AsRef<Path> + Debug,
61 {
62 // Find the workspace id which is uniquely identified by this set of paths
63 // return it if found
64 let mut workspace_id = self.workspace_id(worktree_roots);
65 if workspace_id.is_none() && worktree_roots.len() == 0 {
66 workspace_id = self.last_workspace_id();
67 }
68
69 if let Some(workspace_id) = workspace_id {
70 SerializedWorkspace {
71 workspace_id,
72 dock_pane: self.get_dock_pane(workspace_id),
73 }
74 } else {
75 self.make_new_workspace(worktree_roots)
76 }
77 }
78
79 fn make_new_workspace<P>(&self, worktree_roots: &[P]) -> SerializedWorkspace
80 where
81 P: AsRef<Path> + Debug,
82 {
83 let result = (|| {
84 let tx = self.transaction()?;
85 tx.execute("INSERT INTO workspaces(last_opened_timestamp) VALUES" (?), [current_millis()?])?;
86
87 let id = WorkspaceId(tx.last_insert_rowid());
88
89 update_worktree_roots(&tx, &id, worktree_roots)?;
90
91 tx.commit()?;
92
93 Ok(SerializedWorkspace {
94 workspace_id: id,
95 dock_pane: None,
96 })
97 })();
98
99 match result {
100 Ok(serialized_workspace) => serialized_workspace,
101 Err(err) => {
102 log::error!("Failed to insert new workspace into DB: {}", err);
103 Default::default()
104 }
105 }
106 }
107
108 fn workspace_id<P>(&self, worktree_roots: &[P]) -> Option<WorkspaceId>
109 where
110 P: AsRef<Path> + Debug,
111 {
112 self.real()
113 .map(|db| {
114 let lock = db.connection.lock();
115
116 match get_workspace_id(worktree_roots, &lock) {
117 Ok(workspace_id) => workspace_id,
118 Err(err) => {
119 log::error!("Failed to get workspace_id: {}", err);
120 None
121 }
122 }
123 })
124 .unwrap_or(None)
125 }
126
127 // fn get_workspace_row(&self, workspace_id: WorkspaceId) -> WorkspaceRow {
128 // unimplemented!()
129 // }
130
131 /// Updates the open paths for the given workspace id. Will garbage collect items from
132 /// any workspace ids which are no replaced by the new workspace id. Updates the timestamps
133 /// in the workspace id table
134 pub fn update_worktrees<P>(&self, workspace_id: &WorkspaceId, worktree_roots: &[P])
135 where
136 P: AsRef<Path> + Debug,
137 {
138 fn logic<P>(
139 connection: &mut Connection,
140 workspace_id: &WorkspaceId,
141 worktree_roots: &[P],
142 ) -> Result<()>
143 where
144 P: AsRef<Path> + Debug,
145 {
146 let tx = connection.transaction()?;
147 update_worktree_roots(&tx, workspace_id, worktree_roots)?;
148 tx.commit()?;
149 Ok(())
150 }
151
152 self.real().map(|db| {
153 let mut lock = db.connection.lock();
154
155 match logic(&mut lock, workspace_id, worktree_roots) {
156 Ok(_) => {}
157 Err(err) => {
158 log::error!(
159 "Failed to update the worktree roots for {:?}, roots: {:?}, error: {}",
160 workspace_id,
161 worktree_roots,
162 err
163 );
164 }
165 }
166 });
167 }
168
169 fn last_workspace_id(&self) -> Option<WorkspaceId> {
170 fn logic(connection: &mut Connection) -> Result<Option<WorkspaceId>> {
171 let mut stmt = connection.prepare(
172 "SELECT workspace_id FROM workspaces ORDER BY last_opened_timestamp DESC LIMIT 1",
173 )?;
174
175 Ok(stmt
176 .query_row([], |row| Ok(WorkspaceId(row.get(0)?)))
177 .optional()?)
178 }
179
180 self.real()
181 .map(|db| {
182 let mut lock = db.connection.lock();
183
184 match logic(&mut lock) {
185 Ok(result) => result,
186 Err(err) => {
187 log::error!("Failed to get last workspace id, err: {}", err);
188 None
189 }
190 }
191 })
192 .unwrap_or(None)
193 }
194
195 /// Returns the previous workspace ids sorted by last modified along with their opened worktree roots
196 pub fn recent_workspaces(&self, limit: usize) -> Vec<(WorkspaceId, Vec<Arc<Path>>)> {
197 fn logic(
198 connection: &mut Connection,
199 limit: usize,
200 ) -> Result<Vec<(WorkspaceId, Vec<Arc<Path>>)>, anyhow::Error> {
201 let tx = connection.transaction()?;
202 let result = {
203 let mut stmt = tx.prepare(
204 "SELECT workspace_id FROM workspaces ORDER BY last_opened_timestamp DESC LIMIT ?",
205 )?;
206
207 let workspace_ids = stmt
208 .query_map([limit], |row| Ok(WorkspaceId(row.get(0)?)))?
209 .collect::<Result<Vec<_>, rusqlite::Error>>()?;
210
211 let mut result = Vec::new();
212 let mut stmt =
213 tx.prepare("SELECT worktree_root FROM worktree_roots WHERE workspace_id = ?")?;
214 for workspace_id in workspace_ids {
215 let roots = stmt
216 .query_map([workspace_id.0], |row| {
217 let row = row.get::<_, Vec<u8>>(0)?;
218 Ok(PathBuf::from(OsStr::from_bytes(&row)).into())
219 // If you need to debug this, here's the string parsing:
220 // let row = row.get::<_, String>(0)?;
221 // Ok(PathBuf::from(row).into())
222 })?
223 .collect::<Result<Vec<_>, rusqlite::Error>>()?;
224 result.push((workspace_id, roots))
225 }
226
227 result
228 };
229 tx.commit()?;
230 return Ok(result);
231 }
232
233 self.real()
234 .map(|db| {
235 let mut lock = db.connection.lock();
236
237 match logic(&mut lock, limit) {
238 Ok(result) => result,
239 Err(err) => {
240 log::error!("Failed to get recent workspaces, err: {}", err);
241 Vec::new()
242 }
243 }
244 })
245 .unwrap_or_else(|| Vec::new())
246 }
247}
248
249fn current_millis() -> Result<u64, anyhow::Error> {
250 // SQLite only supports u64 integers, which means this code will trigger
251 // undefined behavior in 584 million years. It's probably fine.
252 Ok(SystemTime::now().duration_since(UNIX_EPOCH)?.as_millis() as u64)
253}
254
255fn update_worktree_roots<P>(
256 connection: &Connection,
257 workspace_id: &WorkspaceId,
258 worktree_roots: &[P],
259) -> Result<()>
260where
261 P: AsRef<Path> + Debug,
262{
263 // Lookup any old WorkspaceIds which have the same set of roots, and delete them.
264 let preexisting_id = get_workspace_id(worktree_roots, &connection)?;
265 if let Some(preexisting_id) = preexisting_id {
266 if preexisting_id != *workspace_id {
267 // Should also delete fields in other tables with cascading updates
268 connection.execute(
269 "DELETE FROM workspaces WHERE workspace_id = ?",
270 [preexisting_id.0],
271 )?;
272 }
273 }
274
275 connection.execute(
276 "DELETE FROM worktree_roots WHERE workspace_id = ?",
277 [workspace_id.0],
278 )?;
279
280 for root in worktree_roots {
281 let path = root.as_ref().as_os_str().as_bytes();
282 // If you need to debug this, here's the string parsing:
283 // let path = root.as_ref().to_string_lossy().to_string();
284
285 connection.execute(
286 "INSERT INTO worktree_roots(workspace_id, worktree_root) VALUES (?, ?)",
287 params![workspace_id.0, path],
288 )?;
289 }
290
291 connection.execute(
292 "UPDATE workspaces SET last_opened_timestamp = ? WHERE workspace_id = ?",
293 params![current_millis()?, workspace_id.0],
294 )?;
295
296 Ok(())
297}
298
299fn get_workspace_id<P>(worktree_roots: &[P], connection: &Connection) -> Result<Option<WorkspaceId>>
300where
301 P: AsRef<Path> + Debug,
302{
303 // fn logic<P>(
304 // worktree_roots: &[P],
305 // connection: &Connection,
306 // ) -> Result<Option<WorkspaceId>, anyhow::Error>
307 // where
308 // P: AsRef<Path> + Debug,
309 // {
310 // Short circuit if we can
311 if worktree_roots.len() == 0 {
312 return Ok(None);
313 }
314
315 // Prepare the array binding string. SQL doesn't have syntax for this, so
316 // we have to do it ourselves.
317 let mut array_binding_stmt = "(".to_string();
318 for i in 0..worktree_roots.len() {
319 // This uses ?NNN for numbered placeholder syntax
320 array_binding_stmt.push_str(&format!("?{}", (i + 1))); //sqlite is 1-based
321 if i < worktree_roots.len() - 1 {
322 array_binding_stmt.push(',');
323 array_binding_stmt.push(' ');
324 }
325 }
326 array_binding_stmt.push(')');
327 // Any workspace can have multiple independent paths, and these paths
328 // can overlap in the database. Take this test data for example:
329 //
330 // [/tmp, /tmp2] -> 1
331 // [/tmp] -> 2
332 // [/tmp2, /tmp3] -> 3
333 //
334 // This would be stred in the database like so:
335 //
336 // ID PATH
337 // 1 /tmp
338 // 1 /tmp2
339 // 2 /tmp
340 // 3 /tmp2
341 // 3 /tmp3
342 //
343 // Note how both /tmp and /tmp2 are associated with multiple workspace IDs.
344 // So, given an array of worktree roots, how can we find the exactly matching ID?
345 // Let's analyze what happens when querying for [/tmp, /tmp2], from the inside out:
346 // - We start with a join of this table on itself, generating every possible
347 // pair of ((path, ID), (path, ID)), and filtering the join down to just the
348 // *overlapping but non-matching* workspace IDs. For this small data set,
349 // this would look like:
350 //
351 // wt1.ID wt1.PATH | wt2.ID wt2.PATH
352 // 3 /tmp3 3 /tmp2
353 //
354 // - Moving one SELECT out, we use the first pair's ID column to invert the selection,
355 // meaning we now have a list of all the entries for our array, minus overlapping sets,
356 // but including *subsets* of our worktree roots:
357 //
358 // ID PATH
359 // 1 /tmp
360 // 1 /tmp2
361 // 2 /tmp
362 //
363 // - To trim out the subsets, we can to exploit the PRIMARY KEY constraint that there are no
364 // duplicate entries in this table. Using a GROUP BY and a COUNT we can find the subsets of
365 // our keys:
366 //
367 // ID num_matching
368 // 1 2
369 // 2 1
370 //
371 // - And with one final WHERE num_matching = $num_of_worktree_roots, we're done! We've found the
372 // matching ID correctly :D
373 //
374 // Note: due to limitations in SQLite's query binding, we have to generate the prepared
375 // statement with string substitution (the {array_bind}) below, and then bind the
376 // parameters by number.
377 let query = format!(
378 r#"
379 SELECT workspace_id
380 FROM (SELECT count(workspace_id) as num_matching, workspace_id FROM worktree_roots
381 WHERE worktree_root in {array_bind} AND workspace_id NOT IN
382 (SELECT wt1.workspace_id FROM worktree_roots as wt1
383 JOIN worktree_roots as wt2
384 ON wt1.workspace_id = wt2.workspace_id
385 WHERE wt1.worktree_root NOT in {array_bind} AND wt2.worktree_root in {array_bind})
386 GROUP BY workspace_id)
387 WHERE num_matching = ?
388 "#,
389 array_bind = array_binding_stmt
390 );
391
392 // This will only be called on start up and when root workspaces change, no need to waste memory
393 // caching it.
394 let mut stmt = connection.prepare(&query)?;
395 // Make sure we bound the parameters correctly
396 debug_assert!(worktree_roots.len() + 1 == stmt.parameter_count());
397
398 for i in 0..worktree_roots.len() {
399 let path = &worktree_roots[i].as_ref().as_os_str().as_bytes();
400 // If you need to debug this, here's the string parsing:
401 // let path = &worktree_roots[i].as_ref().to_string_lossy().to_string()
402 stmt.raw_bind_parameter(i + 1, path)?
403 }
404 // No -1, because SQLite is 1 based
405 stmt.raw_bind_parameter(worktree_roots.len() + 1, worktree_roots.len())?;
406
407 let mut rows = stmt.raw_query();
408 let row = rows.next();
409 let result = if let Ok(Some(row)) = row {
410 Ok(Some(WorkspaceId(row.get(0)?)))
411 } else {
412 Ok(None)
413 };
414
415 // Ensure that this query only returns one row. The PRIMARY KEY constraint should catch this case
416 // but this is here to catch if someone refactors that constraint out.
417 debug_assert!(matches!(rows.next(), Ok(None)));
418
419 result
420 // }
421
422 // match logic(worktree_roots, connection) {
423 // Ok(result) => result,
424 // Err(err) => {
425 // log::error!(
426 // "Failed to get the workspace ID for paths {:?}, err: {}",
427 // worktree_roots,
428 // err
429 // );
430 // None
431 // }
432 // }
433}
434
435#[cfg(test)]
436mod tests {
437
438 use std::{
439 path::{Path, PathBuf},
440 sync::Arc,
441 thread::sleep,
442 time::Duration,
443 };
444
445 use crate::Db;
446
447 use super::WorkspaceId;
448
449 #[test]
450 fn test_new_worktrees_for_roots() {
451 let db = Db::open_in_memory();
452
453 // Test creation in 0 case
454 let workspace_1 = db.workspace_for_roots::<String>(&[]);
455 assert_eq!(workspace_1.workspace_id, WorkspaceId(1));
456
457 // Test pulling from recent workspaces
458 let workspace_1 = db.workspace_for_roots::<String>(&[]);
459 assert_eq!(workspace_1.workspace_id, WorkspaceId(1));
460
461 // Ensure the timestamps are different
462 sleep(Duration::from_millis(20));
463 db.make_new_workspace::<String>(&[]);
464
465 // Test pulling another value from recent workspaces
466 let workspace_2 = db.workspace_for_roots::<String>(&[]);
467 assert_eq!(workspace_2.workspace_id, WorkspaceId(2));
468
469 // Ensure the timestamps are different
470 sleep(Duration::from_millis(20));
471
472 // Test creating a new workspace that doesn't exist already
473 let workspace_3 = db.workspace_for_roots(&["/tmp", "/tmp2"]);
474 assert_eq!(workspace_3.workspace_id, WorkspaceId(3));
475
476 // Make sure it's in the recent workspaces....
477 let workspace_3 = db.workspace_for_roots::<String>(&[]);
478 assert_eq!(workspace_3.workspace_id, WorkspaceId(3));
479
480 // And that it can be pulled out again
481 let workspace_3 = db.workspace_for_roots(&["/tmp", "/tmp2"]);
482 assert_eq!(workspace_3.workspace_id, WorkspaceId(3));
483 }
484
485 #[test]
486 fn test_empty_worktrees() {
487 let db = Db::open_in_memory();
488
489 assert_eq!(None, db.workspace_id::<String>(&[]));
490
491 db.make_new_workspace::<String>(&[]); //ID 1
492 db.make_new_workspace::<String>(&[]); //ID 2
493 db.update_worktrees(&WorkspaceId(1), &["/tmp", "/tmp2"]);
494
495 db.write_to("test.db").unwrap();
496 // Sanity check
497 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), Some(WorkspaceId(1)));
498
499 db.update_worktrees::<String>(&WorkspaceId(1), &[]);
500
501 // Make sure 'no worktrees' fails correctly. returning [1, 2] from this
502 // call would be semantically correct (as those are the workspaces that
503 // don't have roots) but I'd prefer that this API to either return exactly one
504 // workspace, and None otherwise
505 assert_eq!(db.workspace_id::<String>(&[]), None,);
506
507 assert_eq!(db.last_workspace_id(), Some(WorkspaceId(1)));
508
509 assert_eq!(
510 db.recent_workspaces(2),
511 vec![(WorkspaceId(1), vec![]), (WorkspaceId(2), vec![]),],
512 )
513 }
514
515 #[test]
516 fn test_more_workspace_ids() {
517 let data = &[
518 (WorkspaceId(1), vec!["/tmp1"]),
519 (WorkspaceId(2), vec!["/tmp1", "/tmp2"]),
520 (WorkspaceId(3), vec!["/tmp1", "/tmp2", "/tmp3"]),
521 (WorkspaceId(4), vec!["/tmp2", "/tmp3"]),
522 (WorkspaceId(5), vec!["/tmp2", "/tmp3", "/tmp4"]),
523 (WorkspaceId(6), vec!["/tmp2", "/tmp4"]),
524 (WorkspaceId(7), vec!["/tmp2"]),
525 ];
526
527 let db = Db::open_in_memory();
528
529 for (workspace_id, entries) in data {
530 db.make_new_workspace::<String>(&[]);
531 db.update_worktrees(workspace_id, entries);
532 }
533
534 assert_eq!(Some(WorkspaceId(1)), db.workspace_id(&["/tmp1"]));
535 assert_eq!(db.workspace_id(&["/tmp1", "/tmp2"]), Some(WorkspaceId(2)));
536 assert_eq!(
537 db.workspace_id(&["/tmp1", "/tmp2", "/tmp3"]),
538 Some(WorkspaceId(3))
539 );
540 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), Some(WorkspaceId(4)));
541 assert_eq!(
542 db.workspace_id(&["/tmp2", "/tmp3", "/tmp4"]),
543 Some(WorkspaceId(5))
544 );
545 assert_eq!(db.workspace_id(&["/tmp2", "/tmp4"]), Some(WorkspaceId(6)));
546 assert_eq!(db.workspace_id(&["/tmp2"]), Some(WorkspaceId(7)));
547
548 assert_eq!(db.workspace_id(&["/tmp1", "/tmp5"]), None);
549 assert_eq!(db.workspace_id(&["/tmp5"]), None);
550 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3", "/tmp4", "/tmp5"]), None);
551 }
552
553 #[test]
554 fn test_detect_workspace_id() {
555 let data = &[
556 (WorkspaceId(1), vec!["/tmp"]),
557 (WorkspaceId(2), vec!["/tmp", "/tmp2"]),
558 (WorkspaceId(3), vec!["/tmp", "/tmp2", "/tmp3"]),
559 ];
560
561 let db = Db::open_in_memory();
562
563 for (workspace_id, entries) in data {
564 db.make_new_workspace::<String>(&[]);
565 db.update_worktrees(workspace_id, entries);
566 }
567
568 assert_eq!(db.workspace_id(&["/tmp2"]), None);
569 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), None);
570 assert_eq!(db.workspace_id(&["/tmp"]), Some(WorkspaceId(1)));
571 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), Some(WorkspaceId(2)));
572 assert_eq!(
573 db.workspace_id(&["/tmp", "/tmp2", "/tmp3"]),
574 Some(WorkspaceId(3))
575 );
576 }
577
578 fn arc_path(path: &'static str) -> Arc<Path> {
579 PathBuf::from(path).into()
580 }
581
582 #[test]
583 fn test_tricky_overlapping_updates() {
584 // DB state:
585 // (/tree) -> ID: 1
586 // (/tree, /tree2) -> ID: 2
587 // (/tree2, /tree3) -> ID: 3
588
589 // -> User updates 2 to: (/tree2, /tree3)
590
591 // DB state:
592 // (/tree) -> ID: 1
593 // (/tree2, /tree3) -> ID: 2
594 // Get rid of 3 for garbage collection
595
596 let data = &[
597 (WorkspaceId(1), vec!["/tmp"]),
598 (WorkspaceId(2), vec!["/tmp", "/tmp2"]),
599 (WorkspaceId(3), vec!["/tmp2", "/tmp3"]),
600 ];
601
602 let db = Db::open_in_memory();
603
604 // Load in the test data
605 for (workspace_id, entries) in data {
606 db.make_new_workspace::<String>(&[]);
607 db.update_worktrees(workspace_id, entries);
608 }
609
610 // Execute the update
611 db.update_worktrees(&WorkspaceId(2), &["/tmp2", "/tmp3"]);
612
613 // Make sure that workspace 3 doesn't exist
614 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), Some(WorkspaceId(2)));
615
616 // And that workspace 1 was untouched
617 assert_eq!(db.workspace_id(&["/tmp"]), Some(WorkspaceId(1)));
618
619 // And that workspace 2 is no longer registered under these roots
620 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), None);
621
622 assert_eq!(Some(WorkspaceId(2)), db.last_workspace_id());
623
624 let recent_workspaces = db.recent_workspaces(10);
625 assert_eq!(
626 recent_workspaces.get(0).unwrap(),
627 &(WorkspaceId(2), vec![arc_path("/tmp2"), arc_path("/tmp3")])
628 );
629 assert_eq!(
630 recent_workspaces.get(1).unwrap(),
631 &(WorkspaceId(1), vec![arc_path("/tmp")])
632 );
633 }
634}