1mod access_token;
2mod contact;
3mod follower;
4mod language_server;
5mod project;
6mod project_collaborator;
7mod room;
8mod room_participant;
9mod server;
10mod signup;
11#[cfg(test)]
12mod tests;
13mod user;
14mod worktree;
15mod worktree_diagnostic_summary;
16mod worktree_entry;
17mod worktree_repository;
18mod worktree_repository_statuses;
19mod worktree_settings_file;
20
21use crate::executor::Executor;
22use crate::{Error, Result};
23use anyhow::anyhow;
24use collections::{BTreeMap, HashMap, HashSet};
25pub use contact::Contact;
26use dashmap::DashMap;
27use futures::StreamExt;
28use hyper::StatusCode;
29use rand::prelude::StdRng;
30use rand::{Rng, SeedableRng};
31use rpc::{proto, ConnectionId};
32use sea_orm::Condition;
33pub use sea_orm::ConnectOptions;
34use sea_orm::{
35 entity::prelude::*, ActiveValue, ConnectionTrait, DatabaseConnection, DatabaseTransaction,
36 DbErr, FromQueryResult, IntoActiveModel, IsolationLevel, JoinType, QueryOrder, QuerySelect,
37 Statement, TransactionTrait,
38};
39use sea_query::{Alias, Expr, OnConflict, Query};
40use serde::{Deserialize, Serialize};
41pub use signup::{Invite, NewSignup, WaitlistSummary};
42use sqlx::migrate::{Migrate, Migration, MigrationSource};
43use sqlx::Connection;
44use std::ops::{Deref, DerefMut};
45use std::path::Path;
46use std::time::Duration;
47use std::{future::Future, marker::PhantomData, rc::Rc, sync::Arc};
48use tokio::sync::{Mutex, OwnedMutexGuard};
49pub use user::Model as User;
50
51pub struct Database {
52 options: ConnectOptions,
53 pool: DatabaseConnection,
54 rooms: DashMap<RoomId, Arc<Mutex<()>>>,
55 rng: Mutex<StdRng>,
56 executor: Executor,
57 #[cfg(test)]
58 runtime: Option<tokio::runtime::Runtime>,
59}
60
61impl Database {
62 pub async fn new(options: ConnectOptions, executor: Executor) -> Result<Self> {
63 Ok(Self {
64 options: options.clone(),
65 pool: sea_orm::Database::connect(options).await?,
66 rooms: DashMap::with_capacity(16384),
67 rng: Mutex::new(StdRng::seed_from_u64(0)),
68 executor,
69 #[cfg(test)]
70 runtime: None,
71 })
72 }
73
74 #[cfg(test)]
75 pub fn reset(&self) {
76 self.rooms.clear();
77 }
78
79 pub async fn migrate(
80 &self,
81 migrations_path: &Path,
82 ignore_checksum_mismatch: bool,
83 ) -> anyhow::Result<Vec<(Migration, Duration)>> {
84 let migrations = MigrationSource::resolve(migrations_path)
85 .await
86 .map_err(|err| anyhow!("failed to load migrations: {err:?}"))?;
87
88 let mut connection = sqlx::AnyConnection::connect(self.options.get_url()).await?;
89
90 connection.ensure_migrations_table().await?;
91 let applied_migrations: HashMap<_, _> = connection
92 .list_applied_migrations()
93 .await?
94 .into_iter()
95 .map(|m| (m.version, m))
96 .collect();
97
98 let mut new_migrations = Vec::new();
99 for migration in migrations {
100 match applied_migrations.get(&migration.version) {
101 Some(applied_migration) => {
102 if migration.checksum != applied_migration.checksum && !ignore_checksum_mismatch
103 {
104 Err(anyhow!(
105 "checksum mismatch for applied migration {}",
106 migration.description
107 ))?;
108 }
109 }
110 None => {
111 let elapsed = connection.apply(&migration).await?;
112 new_migrations.push((migration, elapsed));
113 }
114 }
115 }
116
117 Ok(new_migrations)
118 }
119
120 pub async fn create_server(&self, environment: &str) -> Result<ServerId> {
121 self.transaction(|tx| async move {
122 let server = server::ActiveModel {
123 environment: ActiveValue::set(environment.into()),
124 ..Default::default()
125 }
126 .insert(&*tx)
127 .await?;
128 Ok(server.id)
129 })
130 .await
131 }
132
133 pub async fn stale_room_ids(
134 &self,
135 environment: &str,
136 new_server_id: ServerId,
137 ) -> Result<Vec<RoomId>> {
138 self.transaction(|tx| async move {
139 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
140 enum QueryAs {
141 RoomId,
142 }
143
144 let stale_server_epochs = self
145 .stale_server_ids(environment, new_server_id, &tx)
146 .await?;
147 Ok(room_participant::Entity::find()
148 .select_only()
149 .column(room_participant::Column::RoomId)
150 .distinct()
151 .filter(
152 room_participant::Column::AnsweringConnectionServerId
153 .is_in(stale_server_epochs),
154 )
155 .into_values::<_, QueryAs>()
156 .all(&*tx)
157 .await?)
158 })
159 .await
160 }
161
162 pub async fn refresh_room(
163 &self,
164 room_id: RoomId,
165 new_server_id: ServerId,
166 ) -> Result<RoomGuard<RefreshedRoom>> {
167 self.room_transaction(room_id, |tx| async move {
168 let stale_participant_filter = Condition::all()
169 .add(room_participant::Column::RoomId.eq(room_id))
170 .add(room_participant::Column::AnsweringConnectionId.is_not_null())
171 .add(room_participant::Column::AnsweringConnectionServerId.ne(new_server_id));
172
173 let stale_participant_user_ids = room_participant::Entity::find()
174 .filter(stale_participant_filter.clone())
175 .all(&*tx)
176 .await?
177 .into_iter()
178 .map(|participant| participant.user_id)
179 .collect::<Vec<_>>();
180
181 // Delete participants who failed to reconnect and cancel their calls.
182 let mut canceled_calls_to_user_ids = Vec::new();
183 room_participant::Entity::delete_many()
184 .filter(stale_participant_filter)
185 .exec(&*tx)
186 .await?;
187 let called_participants = room_participant::Entity::find()
188 .filter(
189 Condition::all()
190 .add(
191 room_participant::Column::CallingUserId
192 .is_in(stale_participant_user_ids.iter().copied()),
193 )
194 .add(room_participant::Column::AnsweringConnectionId.is_null()),
195 )
196 .all(&*tx)
197 .await?;
198 room_participant::Entity::delete_many()
199 .filter(
200 room_participant::Column::Id
201 .is_in(called_participants.iter().map(|participant| participant.id)),
202 )
203 .exec(&*tx)
204 .await?;
205 canceled_calls_to_user_ids.extend(
206 called_participants
207 .into_iter()
208 .map(|participant| participant.user_id),
209 );
210
211 let room = self.get_room(room_id, &tx).await?;
212 // Delete the room if it becomes empty.
213 if room.participants.is_empty() {
214 project::Entity::delete_many()
215 .filter(project::Column::RoomId.eq(room_id))
216 .exec(&*tx)
217 .await?;
218 room::Entity::delete_by_id(room_id).exec(&*tx).await?;
219 }
220
221 Ok(RefreshedRoom {
222 room,
223 stale_participant_user_ids,
224 canceled_calls_to_user_ids,
225 })
226 })
227 .await
228 }
229
230 pub async fn delete_stale_servers(
231 &self,
232 environment: &str,
233 new_server_id: ServerId,
234 ) -> Result<()> {
235 self.transaction(|tx| async move {
236 server::Entity::delete_many()
237 .filter(
238 Condition::all()
239 .add(server::Column::Environment.eq(environment))
240 .add(server::Column::Id.ne(new_server_id)),
241 )
242 .exec(&*tx)
243 .await?;
244 Ok(())
245 })
246 .await
247 }
248
249 async fn stale_server_ids(
250 &self,
251 environment: &str,
252 new_server_id: ServerId,
253 tx: &DatabaseTransaction,
254 ) -> Result<Vec<ServerId>> {
255 let stale_servers = server::Entity::find()
256 .filter(
257 Condition::all()
258 .add(server::Column::Environment.eq(environment))
259 .add(server::Column::Id.ne(new_server_id)),
260 )
261 .all(&*tx)
262 .await?;
263 Ok(stale_servers.into_iter().map(|server| server.id).collect())
264 }
265
266 // users
267
268 pub async fn create_user(
269 &self,
270 email_address: &str,
271 admin: bool,
272 params: NewUserParams,
273 ) -> Result<NewUserResult> {
274 self.transaction(|tx| async {
275 let tx = tx;
276 let user = user::Entity::insert(user::ActiveModel {
277 email_address: ActiveValue::set(Some(email_address.into())),
278 github_login: ActiveValue::set(params.github_login.clone()),
279 github_user_id: ActiveValue::set(Some(params.github_user_id)),
280 admin: ActiveValue::set(admin),
281 metrics_id: ActiveValue::set(Uuid::new_v4()),
282 ..Default::default()
283 })
284 .on_conflict(
285 OnConflict::column(user::Column::GithubLogin)
286 .update_column(user::Column::GithubLogin)
287 .to_owned(),
288 )
289 .exec_with_returning(&*tx)
290 .await?;
291
292 Ok(NewUserResult {
293 user_id: user.id,
294 metrics_id: user.metrics_id.to_string(),
295 signup_device_id: None,
296 inviting_user_id: None,
297 })
298 })
299 .await
300 }
301
302 pub async fn get_user_by_id(&self, id: UserId) -> Result<Option<user::Model>> {
303 self.transaction(|tx| async move { Ok(user::Entity::find_by_id(id).one(&*tx).await?) })
304 .await
305 }
306
307 pub async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<user::Model>> {
308 self.transaction(|tx| async {
309 let tx = tx;
310 Ok(user::Entity::find()
311 .filter(user::Column::Id.is_in(ids.iter().copied()))
312 .all(&*tx)
313 .await?)
314 })
315 .await
316 }
317
318 pub async fn get_user_by_github_login(&self, github_login: &str) -> Result<Option<User>> {
319 self.transaction(|tx| async move {
320 Ok(user::Entity::find()
321 .filter(user::Column::GithubLogin.eq(github_login))
322 .one(&*tx)
323 .await?)
324 })
325 .await
326 }
327
328 pub async fn get_or_create_user_by_github_account(
329 &self,
330 github_login: &str,
331 github_user_id: Option<i32>,
332 github_email: Option<&str>,
333 ) -> Result<Option<User>> {
334 self.transaction(|tx| async move {
335 let tx = &*tx;
336 if let Some(github_user_id) = github_user_id {
337 if let Some(user_by_github_user_id) = user::Entity::find()
338 .filter(user::Column::GithubUserId.eq(github_user_id))
339 .one(tx)
340 .await?
341 {
342 let mut user_by_github_user_id = user_by_github_user_id.into_active_model();
343 user_by_github_user_id.github_login = ActiveValue::set(github_login.into());
344 Ok(Some(user_by_github_user_id.update(tx).await?))
345 } else if let Some(user_by_github_login) = user::Entity::find()
346 .filter(user::Column::GithubLogin.eq(github_login))
347 .one(tx)
348 .await?
349 {
350 let mut user_by_github_login = user_by_github_login.into_active_model();
351 user_by_github_login.github_user_id = ActiveValue::set(Some(github_user_id));
352 Ok(Some(user_by_github_login.update(tx).await?))
353 } else {
354 let user = user::Entity::insert(user::ActiveModel {
355 email_address: ActiveValue::set(github_email.map(|email| email.into())),
356 github_login: ActiveValue::set(github_login.into()),
357 github_user_id: ActiveValue::set(Some(github_user_id)),
358 admin: ActiveValue::set(false),
359 invite_count: ActiveValue::set(0),
360 invite_code: ActiveValue::set(None),
361 metrics_id: ActiveValue::set(Uuid::new_v4()),
362 ..Default::default()
363 })
364 .exec_with_returning(&*tx)
365 .await?;
366 Ok(Some(user))
367 }
368 } else {
369 Ok(user::Entity::find()
370 .filter(user::Column::GithubLogin.eq(github_login))
371 .one(tx)
372 .await?)
373 }
374 })
375 .await
376 }
377
378 pub async fn get_all_users(&self, page: u32, limit: u32) -> Result<Vec<User>> {
379 self.transaction(|tx| async move {
380 Ok(user::Entity::find()
381 .order_by_asc(user::Column::GithubLogin)
382 .limit(limit as u64)
383 .offset(page as u64 * limit as u64)
384 .all(&*tx)
385 .await?)
386 })
387 .await
388 }
389
390 pub async fn get_users_with_no_invites(
391 &self,
392 invited_by_another_user: bool,
393 ) -> Result<Vec<User>> {
394 self.transaction(|tx| async move {
395 Ok(user::Entity::find()
396 .filter(
397 user::Column::InviteCount
398 .eq(0)
399 .and(if invited_by_another_user {
400 user::Column::InviterId.is_not_null()
401 } else {
402 user::Column::InviterId.is_null()
403 }),
404 )
405 .all(&*tx)
406 .await?)
407 })
408 .await
409 }
410
411 pub async fn get_user_metrics_id(&self, id: UserId) -> Result<String> {
412 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
413 enum QueryAs {
414 MetricsId,
415 }
416
417 self.transaction(|tx| async move {
418 let metrics_id: Uuid = user::Entity::find_by_id(id)
419 .select_only()
420 .column(user::Column::MetricsId)
421 .into_values::<_, QueryAs>()
422 .one(&*tx)
423 .await?
424 .ok_or_else(|| anyhow!("could not find user"))?;
425 Ok(metrics_id.to_string())
426 })
427 .await
428 }
429
430 pub async fn set_user_is_admin(&self, id: UserId, is_admin: bool) -> Result<()> {
431 self.transaction(|tx| async move {
432 user::Entity::update_many()
433 .filter(user::Column::Id.eq(id))
434 .set(user::ActiveModel {
435 admin: ActiveValue::set(is_admin),
436 ..Default::default()
437 })
438 .exec(&*tx)
439 .await?;
440 Ok(())
441 })
442 .await
443 }
444
445 pub async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()> {
446 self.transaction(|tx| async move {
447 user::Entity::update_many()
448 .filter(user::Column::Id.eq(id))
449 .set(user::ActiveModel {
450 connected_once: ActiveValue::set(connected_once),
451 ..Default::default()
452 })
453 .exec(&*tx)
454 .await?;
455 Ok(())
456 })
457 .await
458 }
459
460 pub async fn destroy_user(&self, id: UserId) -> Result<()> {
461 self.transaction(|tx| async move {
462 access_token::Entity::delete_many()
463 .filter(access_token::Column::UserId.eq(id))
464 .exec(&*tx)
465 .await?;
466 user::Entity::delete_by_id(id).exec(&*tx).await?;
467 Ok(())
468 })
469 .await
470 }
471
472 // contacts
473
474 pub async fn get_contacts(&self, user_id: UserId) -> Result<Vec<Contact>> {
475 #[derive(Debug, FromQueryResult)]
476 struct ContactWithUserBusyStatuses {
477 user_id_a: UserId,
478 user_id_b: UserId,
479 a_to_b: bool,
480 accepted: bool,
481 should_notify: bool,
482 user_a_busy: bool,
483 user_b_busy: bool,
484 }
485
486 self.transaction(|tx| async move {
487 let user_a_participant = Alias::new("user_a_participant");
488 let user_b_participant = Alias::new("user_b_participant");
489 let mut db_contacts = contact::Entity::find()
490 .column_as(
491 Expr::tbl(user_a_participant.clone(), room_participant::Column::Id)
492 .is_not_null(),
493 "user_a_busy",
494 )
495 .column_as(
496 Expr::tbl(user_b_participant.clone(), room_participant::Column::Id)
497 .is_not_null(),
498 "user_b_busy",
499 )
500 .filter(
501 contact::Column::UserIdA
502 .eq(user_id)
503 .or(contact::Column::UserIdB.eq(user_id)),
504 )
505 .join_as(
506 JoinType::LeftJoin,
507 contact::Relation::UserARoomParticipant.def(),
508 user_a_participant,
509 )
510 .join_as(
511 JoinType::LeftJoin,
512 contact::Relation::UserBRoomParticipant.def(),
513 user_b_participant,
514 )
515 .into_model::<ContactWithUserBusyStatuses>()
516 .stream(&*tx)
517 .await?;
518
519 let mut contacts = Vec::new();
520 while let Some(db_contact) = db_contacts.next().await {
521 let db_contact = db_contact?;
522 if db_contact.user_id_a == user_id {
523 if db_contact.accepted {
524 contacts.push(Contact::Accepted {
525 user_id: db_contact.user_id_b,
526 should_notify: db_contact.should_notify && db_contact.a_to_b,
527 busy: db_contact.user_b_busy,
528 });
529 } else if db_contact.a_to_b {
530 contacts.push(Contact::Outgoing {
531 user_id: db_contact.user_id_b,
532 })
533 } else {
534 contacts.push(Contact::Incoming {
535 user_id: db_contact.user_id_b,
536 should_notify: db_contact.should_notify,
537 });
538 }
539 } else if db_contact.accepted {
540 contacts.push(Contact::Accepted {
541 user_id: db_contact.user_id_a,
542 should_notify: db_contact.should_notify && !db_contact.a_to_b,
543 busy: db_contact.user_a_busy,
544 });
545 } else if db_contact.a_to_b {
546 contacts.push(Contact::Incoming {
547 user_id: db_contact.user_id_a,
548 should_notify: db_contact.should_notify,
549 });
550 } else {
551 contacts.push(Contact::Outgoing {
552 user_id: db_contact.user_id_a,
553 });
554 }
555 }
556
557 contacts.sort_unstable_by_key(|contact| contact.user_id());
558
559 Ok(contacts)
560 })
561 .await
562 }
563
564 pub async fn is_user_busy(&self, user_id: UserId) -> Result<bool> {
565 self.transaction(|tx| async move {
566 let participant = room_participant::Entity::find()
567 .filter(room_participant::Column::UserId.eq(user_id))
568 .one(&*tx)
569 .await?;
570 Ok(participant.is_some())
571 })
572 .await
573 }
574
575 pub async fn has_contact(&self, user_id_1: UserId, user_id_2: UserId) -> Result<bool> {
576 self.transaction(|tx| async move {
577 let (id_a, id_b) = if user_id_1 < user_id_2 {
578 (user_id_1, user_id_2)
579 } else {
580 (user_id_2, user_id_1)
581 };
582
583 Ok(contact::Entity::find()
584 .filter(
585 contact::Column::UserIdA
586 .eq(id_a)
587 .and(contact::Column::UserIdB.eq(id_b))
588 .and(contact::Column::Accepted.eq(true)),
589 )
590 .one(&*tx)
591 .await?
592 .is_some())
593 })
594 .await
595 }
596
597 pub async fn send_contact_request(&self, sender_id: UserId, receiver_id: UserId) -> Result<()> {
598 self.transaction(|tx| async move {
599 let (id_a, id_b, a_to_b) = if sender_id < receiver_id {
600 (sender_id, receiver_id, true)
601 } else {
602 (receiver_id, sender_id, false)
603 };
604
605 let rows_affected = contact::Entity::insert(contact::ActiveModel {
606 user_id_a: ActiveValue::set(id_a),
607 user_id_b: ActiveValue::set(id_b),
608 a_to_b: ActiveValue::set(a_to_b),
609 accepted: ActiveValue::set(false),
610 should_notify: ActiveValue::set(true),
611 ..Default::default()
612 })
613 .on_conflict(
614 OnConflict::columns([contact::Column::UserIdA, contact::Column::UserIdB])
615 .values([
616 (contact::Column::Accepted, true.into()),
617 (contact::Column::ShouldNotify, false.into()),
618 ])
619 .action_and_where(
620 contact::Column::Accepted.eq(false).and(
621 contact::Column::AToB
622 .eq(a_to_b)
623 .and(contact::Column::UserIdA.eq(id_b))
624 .or(contact::Column::AToB
625 .ne(a_to_b)
626 .and(contact::Column::UserIdA.eq(id_a))),
627 ),
628 )
629 .to_owned(),
630 )
631 .exec_without_returning(&*tx)
632 .await?;
633
634 if rows_affected == 1 {
635 Ok(())
636 } else {
637 Err(anyhow!("contact already requested"))?
638 }
639 })
640 .await
641 }
642
643 /// Returns a bool indicating whether the removed contact had originally accepted or not
644 ///
645 /// Deletes the contact identified by the requester and responder ids, and then returns
646 /// whether the deleted contact had originally accepted or was a pending contact request.
647 ///
648 /// # Arguments
649 ///
650 /// * `requester_id` - The user that initiates this request
651 /// * `responder_id` - The user that will be removed
652 pub async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<bool> {
653 self.transaction(|tx| async move {
654 let (id_a, id_b) = if responder_id < requester_id {
655 (responder_id, requester_id)
656 } else {
657 (requester_id, responder_id)
658 };
659
660 let contact = contact::Entity::find()
661 .filter(
662 contact::Column::UserIdA
663 .eq(id_a)
664 .and(contact::Column::UserIdB.eq(id_b)),
665 )
666 .one(&*tx)
667 .await?
668 .ok_or_else(|| anyhow!("no such contact"))?;
669
670 contact::Entity::delete_by_id(contact.id).exec(&*tx).await?;
671 Ok(contact.accepted)
672 })
673 .await
674 }
675
676 pub async fn dismiss_contact_notification(
677 &self,
678 user_id: UserId,
679 contact_user_id: UserId,
680 ) -> Result<()> {
681 self.transaction(|tx| async move {
682 let (id_a, id_b, a_to_b) = if user_id < contact_user_id {
683 (user_id, contact_user_id, true)
684 } else {
685 (contact_user_id, user_id, false)
686 };
687
688 let result = contact::Entity::update_many()
689 .set(contact::ActiveModel {
690 should_notify: ActiveValue::set(false),
691 ..Default::default()
692 })
693 .filter(
694 contact::Column::UserIdA
695 .eq(id_a)
696 .and(contact::Column::UserIdB.eq(id_b))
697 .and(
698 contact::Column::AToB
699 .eq(a_to_b)
700 .and(contact::Column::Accepted.eq(true))
701 .or(contact::Column::AToB
702 .ne(a_to_b)
703 .and(contact::Column::Accepted.eq(false))),
704 ),
705 )
706 .exec(&*tx)
707 .await?;
708 if result.rows_affected == 0 {
709 Err(anyhow!("no such contact request"))?
710 } else {
711 Ok(())
712 }
713 })
714 .await
715 }
716
717 pub async fn respond_to_contact_request(
718 &self,
719 responder_id: UserId,
720 requester_id: UserId,
721 accept: bool,
722 ) -> Result<()> {
723 self.transaction(|tx| async move {
724 let (id_a, id_b, a_to_b) = if responder_id < requester_id {
725 (responder_id, requester_id, false)
726 } else {
727 (requester_id, responder_id, true)
728 };
729 let rows_affected = if accept {
730 let result = contact::Entity::update_many()
731 .set(contact::ActiveModel {
732 accepted: ActiveValue::set(true),
733 should_notify: ActiveValue::set(true),
734 ..Default::default()
735 })
736 .filter(
737 contact::Column::UserIdA
738 .eq(id_a)
739 .and(contact::Column::UserIdB.eq(id_b))
740 .and(contact::Column::AToB.eq(a_to_b)),
741 )
742 .exec(&*tx)
743 .await?;
744 result.rows_affected
745 } else {
746 let result = contact::Entity::delete_many()
747 .filter(
748 contact::Column::UserIdA
749 .eq(id_a)
750 .and(contact::Column::UserIdB.eq(id_b))
751 .and(contact::Column::AToB.eq(a_to_b))
752 .and(contact::Column::Accepted.eq(false)),
753 )
754 .exec(&*tx)
755 .await?;
756
757 result.rows_affected
758 };
759
760 if rows_affected == 1 {
761 Ok(())
762 } else {
763 Err(anyhow!("no such contact request"))?
764 }
765 })
766 .await
767 }
768
769 pub fn fuzzy_like_string(string: &str) -> String {
770 let mut result = String::with_capacity(string.len() * 2 + 1);
771 for c in string.chars() {
772 if c.is_alphanumeric() {
773 result.push('%');
774 result.push(c);
775 }
776 }
777 result.push('%');
778 result
779 }
780
781 pub async fn fuzzy_search_users(&self, name_query: &str, limit: u32) -> Result<Vec<User>> {
782 self.transaction(|tx| async {
783 let tx = tx;
784 let like_string = Self::fuzzy_like_string(name_query);
785 let query = "
786 SELECT users.*
787 FROM users
788 WHERE github_login ILIKE $1
789 ORDER BY github_login <-> $2
790 LIMIT $3
791 ";
792
793 Ok(user::Entity::find()
794 .from_raw_sql(Statement::from_sql_and_values(
795 self.pool.get_database_backend(),
796 query.into(),
797 vec![like_string.into(), name_query.into(), limit.into()],
798 ))
799 .all(&*tx)
800 .await?)
801 })
802 .await
803 }
804
805 // signups
806
807 pub async fn create_signup(&self, signup: &NewSignup) -> Result<()> {
808 self.transaction(|tx| async move {
809 signup::Entity::insert(signup::ActiveModel {
810 email_address: ActiveValue::set(signup.email_address.clone()),
811 email_confirmation_code: ActiveValue::set(random_email_confirmation_code()),
812 email_confirmation_sent: ActiveValue::set(false),
813 platform_mac: ActiveValue::set(signup.platform_mac),
814 platform_windows: ActiveValue::set(signup.platform_windows),
815 platform_linux: ActiveValue::set(signup.platform_linux),
816 platform_unknown: ActiveValue::set(false),
817 editor_features: ActiveValue::set(Some(signup.editor_features.clone())),
818 programming_languages: ActiveValue::set(Some(signup.programming_languages.clone())),
819 device_id: ActiveValue::set(signup.device_id.clone()),
820 added_to_mailing_list: ActiveValue::set(signup.added_to_mailing_list),
821 ..Default::default()
822 })
823 .on_conflict(
824 OnConflict::column(signup::Column::EmailAddress)
825 .update_columns([
826 signup::Column::PlatformMac,
827 signup::Column::PlatformWindows,
828 signup::Column::PlatformLinux,
829 signup::Column::EditorFeatures,
830 signup::Column::ProgrammingLanguages,
831 signup::Column::DeviceId,
832 signup::Column::AddedToMailingList,
833 ])
834 .to_owned(),
835 )
836 .exec(&*tx)
837 .await?;
838 Ok(())
839 })
840 .await
841 }
842
843 pub async fn get_signup(&self, email_address: &str) -> Result<signup::Model> {
844 self.transaction(|tx| async move {
845 let signup = signup::Entity::find()
846 .filter(signup::Column::EmailAddress.eq(email_address))
847 .one(&*tx)
848 .await?
849 .ok_or_else(|| {
850 anyhow!("signup with email address {} doesn't exist", email_address)
851 })?;
852
853 Ok(signup)
854 })
855 .await
856 }
857
858 pub async fn get_waitlist_summary(&self) -> Result<WaitlistSummary> {
859 self.transaction(|tx| async move {
860 let query = "
861 SELECT
862 COUNT(*) as count,
863 COALESCE(SUM(CASE WHEN platform_linux THEN 1 ELSE 0 END), 0) as linux_count,
864 COALESCE(SUM(CASE WHEN platform_mac THEN 1 ELSE 0 END), 0) as mac_count,
865 COALESCE(SUM(CASE WHEN platform_windows THEN 1 ELSE 0 END), 0) as windows_count,
866 COALESCE(SUM(CASE WHEN platform_unknown THEN 1 ELSE 0 END), 0) as unknown_count
867 FROM (
868 SELECT *
869 FROM signups
870 WHERE
871 NOT email_confirmation_sent
872 ) AS unsent
873 ";
874 Ok(
875 WaitlistSummary::find_by_statement(Statement::from_sql_and_values(
876 self.pool.get_database_backend(),
877 query.into(),
878 vec![],
879 ))
880 .one(&*tx)
881 .await?
882 .ok_or_else(|| anyhow!("invalid result"))?,
883 )
884 })
885 .await
886 }
887
888 pub async fn record_sent_invites(&self, invites: &[Invite]) -> Result<()> {
889 let emails = invites
890 .iter()
891 .map(|s| s.email_address.as_str())
892 .collect::<Vec<_>>();
893 self.transaction(|tx| async {
894 let tx = tx;
895 signup::Entity::update_many()
896 .filter(signup::Column::EmailAddress.is_in(emails.iter().copied()))
897 .set(signup::ActiveModel {
898 email_confirmation_sent: ActiveValue::set(true),
899 ..Default::default()
900 })
901 .exec(&*tx)
902 .await?;
903 Ok(())
904 })
905 .await
906 }
907
908 pub async fn get_unsent_invites(&self, count: usize) -> Result<Vec<Invite>> {
909 self.transaction(|tx| async move {
910 Ok(signup::Entity::find()
911 .select_only()
912 .column(signup::Column::EmailAddress)
913 .column(signup::Column::EmailConfirmationCode)
914 .filter(
915 signup::Column::EmailConfirmationSent.eq(false).and(
916 signup::Column::PlatformMac
917 .eq(true)
918 .or(signup::Column::PlatformUnknown.eq(true)),
919 ),
920 )
921 .order_by_asc(signup::Column::CreatedAt)
922 .limit(count as u64)
923 .into_model()
924 .all(&*tx)
925 .await?)
926 })
927 .await
928 }
929
930 // invite codes
931
932 pub async fn create_invite_from_code(
933 &self,
934 code: &str,
935 email_address: &str,
936 device_id: Option<&str>,
937 added_to_mailing_list: bool,
938 ) -> Result<Invite> {
939 self.transaction(|tx| async move {
940 let existing_user = user::Entity::find()
941 .filter(user::Column::EmailAddress.eq(email_address))
942 .one(&*tx)
943 .await?;
944
945 if existing_user.is_some() {
946 Err(anyhow!("email address is already in use"))?;
947 }
948
949 let inviting_user_with_invites = match user::Entity::find()
950 .filter(
951 user::Column::InviteCode
952 .eq(code)
953 .and(user::Column::InviteCount.gt(0)),
954 )
955 .one(&*tx)
956 .await?
957 {
958 Some(inviting_user) => inviting_user,
959 None => {
960 return Err(Error::Http(
961 StatusCode::UNAUTHORIZED,
962 "unable to find an invite code with invites remaining".to_string(),
963 ))?
964 }
965 };
966 user::Entity::update_many()
967 .filter(
968 user::Column::Id
969 .eq(inviting_user_with_invites.id)
970 .and(user::Column::InviteCount.gt(0)),
971 )
972 .col_expr(
973 user::Column::InviteCount,
974 Expr::col(user::Column::InviteCount).sub(1),
975 )
976 .exec(&*tx)
977 .await?;
978
979 let signup = signup::Entity::insert(signup::ActiveModel {
980 email_address: ActiveValue::set(email_address.into()),
981 email_confirmation_code: ActiveValue::set(random_email_confirmation_code()),
982 email_confirmation_sent: ActiveValue::set(false),
983 inviting_user_id: ActiveValue::set(Some(inviting_user_with_invites.id)),
984 platform_linux: ActiveValue::set(false),
985 platform_mac: ActiveValue::set(false),
986 platform_windows: ActiveValue::set(false),
987 platform_unknown: ActiveValue::set(true),
988 device_id: ActiveValue::set(device_id.map(|device_id| device_id.into())),
989 added_to_mailing_list: ActiveValue::set(added_to_mailing_list),
990 ..Default::default()
991 })
992 .on_conflict(
993 OnConflict::column(signup::Column::EmailAddress)
994 .update_column(signup::Column::InvitingUserId)
995 .to_owned(),
996 )
997 .exec_with_returning(&*tx)
998 .await?;
999
1000 Ok(Invite {
1001 email_address: signup.email_address,
1002 email_confirmation_code: signup.email_confirmation_code,
1003 })
1004 })
1005 .await
1006 }
1007
1008 pub async fn create_user_from_invite(
1009 &self,
1010 invite: &Invite,
1011 user: NewUserParams,
1012 ) -> Result<Option<NewUserResult>> {
1013 self.transaction(|tx| async {
1014 let tx = tx;
1015 let signup = signup::Entity::find()
1016 .filter(
1017 signup::Column::EmailAddress
1018 .eq(invite.email_address.as_str())
1019 .and(
1020 signup::Column::EmailConfirmationCode
1021 .eq(invite.email_confirmation_code.as_str()),
1022 ),
1023 )
1024 .one(&*tx)
1025 .await?
1026 .ok_or_else(|| Error::Http(StatusCode::NOT_FOUND, "no such invite".to_string()))?;
1027
1028 if signup.user_id.is_some() {
1029 return Ok(None);
1030 }
1031
1032 let user = user::Entity::insert(user::ActiveModel {
1033 email_address: ActiveValue::set(Some(invite.email_address.clone())),
1034 github_login: ActiveValue::set(user.github_login.clone()),
1035 github_user_id: ActiveValue::set(Some(user.github_user_id)),
1036 admin: ActiveValue::set(false),
1037 invite_count: ActiveValue::set(user.invite_count),
1038 invite_code: ActiveValue::set(Some(random_invite_code())),
1039 metrics_id: ActiveValue::set(Uuid::new_v4()),
1040 ..Default::default()
1041 })
1042 .on_conflict(
1043 OnConflict::column(user::Column::GithubLogin)
1044 .update_columns([
1045 user::Column::EmailAddress,
1046 user::Column::GithubUserId,
1047 user::Column::Admin,
1048 ])
1049 .to_owned(),
1050 )
1051 .exec_with_returning(&*tx)
1052 .await?;
1053
1054 let mut signup = signup.into_active_model();
1055 signup.user_id = ActiveValue::set(Some(user.id));
1056 let signup = signup.update(&*tx).await?;
1057
1058 if let Some(inviting_user_id) = signup.inviting_user_id {
1059 let (user_id_a, user_id_b, a_to_b) = if inviting_user_id < user.id {
1060 (inviting_user_id, user.id, true)
1061 } else {
1062 (user.id, inviting_user_id, false)
1063 };
1064
1065 contact::Entity::insert(contact::ActiveModel {
1066 user_id_a: ActiveValue::set(user_id_a),
1067 user_id_b: ActiveValue::set(user_id_b),
1068 a_to_b: ActiveValue::set(a_to_b),
1069 should_notify: ActiveValue::set(true),
1070 accepted: ActiveValue::set(true),
1071 ..Default::default()
1072 })
1073 .on_conflict(OnConflict::new().do_nothing().to_owned())
1074 .exec_without_returning(&*tx)
1075 .await?;
1076 }
1077
1078 Ok(Some(NewUserResult {
1079 user_id: user.id,
1080 metrics_id: user.metrics_id.to_string(),
1081 inviting_user_id: signup.inviting_user_id,
1082 signup_device_id: signup.device_id,
1083 }))
1084 })
1085 .await
1086 }
1087
1088 pub async fn set_invite_count_for_user(&self, id: UserId, count: i32) -> Result<()> {
1089 self.transaction(|tx| async move {
1090 if count > 0 {
1091 user::Entity::update_many()
1092 .filter(
1093 user::Column::Id
1094 .eq(id)
1095 .and(user::Column::InviteCode.is_null()),
1096 )
1097 .set(user::ActiveModel {
1098 invite_code: ActiveValue::set(Some(random_invite_code())),
1099 ..Default::default()
1100 })
1101 .exec(&*tx)
1102 .await?;
1103 }
1104
1105 user::Entity::update_many()
1106 .filter(user::Column::Id.eq(id))
1107 .set(user::ActiveModel {
1108 invite_count: ActiveValue::set(count),
1109 ..Default::default()
1110 })
1111 .exec(&*tx)
1112 .await?;
1113 Ok(())
1114 })
1115 .await
1116 }
1117
1118 pub async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, i32)>> {
1119 self.transaction(|tx| async move {
1120 match user::Entity::find_by_id(id).one(&*tx).await? {
1121 Some(user) if user.invite_code.is_some() => {
1122 Ok(Some((user.invite_code.unwrap(), user.invite_count)))
1123 }
1124 _ => Ok(None),
1125 }
1126 })
1127 .await
1128 }
1129
1130 pub async fn get_user_for_invite_code(&self, code: &str) -> Result<User> {
1131 self.transaction(|tx| async move {
1132 user::Entity::find()
1133 .filter(user::Column::InviteCode.eq(code))
1134 .one(&*tx)
1135 .await?
1136 .ok_or_else(|| {
1137 Error::Http(
1138 StatusCode::NOT_FOUND,
1139 "that invite code does not exist".to_string(),
1140 )
1141 })
1142 })
1143 .await
1144 }
1145
1146 // rooms
1147
1148 pub async fn incoming_call_for_user(
1149 &self,
1150 user_id: UserId,
1151 ) -> Result<Option<proto::IncomingCall>> {
1152 self.transaction(|tx| async move {
1153 let pending_participant = room_participant::Entity::find()
1154 .filter(
1155 room_participant::Column::UserId
1156 .eq(user_id)
1157 .and(room_participant::Column::AnsweringConnectionId.is_null()),
1158 )
1159 .one(&*tx)
1160 .await?;
1161
1162 if let Some(pending_participant) = pending_participant {
1163 let room = self.get_room(pending_participant.room_id, &tx).await?;
1164 Ok(Self::build_incoming_call(&room, user_id))
1165 } else {
1166 Ok(None)
1167 }
1168 })
1169 .await
1170 }
1171
1172 pub async fn create_room(
1173 &self,
1174 user_id: UserId,
1175 connection: ConnectionId,
1176 live_kit_room: &str,
1177 ) -> Result<proto::Room> {
1178 self.transaction(|tx| async move {
1179 let room = room::ActiveModel {
1180 live_kit_room: ActiveValue::set(live_kit_room.into()),
1181 ..Default::default()
1182 }
1183 .insert(&*tx)
1184 .await?;
1185 room_participant::ActiveModel {
1186 room_id: ActiveValue::set(room.id),
1187 user_id: ActiveValue::set(user_id),
1188 answering_connection_id: ActiveValue::set(Some(connection.id as i32)),
1189 answering_connection_server_id: ActiveValue::set(Some(ServerId(
1190 connection.owner_id as i32,
1191 ))),
1192 answering_connection_lost: ActiveValue::set(false),
1193 calling_user_id: ActiveValue::set(user_id),
1194 calling_connection_id: ActiveValue::set(connection.id as i32),
1195 calling_connection_server_id: ActiveValue::set(Some(ServerId(
1196 connection.owner_id as i32,
1197 ))),
1198 ..Default::default()
1199 }
1200 .insert(&*tx)
1201 .await?;
1202
1203 let room = self.get_room(room.id, &tx).await?;
1204 Ok(room)
1205 })
1206 .await
1207 }
1208
1209 pub async fn call(
1210 &self,
1211 room_id: RoomId,
1212 calling_user_id: UserId,
1213 calling_connection: ConnectionId,
1214 called_user_id: UserId,
1215 initial_project_id: Option<ProjectId>,
1216 ) -> Result<RoomGuard<(proto::Room, proto::IncomingCall)>> {
1217 self.room_transaction(room_id, |tx| async move {
1218 room_participant::ActiveModel {
1219 room_id: ActiveValue::set(room_id),
1220 user_id: ActiveValue::set(called_user_id),
1221 answering_connection_lost: ActiveValue::set(false),
1222 calling_user_id: ActiveValue::set(calling_user_id),
1223 calling_connection_id: ActiveValue::set(calling_connection.id as i32),
1224 calling_connection_server_id: ActiveValue::set(Some(ServerId(
1225 calling_connection.owner_id as i32,
1226 ))),
1227 initial_project_id: ActiveValue::set(initial_project_id),
1228 ..Default::default()
1229 }
1230 .insert(&*tx)
1231 .await?;
1232
1233 let room = self.get_room(room_id, &tx).await?;
1234 let incoming_call = Self::build_incoming_call(&room, called_user_id)
1235 .ok_or_else(|| anyhow!("failed to build incoming call"))?;
1236 Ok((room, incoming_call))
1237 })
1238 .await
1239 }
1240
1241 pub async fn call_failed(
1242 &self,
1243 room_id: RoomId,
1244 called_user_id: UserId,
1245 ) -> Result<RoomGuard<proto::Room>> {
1246 self.room_transaction(room_id, |tx| async move {
1247 room_participant::Entity::delete_many()
1248 .filter(
1249 room_participant::Column::RoomId
1250 .eq(room_id)
1251 .and(room_participant::Column::UserId.eq(called_user_id)),
1252 )
1253 .exec(&*tx)
1254 .await?;
1255 let room = self.get_room(room_id, &tx).await?;
1256 Ok(room)
1257 })
1258 .await
1259 }
1260
1261 pub async fn decline_call(
1262 &self,
1263 expected_room_id: Option<RoomId>,
1264 user_id: UserId,
1265 ) -> Result<Option<RoomGuard<proto::Room>>> {
1266 self.optional_room_transaction(|tx| async move {
1267 let mut filter = Condition::all()
1268 .add(room_participant::Column::UserId.eq(user_id))
1269 .add(room_participant::Column::AnsweringConnectionId.is_null());
1270 if let Some(room_id) = expected_room_id {
1271 filter = filter.add(room_participant::Column::RoomId.eq(room_id));
1272 }
1273 let participant = room_participant::Entity::find()
1274 .filter(filter)
1275 .one(&*tx)
1276 .await?;
1277
1278 let participant = if let Some(participant) = participant {
1279 participant
1280 } else if expected_room_id.is_some() {
1281 return Err(anyhow!("could not find call to decline"))?;
1282 } else {
1283 return Ok(None);
1284 };
1285
1286 let room_id = participant.room_id;
1287 room_participant::Entity::delete(participant.into_active_model())
1288 .exec(&*tx)
1289 .await?;
1290
1291 let room = self.get_room(room_id, &tx).await?;
1292 Ok(Some((room_id, room)))
1293 })
1294 .await
1295 }
1296
1297 pub async fn cancel_call(
1298 &self,
1299 room_id: RoomId,
1300 calling_connection: ConnectionId,
1301 called_user_id: UserId,
1302 ) -> Result<RoomGuard<proto::Room>> {
1303 self.room_transaction(room_id, |tx| async move {
1304 let participant = room_participant::Entity::find()
1305 .filter(
1306 Condition::all()
1307 .add(room_participant::Column::UserId.eq(called_user_id))
1308 .add(room_participant::Column::RoomId.eq(room_id))
1309 .add(
1310 room_participant::Column::CallingConnectionId
1311 .eq(calling_connection.id as i32),
1312 )
1313 .add(
1314 room_participant::Column::CallingConnectionServerId
1315 .eq(calling_connection.owner_id as i32),
1316 )
1317 .add(room_participant::Column::AnsweringConnectionId.is_null()),
1318 )
1319 .one(&*tx)
1320 .await?
1321 .ok_or_else(|| anyhow!("no call to cancel"))?;
1322
1323 room_participant::Entity::delete(participant.into_active_model())
1324 .exec(&*tx)
1325 .await?;
1326
1327 let room = self.get_room(room_id, &tx).await?;
1328 Ok(room)
1329 })
1330 .await
1331 }
1332
1333 pub async fn join_room(
1334 &self,
1335 room_id: RoomId,
1336 user_id: UserId,
1337 connection: ConnectionId,
1338 ) -> Result<RoomGuard<proto::Room>> {
1339 self.room_transaction(room_id, |tx| async move {
1340 let result = room_participant::Entity::update_many()
1341 .filter(
1342 Condition::all()
1343 .add(room_participant::Column::RoomId.eq(room_id))
1344 .add(room_participant::Column::UserId.eq(user_id))
1345 .add(room_participant::Column::AnsweringConnectionId.is_null()),
1346 )
1347 .set(room_participant::ActiveModel {
1348 answering_connection_id: ActiveValue::set(Some(connection.id as i32)),
1349 answering_connection_server_id: ActiveValue::set(Some(ServerId(
1350 connection.owner_id as i32,
1351 ))),
1352 answering_connection_lost: ActiveValue::set(false),
1353 ..Default::default()
1354 })
1355 .exec(&*tx)
1356 .await?;
1357 if result.rows_affected == 0 {
1358 Err(anyhow!("room does not exist or was already joined"))?
1359 } else {
1360 let room = self.get_room(room_id, &tx).await?;
1361 Ok(room)
1362 }
1363 })
1364 .await
1365 }
1366
1367 pub async fn rejoin_room(
1368 &self,
1369 rejoin_room: proto::RejoinRoom,
1370 user_id: UserId,
1371 connection: ConnectionId,
1372 ) -> Result<RoomGuard<RejoinedRoom>> {
1373 let room_id = RoomId::from_proto(rejoin_room.id);
1374 self.room_transaction(room_id, |tx| async {
1375 let tx = tx;
1376 let participant_update = room_participant::Entity::update_many()
1377 .filter(
1378 Condition::all()
1379 .add(room_participant::Column::RoomId.eq(room_id))
1380 .add(room_participant::Column::UserId.eq(user_id))
1381 .add(room_participant::Column::AnsweringConnectionId.is_not_null())
1382 .add(
1383 Condition::any()
1384 .add(room_participant::Column::AnsweringConnectionLost.eq(true))
1385 .add(
1386 room_participant::Column::AnsweringConnectionServerId
1387 .ne(connection.owner_id as i32),
1388 ),
1389 ),
1390 )
1391 .set(room_participant::ActiveModel {
1392 answering_connection_id: ActiveValue::set(Some(connection.id as i32)),
1393 answering_connection_server_id: ActiveValue::set(Some(ServerId(
1394 connection.owner_id as i32,
1395 ))),
1396 answering_connection_lost: ActiveValue::set(false),
1397 ..Default::default()
1398 })
1399 .exec(&*tx)
1400 .await?;
1401 if participant_update.rows_affected == 0 {
1402 return Err(anyhow!("room does not exist or was already joined"))?;
1403 }
1404
1405 let mut reshared_projects = Vec::new();
1406 for reshared_project in &rejoin_room.reshared_projects {
1407 let project_id = ProjectId::from_proto(reshared_project.project_id);
1408 let project = project::Entity::find_by_id(project_id)
1409 .one(&*tx)
1410 .await?
1411 .ok_or_else(|| anyhow!("project does not exist"))?;
1412 if project.host_user_id != user_id {
1413 return Err(anyhow!("no such project"))?;
1414 }
1415
1416 let mut collaborators = project
1417 .find_related(project_collaborator::Entity)
1418 .all(&*tx)
1419 .await?;
1420 let host_ix = collaborators
1421 .iter()
1422 .position(|collaborator| {
1423 collaborator.user_id == user_id && collaborator.is_host
1424 })
1425 .ok_or_else(|| anyhow!("host not found among collaborators"))?;
1426 let host = collaborators.swap_remove(host_ix);
1427 let old_connection_id = host.connection();
1428
1429 project::Entity::update(project::ActiveModel {
1430 host_connection_id: ActiveValue::set(Some(connection.id as i32)),
1431 host_connection_server_id: ActiveValue::set(Some(ServerId(
1432 connection.owner_id as i32,
1433 ))),
1434 ..project.into_active_model()
1435 })
1436 .exec(&*tx)
1437 .await?;
1438 project_collaborator::Entity::update(project_collaborator::ActiveModel {
1439 connection_id: ActiveValue::set(connection.id as i32),
1440 connection_server_id: ActiveValue::set(ServerId(connection.owner_id as i32)),
1441 ..host.into_active_model()
1442 })
1443 .exec(&*tx)
1444 .await?;
1445
1446 self.update_project_worktrees(project_id, &reshared_project.worktrees, &tx)
1447 .await?;
1448
1449 reshared_projects.push(ResharedProject {
1450 id: project_id,
1451 old_connection_id,
1452 collaborators: collaborators
1453 .iter()
1454 .map(|collaborator| ProjectCollaborator {
1455 connection_id: collaborator.connection(),
1456 user_id: collaborator.user_id,
1457 replica_id: collaborator.replica_id,
1458 is_host: collaborator.is_host,
1459 })
1460 .collect(),
1461 worktrees: reshared_project.worktrees.clone(),
1462 });
1463 }
1464
1465 project::Entity::delete_many()
1466 .filter(
1467 Condition::all()
1468 .add(project::Column::RoomId.eq(room_id))
1469 .add(project::Column::HostUserId.eq(user_id))
1470 .add(
1471 project::Column::Id
1472 .is_not_in(reshared_projects.iter().map(|project| project.id)),
1473 ),
1474 )
1475 .exec(&*tx)
1476 .await?;
1477
1478 let mut rejoined_projects = Vec::new();
1479 for rejoined_project in &rejoin_room.rejoined_projects {
1480 let project_id = ProjectId::from_proto(rejoined_project.id);
1481 let Some(project) = project::Entity::find_by_id(project_id)
1482 .one(&*tx)
1483 .await? else { continue };
1484
1485 let mut worktrees = Vec::new();
1486 let db_worktrees = project.find_related(worktree::Entity).all(&*tx).await?;
1487 for db_worktree in db_worktrees {
1488 let mut worktree = RejoinedWorktree {
1489 id: db_worktree.id as u64,
1490 abs_path: db_worktree.abs_path,
1491 root_name: db_worktree.root_name,
1492 visible: db_worktree.visible,
1493 updated_entries: Default::default(),
1494 removed_entries: Default::default(),
1495 updated_repositories: Default::default(),
1496 removed_repositories: Default::default(),
1497 diagnostic_summaries: Default::default(),
1498 settings_files: Default::default(),
1499 scan_id: db_worktree.scan_id as u64,
1500 completed_scan_id: db_worktree.completed_scan_id as u64,
1501 };
1502
1503 let rejoined_worktree = rejoined_project
1504 .worktrees
1505 .iter()
1506 .find(|worktree| worktree.id == db_worktree.id as u64);
1507
1508 // File entries
1509 {
1510 let entry_filter = if let Some(rejoined_worktree) = rejoined_worktree {
1511 worktree_entry::Column::ScanId.gt(rejoined_worktree.scan_id)
1512 } else {
1513 worktree_entry::Column::IsDeleted.eq(false)
1514 };
1515
1516 let mut db_entries = worktree_entry::Entity::find()
1517 .filter(
1518 Condition::all()
1519 .add(worktree_entry::Column::ProjectId.eq(project.id))
1520 .add(worktree_entry::Column::WorktreeId.eq(worktree.id))
1521 .add(entry_filter),
1522 )
1523 .stream(&*tx)
1524 .await?;
1525
1526 while let Some(db_entry) = db_entries.next().await {
1527 let db_entry = db_entry?;
1528 if db_entry.is_deleted {
1529 worktree.removed_entries.push(db_entry.id as u64);
1530 } else {
1531 worktree.updated_entries.push(proto::Entry {
1532 id: db_entry.id as u64,
1533 is_dir: db_entry.is_dir,
1534 path: db_entry.path,
1535 inode: db_entry.inode as u64,
1536 mtime: Some(proto::Timestamp {
1537 seconds: db_entry.mtime_seconds as u64,
1538 nanos: db_entry.mtime_nanos as u32,
1539 }),
1540 is_symlink: db_entry.is_symlink,
1541 is_ignored: db_entry.is_ignored,
1542 is_external: db_entry.is_external,
1543 git_status: db_entry.git_status.map(|status| status as i32),
1544 });
1545 }
1546 }
1547 }
1548
1549 // Repository Entries
1550 {
1551 let repository_entry_filter =
1552 if let Some(rejoined_worktree) = rejoined_worktree {
1553 worktree_repository::Column::ScanId.gt(rejoined_worktree.scan_id)
1554 } else {
1555 worktree_repository::Column::IsDeleted.eq(false)
1556 };
1557
1558 let mut db_repositories = worktree_repository::Entity::find()
1559 .filter(
1560 Condition::all()
1561 .add(worktree_repository::Column::ProjectId.eq(project.id))
1562 .add(worktree_repository::Column::WorktreeId.eq(worktree.id))
1563 .add(repository_entry_filter),
1564 )
1565 .stream(&*tx)
1566 .await?;
1567
1568 while let Some(db_repository) = db_repositories.next().await {
1569 let db_repository = db_repository?;
1570 if db_repository.is_deleted {
1571 worktree
1572 .removed_repositories
1573 .push(db_repository.work_directory_id as u64);
1574 } else {
1575 worktree.updated_repositories.push(proto::RepositoryEntry {
1576 work_directory_id: db_repository.work_directory_id as u64,
1577 branch: db_repository.branch,
1578 });
1579 }
1580 }
1581 }
1582
1583 worktrees.push(worktree);
1584 }
1585
1586 let language_servers = project
1587 .find_related(language_server::Entity)
1588 .all(&*tx)
1589 .await?
1590 .into_iter()
1591 .map(|language_server| proto::LanguageServer {
1592 id: language_server.id as u64,
1593 name: language_server.name,
1594 })
1595 .collect::<Vec<_>>();
1596
1597 {
1598 let mut db_settings_files = worktree_settings_file::Entity::find()
1599 .filter(worktree_settings_file::Column::ProjectId.eq(project_id))
1600 .stream(&*tx)
1601 .await?;
1602 while let Some(db_settings_file) = db_settings_files.next().await {
1603 let db_settings_file = db_settings_file?;
1604 if let Some(worktree) = worktrees
1605 .iter_mut()
1606 .find(|w| w.id == db_settings_file.worktree_id as u64)
1607 {
1608 worktree.settings_files.push(WorktreeSettingsFile {
1609 path: db_settings_file.path,
1610 content: db_settings_file.content,
1611 });
1612 }
1613 }
1614 }
1615
1616 let mut collaborators = project
1617 .find_related(project_collaborator::Entity)
1618 .all(&*tx)
1619 .await?;
1620 let self_collaborator = if let Some(self_collaborator_ix) = collaborators
1621 .iter()
1622 .position(|collaborator| collaborator.user_id == user_id)
1623 {
1624 collaborators.swap_remove(self_collaborator_ix)
1625 } else {
1626 continue;
1627 };
1628 let old_connection_id = self_collaborator.connection();
1629 project_collaborator::Entity::update(project_collaborator::ActiveModel {
1630 connection_id: ActiveValue::set(connection.id as i32),
1631 connection_server_id: ActiveValue::set(ServerId(connection.owner_id as i32)),
1632 ..self_collaborator.into_active_model()
1633 })
1634 .exec(&*tx)
1635 .await?;
1636
1637 let collaborators = collaborators
1638 .into_iter()
1639 .map(|collaborator| ProjectCollaborator {
1640 connection_id: collaborator.connection(),
1641 user_id: collaborator.user_id,
1642 replica_id: collaborator.replica_id,
1643 is_host: collaborator.is_host,
1644 })
1645 .collect::<Vec<_>>();
1646
1647 rejoined_projects.push(RejoinedProject {
1648 id: project_id,
1649 old_connection_id,
1650 collaborators,
1651 worktrees,
1652 language_servers,
1653 });
1654 }
1655
1656 let room = self.get_room(room_id, &tx).await?;
1657 Ok(RejoinedRoom {
1658 room,
1659 rejoined_projects,
1660 reshared_projects,
1661 })
1662 })
1663 .await
1664 }
1665
1666 pub async fn leave_room(
1667 &self,
1668 connection: ConnectionId,
1669 ) -> Result<Option<RoomGuard<LeftRoom>>> {
1670 self.optional_room_transaction(|tx| async move {
1671 let leaving_participant = room_participant::Entity::find()
1672 .filter(
1673 Condition::all()
1674 .add(
1675 room_participant::Column::AnsweringConnectionId
1676 .eq(connection.id as i32),
1677 )
1678 .add(
1679 room_participant::Column::AnsweringConnectionServerId
1680 .eq(connection.owner_id as i32),
1681 ),
1682 )
1683 .one(&*tx)
1684 .await?;
1685
1686 if let Some(leaving_participant) = leaving_participant {
1687 // Leave room.
1688 let room_id = leaving_participant.room_id;
1689 room_participant::Entity::delete_by_id(leaving_participant.id)
1690 .exec(&*tx)
1691 .await?;
1692
1693 // Cancel pending calls initiated by the leaving user.
1694 let called_participants = room_participant::Entity::find()
1695 .filter(
1696 Condition::all()
1697 .add(
1698 room_participant::Column::CallingUserId
1699 .eq(leaving_participant.user_id),
1700 )
1701 .add(room_participant::Column::AnsweringConnectionId.is_null()),
1702 )
1703 .all(&*tx)
1704 .await?;
1705 room_participant::Entity::delete_many()
1706 .filter(
1707 room_participant::Column::Id
1708 .is_in(called_participants.iter().map(|participant| participant.id)),
1709 )
1710 .exec(&*tx)
1711 .await?;
1712 let canceled_calls_to_user_ids = called_participants
1713 .into_iter()
1714 .map(|participant| participant.user_id)
1715 .collect();
1716
1717 // Detect left projects.
1718 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
1719 enum QueryProjectIds {
1720 ProjectId,
1721 }
1722 let project_ids: Vec<ProjectId> = project_collaborator::Entity::find()
1723 .select_only()
1724 .column_as(
1725 project_collaborator::Column::ProjectId,
1726 QueryProjectIds::ProjectId,
1727 )
1728 .filter(
1729 Condition::all()
1730 .add(
1731 project_collaborator::Column::ConnectionId.eq(connection.id as i32),
1732 )
1733 .add(
1734 project_collaborator::Column::ConnectionServerId
1735 .eq(connection.owner_id as i32),
1736 ),
1737 )
1738 .into_values::<_, QueryProjectIds>()
1739 .all(&*tx)
1740 .await?;
1741 let mut left_projects = HashMap::default();
1742 let mut collaborators = project_collaborator::Entity::find()
1743 .filter(project_collaborator::Column::ProjectId.is_in(project_ids))
1744 .stream(&*tx)
1745 .await?;
1746 while let Some(collaborator) = collaborators.next().await {
1747 let collaborator = collaborator?;
1748 let left_project =
1749 left_projects
1750 .entry(collaborator.project_id)
1751 .or_insert(LeftProject {
1752 id: collaborator.project_id,
1753 host_user_id: Default::default(),
1754 connection_ids: Default::default(),
1755 host_connection_id: Default::default(),
1756 });
1757
1758 let collaborator_connection_id = collaborator.connection();
1759 if collaborator_connection_id != connection {
1760 left_project.connection_ids.push(collaborator_connection_id);
1761 }
1762
1763 if collaborator.is_host {
1764 left_project.host_user_id = collaborator.user_id;
1765 left_project.host_connection_id = collaborator_connection_id;
1766 }
1767 }
1768 drop(collaborators);
1769
1770 // Leave projects.
1771 project_collaborator::Entity::delete_many()
1772 .filter(
1773 Condition::all()
1774 .add(
1775 project_collaborator::Column::ConnectionId.eq(connection.id as i32),
1776 )
1777 .add(
1778 project_collaborator::Column::ConnectionServerId
1779 .eq(connection.owner_id as i32),
1780 ),
1781 )
1782 .exec(&*tx)
1783 .await?;
1784
1785 // Unshare projects.
1786 project::Entity::delete_many()
1787 .filter(
1788 Condition::all()
1789 .add(project::Column::RoomId.eq(room_id))
1790 .add(project::Column::HostConnectionId.eq(connection.id as i32))
1791 .add(
1792 project::Column::HostConnectionServerId
1793 .eq(connection.owner_id as i32),
1794 ),
1795 )
1796 .exec(&*tx)
1797 .await?;
1798
1799 let room = self.get_room(room_id, &tx).await?;
1800 if room.participants.is_empty() {
1801 room::Entity::delete_by_id(room_id).exec(&*tx).await?;
1802 }
1803
1804 let left_room = LeftRoom {
1805 room,
1806 left_projects,
1807 canceled_calls_to_user_ids,
1808 };
1809
1810 if left_room.room.participants.is_empty() {
1811 self.rooms.remove(&room_id);
1812 }
1813
1814 Ok(Some((room_id, left_room)))
1815 } else {
1816 Ok(None)
1817 }
1818 })
1819 .await
1820 }
1821
1822 pub async fn follow(
1823 &self,
1824 project_id: ProjectId,
1825 leader_connection: ConnectionId,
1826 follower_connection: ConnectionId,
1827 ) -> Result<RoomGuard<proto::Room>> {
1828 let room_id = self.room_id_for_project(project_id).await?;
1829 self.room_transaction(room_id, |tx| async move {
1830 follower::ActiveModel {
1831 room_id: ActiveValue::set(room_id),
1832 project_id: ActiveValue::set(project_id),
1833 leader_connection_server_id: ActiveValue::set(ServerId(
1834 leader_connection.owner_id as i32,
1835 )),
1836 leader_connection_id: ActiveValue::set(leader_connection.id as i32),
1837 follower_connection_server_id: ActiveValue::set(ServerId(
1838 follower_connection.owner_id as i32,
1839 )),
1840 follower_connection_id: ActiveValue::set(follower_connection.id as i32),
1841 ..Default::default()
1842 }
1843 .insert(&*tx)
1844 .await?;
1845
1846 let room = self.get_room(room_id, &*tx).await?;
1847 Ok(room)
1848 })
1849 .await
1850 }
1851
1852 pub async fn unfollow(
1853 &self,
1854 project_id: ProjectId,
1855 leader_connection: ConnectionId,
1856 follower_connection: ConnectionId,
1857 ) -> Result<RoomGuard<proto::Room>> {
1858 let room_id = self.room_id_for_project(project_id).await?;
1859 self.room_transaction(room_id, |tx| async move {
1860 follower::Entity::delete_many()
1861 .filter(
1862 Condition::all()
1863 .add(follower::Column::ProjectId.eq(project_id))
1864 .add(
1865 follower::Column::LeaderConnectionServerId
1866 .eq(leader_connection.owner_id),
1867 )
1868 .add(follower::Column::LeaderConnectionId.eq(leader_connection.id))
1869 .add(
1870 follower::Column::FollowerConnectionServerId
1871 .eq(follower_connection.owner_id),
1872 )
1873 .add(follower::Column::FollowerConnectionId.eq(follower_connection.id)),
1874 )
1875 .exec(&*tx)
1876 .await?;
1877
1878 let room = self.get_room(room_id, &*tx).await?;
1879 Ok(room)
1880 })
1881 .await
1882 }
1883
1884 pub async fn update_room_participant_location(
1885 &self,
1886 room_id: RoomId,
1887 connection: ConnectionId,
1888 location: proto::ParticipantLocation,
1889 ) -> Result<RoomGuard<proto::Room>> {
1890 self.room_transaction(room_id, |tx| async {
1891 let tx = tx;
1892 let location_kind;
1893 let location_project_id;
1894 match location
1895 .variant
1896 .as_ref()
1897 .ok_or_else(|| anyhow!("invalid location"))?
1898 {
1899 proto::participant_location::Variant::SharedProject(project) => {
1900 location_kind = 0;
1901 location_project_id = Some(ProjectId::from_proto(project.id));
1902 }
1903 proto::participant_location::Variant::UnsharedProject(_) => {
1904 location_kind = 1;
1905 location_project_id = None;
1906 }
1907 proto::participant_location::Variant::External(_) => {
1908 location_kind = 2;
1909 location_project_id = None;
1910 }
1911 }
1912
1913 let result = room_participant::Entity::update_many()
1914 .filter(
1915 Condition::all()
1916 .add(room_participant::Column::RoomId.eq(room_id))
1917 .add(
1918 room_participant::Column::AnsweringConnectionId
1919 .eq(connection.id as i32),
1920 )
1921 .add(
1922 room_participant::Column::AnsweringConnectionServerId
1923 .eq(connection.owner_id as i32),
1924 ),
1925 )
1926 .set(room_participant::ActiveModel {
1927 location_kind: ActiveValue::set(Some(location_kind)),
1928 location_project_id: ActiveValue::set(location_project_id),
1929 ..Default::default()
1930 })
1931 .exec(&*tx)
1932 .await?;
1933
1934 if result.rows_affected == 1 {
1935 let room = self.get_room(room_id, &tx).await?;
1936 Ok(room)
1937 } else {
1938 Err(anyhow!("could not update room participant location"))?
1939 }
1940 })
1941 .await
1942 }
1943
1944 pub async fn connection_lost(&self, connection: ConnectionId) -> Result<()> {
1945 self.transaction(|tx| async move {
1946 let participant = room_participant::Entity::find()
1947 .filter(
1948 Condition::all()
1949 .add(
1950 room_participant::Column::AnsweringConnectionId
1951 .eq(connection.id as i32),
1952 )
1953 .add(
1954 room_participant::Column::AnsweringConnectionServerId
1955 .eq(connection.owner_id as i32),
1956 ),
1957 )
1958 .one(&*tx)
1959 .await?
1960 .ok_or_else(|| anyhow!("not a participant in any room"))?;
1961
1962 room_participant::Entity::update(room_participant::ActiveModel {
1963 answering_connection_lost: ActiveValue::set(true),
1964 ..participant.into_active_model()
1965 })
1966 .exec(&*tx)
1967 .await?;
1968
1969 Ok(())
1970 })
1971 .await
1972 }
1973
1974 fn build_incoming_call(
1975 room: &proto::Room,
1976 called_user_id: UserId,
1977 ) -> Option<proto::IncomingCall> {
1978 let pending_participant = room
1979 .pending_participants
1980 .iter()
1981 .find(|participant| participant.user_id == called_user_id.to_proto())?;
1982
1983 Some(proto::IncomingCall {
1984 room_id: room.id,
1985 calling_user_id: pending_participant.calling_user_id,
1986 participant_user_ids: room
1987 .participants
1988 .iter()
1989 .map(|participant| participant.user_id)
1990 .collect(),
1991 initial_project: room.participants.iter().find_map(|participant| {
1992 let initial_project_id = pending_participant.initial_project_id?;
1993 participant
1994 .projects
1995 .iter()
1996 .find(|project| project.id == initial_project_id)
1997 .cloned()
1998 }),
1999 })
2000 }
2001
2002 async fn get_room(&self, room_id: RoomId, tx: &DatabaseTransaction) -> Result<proto::Room> {
2003 let db_room = room::Entity::find_by_id(room_id)
2004 .one(tx)
2005 .await?
2006 .ok_or_else(|| anyhow!("could not find room"))?;
2007
2008 let mut db_participants = db_room
2009 .find_related(room_participant::Entity)
2010 .stream(tx)
2011 .await?;
2012 let mut participants = HashMap::default();
2013 let mut pending_participants = Vec::new();
2014 while let Some(db_participant) = db_participants.next().await {
2015 let db_participant = db_participant?;
2016 if let Some((answering_connection_id, answering_connection_server_id)) = db_participant
2017 .answering_connection_id
2018 .zip(db_participant.answering_connection_server_id)
2019 {
2020 let location = match (
2021 db_participant.location_kind,
2022 db_participant.location_project_id,
2023 ) {
2024 (Some(0), Some(project_id)) => {
2025 Some(proto::participant_location::Variant::SharedProject(
2026 proto::participant_location::SharedProject {
2027 id: project_id.to_proto(),
2028 },
2029 ))
2030 }
2031 (Some(1), _) => Some(proto::participant_location::Variant::UnsharedProject(
2032 Default::default(),
2033 )),
2034 _ => Some(proto::participant_location::Variant::External(
2035 Default::default(),
2036 )),
2037 };
2038
2039 let answering_connection = ConnectionId {
2040 owner_id: answering_connection_server_id.0 as u32,
2041 id: answering_connection_id as u32,
2042 };
2043 participants.insert(
2044 answering_connection,
2045 proto::Participant {
2046 user_id: db_participant.user_id.to_proto(),
2047 peer_id: Some(answering_connection.into()),
2048 projects: Default::default(),
2049 location: Some(proto::ParticipantLocation { variant: location }),
2050 },
2051 );
2052 } else {
2053 pending_participants.push(proto::PendingParticipant {
2054 user_id: db_participant.user_id.to_proto(),
2055 calling_user_id: db_participant.calling_user_id.to_proto(),
2056 initial_project_id: db_participant.initial_project_id.map(|id| id.to_proto()),
2057 });
2058 }
2059 }
2060 drop(db_participants);
2061
2062 let mut db_projects = db_room
2063 .find_related(project::Entity)
2064 .find_with_related(worktree::Entity)
2065 .stream(tx)
2066 .await?;
2067
2068 while let Some(row) = db_projects.next().await {
2069 let (db_project, db_worktree) = row?;
2070 let host_connection = db_project.host_connection()?;
2071 if let Some(participant) = participants.get_mut(&host_connection) {
2072 let project = if let Some(project) = participant
2073 .projects
2074 .iter_mut()
2075 .find(|project| project.id == db_project.id.to_proto())
2076 {
2077 project
2078 } else {
2079 participant.projects.push(proto::ParticipantProject {
2080 id: db_project.id.to_proto(),
2081 worktree_root_names: Default::default(),
2082 });
2083 participant.projects.last_mut().unwrap()
2084 };
2085
2086 if let Some(db_worktree) = db_worktree {
2087 if db_worktree.visible {
2088 project.worktree_root_names.push(db_worktree.root_name);
2089 }
2090 }
2091 }
2092 }
2093 drop(db_projects);
2094
2095 let mut db_followers = db_room.find_related(follower::Entity).stream(tx).await?;
2096 let mut followers = Vec::new();
2097 while let Some(db_follower) = db_followers.next().await {
2098 let db_follower = db_follower?;
2099 followers.push(proto::Follower {
2100 leader_id: Some(db_follower.leader_connection().into()),
2101 follower_id: Some(db_follower.follower_connection().into()),
2102 project_id: db_follower.project_id.to_proto(),
2103 });
2104 }
2105
2106 Ok(proto::Room {
2107 id: db_room.id.to_proto(),
2108 live_kit_room: db_room.live_kit_room,
2109 participants: participants.into_values().collect(),
2110 pending_participants,
2111 followers,
2112 })
2113 }
2114
2115 // projects
2116
2117 pub async fn project_count_excluding_admins(&self) -> Result<usize> {
2118 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
2119 enum QueryAs {
2120 Count,
2121 }
2122
2123 self.transaction(|tx| async move {
2124 Ok(project::Entity::find()
2125 .select_only()
2126 .column_as(project::Column::Id.count(), QueryAs::Count)
2127 .inner_join(user::Entity)
2128 .filter(user::Column::Admin.eq(false))
2129 .into_values::<_, QueryAs>()
2130 .one(&*tx)
2131 .await?
2132 .unwrap_or(0i64) as usize)
2133 })
2134 .await
2135 }
2136
2137 pub async fn share_project(
2138 &self,
2139 room_id: RoomId,
2140 connection: ConnectionId,
2141 worktrees: &[proto::WorktreeMetadata],
2142 ) -> Result<RoomGuard<(ProjectId, proto::Room)>> {
2143 self.room_transaction(room_id, |tx| async move {
2144 let participant = room_participant::Entity::find()
2145 .filter(
2146 Condition::all()
2147 .add(
2148 room_participant::Column::AnsweringConnectionId
2149 .eq(connection.id as i32),
2150 )
2151 .add(
2152 room_participant::Column::AnsweringConnectionServerId
2153 .eq(connection.owner_id as i32),
2154 ),
2155 )
2156 .one(&*tx)
2157 .await?
2158 .ok_or_else(|| anyhow!("could not find participant"))?;
2159 if participant.room_id != room_id {
2160 return Err(anyhow!("shared project on unexpected room"))?;
2161 }
2162
2163 let project = project::ActiveModel {
2164 room_id: ActiveValue::set(participant.room_id),
2165 host_user_id: ActiveValue::set(participant.user_id),
2166 host_connection_id: ActiveValue::set(Some(connection.id as i32)),
2167 host_connection_server_id: ActiveValue::set(Some(ServerId(
2168 connection.owner_id as i32,
2169 ))),
2170 ..Default::default()
2171 }
2172 .insert(&*tx)
2173 .await?;
2174
2175 if !worktrees.is_empty() {
2176 worktree::Entity::insert_many(worktrees.iter().map(|worktree| {
2177 worktree::ActiveModel {
2178 id: ActiveValue::set(worktree.id as i64),
2179 project_id: ActiveValue::set(project.id),
2180 abs_path: ActiveValue::set(worktree.abs_path.clone()),
2181 root_name: ActiveValue::set(worktree.root_name.clone()),
2182 visible: ActiveValue::set(worktree.visible),
2183 scan_id: ActiveValue::set(0),
2184 completed_scan_id: ActiveValue::set(0),
2185 }
2186 }))
2187 .exec(&*tx)
2188 .await?;
2189 }
2190
2191 project_collaborator::ActiveModel {
2192 project_id: ActiveValue::set(project.id),
2193 connection_id: ActiveValue::set(connection.id as i32),
2194 connection_server_id: ActiveValue::set(ServerId(connection.owner_id as i32)),
2195 user_id: ActiveValue::set(participant.user_id),
2196 replica_id: ActiveValue::set(ReplicaId(0)),
2197 is_host: ActiveValue::set(true),
2198 ..Default::default()
2199 }
2200 .insert(&*tx)
2201 .await?;
2202
2203 let room = self.get_room(room_id, &tx).await?;
2204 Ok((project.id, room))
2205 })
2206 .await
2207 }
2208
2209 pub async fn unshare_project(
2210 &self,
2211 project_id: ProjectId,
2212 connection: ConnectionId,
2213 ) -> Result<RoomGuard<(proto::Room, Vec<ConnectionId>)>> {
2214 let room_id = self.room_id_for_project(project_id).await?;
2215 self.room_transaction(room_id, |tx| async move {
2216 let guest_connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
2217
2218 let project = project::Entity::find_by_id(project_id)
2219 .one(&*tx)
2220 .await?
2221 .ok_or_else(|| anyhow!("project not found"))?;
2222 if project.host_connection()? == connection {
2223 project::Entity::delete(project.into_active_model())
2224 .exec(&*tx)
2225 .await?;
2226 let room = self.get_room(room_id, &tx).await?;
2227 Ok((room, guest_connection_ids))
2228 } else {
2229 Err(anyhow!("cannot unshare a project hosted by another user"))?
2230 }
2231 })
2232 .await
2233 }
2234
2235 pub async fn update_project(
2236 &self,
2237 project_id: ProjectId,
2238 connection: ConnectionId,
2239 worktrees: &[proto::WorktreeMetadata],
2240 ) -> Result<RoomGuard<(proto::Room, Vec<ConnectionId>)>> {
2241 let room_id = self.room_id_for_project(project_id).await?;
2242 self.room_transaction(room_id, |tx| async move {
2243 let project = project::Entity::find_by_id(project_id)
2244 .filter(
2245 Condition::all()
2246 .add(project::Column::HostConnectionId.eq(connection.id as i32))
2247 .add(
2248 project::Column::HostConnectionServerId.eq(connection.owner_id as i32),
2249 ),
2250 )
2251 .one(&*tx)
2252 .await?
2253 .ok_or_else(|| anyhow!("no such project"))?;
2254
2255 self.update_project_worktrees(project.id, worktrees, &tx)
2256 .await?;
2257
2258 let guest_connection_ids = self.project_guest_connection_ids(project.id, &tx).await?;
2259 let room = self.get_room(project.room_id, &tx).await?;
2260 Ok((room, guest_connection_ids))
2261 })
2262 .await
2263 }
2264
2265 async fn update_project_worktrees(
2266 &self,
2267 project_id: ProjectId,
2268 worktrees: &[proto::WorktreeMetadata],
2269 tx: &DatabaseTransaction,
2270 ) -> Result<()> {
2271 if !worktrees.is_empty() {
2272 worktree::Entity::insert_many(worktrees.iter().map(|worktree| worktree::ActiveModel {
2273 id: ActiveValue::set(worktree.id as i64),
2274 project_id: ActiveValue::set(project_id),
2275 abs_path: ActiveValue::set(worktree.abs_path.clone()),
2276 root_name: ActiveValue::set(worktree.root_name.clone()),
2277 visible: ActiveValue::set(worktree.visible),
2278 scan_id: ActiveValue::set(0),
2279 completed_scan_id: ActiveValue::set(0),
2280 }))
2281 .on_conflict(
2282 OnConflict::columns([worktree::Column::ProjectId, worktree::Column::Id])
2283 .update_column(worktree::Column::RootName)
2284 .to_owned(),
2285 )
2286 .exec(&*tx)
2287 .await?;
2288 }
2289
2290 worktree::Entity::delete_many()
2291 .filter(worktree::Column::ProjectId.eq(project_id).and(
2292 worktree::Column::Id.is_not_in(worktrees.iter().map(|worktree| worktree.id as i64)),
2293 ))
2294 .exec(&*tx)
2295 .await?;
2296
2297 Ok(())
2298 }
2299
2300 pub async fn update_worktree(
2301 &self,
2302 update: &proto::UpdateWorktree,
2303 connection: ConnectionId,
2304 ) -> Result<RoomGuard<Vec<ConnectionId>>> {
2305 let project_id = ProjectId::from_proto(update.project_id);
2306 let worktree_id = update.worktree_id as i64;
2307 let room_id = self.room_id_for_project(project_id).await?;
2308 self.room_transaction(room_id, |tx| async move {
2309 // Ensure the update comes from the host.
2310 let _project = project::Entity::find_by_id(project_id)
2311 .filter(
2312 Condition::all()
2313 .add(project::Column::HostConnectionId.eq(connection.id as i32))
2314 .add(
2315 project::Column::HostConnectionServerId.eq(connection.owner_id as i32),
2316 ),
2317 )
2318 .one(&*tx)
2319 .await?
2320 .ok_or_else(|| anyhow!("no such project"))?;
2321
2322 // Update metadata.
2323 worktree::Entity::update(worktree::ActiveModel {
2324 id: ActiveValue::set(worktree_id),
2325 project_id: ActiveValue::set(project_id),
2326 root_name: ActiveValue::set(update.root_name.clone()),
2327 scan_id: ActiveValue::set(update.scan_id as i64),
2328 completed_scan_id: if update.is_last_update {
2329 ActiveValue::set(update.scan_id as i64)
2330 } else {
2331 ActiveValue::default()
2332 },
2333 abs_path: ActiveValue::set(update.abs_path.clone()),
2334 ..Default::default()
2335 })
2336 .exec(&*tx)
2337 .await?;
2338
2339 if !update.updated_entries.is_empty() {
2340 worktree_entry::Entity::insert_many(update.updated_entries.iter().map(|entry| {
2341 let mtime = entry.mtime.clone().unwrap_or_default();
2342 worktree_entry::ActiveModel {
2343 project_id: ActiveValue::set(project_id),
2344 worktree_id: ActiveValue::set(worktree_id),
2345 id: ActiveValue::set(entry.id as i64),
2346 is_dir: ActiveValue::set(entry.is_dir),
2347 path: ActiveValue::set(entry.path.clone()),
2348 inode: ActiveValue::set(entry.inode as i64),
2349 mtime_seconds: ActiveValue::set(mtime.seconds as i64),
2350 mtime_nanos: ActiveValue::set(mtime.nanos as i32),
2351 is_symlink: ActiveValue::set(entry.is_symlink),
2352 is_ignored: ActiveValue::set(entry.is_ignored),
2353 is_external: ActiveValue::set(entry.is_external),
2354 git_status: ActiveValue::set(entry.git_status.map(|status| status as i64)),
2355 is_deleted: ActiveValue::set(false),
2356 scan_id: ActiveValue::set(update.scan_id as i64),
2357 }
2358 }))
2359 .on_conflict(
2360 OnConflict::columns([
2361 worktree_entry::Column::ProjectId,
2362 worktree_entry::Column::WorktreeId,
2363 worktree_entry::Column::Id,
2364 ])
2365 .update_columns([
2366 worktree_entry::Column::IsDir,
2367 worktree_entry::Column::Path,
2368 worktree_entry::Column::Inode,
2369 worktree_entry::Column::MtimeSeconds,
2370 worktree_entry::Column::MtimeNanos,
2371 worktree_entry::Column::IsSymlink,
2372 worktree_entry::Column::IsIgnored,
2373 worktree_entry::Column::GitStatus,
2374 worktree_entry::Column::ScanId,
2375 ])
2376 .to_owned(),
2377 )
2378 .exec(&*tx)
2379 .await?;
2380 }
2381
2382 if !update.removed_entries.is_empty() {
2383 worktree_entry::Entity::update_many()
2384 .filter(
2385 worktree_entry::Column::ProjectId
2386 .eq(project_id)
2387 .and(worktree_entry::Column::WorktreeId.eq(worktree_id))
2388 .and(
2389 worktree_entry::Column::Id
2390 .is_in(update.removed_entries.iter().map(|id| *id as i64)),
2391 ),
2392 )
2393 .set(worktree_entry::ActiveModel {
2394 is_deleted: ActiveValue::Set(true),
2395 scan_id: ActiveValue::Set(update.scan_id as i64),
2396 ..Default::default()
2397 })
2398 .exec(&*tx)
2399 .await?;
2400 }
2401
2402 if !update.updated_repositories.is_empty() {
2403 worktree_repository::Entity::insert_many(update.updated_repositories.iter().map(
2404 |repository| worktree_repository::ActiveModel {
2405 project_id: ActiveValue::set(project_id),
2406 worktree_id: ActiveValue::set(worktree_id),
2407 work_directory_id: ActiveValue::set(repository.work_directory_id as i64),
2408 scan_id: ActiveValue::set(update.scan_id as i64),
2409 branch: ActiveValue::set(repository.branch.clone()),
2410 is_deleted: ActiveValue::set(false),
2411 },
2412 ))
2413 .on_conflict(
2414 OnConflict::columns([
2415 worktree_repository::Column::ProjectId,
2416 worktree_repository::Column::WorktreeId,
2417 worktree_repository::Column::WorkDirectoryId,
2418 ])
2419 .update_columns([
2420 worktree_repository::Column::ScanId,
2421 worktree_repository::Column::Branch,
2422 ])
2423 .to_owned(),
2424 )
2425 .exec(&*tx)
2426 .await?;
2427 }
2428
2429 if !update.removed_repositories.is_empty() {
2430 worktree_repository::Entity::update_many()
2431 .filter(
2432 worktree_repository::Column::ProjectId
2433 .eq(project_id)
2434 .and(worktree_repository::Column::WorktreeId.eq(worktree_id))
2435 .and(
2436 worktree_repository::Column::WorkDirectoryId
2437 .is_in(update.removed_repositories.iter().map(|id| *id as i64)),
2438 ),
2439 )
2440 .set(worktree_repository::ActiveModel {
2441 is_deleted: ActiveValue::Set(true),
2442 scan_id: ActiveValue::Set(update.scan_id as i64),
2443 ..Default::default()
2444 })
2445 .exec(&*tx)
2446 .await?;
2447 }
2448
2449 let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
2450 Ok(connection_ids)
2451 })
2452 .await
2453 }
2454
2455 pub async fn update_diagnostic_summary(
2456 &self,
2457 update: &proto::UpdateDiagnosticSummary,
2458 connection: ConnectionId,
2459 ) -> Result<RoomGuard<Vec<ConnectionId>>> {
2460 let project_id = ProjectId::from_proto(update.project_id);
2461 let worktree_id = update.worktree_id as i64;
2462 let room_id = self.room_id_for_project(project_id).await?;
2463 self.room_transaction(room_id, |tx| async move {
2464 let summary = update
2465 .summary
2466 .as_ref()
2467 .ok_or_else(|| anyhow!("invalid summary"))?;
2468
2469 // Ensure the update comes from the host.
2470 let project = project::Entity::find_by_id(project_id)
2471 .one(&*tx)
2472 .await?
2473 .ok_or_else(|| anyhow!("no such project"))?;
2474 if project.host_connection()? != connection {
2475 return Err(anyhow!("can't update a project hosted by someone else"))?;
2476 }
2477
2478 // Update summary.
2479 worktree_diagnostic_summary::Entity::insert(worktree_diagnostic_summary::ActiveModel {
2480 project_id: ActiveValue::set(project_id),
2481 worktree_id: ActiveValue::set(worktree_id),
2482 path: ActiveValue::set(summary.path.clone()),
2483 language_server_id: ActiveValue::set(summary.language_server_id as i64),
2484 error_count: ActiveValue::set(summary.error_count as i32),
2485 warning_count: ActiveValue::set(summary.warning_count as i32),
2486 ..Default::default()
2487 })
2488 .on_conflict(
2489 OnConflict::columns([
2490 worktree_diagnostic_summary::Column::ProjectId,
2491 worktree_diagnostic_summary::Column::WorktreeId,
2492 worktree_diagnostic_summary::Column::Path,
2493 ])
2494 .update_columns([
2495 worktree_diagnostic_summary::Column::LanguageServerId,
2496 worktree_diagnostic_summary::Column::ErrorCount,
2497 worktree_diagnostic_summary::Column::WarningCount,
2498 ])
2499 .to_owned(),
2500 )
2501 .exec(&*tx)
2502 .await?;
2503
2504 let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
2505 Ok(connection_ids)
2506 })
2507 .await
2508 }
2509
2510 pub async fn start_language_server(
2511 &self,
2512 update: &proto::StartLanguageServer,
2513 connection: ConnectionId,
2514 ) -> Result<RoomGuard<Vec<ConnectionId>>> {
2515 let project_id = ProjectId::from_proto(update.project_id);
2516 let room_id = self.room_id_for_project(project_id).await?;
2517 self.room_transaction(room_id, |tx| async move {
2518 let server = update
2519 .server
2520 .as_ref()
2521 .ok_or_else(|| anyhow!("invalid language server"))?;
2522
2523 // Ensure the update comes from the host.
2524 let project = project::Entity::find_by_id(project_id)
2525 .one(&*tx)
2526 .await?
2527 .ok_or_else(|| anyhow!("no such project"))?;
2528 if project.host_connection()? != connection {
2529 return Err(anyhow!("can't update a project hosted by someone else"))?;
2530 }
2531
2532 // Add the newly-started language server.
2533 language_server::Entity::insert(language_server::ActiveModel {
2534 project_id: ActiveValue::set(project_id),
2535 id: ActiveValue::set(server.id as i64),
2536 name: ActiveValue::set(server.name.clone()),
2537 ..Default::default()
2538 })
2539 .on_conflict(
2540 OnConflict::columns([
2541 language_server::Column::ProjectId,
2542 language_server::Column::Id,
2543 ])
2544 .update_column(language_server::Column::Name)
2545 .to_owned(),
2546 )
2547 .exec(&*tx)
2548 .await?;
2549
2550 let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
2551 Ok(connection_ids)
2552 })
2553 .await
2554 }
2555
2556 pub async fn update_worktree_settings(
2557 &self,
2558 update: &proto::UpdateWorktreeSettings,
2559 connection: ConnectionId,
2560 ) -> Result<RoomGuard<Vec<ConnectionId>>> {
2561 let project_id = ProjectId::from_proto(update.project_id);
2562 let room_id = self.room_id_for_project(project_id).await?;
2563 self.room_transaction(room_id, |tx| async move {
2564 // Ensure the update comes from the host.
2565 let project = project::Entity::find_by_id(project_id)
2566 .one(&*tx)
2567 .await?
2568 .ok_or_else(|| anyhow!("no such project"))?;
2569 if project.host_connection()? != connection {
2570 return Err(anyhow!("can't update a project hosted by someone else"))?;
2571 }
2572
2573 if let Some(content) = &update.content {
2574 worktree_settings_file::Entity::insert(worktree_settings_file::ActiveModel {
2575 project_id: ActiveValue::Set(project_id),
2576 worktree_id: ActiveValue::Set(update.worktree_id as i64),
2577 path: ActiveValue::Set(update.path.clone()),
2578 content: ActiveValue::Set(content.clone()),
2579 })
2580 .on_conflict(
2581 OnConflict::columns([
2582 worktree_settings_file::Column::ProjectId,
2583 worktree_settings_file::Column::WorktreeId,
2584 worktree_settings_file::Column::Path,
2585 ])
2586 .update_column(worktree_settings_file::Column::Content)
2587 .to_owned(),
2588 )
2589 .exec(&*tx)
2590 .await?;
2591 } else {
2592 worktree_settings_file::Entity::delete(worktree_settings_file::ActiveModel {
2593 project_id: ActiveValue::Set(project_id),
2594 worktree_id: ActiveValue::Set(update.worktree_id as i64),
2595 path: ActiveValue::Set(update.path.clone()),
2596 ..Default::default()
2597 })
2598 .exec(&*tx)
2599 .await?;
2600 }
2601
2602 let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
2603 Ok(connection_ids)
2604 })
2605 .await
2606 }
2607
2608 pub async fn join_project(
2609 &self,
2610 project_id: ProjectId,
2611 connection: ConnectionId,
2612 ) -> Result<RoomGuard<(Project, ReplicaId)>> {
2613 let room_id = self.room_id_for_project(project_id).await?;
2614 self.room_transaction(room_id, |tx| async move {
2615 let participant = room_participant::Entity::find()
2616 .filter(
2617 Condition::all()
2618 .add(
2619 room_participant::Column::AnsweringConnectionId
2620 .eq(connection.id as i32),
2621 )
2622 .add(
2623 room_participant::Column::AnsweringConnectionServerId
2624 .eq(connection.owner_id as i32),
2625 ),
2626 )
2627 .one(&*tx)
2628 .await?
2629 .ok_or_else(|| anyhow!("must join a room first"))?;
2630
2631 let project = project::Entity::find_by_id(project_id)
2632 .one(&*tx)
2633 .await?
2634 .ok_or_else(|| anyhow!("no such project"))?;
2635 if project.room_id != participant.room_id {
2636 return Err(anyhow!("no such project"))?;
2637 }
2638
2639 let mut collaborators = project
2640 .find_related(project_collaborator::Entity)
2641 .all(&*tx)
2642 .await?;
2643 let replica_ids = collaborators
2644 .iter()
2645 .map(|c| c.replica_id)
2646 .collect::<HashSet<_>>();
2647 let mut replica_id = ReplicaId(1);
2648 while replica_ids.contains(&replica_id) {
2649 replica_id.0 += 1;
2650 }
2651 let new_collaborator = project_collaborator::ActiveModel {
2652 project_id: ActiveValue::set(project_id),
2653 connection_id: ActiveValue::set(connection.id as i32),
2654 connection_server_id: ActiveValue::set(ServerId(connection.owner_id as i32)),
2655 user_id: ActiveValue::set(participant.user_id),
2656 replica_id: ActiveValue::set(replica_id),
2657 is_host: ActiveValue::set(false),
2658 ..Default::default()
2659 }
2660 .insert(&*tx)
2661 .await?;
2662 collaborators.push(new_collaborator);
2663
2664 let db_worktrees = project.find_related(worktree::Entity).all(&*tx).await?;
2665 let mut worktrees = db_worktrees
2666 .into_iter()
2667 .map(|db_worktree| {
2668 (
2669 db_worktree.id as u64,
2670 Worktree {
2671 id: db_worktree.id as u64,
2672 abs_path: db_worktree.abs_path,
2673 root_name: db_worktree.root_name,
2674 visible: db_worktree.visible,
2675 entries: Default::default(),
2676 repository_entries: Default::default(),
2677 diagnostic_summaries: Default::default(),
2678 settings_files: Default::default(),
2679 scan_id: db_worktree.scan_id as u64,
2680 completed_scan_id: db_worktree.completed_scan_id as u64,
2681 },
2682 )
2683 })
2684 .collect::<BTreeMap<_, _>>();
2685
2686 // Populate worktree entries.
2687 {
2688 let mut db_entries = worktree_entry::Entity::find()
2689 .filter(
2690 Condition::all()
2691 .add(worktree_entry::Column::ProjectId.eq(project_id))
2692 .add(worktree_entry::Column::IsDeleted.eq(false)),
2693 )
2694 .stream(&*tx)
2695 .await?;
2696 while let Some(db_entry) = db_entries.next().await {
2697 let db_entry = db_entry?;
2698 if let Some(worktree) = worktrees.get_mut(&(db_entry.worktree_id as u64)) {
2699 worktree.entries.push(proto::Entry {
2700 id: db_entry.id as u64,
2701 is_dir: db_entry.is_dir,
2702 path: db_entry.path,
2703 inode: db_entry.inode as u64,
2704 mtime: Some(proto::Timestamp {
2705 seconds: db_entry.mtime_seconds as u64,
2706 nanos: db_entry.mtime_nanos as u32,
2707 }),
2708 is_symlink: db_entry.is_symlink,
2709 is_ignored: db_entry.is_ignored,
2710 is_external: db_entry.is_external,
2711 git_status: db_entry.git_status.map(|status| status as i32),
2712 });
2713 }
2714 }
2715 }
2716
2717 // Populate repository entries.
2718 {
2719 let mut db_repository_entries = worktree_repository::Entity::find()
2720 .filter(
2721 Condition::all()
2722 .add(worktree_repository::Column::ProjectId.eq(project_id))
2723 .add(worktree_repository::Column::IsDeleted.eq(false)),
2724 )
2725 .stream(&*tx)
2726 .await?;
2727 while let Some(db_repository_entry) = db_repository_entries.next().await {
2728 let db_repository_entry = db_repository_entry?;
2729 if let Some(worktree) =
2730 worktrees.get_mut(&(db_repository_entry.worktree_id as u64))
2731 {
2732 worktree.repository_entries.insert(
2733 db_repository_entry.work_directory_id as u64,
2734 proto::RepositoryEntry {
2735 work_directory_id: db_repository_entry.work_directory_id as u64,
2736 branch: db_repository_entry.branch,
2737 },
2738 );
2739 }
2740 }
2741 }
2742
2743 // Populate worktree diagnostic summaries.
2744 {
2745 let mut db_summaries = worktree_diagnostic_summary::Entity::find()
2746 .filter(worktree_diagnostic_summary::Column::ProjectId.eq(project_id))
2747 .stream(&*tx)
2748 .await?;
2749 while let Some(db_summary) = db_summaries.next().await {
2750 let db_summary = db_summary?;
2751 if let Some(worktree) = worktrees.get_mut(&(db_summary.worktree_id as u64)) {
2752 worktree
2753 .diagnostic_summaries
2754 .push(proto::DiagnosticSummary {
2755 path: db_summary.path,
2756 language_server_id: db_summary.language_server_id as u64,
2757 error_count: db_summary.error_count as u32,
2758 warning_count: db_summary.warning_count as u32,
2759 });
2760 }
2761 }
2762 }
2763
2764 // Populate worktree settings files
2765 {
2766 let mut db_settings_files = worktree_settings_file::Entity::find()
2767 .filter(worktree_settings_file::Column::ProjectId.eq(project_id))
2768 .stream(&*tx)
2769 .await?;
2770 while let Some(db_settings_file) = db_settings_files.next().await {
2771 let db_settings_file = db_settings_file?;
2772 if let Some(worktree) =
2773 worktrees.get_mut(&(db_settings_file.worktree_id as u64))
2774 {
2775 worktree.settings_files.push(WorktreeSettingsFile {
2776 path: db_settings_file.path,
2777 content: db_settings_file.content,
2778 });
2779 }
2780 }
2781 }
2782
2783 // Populate language servers.
2784 let language_servers = project
2785 .find_related(language_server::Entity)
2786 .all(&*tx)
2787 .await?;
2788
2789 let project = Project {
2790 collaborators: collaborators
2791 .into_iter()
2792 .map(|collaborator| ProjectCollaborator {
2793 connection_id: collaborator.connection(),
2794 user_id: collaborator.user_id,
2795 replica_id: collaborator.replica_id,
2796 is_host: collaborator.is_host,
2797 })
2798 .collect(),
2799 worktrees,
2800 language_servers: language_servers
2801 .into_iter()
2802 .map(|language_server| proto::LanguageServer {
2803 id: language_server.id as u64,
2804 name: language_server.name,
2805 })
2806 .collect(),
2807 };
2808 Ok((project, replica_id as ReplicaId))
2809 })
2810 .await
2811 }
2812
2813 pub async fn leave_project(
2814 &self,
2815 project_id: ProjectId,
2816 connection: ConnectionId,
2817 ) -> Result<RoomGuard<(proto::Room, LeftProject)>> {
2818 let room_id = self.room_id_for_project(project_id).await?;
2819 self.room_transaction(room_id, |tx| async move {
2820 let result = project_collaborator::Entity::delete_many()
2821 .filter(
2822 Condition::all()
2823 .add(project_collaborator::Column::ProjectId.eq(project_id))
2824 .add(project_collaborator::Column::ConnectionId.eq(connection.id as i32))
2825 .add(
2826 project_collaborator::Column::ConnectionServerId
2827 .eq(connection.owner_id as i32),
2828 ),
2829 )
2830 .exec(&*tx)
2831 .await?;
2832 if result.rows_affected == 0 {
2833 Err(anyhow!("not a collaborator on this project"))?;
2834 }
2835
2836 let project = project::Entity::find_by_id(project_id)
2837 .one(&*tx)
2838 .await?
2839 .ok_or_else(|| anyhow!("no such project"))?;
2840 let collaborators = project
2841 .find_related(project_collaborator::Entity)
2842 .all(&*tx)
2843 .await?;
2844 let connection_ids = collaborators
2845 .into_iter()
2846 .map(|collaborator| collaborator.connection())
2847 .collect();
2848
2849 follower::Entity::delete_many()
2850 .filter(
2851 Condition::any()
2852 .add(
2853 Condition::all()
2854 .add(follower::Column::ProjectId.eq(project_id))
2855 .add(
2856 follower::Column::LeaderConnectionServerId
2857 .eq(connection.owner_id),
2858 )
2859 .add(follower::Column::LeaderConnectionId.eq(connection.id)),
2860 )
2861 .add(
2862 Condition::all()
2863 .add(follower::Column::ProjectId.eq(project_id))
2864 .add(
2865 follower::Column::FollowerConnectionServerId
2866 .eq(connection.owner_id),
2867 )
2868 .add(follower::Column::FollowerConnectionId.eq(connection.id)),
2869 ),
2870 )
2871 .exec(&*tx)
2872 .await?;
2873
2874 let room = self.get_room(project.room_id, &tx).await?;
2875 let left_project = LeftProject {
2876 id: project_id,
2877 host_user_id: project.host_user_id,
2878 host_connection_id: project.host_connection()?,
2879 connection_ids,
2880 };
2881 Ok((room, left_project))
2882 })
2883 .await
2884 }
2885
2886 pub async fn project_collaborators(
2887 &self,
2888 project_id: ProjectId,
2889 connection_id: ConnectionId,
2890 ) -> Result<RoomGuard<Vec<ProjectCollaborator>>> {
2891 let room_id = self.room_id_for_project(project_id).await?;
2892 self.room_transaction(room_id, |tx| async move {
2893 let collaborators = project_collaborator::Entity::find()
2894 .filter(project_collaborator::Column::ProjectId.eq(project_id))
2895 .all(&*tx)
2896 .await?
2897 .into_iter()
2898 .map(|collaborator| ProjectCollaborator {
2899 connection_id: collaborator.connection(),
2900 user_id: collaborator.user_id,
2901 replica_id: collaborator.replica_id,
2902 is_host: collaborator.is_host,
2903 })
2904 .collect::<Vec<_>>();
2905
2906 if collaborators
2907 .iter()
2908 .any(|collaborator| collaborator.connection_id == connection_id)
2909 {
2910 Ok(collaborators)
2911 } else {
2912 Err(anyhow!("no such project"))?
2913 }
2914 })
2915 .await
2916 }
2917
2918 pub async fn project_connection_ids(
2919 &self,
2920 project_id: ProjectId,
2921 connection_id: ConnectionId,
2922 ) -> Result<RoomGuard<HashSet<ConnectionId>>> {
2923 let room_id = self.room_id_for_project(project_id).await?;
2924 self.room_transaction(room_id, |tx| async move {
2925 let mut collaborators = project_collaborator::Entity::find()
2926 .filter(project_collaborator::Column::ProjectId.eq(project_id))
2927 .stream(&*tx)
2928 .await?;
2929
2930 let mut connection_ids = HashSet::default();
2931 while let Some(collaborator) = collaborators.next().await {
2932 let collaborator = collaborator?;
2933 connection_ids.insert(collaborator.connection());
2934 }
2935
2936 if connection_ids.contains(&connection_id) {
2937 Ok(connection_ids)
2938 } else {
2939 Err(anyhow!("no such project"))?
2940 }
2941 })
2942 .await
2943 }
2944
2945 async fn project_guest_connection_ids(
2946 &self,
2947 project_id: ProjectId,
2948 tx: &DatabaseTransaction,
2949 ) -> Result<Vec<ConnectionId>> {
2950 let mut collaborators = project_collaborator::Entity::find()
2951 .filter(
2952 project_collaborator::Column::ProjectId
2953 .eq(project_id)
2954 .and(project_collaborator::Column::IsHost.eq(false)),
2955 )
2956 .stream(tx)
2957 .await?;
2958
2959 let mut guest_connection_ids = Vec::new();
2960 while let Some(collaborator) = collaborators.next().await {
2961 let collaborator = collaborator?;
2962 guest_connection_ids.push(collaborator.connection());
2963 }
2964 Ok(guest_connection_ids)
2965 }
2966
2967 async fn room_id_for_project(&self, project_id: ProjectId) -> Result<RoomId> {
2968 self.transaction(|tx| async move {
2969 let project = project::Entity::find_by_id(project_id)
2970 .one(&*tx)
2971 .await?
2972 .ok_or_else(|| anyhow!("project {} not found", project_id))?;
2973 Ok(project.room_id)
2974 })
2975 .await
2976 }
2977
2978 // access tokens
2979
2980 pub async fn create_access_token(
2981 &self,
2982 user_id: UserId,
2983 access_token_hash: &str,
2984 max_access_token_count: usize,
2985 ) -> Result<AccessTokenId> {
2986 self.transaction(|tx| async {
2987 let tx = tx;
2988
2989 let token = access_token::ActiveModel {
2990 user_id: ActiveValue::set(user_id),
2991 hash: ActiveValue::set(access_token_hash.into()),
2992 ..Default::default()
2993 }
2994 .insert(&*tx)
2995 .await?;
2996
2997 access_token::Entity::delete_many()
2998 .filter(
2999 access_token::Column::Id.in_subquery(
3000 Query::select()
3001 .column(access_token::Column::Id)
3002 .from(access_token::Entity)
3003 .and_where(access_token::Column::UserId.eq(user_id))
3004 .order_by(access_token::Column::Id, sea_orm::Order::Desc)
3005 .limit(10000)
3006 .offset(max_access_token_count as u64)
3007 .to_owned(),
3008 ),
3009 )
3010 .exec(&*tx)
3011 .await?;
3012 Ok(token.id)
3013 })
3014 .await
3015 }
3016
3017 pub async fn get_access_token(
3018 &self,
3019 access_token_id: AccessTokenId,
3020 ) -> Result<access_token::Model> {
3021 self.transaction(|tx| async move {
3022 Ok(access_token::Entity::find_by_id(access_token_id)
3023 .one(&*tx)
3024 .await?
3025 .ok_or_else(|| anyhow!("no such access token"))?)
3026 })
3027 .await
3028 }
3029
3030 async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
3031 where
3032 F: Send + Fn(TransactionHandle) -> Fut,
3033 Fut: Send + Future<Output = Result<T>>,
3034 {
3035 let body = async {
3036 let mut i = 0;
3037 loop {
3038 let (tx, result) = self.with_transaction(&f).await?;
3039 match result {
3040 Ok(result) => match tx.commit().await.map_err(Into::into) {
3041 Ok(()) => return Ok(result),
3042 Err(error) => {
3043 if !self.retry_on_serialization_error(&error, i).await {
3044 return Err(error);
3045 }
3046 }
3047 },
3048 Err(error) => {
3049 tx.rollback().await?;
3050 if !self.retry_on_serialization_error(&error, i).await {
3051 return Err(error);
3052 }
3053 }
3054 }
3055 i += 1;
3056 }
3057 };
3058
3059 self.run(body).await
3060 }
3061
3062 async fn optional_room_transaction<F, Fut, T>(&self, f: F) -> Result<Option<RoomGuard<T>>>
3063 where
3064 F: Send + Fn(TransactionHandle) -> Fut,
3065 Fut: Send + Future<Output = Result<Option<(RoomId, T)>>>,
3066 {
3067 let body = async {
3068 let mut i = 0;
3069 loop {
3070 let (tx, result) = self.with_transaction(&f).await?;
3071 match result {
3072 Ok(Some((room_id, data))) => {
3073 let lock = self.rooms.entry(room_id).or_default().clone();
3074 let _guard = lock.lock_owned().await;
3075 match tx.commit().await.map_err(Into::into) {
3076 Ok(()) => {
3077 return Ok(Some(RoomGuard {
3078 data,
3079 _guard,
3080 _not_send: PhantomData,
3081 }));
3082 }
3083 Err(error) => {
3084 if !self.retry_on_serialization_error(&error, i).await {
3085 return Err(error);
3086 }
3087 }
3088 }
3089 }
3090 Ok(None) => match tx.commit().await.map_err(Into::into) {
3091 Ok(()) => return Ok(None),
3092 Err(error) => {
3093 if !self.retry_on_serialization_error(&error, i).await {
3094 return Err(error);
3095 }
3096 }
3097 },
3098 Err(error) => {
3099 tx.rollback().await?;
3100 if !self.retry_on_serialization_error(&error, i).await {
3101 return Err(error);
3102 }
3103 }
3104 }
3105 i += 1;
3106 }
3107 };
3108
3109 self.run(body).await
3110 }
3111
3112 async fn room_transaction<F, Fut, T>(&self, room_id: RoomId, f: F) -> Result<RoomGuard<T>>
3113 where
3114 F: Send + Fn(TransactionHandle) -> Fut,
3115 Fut: Send + Future<Output = Result<T>>,
3116 {
3117 let body = async {
3118 let mut i = 0;
3119 loop {
3120 let lock = self.rooms.entry(room_id).or_default().clone();
3121 let _guard = lock.lock_owned().await;
3122 let (tx, result) = self.with_transaction(&f).await?;
3123 match result {
3124 Ok(data) => match tx.commit().await.map_err(Into::into) {
3125 Ok(()) => {
3126 return Ok(RoomGuard {
3127 data,
3128 _guard,
3129 _not_send: PhantomData,
3130 });
3131 }
3132 Err(error) => {
3133 if !self.retry_on_serialization_error(&error, i).await {
3134 return Err(error);
3135 }
3136 }
3137 },
3138 Err(error) => {
3139 tx.rollback().await?;
3140 if !self.retry_on_serialization_error(&error, i).await {
3141 return Err(error);
3142 }
3143 }
3144 }
3145 i += 1;
3146 }
3147 };
3148
3149 self.run(body).await
3150 }
3151
3152 async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
3153 where
3154 F: Send + Fn(TransactionHandle) -> Fut,
3155 Fut: Send + Future<Output = Result<T>>,
3156 {
3157 let tx = self
3158 .pool
3159 .begin_with_config(Some(IsolationLevel::Serializable), None)
3160 .await?;
3161
3162 let mut tx = Arc::new(Some(tx));
3163 let result = f(TransactionHandle(tx.clone())).await;
3164 let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
3165 return Err(anyhow!("couldn't complete transaction because it's still in use"))?;
3166 };
3167
3168 Ok((tx, result))
3169 }
3170
3171 async fn run<F, T>(&self, future: F) -> Result<T>
3172 where
3173 F: Future<Output = Result<T>>,
3174 {
3175 #[cfg(test)]
3176 {
3177 if let Executor::Deterministic(executor) = &self.executor {
3178 executor.simulate_random_delay().await;
3179 }
3180
3181 self.runtime.as_ref().unwrap().block_on(future)
3182 }
3183
3184 #[cfg(not(test))]
3185 {
3186 future.await
3187 }
3188 }
3189
3190 async fn retry_on_serialization_error(&self, error: &Error, prev_attempt_count: u32) -> bool {
3191 // If the error is due to a failure to serialize concurrent transactions, then retry
3192 // this transaction after a delay. With each subsequent retry, double the delay duration.
3193 // Also vary the delay randomly in order to ensure different database connections retry
3194 // at different times.
3195 if is_serialization_error(error) {
3196 let base_delay = 4_u64 << prev_attempt_count.min(16);
3197 let randomized_delay = base_delay as f32 * self.rng.lock().await.gen_range(0.5..=2.0);
3198 log::info!(
3199 "retrying transaction after serialization error. delay: {} ms.",
3200 randomized_delay
3201 );
3202 self.executor
3203 .sleep(Duration::from_millis(randomized_delay as u64))
3204 .await;
3205 true
3206 } else {
3207 false
3208 }
3209 }
3210}
3211
3212fn is_serialization_error(error: &Error) -> bool {
3213 const SERIALIZATION_FAILURE_CODE: &'static str = "40001";
3214 match error {
3215 Error::Database(
3216 DbErr::Exec(sea_orm::RuntimeErr::SqlxError(error))
3217 | DbErr::Query(sea_orm::RuntimeErr::SqlxError(error)),
3218 ) if error
3219 .as_database_error()
3220 .and_then(|error| error.code())
3221 .as_deref()
3222 == Some(SERIALIZATION_FAILURE_CODE) =>
3223 {
3224 true
3225 }
3226 _ => false,
3227 }
3228}
3229
3230struct TransactionHandle(Arc<Option<DatabaseTransaction>>);
3231
3232impl Deref for TransactionHandle {
3233 type Target = DatabaseTransaction;
3234
3235 fn deref(&self) -> &Self::Target {
3236 self.0.as_ref().as_ref().unwrap()
3237 }
3238}
3239
3240pub struct RoomGuard<T> {
3241 data: T,
3242 _guard: OwnedMutexGuard<()>,
3243 _not_send: PhantomData<Rc<()>>,
3244}
3245
3246impl<T> Deref for RoomGuard<T> {
3247 type Target = T;
3248
3249 fn deref(&self) -> &T {
3250 &self.data
3251 }
3252}
3253
3254impl<T> DerefMut for RoomGuard<T> {
3255 fn deref_mut(&mut self) -> &mut T {
3256 &mut self.data
3257 }
3258}
3259
3260#[derive(Debug, Serialize, Deserialize)]
3261pub struct NewUserParams {
3262 pub github_login: String,
3263 pub github_user_id: i32,
3264 pub invite_count: i32,
3265}
3266
3267#[derive(Debug)]
3268pub struct NewUserResult {
3269 pub user_id: UserId,
3270 pub metrics_id: String,
3271 pub inviting_user_id: Option<UserId>,
3272 pub signup_device_id: Option<String>,
3273}
3274
3275fn random_invite_code() -> String {
3276 nanoid::nanoid!(16)
3277}
3278
3279fn random_email_confirmation_code() -> String {
3280 nanoid::nanoid!(64)
3281}
3282
3283macro_rules! id_type {
3284 ($name:ident) => {
3285 #[derive(
3286 Clone,
3287 Copy,
3288 Debug,
3289 Default,
3290 PartialEq,
3291 Eq,
3292 PartialOrd,
3293 Ord,
3294 Hash,
3295 Serialize,
3296 Deserialize,
3297 )]
3298 #[serde(transparent)]
3299 pub struct $name(pub i32);
3300
3301 impl $name {
3302 #[allow(unused)]
3303 pub const MAX: Self = Self(i32::MAX);
3304
3305 #[allow(unused)]
3306 pub fn from_proto(value: u64) -> Self {
3307 Self(value as i32)
3308 }
3309
3310 #[allow(unused)]
3311 pub fn to_proto(self) -> u64 {
3312 self.0 as u64
3313 }
3314 }
3315
3316 impl std::fmt::Display for $name {
3317 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
3318 self.0.fmt(f)
3319 }
3320 }
3321
3322 impl From<$name> for sea_query::Value {
3323 fn from(value: $name) -> Self {
3324 sea_query::Value::Int(Some(value.0))
3325 }
3326 }
3327
3328 impl sea_orm::TryGetable for $name {
3329 fn try_get(
3330 res: &sea_orm::QueryResult,
3331 pre: &str,
3332 col: &str,
3333 ) -> Result<Self, sea_orm::TryGetError> {
3334 Ok(Self(i32::try_get(res, pre, col)?))
3335 }
3336 }
3337
3338 impl sea_query::ValueType for $name {
3339 fn try_from(v: Value) -> Result<Self, sea_query::ValueTypeErr> {
3340 match v {
3341 Value::TinyInt(Some(int)) => {
3342 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3343 }
3344 Value::SmallInt(Some(int)) => {
3345 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3346 }
3347 Value::Int(Some(int)) => {
3348 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3349 }
3350 Value::BigInt(Some(int)) => {
3351 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3352 }
3353 Value::TinyUnsigned(Some(int)) => {
3354 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3355 }
3356 Value::SmallUnsigned(Some(int)) => {
3357 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3358 }
3359 Value::Unsigned(Some(int)) => {
3360 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3361 }
3362 Value::BigUnsigned(Some(int)) => {
3363 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
3364 }
3365 _ => Err(sea_query::ValueTypeErr),
3366 }
3367 }
3368
3369 fn type_name() -> String {
3370 stringify!($name).into()
3371 }
3372
3373 fn array_type() -> sea_query::ArrayType {
3374 sea_query::ArrayType::Int
3375 }
3376
3377 fn column_type() -> sea_query::ColumnType {
3378 sea_query::ColumnType::Integer(None)
3379 }
3380 }
3381
3382 impl sea_orm::TryFromU64 for $name {
3383 fn try_from_u64(n: u64) -> Result<Self, DbErr> {
3384 Ok(Self(n.try_into().map_err(|_| {
3385 DbErr::ConvertFromU64(concat!(
3386 "error converting ",
3387 stringify!($name),
3388 " to u64"
3389 ))
3390 })?))
3391 }
3392 }
3393
3394 impl sea_query::Nullable for $name {
3395 fn null() -> Value {
3396 Value::Int(None)
3397 }
3398 }
3399 };
3400}
3401
3402id_type!(AccessTokenId);
3403id_type!(ContactId);
3404id_type!(FollowerId);
3405id_type!(RoomId);
3406id_type!(RoomParticipantId);
3407id_type!(ProjectId);
3408id_type!(ProjectCollaboratorId);
3409id_type!(ReplicaId);
3410id_type!(ServerId);
3411id_type!(SignupId);
3412id_type!(UserId);
3413
3414pub struct RejoinedRoom {
3415 pub room: proto::Room,
3416 pub rejoined_projects: Vec<RejoinedProject>,
3417 pub reshared_projects: Vec<ResharedProject>,
3418}
3419
3420pub struct ResharedProject {
3421 pub id: ProjectId,
3422 pub old_connection_id: ConnectionId,
3423 pub collaborators: Vec<ProjectCollaborator>,
3424 pub worktrees: Vec<proto::WorktreeMetadata>,
3425}
3426
3427pub struct RejoinedProject {
3428 pub id: ProjectId,
3429 pub old_connection_id: ConnectionId,
3430 pub collaborators: Vec<ProjectCollaborator>,
3431 pub worktrees: Vec<RejoinedWorktree>,
3432 pub language_servers: Vec<proto::LanguageServer>,
3433}
3434
3435#[derive(Debug)]
3436pub struct RejoinedWorktree {
3437 pub id: u64,
3438 pub abs_path: String,
3439 pub root_name: String,
3440 pub visible: bool,
3441 pub updated_entries: Vec<proto::Entry>,
3442 pub removed_entries: Vec<u64>,
3443 pub updated_repositories: Vec<proto::RepositoryEntry>,
3444 pub removed_repositories: Vec<u64>,
3445 pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
3446 pub settings_files: Vec<WorktreeSettingsFile>,
3447 pub scan_id: u64,
3448 pub completed_scan_id: u64,
3449}
3450
3451pub struct LeftRoom {
3452 pub room: proto::Room,
3453 pub left_projects: HashMap<ProjectId, LeftProject>,
3454 pub canceled_calls_to_user_ids: Vec<UserId>,
3455}
3456
3457pub struct RefreshedRoom {
3458 pub room: proto::Room,
3459 pub stale_participant_user_ids: Vec<UserId>,
3460 pub canceled_calls_to_user_ids: Vec<UserId>,
3461}
3462
3463pub struct Project {
3464 pub collaborators: Vec<ProjectCollaborator>,
3465 pub worktrees: BTreeMap<u64, Worktree>,
3466 pub language_servers: Vec<proto::LanguageServer>,
3467}
3468
3469pub struct ProjectCollaborator {
3470 pub connection_id: ConnectionId,
3471 pub user_id: UserId,
3472 pub replica_id: ReplicaId,
3473 pub is_host: bool,
3474}
3475
3476impl ProjectCollaborator {
3477 pub fn to_proto(&self) -> proto::Collaborator {
3478 proto::Collaborator {
3479 peer_id: Some(self.connection_id.into()),
3480 replica_id: self.replica_id.0 as u32,
3481 user_id: self.user_id.to_proto(),
3482 }
3483 }
3484}
3485
3486#[derive(Debug)]
3487pub struct LeftProject {
3488 pub id: ProjectId,
3489 pub host_user_id: UserId,
3490 pub host_connection_id: ConnectionId,
3491 pub connection_ids: Vec<ConnectionId>,
3492}
3493
3494pub struct Worktree {
3495 pub id: u64,
3496 pub abs_path: String,
3497 pub root_name: String,
3498 pub visible: bool,
3499 pub entries: Vec<proto::Entry>,
3500 pub repository_entries: BTreeMap<u64, proto::RepositoryEntry>,
3501 pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
3502 pub settings_files: Vec<WorktreeSettingsFile>,
3503 pub scan_id: u64,
3504 pub completed_scan_id: u64,
3505}
3506
3507#[derive(Debug)]
3508pub struct WorktreeSettingsFile {
3509 pub path: String,
3510 pub content: String,
3511}
3512
3513#[cfg(test)]
3514pub use test::*;
3515
3516#[cfg(test)]
3517mod test {
3518 use super::*;
3519 use gpui::executor::Background;
3520 use parking_lot::Mutex;
3521 use sea_orm::ConnectionTrait;
3522 use sqlx::migrate::MigrateDatabase;
3523 use std::sync::Arc;
3524
3525 pub struct TestDb {
3526 pub db: Option<Arc<Database>>,
3527 pub connection: Option<sqlx::AnyConnection>,
3528 }
3529
3530 impl TestDb {
3531 pub fn sqlite(background: Arc<Background>) -> Self {
3532 let url = format!("sqlite::memory:");
3533 let runtime = tokio::runtime::Builder::new_current_thread()
3534 .enable_io()
3535 .enable_time()
3536 .build()
3537 .unwrap();
3538
3539 let mut db = runtime.block_on(async {
3540 let mut options = ConnectOptions::new(url);
3541 options.max_connections(5);
3542 let db = Database::new(options, Executor::Deterministic(background))
3543 .await
3544 .unwrap();
3545 let sql = include_str!(concat!(
3546 env!("CARGO_MANIFEST_DIR"),
3547 "/migrations.sqlite/20221109000000_test_schema.sql"
3548 ));
3549 db.pool
3550 .execute(sea_orm::Statement::from_string(
3551 db.pool.get_database_backend(),
3552 sql.into(),
3553 ))
3554 .await
3555 .unwrap();
3556 db
3557 });
3558
3559 db.runtime = Some(runtime);
3560
3561 Self {
3562 db: Some(Arc::new(db)),
3563 connection: None,
3564 }
3565 }
3566
3567 pub fn postgres(background: Arc<Background>) -> Self {
3568 static LOCK: Mutex<()> = Mutex::new(());
3569
3570 let _guard = LOCK.lock();
3571 let mut rng = StdRng::from_entropy();
3572 let url = format!(
3573 "postgres://postgres@localhost/zed-test-{}",
3574 rng.gen::<u128>()
3575 );
3576 let runtime = tokio::runtime::Builder::new_current_thread()
3577 .enable_io()
3578 .enable_time()
3579 .build()
3580 .unwrap();
3581
3582 let mut db = runtime.block_on(async {
3583 sqlx::Postgres::create_database(&url)
3584 .await
3585 .expect("failed to create test db");
3586 let mut options = ConnectOptions::new(url);
3587 options
3588 .max_connections(5)
3589 .idle_timeout(Duration::from_secs(0));
3590 let db = Database::new(options, Executor::Deterministic(background))
3591 .await
3592 .unwrap();
3593 let migrations_path = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
3594 db.migrate(Path::new(migrations_path), false).await.unwrap();
3595 db
3596 });
3597
3598 db.runtime = Some(runtime);
3599
3600 Self {
3601 db: Some(Arc::new(db)),
3602 connection: None,
3603 }
3604 }
3605
3606 pub fn db(&self) -> &Arc<Database> {
3607 self.db.as_ref().unwrap()
3608 }
3609 }
3610
3611 impl Drop for TestDb {
3612 fn drop(&mut self) {
3613 let db = self.db.take().unwrap();
3614 if let sea_orm::DatabaseBackend::Postgres = db.pool.get_database_backend() {
3615 db.runtime.as_ref().unwrap().block_on(async {
3616 use util::ResultExt;
3617 let query = "
3618 SELECT pg_terminate_backend(pg_stat_activity.pid)
3619 FROM pg_stat_activity
3620 WHERE
3621 pg_stat_activity.datname = current_database() AND
3622 pid <> pg_backend_pid();
3623 ";
3624 db.pool
3625 .execute(sea_orm::Statement::from_string(
3626 db.pool.get_database_backend(),
3627 query.into(),
3628 ))
3629 .await
3630 .log_err();
3631 sqlx::Postgres::drop_database(db.options.get_url())
3632 .await
3633 .log_err();
3634 })
3635 }
3636 }
3637 }
3638}