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