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