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};
10
11use indoc::indoc;
12use sqlez::{
13 connection::Connection, migrations::Migration,
14};
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
23pub(crate) const WORKSPACES_MIGRATION: Migration = Migration::new(
24 "workspace",
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,
57 /// returns the last workspace which was updated
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 res = self.with_savepoint("make_new_workspace", |conn| {
84 let workspace_id = WorkspaceId(
85 conn.prepare("INSERT INTO workspaces DEFAULT VALUES")?
86 .insert()?,
87 );
88
89 update_worktree_roots(conn, &workspace_id, worktree_roots)?;
90
91 Ok(SerializedWorkspace {
92 workspace_id,
93 dock_pane: None,
94 })
95 });
96
97 match res {
98 Ok(serialized_workspace) => serialized_workspace,
99 Err(err) => {
100 log::error!("Failed to insert new workspace into DB: {}", err);
101 Default::default()
102 }
103 }
104 }
105
106 fn workspace_id<P>(&self, worktree_roots: &[P]) -> Option<WorkspaceId>
107 where
108 P: AsRef<Path> + Debug,
109 {
110 match get_workspace_id(worktree_roots, &self) {
111 Ok(workspace_id) => workspace_id,
112 Err(err) => {
113 log::error!("Failed to get workspace_id: {}", err);
114 None
115 }
116 }
117 }
118
119 // fn get_workspace_row(&self, workspace_id: WorkspaceId) -> WorkspaceRow {
120 // unimplemented!()
121 // }
122
123 /// Updates the open paths for the given workspace id. Will garbage collect items from
124 /// any workspace ids which are no replaced by the new workspace id. Updates the timestamps
125 /// in the workspace id table
126 pub fn update_worktrees<P>(&self, workspace_id: &WorkspaceId, worktree_roots: &[P])
127 where
128 P: AsRef<Path> + Debug,
129 {
130 match self.with_savepoint("update_worktrees", |conn| {
131 update_worktree_roots(conn, workspace_id, worktree_roots)
132 }) {
133 Ok(_) => {}
134 Err(err) => log::error!(
135 "Failed to update workspace {:?} with roots {:?}, error: {}",
136 workspace_id,
137 worktree_roots,
138 err
139 ),
140 }
141 }
142
143 fn last_workspace_id(&self) -> Option<WorkspaceId> {
144 let res = self
145 .prepare(
146 "SELECT workspace_id FROM workspaces ORDER BY last_opened_timestamp DESC LIMIT 1",
147 )
148 .and_then(|stmt| stmt.maybe_row())
149 .map(|row| row.map(|id| WorkspaceId(id)));
150
151 match res {
152 Ok(result) => result,
153 Err(err) => {
154 log::error!("Failed to get last workspace id, err: {}", err);
155 return None;
156 }
157 }
158 }
159
160 /// Returns the previous workspace ids sorted by last modified along with their opened worktree roots
161 pub fn recent_workspaces(&self, limit: usize) -> Vec<(WorkspaceId, Vec<Arc<Path>>)> {
162 let res = self.with_savepoint("recent_workspaces", |conn| {
163 let ids = conn.prepare("SELECT workspace_id FROM workspaces ORDER BY last_opened_timestamp DESC LIMIT ?")?
164 .bind(limit)?
165 .rows::<i64>()?
166 .iter()
167 .map(|row| WorkspaceId(*row));
168
169 let result = Vec::new();
170
171 let stmt = conn.prepare("SELECT worktree_root FROM worktree_roots WHERE workspace_id = ?")?;
172 for workspace_id in ids {
173 let roots = stmt.bind(workspace_id.0)?
174 .rows::<Vec<u8>>()?
175 .iter()
176 .map(|row| {
177 PathBuf::from(OsStr::from_bytes(&row)).into()
178 })
179 .collect();
180 result.push((workspace_id, roots))
181 }
182
183
184 Ok(result)
185 });
186
187 match res {
188 Ok(result) => result,
189 Err(err) => {
190 log::error!("Failed to get recent workspaces, err: {}", err);
191 Vec::new()
192 }
193 }
194 }
195}
196
197fn update_worktree_roots<P>(
198 connection: &Connection,
199 workspace_id: &WorkspaceId,
200 worktree_roots: &[P],
201) -> Result<()>
202where
203 P: AsRef<Path> + Debug,
204{
205 // Lookup any old WorkspaceIds which have the same set of roots, and delete them.
206 let preexisting_id = get_workspace_id(worktree_roots, &connection)?;
207 if let Some(preexisting_id) = preexisting_id {
208 if preexisting_id != *workspace_id {
209 // Should also delete fields in other tables with cascading updates
210 connection.prepare(
211 "DELETE FROM workspaces WHERE workspace_id = ?",
212 )?
213 .bind(preexisting_id.0)?
214 .exec()?;
215 }
216 }
217
218 connection
219 .prepare("DELETE FROM worktree_roots WHERE workspace_id = ?")?
220 .bind(workspace_id.0)?
221 .exec()?;
222
223 for root in worktree_roots {
224 let path = root.as_ref().as_os_str().as_bytes();
225 // If you need to debug this, here's the string parsing:
226 // let path = root.as_ref().to_string_lossy().to_string();
227
228 connection.prepare("INSERT INTO worktree_roots(workspace_id, worktree_root) VALUES (?, ?)")?
229 .bind((workspace_id.0, path))?
230 .exec()?;
231 }
232
233 connection.prepare("UPDATE workspaces SET last_opened_timestamp = CURRENT_TIMESTAMP WHERE workspace_id = ?")?
234 .bind(workspace_id.0)?
235 .exec()?;
236
237 Ok(())
238}
239
240fn get_workspace_id<P>(worktree_roots: &[P], connection: &Connection) -> Result<Option<WorkspaceId>>
241where
242 P: AsRef<Path> + Debug,
243{
244 // Short circuit if we can
245 if worktree_roots.len() == 0 {
246 return Ok(None);
247 }
248
249 // Prepare the array binding string. SQL doesn't have syntax for this, so
250 // we have to do it ourselves.
251 let mut array_binding_stmt = "(".to_string();
252 for i in 0..worktree_roots.len() {
253 // This uses ?NNN for numbered placeholder syntax
254 array_binding_stmt.push_str(&format!("?{}", (i + 1))); //sqlite is 1-based
255 if i < worktree_roots.len() - 1 {
256 array_binding_stmt.push(',');
257 array_binding_stmt.push(' ');
258 }
259 }
260 array_binding_stmt.push(')');
261
262 // Any workspace can have multiple independent paths, and these paths
263 // can overlap in the database. Take this test data for example:
264 //
265 // [/tmp, /tmp2] -> 1
266 // [/tmp] -> 2
267 // [/tmp2, /tmp3] -> 3
268 //
269 // This would be stred in the database like so:
270 //
271 // ID PATH
272 // 1 /tmp
273 // 1 /tmp2
274 // 2 /tmp
275 // 3 /tmp2
276 // 3 /tmp3
277 //
278 // Note how both /tmp and /tmp2 are associated with multiple workspace IDs.
279 // So, given an array of worktree roots, how can we find the exactly matching ID?
280 // Let's analyze what happens when querying for [/tmp, /tmp2], from the inside out:
281 // - We start with a join of this table on itself, generating every possible
282 // pair of ((path, ID), (path, ID)), and filtering the join down to just the
283 // *overlapping but non-matching* workspace IDs. For this small data set,
284 // this would look like:
285 //
286 // wt1.ID wt1.PATH | wt2.ID wt2.PATH
287 // 3 /tmp3 3 /tmp2
288 //
289 // - Moving one SELECT out, we use the first pair's ID column to invert the selection,
290 // meaning we now have a list of all the entries for our array, minus overlapping sets,
291 // but including *subsets* of our worktree roots:
292 //
293 // ID PATH
294 // 1 /tmp
295 // 1 /tmp2
296 // 2 /tmp
297 //
298 // - To trim out the subsets, we can to exploit the PRIMARY KEY constraint that there are no
299 // duplicate entries in this table. Using a GROUP BY and a COUNT we can find the subsets of
300 // our keys:
301 //
302 // ID num_matching
303 // 1 2
304 // 2 1
305 //
306 // - And with one final WHERE num_matching = $num_of_worktree_roots, we're done! We've found the
307 // matching ID correctly :D
308 //
309 // Note: due to limitations in SQLite's query binding, we have to generate the prepared
310 // statement with string substitution (the {array_bind}) below, and then bind the
311 // parameters by number.
312 let query = format!(
313 r#"
314 SELECT workspace_id
315 FROM (SELECT count(workspace_id) as num_matching, workspace_id FROM worktree_roots
316 WHERE worktree_root in {array_bind} AND workspace_id NOT IN
317 (SELECT wt1.workspace_id FROM worktree_roots as wt1
318 JOIN worktree_roots as wt2
319 ON wt1.workspace_id = wt2.workspace_id
320 WHERE wt1.worktree_root NOT in {array_bind} AND wt2.worktree_root in {array_bind})
321 GROUP BY workspace_id)
322 WHERE num_matching = ?
323 "#,
324 array_bind = array_binding_stmt
325 );
326
327 // This will only be called on start up and when root workspaces change, no need to waste memory
328 // caching it.
329 let mut stmt = connection.prepare(&query)?;
330 // Make sure we bound the parameters correctly
331 debug_assert!(worktree_roots.len() as i32 + 1 == stmt.parameter_count());
332
333 for i in 0..worktree_roots.len() {
334 let path = &worktree_roots[i].as_ref().as_os_str().as_bytes();
335 // If you need to debug this, here's the string parsing:
336 // let path = &worktree_roots[i].as_ref().to_string_lossy().to_string()
337 stmt.bind_value(*path, i as i32 + 1);
338 }
339 // No -1, because SQLite is 1 based
340 stmt.bind_value(worktree_roots.len(), worktree_roots.len() as i32 + 1)?;
341
342 stmt.maybe_row()
343 .map(|row| row.map(|id| WorkspaceId(id)))
344}
345
346#[cfg(test)]
347mod tests {
348
349 use std::{
350 path::{Path, PathBuf},
351 sync::Arc,
352 thread::sleep,
353 time::Duration,
354 };
355
356 use crate::Db;
357
358 use super::WorkspaceId;
359
360 #[test]
361 fn test_new_worktrees_for_roots() {
362 let db = Db::open_in_memory();
363
364 // Test creation in 0 case
365 let workspace_1 = db.workspace_for_roots::<String>(&[]);
366 assert_eq!(workspace_1.workspace_id, WorkspaceId(1));
367
368 // Test pulling from recent workspaces
369 let workspace_1 = db.workspace_for_roots::<String>(&[]);
370 assert_eq!(workspace_1.workspace_id, WorkspaceId(1));
371
372 // Ensure the timestamps are different
373 sleep(Duration::from_millis(20));
374 db.make_new_workspace::<String>(&[]);
375
376 // Test pulling another value from recent workspaces
377 let workspace_2 = db.workspace_for_roots::<String>(&[]);
378 assert_eq!(workspace_2.workspace_id, WorkspaceId(2));
379
380 // Ensure the timestamps are different
381 sleep(Duration::from_millis(20));
382
383 // Test creating a new workspace that doesn't exist already
384 let workspace_3 = db.workspace_for_roots(&["/tmp", "/tmp2"]);
385 assert_eq!(workspace_3.workspace_id, WorkspaceId(3));
386
387 // Make sure it's in the recent workspaces....
388 let workspace_3 = db.workspace_for_roots::<String>(&[]);
389 assert_eq!(workspace_3.workspace_id, WorkspaceId(3));
390
391 // And that it can be pulled out again
392 let workspace_3 = db.workspace_for_roots(&["/tmp", "/tmp2"]);
393 assert_eq!(workspace_3.workspace_id, WorkspaceId(3));
394 }
395
396 #[test]
397 fn test_empty_worktrees() {
398 let db = Db::open_in_memory();
399
400 assert_eq!(None, db.workspace_id::<String>(&[]));
401
402 db.make_new_workspace::<String>(&[]); //ID 1
403 db.make_new_workspace::<String>(&[]); //ID 2
404 db.update_worktrees(&WorkspaceId(1), &["/tmp", "/tmp2"]);
405
406 db.write_to("test.db").unwrap();
407 // Sanity check
408 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), Some(WorkspaceId(1)));
409
410 db.update_worktrees::<String>(&WorkspaceId(1), &[]);
411
412 // Make sure 'no worktrees' fails correctly. returning [1, 2] from this
413 // call would be semantically correct (as those are the workspaces that
414 // don't have roots) but I'd prefer that this API to either return exactly one
415 // workspace, and None otherwise
416 assert_eq!(db.workspace_id::<String>(&[]), None,);
417
418 assert_eq!(db.last_workspace_id(), Some(WorkspaceId(1)));
419
420 assert_eq!(
421 db.recent_workspaces(2),
422 vec![(WorkspaceId(1), vec![]), (WorkspaceId(2), vec![]),],
423 )
424 }
425
426 #[test]
427 fn test_more_workspace_ids() {
428 let data = &[
429 (WorkspaceId(1), vec!["/tmp1"]),
430 (WorkspaceId(2), vec!["/tmp1", "/tmp2"]),
431 (WorkspaceId(3), vec!["/tmp1", "/tmp2", "/tmp3"]),
432 (WorkspaceId(4), vec!["/tmp2", "/tmp3"]),
433 (WorkspaceId(5), vec!["/tmp2", "/tmp3", "/tmp4"]),
434 (WorkspaceId(6), vec!["/tmp2", "/tmp4"]),
435 (WorkspaceId(7), vec!["/tmp2"]),
436 ];
437
438 let db = Db::open_in_memory();
439
440 for (workspace_id, entries) in data {
441 db.make_new_workspace::<String>(&[]);
442 db.update_worktrees(workspace_id, entries);
443 }
444
445 assert_eq!(Some(WorkspaceId(1)), db.workspace_id(&["/tmp1"]));
446 assert_eq!(db.workspace_id(&["/tmp1", "/tmp2"]), Some(WorkspaceId(2)));
447 assert_eq!(
448 db.workspace_id(&["/tmp1", "/tmp2", "/tmp3"]),
449 Some(WorkspaceId(3))
450 );
451 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), Some(WorkspaceId(4)));
452 assert_eq!(
453 db.workspace_id(&["/tmp2", "/tmp3", "/tmp4"]),
454 Some(WorkspaceId(5))
455 );
456 assert_eq!(db.workspace_id(&["/tmp2", "/tmp4"]), Some(WorkspaceId(6)));
457 assert_eq!(db.workspace_id(&["/tmp2"]), Some(WorkspaceId(7)));
458
459 assert_eq!(db.workspace_id(&["/tmp1", "/tmp5"]), None);
460 assert_eq!(db.workspace_id(&["/tmp5"]), None);
461 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3", "/tmp4", "/tmp5"]), None);
462 }
463
464 #[test]
465 fn test_detect_workspace_id() {
466 let data = &[
467 (WorkspaceId(1), vec!["/tmp"]),
468 (WorkspaceId(2), vec!["/tmp", "/tmp2"]),
469 (WorkspaceId(3), vec!["/tmp", "/tmp2", "/tmp3"]),
470 ];
471
472 let db = Db::open_in_memory();
473
474 for (workspace_id, entries) in data {
475 db.make_new_workspace::<String>(&[]);
476 db.update_worktrees(workspace_id, entries);
477 }
478
479 assert_eq!(db.workspace_id(&["/tmp2"]), None);
480 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), None);
481 assert_eq!(db.workspace_id(&["/tmp"]), Some(WorkspaceId(1)));
482 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), Some(WorkspaceId(2)));
483 assert_eq!(
484 db.workspace_id(&["/tmp", "/tmp2", "/tmp3"]),
485 Some(WorkspaceId(3))
486 );
487 }
488
489 fn arc_path(path: &'static str) -> Arc<Path> {
490 PathBuf::from(path).into()
491 }
492
493 #[test]
494 fn test_tricky_overlapping_updates() {
495 // DB state:
496 // (/tree) -> ID: 1
497 // (/tree, /tree2) -> ID: 2
498 // (/tree2, /tree3) -> ID: 3
499
500 // -> User updates 2 to: (/tree2, /tree3)
501
502 // DB state:
503 // (/tree) -> ID: 1
504 // (/tree2, /tree3) -> ID: 2
505 // Get rid of 3 for garbage collection
506
507 let data = &[
508 (WorkspaceId(1), vec!["/tmp"]),
509 (WorkspaceId(2), vec!["/tmp", "/tmp2"]),
510 (WorkspaceId(3), vec!["/tmp2", "/tmp3"]),
511 ];
512
513 let db = Db::open_in_memory();
514
515 // Load in the test data
516 for (workspace_id, entries) in data {
517 db.make_new_workspace::<String>(&[]);
518 db.update_worktrees(workspace_id, entries);
519 }
520
521 // Execute the update
522 db.update_worktrees(&WorkspaceId(2), &["/tmp2", "/tmp3"]);
523
524 // Make sure that workspace 3 doesn't exist
525 assert_eq!(db.workspace_id(&["/tmp2", "/tmp3"]), Some(WorkspaceId(2)));
526
527 // And that workspace 1 was untouched
528 assert_eq!(db.workspace_id(&["/tmp"]), Some(WorkspaceId(1)));
529
530 // And that workspace 2 is no longer registered under these roots
531 assert_eq!(db.workspace_id(&["/tmp", "/tmp2"]), None);
532
533 assert_eq!(Some(WorkspaceId(2)), db.last_workspace_id());
534
535 let recent_workspaces = db.recent_workspaces(10);
536 assert_eq!(
537 recent_workspaces.get(0).unwrap(),
538 &(WorkspaceId(2), vec![arc_path("/tmp2"), arc_path("/tmp3")])
539 );
540 assert_eq!(
541 recent_workspaces.get(1).unwrap(),
542 &(WorkspaceId(1), vec![arc_path("/tmp")])
543 );
544 }
545}