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