1mod access_token;
2mod contact;
3mod language_server;
4mod project;
5mod project_collaborator;
6mod room;
7mod room_participant;
8mod signup;
9#[cfg(test)]
10mod tests;
11mod user;
12mod worktree;
13mod worktree_diagnostic_summary;
14mod worktree_entry;
15
16use crate::{Error, Result};
17use anyhow::anyhow;
18use collections::{BTreeMap, HashMap, HashSet};
19pub use contact::Contact;
20use dashmap::DashMap;
21use futures::StreamExt;
22use hyper::StatusCode;
23use rpc::{proto, ConnectionId};
24pub use sea_orm::ConnectOptions;
25use sea_orm::{
26 entity::prelude::*, ActiveValue, ConnectionTrait, DatabaseConnection, DatabaseTransaction,
27 DbErr, FromQueryResult, IntoActiveModel, IsolationLevel, JoinType, QueryOrder, QuerySelect,
28 Statement, TransactionTrait,
29};
30use sea_query::{Alias, Expr, OnConflict, Query};
31use serde::{Deserialize, Serialize};
32pub use signup::{Invite, NewSignup, WaitlistSummary};
33use sqlx::migrate::{Migrate, Migration, MigrationSource};
34use sqlx::Connection;
35use std::ops::{Deref, DerefMut};
36use std::path::Path;
37use std::time::Duration;
38use std::{future::Future, marker::PhantomData, rc::Rc, sync::Arc};
39use tokio::sync::{Mutex, OwnedMutexGuard};
40pub use user::Model as User;
41
42pub struct Database {
43 options: ConnectOptions,
44 pool: DatabaseConnection,
45 rooms: DashMap<RoomId, Arc<Mutex<()>>>,
46 #[cfg(test)]
47 background: Option<std::sync::Arc<gpui::executor::Background>>,
48 #[cfg(test)]
49 runtime: Option<tokio::runtime::Runtime>,
50 epoch: Uuid,
51}
52
53impl Database {
54 pub async fn new(options: ConnectOptions) -> Result<Self> {
55 Ok(Self {
56 options: options.clone(),
57 pool: sea_orm::Database::connect(options).await?,
58 rooms: DashMap::with_capacity(16384),
59 #[cfg(test)]
60 background: None,
61 #[cfg(test)]
62 runtime: None,
63 epoch: Uuid::new_v4(),
64 })
65 }
66
67 pub async fn migrate(
68 &self,
69 migrations_path: &Path,
70 ignore_checksum_mismatch: bool,
71 ) -> anyhow::Result<Vec<(Migration, Duration)>> {
72 let migrations = MigrationSource::resolve(migrations_path)
73 .await
74 .map_err(|err| anyhow!("failed to load migrations: {err:?}"))?;
75
76 let mut connection = sqlx::AnyConnection::connect(self.options.get_url()).await?;
77
78 connection.ensure_migrations_table().await?;
79 let applied_migrations: HashMap<_, _> = connection
80 .list_applied_migrations()
81 .await?
82 .into_iter()
83 .map(|m| (m.version, m))
84 .collect();
85
86 let mut new_migrations = Vec::new();
87 for migration in migrations {
88 match applied_migrations.get(&migration.version) {
89 Some(applied_migration) => {
90 if migration.checksum != applied_migration.checksum && !ignore_checksum_mismatch
91 {
92 Err(anyhow!(
93 "checksum mismatch for applied migration {}",
94 migration.description
95 ))?;
96 }
97 }
98 None => {
99 let elapsed = connection.apply(&migration).await?;
100 new_migrations.push((migration, elapsed));
101 }
102 }
103 }
104
105 Ok(new_migrations)
106 }
107
108 pub async fn clear_stale_data(&self) -> Result<()> {
109 self.transaction(|tx| async move {
110 project_collaborator::Entity::delete_many()
111 .filter(project_collaborator::Column::ConnectionEpoch.ne(self.epoch))
112 .exec(&*tx)
113 .await?;
114 room_participant::Entity::delete_many()
115 .filter(
116 room_participant::Column::AnsweringConnectionEpoch
117 .ne(self.epoch)
118 .or(room_participant::Column::CallingConnectionEpoch.ne(self.epoch)),
119 )
120 .exec(&*tx)
121 .await?;
122 project::Entity::delete_many()
123 .filter(project::Column::HostConnectionEpoch.ne(self.epoch))
124 .exec(&*tx)
125 .await?;
126 room::Entity::delete_many()
127 .filter(
128 room::Column::Id.not_in_subquery(
129 Query::select()
130 .column(room_participant::Column::RoomId)
131 .from(room_participant::Entity)
132 .distinct()
133 .to_owned(),
134 ),
135 )
136 .exec(&*tx)
137 .await?;
138 Ok(())
139 })
140 .await
141 }
142
143 // users
144
145 pub async fn create_user(
146 &self,
147 email_address: &str,
148 admin: bool,
149 params: NewUserParams,
150 ) -> Result<NewUserResult> {
151 self.transaction(|tx| async {
152 let tx = tx;
153 let user = user::Entity::insert(user::ActiveModel {
154 email_address: ActiveValue::set(Some(email_address.into())),
155 github_login: ActiveValue::set(params.github_login.clone()),
156 github_user_id: ActiveValue::set(Some(params.github_user_id)),
157 admin: ActiveValue::set(admin),
158 metrics_id: ActiveValue::set(Uuid::new_v4()),
159 ..Default::default()
160 })
161 .on_conflict(
162 OnConflict::column(user::Column::GithubLogin)
163 .update_column(user::Column::GithubLogin)
164 .to_owned(),
165 )
166 .exec_with_returning(&*tx)
167 .await?;
168
169 Ok(NewUserResult {
170 user_id: user.id,
171 metrics_id: user.metrics_id.to_string(),
172 signup_device_id: None,
173 inviting_user_id: None,
174 })
175 })
176 .await
177 }
178
179 pub async fn get_user_by_id(&self, id: UserId) -> Result<Option<user::Model>> {
180 self.transaction(|tx| async move { Ok(user::Entity::find_by_id(id).one(&*tx).await?) })
181 .await
182 }
183
184 pub async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<user::Model>> {
185 self.transaction(|tx| async {
186 let tx = tx;
187 Ok(user::Entity::find()
188 .filter(user::Column::Id.is_in(ids.iter().copied()))
189 .all(&*tx)
190 .await?)
191 })
192 .await
193 }
194
195 pub async fn get_user_by_github_account(
196 &self,
197 github_login: &str,
198 github_user_id: Option<i32>,
199 ) -> Result<Option<User>> {
200 self.transaction(|tx| async move {
201 let tx = &*tx;
202 if let Some(github_user_id) = github_user_id {
203 if let Some(user_by_github_user_id) = user::Entity::find()
204 .filter(user::Column::GithubUserId.eq(github_user_id))
205 .one(tx)
206 .await?
207 {
208 let mut user_by_github_user_id = user_by_github_user_id.into_active_model();
209 user_by_github_user_id.github_login = ActiveValue::set(github_login.into());
210 Ok(Some(user_by_github_user_id.update(tx).await?))
211 } else if let Some(user_by_github_login) = user::Entity::find()
212 .filter(user::Column::GithubLogin.eq(github_login))
213 .one(tx)
214 .await?
215 {
216 let mut user_by_github_login = user_by_github_login.into_active_model();
217 user_by_github_login.github_user_id = ActiveValue::set(Some(github_user_id));
218 Ok(Some(user_by_github_login.update(tx).await?))
219 } else {
220 Ok(None)
221 }
222 } else {
223 Ok(user::Entity::find()
224 .filter(user::Column::GithubLogin.eq(github_login))
225 .one(tx)
226 .await?)
227 }
228 })
229 .await
230 }
231
232 pub async fn get_all_users(&self, page: u32, limit: u32) -> Result<Vec<User>> {
233 self.transaction(|tx| async move {
234 Ok(user::Entity::find()
235 .order_by_asc(user::Column::GithubLogin)
236 .limit(limit as u64)
237 .offset(page as u64 * limit as u64)
238 .all(&*tx)
239 .await?)
240 })
241 .await
242 }
243
244 pub async fn get_users_with_no_invites(
245 &self,
246 invited_by_another_user: bool,
247 ) -> Result<Vec<User>> {
248 self.transaction(|tx| async move {
249 Ok(user::Entity::find()
250 .filter(
251 user::Column::InviteCount
252 .eq(0)
253 .and(if invited_by_another_user {
254 user::Column::InviterId.is_not_null()
255 } else {
256 user::Column::InviterId.is_null()
257 }),
258 )
259 .all(&*tx)
260 .await?)
261 })
262 .await
263 }
264
265 pub async fn get_user_metrics_id(&self, id: UserId) -> Result<String> {
266 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
267 enum QueryAs {
268 MetricsId,
269 }
270
271 self.transaction(|tx| async move {
272 let metrics_id: Uuid = user::Entity::find_by_id(id)
273 .select_only()
274 .column(user::Column::MetricsId)
275 .into_values::<_, QueryAs>()
276 .one(&*tx)
277 .await?
278 .ok_or_else(|| anyhow!("could not find user"))?;
279 Ok(metrics_id.to_string())
280 })
281 .await
282 }
283
284 pub async fn set_user_is_admin(&self, id: UserId, is_admin: bool) -> Result<()> {
285 self.transaction(|tx| async move {
286 user::Entity::update_many()
287 .filter(user::Column::Id.eq(id))
288 .set(user::ActiveModel {
289 admin: ActiveValue::set(is_admin),
290 ..Default::default()
291 })
292 .exec(&*tx)
293 .await?;
294 Ok(())
295 })
296 .await
297 }
298
299 pub async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()> {
300 self.transaction(|tx| async move {
301 user::Entity::update_many()
302 .filter(user::Column::Id.eq(id))
303 .set(user::ActiveModel {
304 connected_once: ActiveValue::set(connected_once),
305 ..Default::default()
306 })
307 .exec(&*tx)
308 .await?;
309 Ok(())
310 })
311 .await
312 }
313
314 pub async fn destroy_user(&self, id: UserId) -> Result<()> {
315 self.transaction(|tx| async move {
316 access_token::Entity::delete_many()
317 .filter(access_token::Column::UserId.eq(id))
318 .exec(&*tx)
319 .await?;
320 user::Entity::delete_by_id(id).exec(&*tx).await?;
321 Ok(())
322 })
323 .await
324 }
325
326 // contacts
327
328 pub async fn get_contacts(&self, user_id: UserId) -> Result<Vec<Contact>> {
329 #[derive(Debug, FromQueryResult)]
330 struct ContactWithUserBusyStatuses {
331 user_id_a: UserId,
332 user_id_b: UserId,
333 a_to_b: bool,
334 accepted: bool,
335 should_notify: bool,
336 user_a_busy: bool,
337 user_b_busy: bool,
338 }
339
340 self.transaction(|tx| async move {
341 let user_a_participant = Alias::new("user_a_participant");
342 let user_b_participant = Alias::new("user_b_participant");
343 let mut db_contacts = contact::Entity::find()
344 .column_as(
345 Expr::tbl(user_a_participant.clone(), room_participant::Column::Id)
346 .is_not_null(),
347 "user_a_busy",
348 )
349 .column_as(
350 Expr::tbl(user_b_participant.clone(), room_participant::Column::Id)
351 .is_not_null(),
352 "user_b_busy",
353 )
354 .filter(
355 contact::Column::UserIdA
356 .eq(user_id)
357 .or(contact::Column::UserIdB.eq(user_id)),
358 )
359 .join_as(
360 JoinType::LeftJoin,
361 contact::Relation::UserARoomParticipant.def(),
362 user_a_participant,
363 )
364 .join_as(
365 JoinType::LeftJoin,
366 contact::Relation::UserBRoomParticipant.def(),
367 user_b_participant,
368 )
369 .into_model::<ContactWithUserBusyStatuses>()
370 .stream(&*tx)
371 .await?;
372
373 let mut contacts = Vec::new();
374 while let Some(db_contact) = db_contacts.next().await {
375 let db_contact = db_contact?;
376 if db_contact.user_id_a == user_id {
377 if db_contact.accepted {
378 contacts.push(Contact::Accepted {
379 user_id: db_contact.user_id_b,
380 should_notify: db_contact.should_notify && db_contact.a_to_b,
381 busy: db_contact.user_b_busy,
382 });
383 } else if db_contact.a_to_b {
384 contacts.push(Contact::Outgoing {
385 user_id: db_contact.user_id_b,
386 })
387 } else {
388 contacts.push(Contact::Incoming {
389 user_id: db_contact.user_id_b,
390 should_notify: db_contact.should_notify,
391 });
392 }
393 } else if db_contact.accepted {
394 contacts.push(Contact::Accepted {
395 user_id: db_contact.user_id_a,
396 should_notify: db_contact.should_notify && !db_contact.a_to_b,
397 busy: db_contact.user_a_busy,
398 });
399 } else if db_contact.a_to_b {
400 contacts.push(Contact::Incoming {
401 user_id: db_contact.user_id_a,
402 should_notify: db_contact.should_notify,
403 });
404 } else {
405 contacts.push(Contact::Outgoing {
406 user_id: db_contact.user_id_a,
407 });
408 }
409 }
410
411 contacts.sort_unstable_by_key(|contact| contact.user_id());
412
413 Ok(contacts)
414 })
415 .await
416 }
417
418 pub async fn is_user_busy(&self, user_id: UserId) -> Result<bool> {
419 self.transaction(|tx| async move {
420 let participant = room_participant::Entity::find()
421 .filter(room_participant::Column::UserId.eq(user_id))
422 .one(&*tx)
423 .await?;
424 Ok(participant.is_some())
425 })
426 .await
427 }
428
429 pub async fn has_contact(&self, user_id_1: UserId, user_id_2: UserId) -> Result<bool> {
430 self.transaction(|tx| async move {
431 let (id_a, id_b) = if user_id_1 < user_id_2 {
432 (user_id_1, user_id_2)
433 } else {
434 (user_id_2, user_id_1)
435 };
436
437 Ok(contact::Entity::find()
438 .filter(
439 contact::Column::UserIdA
440 .eq(id_a)
441 .and(contact::Column::UserIdB.eq(id_b))
442 .and(contact::Column::Accepted.eq(true)),
443 )
444 .one(&*tx)
445 .await?
446 .is_some())
447 })
448 .await
449 }
450
451 pub async fn send_contact_request(&self, sender_id: UserId, receiver_id: UserId) -> Result<()> {
452 self.transaction(|tx| async move {
453 let (id_a, id_b, a_to_b) = if sender_id < receiver_id {
454 (sender_id, receiver_id, true)
455 } else {
456 (receiver_id, sender_id, false)
457 };
458
459 let rows_affected = contact::Entity::insert(contact::ActiveModel {
460 user_id_a: ActiveValue::set(id_a),
461 user_id_b: ActiveValue::set(id_b),
462 a_to_b: ActiveValue::set(a_to_b),
463 accepted: ActiveValue::set(false),
464 should_notify: ActiveValue::set(true),
465 ..Default::default()
466 })
467 .on_conflict(
468 OnConflict::columns([contact::Column::UserIdA, contact::Column::UserIdB])
469 .values([
470 (contact::Column::Accepted, true.into()),
471 (contact::Column::ShouldNotify, false.into()),
472 ])
473 .action_and_where(
474 contact::Column::Accepted.eq(false).and(
475 contact::Column::AToB
476 .eq(a_to_b)
477 .and(contact::Column::UserIdA.eq(id_b))
478 .or(contact::Column::AToB
479 .ne(a_to_b)
480 .and(contact::Column::UserIdA.eq(id_a))),
481 ),
482 )
483 .to_owned(),
484 )
485 .exec_without_returning(&*tx)
486 .await?;
487
488 if rows_affected == 1 {
489 Ok(())
490 } else {
491 Err(anyhow!("contact already requested"))?
492 }
493 })
494 .await
495 }
496
497 pub async fn remove_contact(&self, requester_id: UserId, responder_id: UserId) -> Result<()> {
498 self.transaction(|tx| async move {
499 let (id_a, id_b) = if responder_id < requester_id {
500 (responder_id, requester_id)
501 } else {
502 (requester_id, responder_id)
503 };
504
505 let result = contact::Entity::delete_many()
506 .filter(
507 contact::Column::UserIdA
508 .eq(id_a)
509 .and(contact::Column::UserIdB.eq(id_b)),
510 )
511 .exec(&*tx)
512 .await?;
513
514 if result.rows_affected == 1 {
515 Ok(())
516 } else {
517 Err(anyhow!("no such contact"))?
518 }
519 })
520 .await
521 }
522
523 pub async fn dismiss_contact_notification(
524 &self,
525 user_id: UserId,
526 contact_user_id: UserId,
527 ) -> Result<()> {
528 self.transaction(|tx| async move {
529 let (id_a, id_b, a_to_b) = if user_id < contact_user_id {
530 (user_id, contact_user_id, true)
531 } else {
532 (contact_user_id, user_id, false)
533 };
534
535 let result = contact::Entity::update_many()
536 .set(contact::ActiveModel {
537 should_notify: ActiveValue::set(false),
538 ..Default::default()
539 })
540 .filter(
541 contact::Column::UserIdA
542 .eq(id_a)
543 .and(contact::Column::UserIdB.eq(id_b))
544 .and(
545 contact::Column::AToB
546 .eq(a_to_b)
547 .and(contact::Column::Accepted.eq(true))
548 .or(contact::Column::AToB
549 .ne(a_to_b)
550 .and(contact::Column::Accepted.eq(false))),
551 ),
552 )
553 .exec(&*tx)
554 .await?;
555 if result.rows_affected == 0 {
556 Err(anyhow!("no such contact request"))?
557 } else {
558 Ok(())
559 }
560 })
561 .await
562 }
563
564 pub async fn respond_to_contact_request(
565 &self,
566 responder_id: UserId,
567 requester_id: UserId,
568 accept: bool,
569 ) -> Result<()> {
570 self.transaction(|tx| async move {
571 let (id_a, id_b, a_to_b) = if responder_id < requester_id {
572 (responder_id, requester_id, false)
573 } else {
574 (requester_id, responder_id, true)
575 };
576 let rows_affected = if accept {
577 let result = contact::Entity::update_many()
578 .set(contact::ActiveModel {
579 accepted: ActiveValue::set(true),
580 should_notify: ActiveValue::set(true),
581 ..Default::default()
582 })
583 .filter(
584 contact::Column::UserIdA
585 .eq(id_a)
586 .and(contact::Column::UserIdB.eq(id_b))
587 .and(contact::Column::AToB.eq(a_to_b)),
588 )
589 .exec(&*tx)
590 .await?;
591 result.rows_affected
592 } else {
593 let result = contact::Entity::delete_many()
594 .filter(
595 contact::Column::UserIdA
596 .eq(id_a)
597 .and(contact::Column::UserIdB.eq(id_b))
598 .and(contact::Column::AToB.eq(a_to_b))
599 .and(contact::Column::Accepted.eq(false)),
600 )
601 .exec(&*tx)
602 .await?;
603
604 result.rows_affected
605 };
606
607 if rows_affected == 1 {
608 Ok(())
609 } else {
610 Err(anyhow!("no such contact request"))?
611 }
612 })
613 .await
614 }
615
616 pub fn fuzzy_like_string(string: &str) -> String {
617 let mut result = String::with_capacity(string.len() * 2 + 1);
618 for c in string.chars() {
619 if c.is_alphanumeric() {
620 result.push('%');
621 result.push(c);
622 }
623 }
624 result.push('%');
625 result
626 }
627
628 pub async fn fuzzy_search_users(&self, name_query: &str, limit: u32) -> Result<Vec<User>> {
629 self.transaction(|tx| async {
630 let tx = tx;
631 let like_string = Self::fuzzy_like_string(name_query);
632 let query = "
633 SELECT users.*
634 FROM users
635 WHERE github_login ILIKE $1
636 ORDER BY github_login <-> $2
637 LIMIT $3
638 ";
639
640 Ok(user::Entity::find()
641 .from_raw_sql(Statement::from_sql_and_values(
642 self.pool.get_database_backend(),
643 query.into(),
644 vec![like_string.into(), name_query.into(), limit.into()],
645 ))
646 .all(&*tx)
647 .await?)
648 })
649 .await
650 }
651
652 // signups
653
654 pub async fn create_signup(&self, signup: &NewSignup) -> Result<()> {
655 self.transaction(|tx| async move {
656 signup::Entity::insert(signup::ActiveModel {
657 email_address: ActiveValue::set(signup.email_address.clone()),
658 email_confirmation_code: ActiveValue::set(random_email_confirmation_code()),
659 email_confirmation_sent: ActiveValue::set(false),
660 platform_mac: ActiveValue::set(signup.platform_mac),
661 platform_windows: ActiveValue::set(signup.platform_windows),
662 platform_linux: ActiveValue::set(signup.platform_linux),
663 platform_unknown: ActiveValue::set(false),
664 editor_features: ActiveValue::set(Some(signup.editor_features.clone())),
665 programming_languages: ActiveValue::set(Some(signup.programming_languages.clone())),
666 device_id: ActiveValue::set(signup.device_id.clone()),
667 added_to_mailing_list: ActiveValue::set(signup.added_to_mailing_list),
668 ..Default::default()
669 })
670 .on_conflict(
671 OnConflict::column(signup::Column::EmailAddress)
672 .update_column(signup::Column::EmailAddress)
673 .to_owned(),
674 )
675 .exec(&*tx)
676 .await?;
677 Ok(())
678 })
679 .await
680 }
681
682 pub async fn get_waitlist_summary(&self) -> Result<WaitlistSummary> {
683 self.transaction(|tx| async move {
684 let query = "
685 SELECT
686 COUNT(*) as count,
687 COALESCE(SUM(CASE WHEN platform_linux THEN 1 ELSE 0 END), 0) as linux_count,
688 COALESCE(SUM(CASE WHEN platform_mac THEN 1 ELSE 0 END), 0) as mac_count,
689 COALESCE(SUM(CASE WHEN platform_windows THEN 1 ELSE 0 END), 0) as windows_count,
690 COALESCE(SUM(CASE WHEN platform_unknown THEN 1 ELSE 0 END), 0) as unknown_count
691 FROM (
692 SELECT *
693 FROM signups
694 WHERE
695 NOT email_confirmation_sent
696 ) AS unsent
697 ";
698 Ok(
699 WaitlistSummary::find_by_statement(Statement::from_sql_and_values(
700 self.pool.get_database_backend(),
701 query.into(),
702 vec![],
703 ))
704 .one(&*tx)
705 .await?
706 .ok_or_else(|| anyhow!("invalid result"))?,
707 )
708 })
709 .await
710 }
711
712 pub async fn record_sent_invites(&self, invites: &[Invite]) -> Result<()> {
713 let emails = invites
714 .iter()
715 .map(|s| s.email_address.as_str())
716 .collect::<Vec<_>>();
717 self.transaction(|tx| async {
718 let tx = tx;
719 signup::Entity::update_many()
720 .filter(signup::Column::EmailAddress.is_in(emails.iter().copied()))
721 .set(signup::ActiveModel {
722 email_confirmation_sent: ActiveValue::set(true),
723 ..Default::default()
724 })
725 .exec(&*tx)
726 .await?;
727 Ok(())
728 })
729 .await
730 }
731
732 pub async fn get_unsent_invites(&self, count: usize) -> Result<Vec<Invite>> {
733 self.transaction(|tx| async move {
734 Ok(signup::Entity::find()
735 .select_only()
736 .column(signup::Column::EmailAddress)
737 .column(signup::Column::EmailConfirmationCode)
738 .filter(
739 signup::Column::EmailConfirmationSent.eq(false).and(
740 signup::Column::PlatformMac
741 .eq(true)
742 .or(signup::Column::PlatformUnknown.eq(true)),
743 ),
744 )
745 .order_by_asc(signup::Column::CreatedAt)
746 .limit(count as u64)
747 .into_model()
748 .all(&*tx)
749 .await?)
750 })
751 .await
752 }
753
754 // invite codes
755
756 pub async fn create_invite_from_code(
757 &self,
758 code: &str,
759 email_address: &str,
760 device_id: Option<&str>,
761 ) -> Result<Invite> {
762 self.transaction(|tx| async move {
763 let existing_user = user::Entity::find()
764 .filter(user::Column::EmailAddress.eq(email_address))
765 .one(&*tx)
766 .await?;
767
768 if existing_user.is_some() {
769 Err(anyhow!("email address is already in use"))?;
770 }
771
772 let inviting_user_with_invites = match user::Entity::find()
773 .filter(
774 user::Column::InviteCode
775 .eq(code)
776 .and(user::Column::InviteCount.gt(0)),
777 )
778 .one(&*tx)
779 .await?
780 {
781 Some(inviting_user) => inviting_user,
782 None => {
783 return Err(Error::Http(
784 StatusCode::UNAUTHORIZED,
785 "unable to find an invite code with invites remaining".to_string(),
786 ))?
787 }
788 };
789 user::Entity::update_many()
790 .filter(
791 user::Column::Id
792 .eq(inviting_user_with_invites.id)
793 .and(user::Column::InviteCount.gt(0)),
794 )
795 .col_expr(
796 user::Column::InviteCount,
797 Expr::col(user::Column::InviteCount).sub(1),
798 )
799 .exec(&*tx)
800 .await?;
801
802 let signup = signup::Entity::insert(signup::ActiveModel {
803 email_address: ActiveValue::set(email_address.into()),
804 email_confirmation_code: ActiveValue::set(random_email_confirmation_code()),
805 email_confirmation_sent: ActiveValue::set(false),
806 inviting_user_id: ActiveValue::set(Some(inviting_user_with_invites.id)),
807 platform_linux: ActiveValue::set(false),
808 platform_mac: ActiveValue::set(false),
809 platform_windows: ActiveValue::set(false),
810 platform_unknown: ActiveValue::set(true),
811 device_id: ActiveValue::set(device_id.map(|device_id| device_id.into())),
812 ..Default::default()
813 })
814 .on_conflict(
815 OnConflict::column(signup::Column::EmailAddress)
816 .update_column(signup::Column::InvitingUserId)
817 .to_owned(),
818 )
819 .exec_with_returning(&*tx)
820 .await?;
821
822 Ok(Invite {
823 email_address: signup.email_address,
824 email_confirmation_code: signup.email_confirmation_code,
825 })
826 })
827 .await
828 }
829
830 pub async fn create_user_from_invite(
831 &self,
832 invite: &Invite,
833 user: NewUserParams,
834 ) -> Result<Option<NewUserResult>> {
835 self.transaction(|tx| async {
836 let tx = tx;
837 let signup = signup::Entity::find()
838 .filter(
839 signup::Column::EmailAddress
840 .eq(invite.email_address.as_str())
841 .and(
842 signup::Column::EmailConfirmationCode
843 .eq(invite.email_confirmation_code.as_str()),
844 ),
845 )
846 .one(&*tx)
847 .await?
848 .ok_or_else(|| Error::Http(StatusCode::NOT_FOUND, "no such invite".to_string()))?;
849
850 if signup.user_id.is_some() {
851 return Ok(None);
852 }
853
854 let user = user::Entity::insert(user::ActiveModel {
855 email_address: ActiveValue::set(Some(invite.email_address.clone())),
856 github_login: ActiveValue::set(user.github_login.clone()),
857 github_user_id: ActiveValue::set(Some(user.github_user_id)),
858 admin: ActiveValue::set(false),
859 invite_count: ActiveValue::set(user.invite_count),
860 invite_code: ActiveValue::set(Some(random_invite_code())),
861 metrics_id: ActiveValue::set(Uuid::new_v4()),
862 ..Default::default()
863 })
864 .on_conflict(
865 OnConflict::column(user::Column::GithubLogin)
866 .update_columns([
867 user::Column::EmailAddress,
868 user::Column::GithubUserId,
869 user::Column::Admin,
870 ])
871 .to_owned(),
872 )
873 .exec_with_returning(&*tx)
874 .await?;
875
876 let mut signup = signup.into_active_model();
877 signup.user_id = ActiveValue::set(Some(user.id));
878 let signup = signup.update(&*tx).await?;
879
880 if let Some(inviting_user_id) = signup.inviting_user_id {
881 let (user_id_a, user_id_b, a_to_b) = if inviting_user_id < user.id {
882 (inviting_user_id, user.id, true)
883 } else {
884 (user.id, inviting_user_id, false)
885 };
886
887 contact::Entity::insert(contact::ActiveModel {
888 user_id_a: ActiveValue::set(user_id_a),
889 user_id_b: ActiveValue::set(user_id_b),
890 a_to_b: ActiveValue::set(a_to_b),
891 should_notify: ActiveValue::set(true),
892 accepted: ActiveValue::set(true),
893 ..Default::default()
894 })
895 .on_conflict(OnConflict::new().do_nothing().to_owned())
896 .exec_without_returning(&*tx)
897 .await?;
898 }
899
900 Ok(Some(NewUserResult {
901 user_id: user.id,
902 metrics_id: user.metrics_id.to_string(),
903 inviting_user_id: signup.inviting_user_id,
904 signup_device_id: signup.device_id,
905 }))
906 })
907 .await
908 }
909
910 pub async fn set_invite_count_for_user(&self, id: UserId, count: i32) -> Result<()> {
911 self.transaction(|tx| async move {
912 if count > 0 {
913 user::Entity::update_many()
914 .filter(
915 user::Column::Id
916 .eq(id)
917 .and(user::Column::InviteCode.is_null()),
918 )
919 .set(user::ActiveModel {
920 invite_code: ActiveValue::set(Some(random_invite_code())),
921 ..Default::default()
922 })
923 .exec(&*tx)
924 .await?;
925 }
926
927 user::Entity::update_many()
928 .filter(user::Column::Id.eq(id))
929 .set(user::ActiveModel {
930 invite_count: ActiveValue::set(count),
931 ..Default::default()
932 })
933 .exec(&*tx)
934 .await?;
935 Ok(())
936 })
937 .await
938 }
939
940 pub async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, i32)>> {
941 self.transaction(|tx| async move {
942 match user::Entity::find_by_id(id).one(&*tx).await? {
943 Some(user) if user.invite_code.is_some() => {
944 Ok(Some((user.invite_code.unwrap(), user.invite_count)))
945 }
946 _ => Ok(None),
947 }
948 })
949 .await
950 }
951
952 pub async fn get_user_for_invite_code(&self, code: &str) -> Result<User> {
953 self.transaction(|tx| async move {
954 user::Entity::find()
955 .filter(user::Column::InviteCode.eq(code))
956 .one(&*tx)
957 .await?
958 .ok_or_else(|| {
959 Error::Http(
960 StatusCode::NOT_FOUND,
961 "that invite code does not exist".to_string(),
962 )
963 })
964 })
965 .await
966 }
967
968 // rooms
969
970 pub async fn incoming_call_for_user(
971 &self,
972 user_id: UserId,
973 ) -> Result<Option<proto::IncomingCall>> {
974 self.transaction(|tx| async move {
975 let pending_participant = room_participant::Entity::find()
976 .filter(
977 room_participant::Column::UserId
978 .eq(user_id)
979 .and(room_participant::Column::AnsweringConnectionId.is_null()),
980 )
981 .one(&*tx)
982 .await?;
983
984 if let Some(pending_participant) = pending_participant {
985 let room = self.get_room(pending_participant.room_id, &tx).await?;
986 Ok(Self::build_incoming_call(&room, user_id))
987 } else {
988 Ok(None)
989 }
990 })
991 .await
992 }
993
994 pub async fn create_room(
995 &self,
996 user_id: UserId,
997 connection_id: ConnectionId,
998 live_kit_room: &str,
999 ) -> Result<RoomGuard<proto::Room>> {
1000 self.room_transaction(|tx| async move {
1001 let room = room::ActiveModel {
1002 live_kit_room: ActiveValue::set(live_kit_room.into()),
1003 ..Default::default()
1004 }
1005 .insert(&*tx)
1006 .await?;
1007 let room_id = room.id;
1008
1009 room_participant::ActiveModel {
1010 room_id: ActiveValue::set(room_id),
1011 user_id: ActiveValue::set(user_id),
1012 answering_connection_id: ActiveValue::set(Some(connection_id.0 as i32)),
1013 answering_connection_epoch: ActiveValue::set(Some(self.epoch)),
1014 calling_user_id: ActiveValue::set(user_id),
1015 calling_connection_id: ActiveValue::set(connection_id.0 as i32),
1016 calling_connection_epoch: ActiveValue::set(self.epoch),
1017 ..Default::default()
1018 }
1019 .insert(&*tx)
1020 .await?;
1021
1022 let room = self.get_room(room_id, &tx).await?;
1023 Ok((room_id, room))
1024 })
1025 .await
1026 }
1027
1028 pub async fn call(
1029 &self,
1030 room_id: RoomId,
1031 calling_user_id: UserId,
1032 calling_connection_id: ConnectionId,
1033 called_user_id: UserId,
1034 initial_project_id: Option<ProjectId>,
1035 ) -> Result<RoomGuard<(proto::Room, proto::IncomingCall)>> {
1036 self.room_transaction(|tx| async move {
1037 room_participant::ActiveModel {
1038 room_id: ActiveValue::set(room_id),
1039 user_id: ActiveValue::set(called_user_id),
1040 calling_user_id: ActiveValue::set(calling_user_id),
1041 calling_connection_id: ActiveValue::set(calling_connection_id.0 as i32),
1042 calling_connection_epoch: ActiveValue::set(self.epoch),
1043 initial_project_id: ActiveValue::set(initial_project_id),
1044 ..Default::default()
1045 }
1046 .insert(&*tx)
1047 .await?;
1048
1049 let room = self.get_room(room_id, &tx).await?;
1050 let incoming_call = Self::build_incoming_call(&room, called_user_id)
1051 .ok_or_else(|| anyhow!("failed to build incoming call"))?;
1052 Ok((room_id, (room, incoming_call)))
1053 })
1054 .await
1055 }
1056
1057 pub async fn call_failed(
1058 &self,
1059 room_id: RoomId,
1060 called_user_id: UserId,
1061 ) -> Result<RoomGuard<proto::Room>> {
1062 self.room_transaction(|tx| async move {
1063 room_participant::Entity::delete_many()
1064 .filter(
1065 room_participant::Column::RoomId
1066 .eq(room_id)
1067 .and(room_participant::Column::UserId.eq(called_user_id)),
1068 )
1069 .exec(&*tx)
1070 .await?;
1071 let room = self.get_room(room_id, &tx).await?;
1072 Ok((room_id, room))
1073 })
1074 .await
1075 }
1076
1077 pub async fn decline_call(
1078 &self,
1079 expected_room_id: Option<RoomId>,
1080 user_id: UserId,
1081 ) -> Result<RoomGuard<proto::Room>> {
1082 self.room_transaction(|tx| async move {
1083 let participant = room_participant::Entity::find()
1084 .filter(
1085 room_participant::Column::UserId
1086 .eq(user_id)
1087 .and(room_participant::Column::AnsweringConnectionId.is_null()),
1088 )
1089 .one(&*tx)
1090 .await?
1091 .ok_or_else(|| anyhow!("could not decline call"))?;
1092 let room_id = participant.room_id;
1093
1094 if expected_room_id.map_or(false, |expected_room_id| expected_room_id != room_id) {
1095 return Err(anyhow!("declining call on unexpected room"))?;
1096 }
1097
1098 room_participant::Entity::delete(participant.into_active_model())
1099 .exec(&*tx)
1100 .await?;
1101
1102 let room = self.get_room(room_id, &tx).await?;
1103 Ok((room_id, room))
1104 })
1105 .await
1106 }
1107
1108 pub async fn cancel_call(
1109 &self,
1110 expected_room_id: Option<RoomId>,
1111 calling_connection_id: ConnectionId,
1112 called_user_id: UserId,
1113 ) -> Result<RoomGuard<proto::Room>> {
1114 self.room_transaction(|tx| async move {
1115 let participant = room_participant::Entity::find()
1116 .filter(
1117 room_participant::Column::UserId
1118 .eq(called_user_id)
1119 .and(
1120 room_participant::Column::CallingConnectionId
1121 .eq(calling_connection_id.0 as i32),
1122 )
1123 .and(room_participant::Column::AnsweringConnectionId.is_null()),
1124 )
1125 .one(&*tx)
1126 .await?
1127 .ok_or_else(|| anyhow!("could not cancel call"))?;
1128 let room_id = participant.room_id;
1129 if expected_room_id.map_or(false, |expected_room_id| expected_room_id != room_id) {
1130 return Err(anyhow!("canceling call on unexpected room"))?;
1131 }
1132
1133 room_participant::Entity::delete(participant.into_active_model())
1134 .exec(&*tx)
1135 .await?;
1136
1137 let room = self.get_room(room_id, &tx).await?;
1138 Ok((room_id, room))
1139 })
1140 .await
1141 }
1142
1143 pub async fn join_room(
1144 &self,
1145 room_id: RoomId,
1146 user_id: UserId,
1147 connection_id: ConnectionId,
1148 ) -> Result<RoomGuard<proto::Room>> {
1149 self.room_transaction(|tx| async move {
1150 let result = room_participant::Entity::update_many()
1151 .filter(
1152 room_participant::Column::RoomId
1153 .eq(room_id)
1154 .and(room_participant::Column::UserId.eq(user_id))
1155 .and(room_participant::Column::AnsweringConnectionId.is_null()),
1156 )
1157 .set(room_participant::ActiveModel {
1158 answering_connection_id: ActiveValue::set(Some(connection_id.0 as i32)),
1159 answering_connection_epoch: ActiveValue::set(Some(self.epoch)),
1160 ..Default::default()
1161 })
1162 .exec(&*tx)
1163 .await?;
1164 if result.rows_affected == 0 {
1165 Err(anyhow!("room does not exist or was already joined"))?
1166 } else {
1167 let room = self.get_room(room_id, &tx).await?;
1168 Ok((room_id, room))
1169 }
1170 })
1171 .await
1172 }
1173
1174 pub async fn leave_room(&self, connection_id: ConnectionId) -> Result<RoomGuard<LeftRoom>> {
1175 self.room_transaction(|tx| async move {
1176 let leaving_participant = room_participant::Entity::find()
1177 .filter(room_participant::Column::AnsweringConnectionId.eq(connection_id.0))
1178 .one(&*tx)
1179 .await?;
1180
1181 if let Some(leaving_participant) = leaving_participant {
1182 // Leave room.
1183 let room_id = leaving_participant.room_id;
1184 room_participant::Entity::delete_by_id(leaving_participant.id)
1185 .exec(&*tx)
1186 .await?;
1187
1188 // Cancel pending calls initiated by the leaving user.
1189 let called_participants = room_participant::Entity::find()
1190 .filter(
1191 room_participant::Column::CallingConnectionId
1192 .eq(connection_id.0)
1193 .and(room_participant::Column::AnsweringConnectionId.is_null()),
1194 )
1195 .all(&*tx)
1196 .await?;
1197 room_participant::Entity::delete_many()
1198 .filter(
1199 room_participant::Column::Id
1200 .is_in(called_participants.iter().map(|participant| participant.id)),
1201 )
1202 .exec(&*tx)
1203 .await?;
1204 let canceled_calls_to_user_ids = called_participants
1205 .into_iter()
1206 .map(|participant| participant.user_id)
1207 .collect();
1208
1209 // Detect left projects.
1210 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
1211 enum QueryProjectIds {
1212 ProjectId,
1213 }
1214 let project_ids: Vec<ProjectId> = project_collaborator::Entity::find()
1215 .select_only()
1216 .column_as(
1217 project_collaborator::Column::ProjectId,
1218 QueryProjectIds::ProjectId,
1219 )
1220 .filter(project_collaborator::Column::ConnectionId.eq(connection_id.0))
1221 .into_values::<_, QueryProjectIds>()
1222 .all(&*tx)
1223 .await?;
1224 let mut left_projects = HashMap::default();
1225 let mut collaborators = project_collaborator::Entity::find()
1226 .filter(project_collaborator::Column::ProjectId.is_in(project_ids))
1227 .stream(&*tx)
1228 .await?;
1229 while let Some(collaborator) = collaborators.next().await {
1230 let collaborator = collaborator?;
1231 let left_project =
1232 left_projects
1233 .entry(collaborator.project_id)
1234 .or_insert(LeftProject {
1235 id: collaborator.project_id,
1236 host_user_id: Default::default(),
1237 connection_ids: Default::default(),
1238 host_connection_id: Default::default(),
1239 });
1240
1241 let collaborator_connection_id =
1242 ConnectionId(collaborator.connection_id as u32);
1243 if collaborator_connection_id != connection_id {
1244 left_project.connection_ids.push(collaborator_connection_id);
1245 }
1246
1247 if collaborator.is_host {
1248 left_project.host_user_id = collaborator.user_id;
1249 left_project.host_connection_id =
1250 ConnectionId(collaborator.connection_id as u32);
1251 }
1252 }
1253 drop(collaborators);
1254
1255 // Leave projects.
1256 project_collaborator::Entity::delete_many()
1257 .filter(project_collaborator::Column::ConnectionId.eq(connection_id.0))
1258 .exec(&*tx)
1259 .await?;
1260
1261 // Unshare projects.
1262 project::Entity::delete_many()
1263 .filter(
1264 project::Column::RoomId
1265 .eq(room_id)
1266 .and(project::Column::HostConnectionId.eq(connection_id.0)),
1267 )
1268 .exec(&*tx)
1269 .await?;
1270
1271 let room = self.get_room(room_id, &tx).await?;
1272 if room.participants.is_empty() {
1273 room::Entity::delete_by_id(room_id).exec(&*tx).await?;
1274 }
1275
1276 let left_room = LeftRoom {
1277 room,
1278 left_projects,
1279 canceled_calls_to_user_ids,
1280 };
1281
1282 if left_room.room.participants.is_empty() {
1283 self.rooms.remove(&room_id);
1284 }
1285
1286 Ok((room_id, left_room))
1287 } else {
1288 Err(anyhow!("could not leave room"))?
1289 }
1290 })
1291 .await
1292 }
1293
1294 pub async fn update_room_participant_location(
1295 &self,
1296 room_id: RoomId,
1297 connection_id: ConnectionId,
1298 location: proto::ParticipantLocation,
1299 ) -> Result<RoomGuard<proto::Room>> {
1300 self.room_transaction(|tx| async {
1301 let tx = tx;
1302 let location_kind;
1303 let location_project_id;
1304 match location
1305 .variant
1306 .as_ref()
1307 .ok_or_else(|| anyhow!("invalid location"))?
1308 {
1309 proto::participant_location::Variant::SharedProject(project) => {
1310 location_kind = 0;
1311 location_project_id = Some(ProjectId::from_proto(project.id));
1312 }
1313 proto::participant_location::Variant::UnsharedProject(_) => {
1314 location_kind = 1;
1315 location_project_id = None;
1316 }
1317 proto::participant_location::Variant::External(_) => {
1318 location_kind = 2;
1319 location_project_id = None;
1320 }
1321 }
1322
1323 let result = room_participant::Entity::update_many()
1324 .filter(
1325 room_participant::Column::RoomId
1326 .eq(room_id)
1327 .and(room_participant::Column::AnsweringConnectionId.eq(connection_id.0)),
1328 )
1329 .set(room_participant::ActiveModel {
1330 location_kind: ActiveValue::set(Some(location_kind)),
1331 location_project_id: ActiveValue::set(location_project_id),
1332 ..Default::default()
1333 })
1334 .exec(&*tx)
1335 .await?;
1336
1337 if result.rows_affected == 1 {
1338 let room = self.get_room(room_id, &tx).await?;
1339 Ok((room_id, room))
1340 } else {
1341 Err(anyhow!("could not update room participant location"))?
1342 }
1343 })
1344 .await
1345 }
1346
1347 fn build_incoming_call(
1348 room: &proto::Room,
1349 called_user_id: UserId,
1350 ) -> Option<proto::IncomingCall> {
1351 let pending_participant = room
1352 .pending_participants
1353 .iter()
1354 .find(|participant| participant.user_id == called_user_id.to_proto())?;
1355
1356 Some(proto::IncomingCall {
1357 room_id: room.id,
1358 calling_user_id: pending_participant.calling_user_id,
1359 participant_user_ids: room
1360 .participants
1361 .iter()
1362 .map(|participant| participant.user_id)
1363 .collect(),
1364 initial_project: room.participants.iter().find_map(|participant| {
1365 let initial_project_id = pending_participant.initial_project_id?;
1366 participant
1367 .projects
1368 .iter()
1369 .find(|project| project.id == initial_project_id)
1370 .cloned()
1371 }),
1372 })
1373 }
1374
1375 async fn get_room(&self, room_id: RoomId, tx: &DatabaseTransaction) -> Result<proto::Room> {
1376 let db_room = room::Entity::find_by_id(room_id)
1377 .one(tx)
1378 .await?
1379 .ok_or_else(|| anyhow!("could not find room"))?;
1380
1381 let mut db_participants = db_room
1382 .find_related(room_participant::Entity)
1383 .stream(tx)
1384 .await?;
1385 let mut participants = HashMap::default();
1386 let mut pending_participants = Vec::new();
1387 while let Some(db_participant) = db_participants.next().await {
1388 let db_participant = db_participant?;
1389 if let Some(answering_connection_id) = db_participant.answering_connection_id {
1390 let location = match (
1391 db_participant.location_kind,
1392 db_participant.location_project_id,
1393 ) {
1394 (Some(0), Some(project_id)) => {
1395 Some(proto::participant_location::Variant::SharedProject(
1396 proto::participant_location::SharedProject {
1397 id: project_id.to_proto(),
1398 },
1399 ))
1400 }
1401 (Some(1), _) => Some(proto::participant_location::Variant::UnsharedProject(
1402 Default::default(),
1403 )),
1404 _ => Some(proto::participant_location::Variant::External(
1405 Default::default(),
1406 )),
1407 };
1408 participants.insert(
1409 answering_connection_id,
1410 proto::Participant {
1411 user_id: db_participant.user_id.to_proto(),
1412 peer_id: answering_connection_id as u32,
1413 projects: Default::default(),
1414 location: Some(proto::ParticipantLocation { variant: location }),
1415 },
1416 );
1417 } else {
1418 pending_participants.push(proto::PendingParticipant {
1419 user_id: db_participant.user_id.to_proto(),
1420 calling_user_id: db_participant.calling_user_id.to_proto(),
1421 initial_project_id: db_participant.initial_project_id.map(|id| id.to_proto()),
1422 });
1423 }
1424 }
1425 drop(db_participants);
1426
1427 let mut db_projects = db_room
1428 .find_related(project::Entity)
1429 .find_with_related(worktree::Entity)
1430 .stream(tx)
1431 .await?;
1432
1433 while let Some(row) = db_projects.next().await {
1434 let (db_project, db_worktree) = row?;
1435 if let Some(participant) = participants.get_mut(&db_project.host_connection_id) {
1436 let project = if let Some(project) = participant
1437 .projects
1438 .iter_mut()
1439 .find(|project| project.id == db_project.id.to_proto())
1440 {
1441 project
1442 } else {
1443 participant.projects.push(proto::ParticipantProject {
1444 id: db_project.id.to_proto(),
1445 worktree_root_names: Default::default(),
1446 });
1447 participant.projects.last_mut().unwrap()
1448 };
1449
1450 if let Some(db_worktree) = db_worktree {
1451 project.worktree_root_names.push(db_worktree.root_name);
1452 }
1453 }
1454 }
1455
1456 Ok(proto::Room {
1457 id: db_room.id.to_proto(),
1458 live_kit_room: db_room.live_kit_room,
1459 participants: participants.into_values().collect(),
1460 pending_participants,
1461 })
1462 }
1463
1464 // projects
1465
1466 pub async fn project_count_excluding_admins(&self) -> Result<usize> {
1467 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
1468 enum QueryAs {
1469 Count,
1470 }
1471
1472 self.transaction(|tx| async move {
1473 Ok(project::Entity::find()
1474 .select_only()
1475 .column_as(project::Column::Id.count(), QueryAs::Count)
1476 .inner_join(user::Entity)
1477 .filter(user::Column::Admin.eq(false))
1478 .into_values::<_, QueryAs>()
1479 .one(&*tx)
1480 .await?
1481 .unwrap_or(0i64) as usize)
1482 })
1483 .await
1484 }
1485
1486 pub async fn share_project(
1487 &self,
1488 room_id: RoomId,
1489 connection_id: ConnectionId,
1490 worktrees: &[proto::WorktreeMetadata],
1491 ) -> Result<RoomGuard<(ProjectId, proto::Room)>> {
1492 self.room_transaction(|tx| async move {
1493 let participant = room_participant::Entity::find()
1494 .filter(room_participant::Column::AnsweringConnectionId.eq(connection_id.0))
1495 .one(&*tx)
1496 .await?
1497 .ok_or_else(|| anyhow!("could not find participant"))?;
1498 if participant.room_id != room_id {
1499 return Err(anyhow!("shared project on unexpected room"))?;
1500 }
1501
1502 let project = project::ActiveModel {
1503 room_id: ActiveValue::set(participant.room_id),
1504 host_user_id: ActiveValue::set(participant.user_id),
1505 host_connection_id: ActiveValue::set(connection_id.0 as i32),
1506 host_connection_epoch: ActiveValue::set(self.epoch),
1507 ..Default::default()
1508 }
1509 .insert(&*tx)
1510 .await?;
1511
1512 if !worktrees.is_empty() {
1513 worktree::Entity::insert_many(worktrees.iter().map(|worktree| {
1514 worktree::ActiveModel {
1515 id: ActiveValue::set(worktree.id as i64),
1516 project_id: ActiveValue::set(project.id),
1517 abs_path: ActiveValue::set(worktree.abs_path.clone()),
1518 root_name: ActiveValue::set(worktree.root_name.clone()),
1519 visible: ActiveValue::set(worktree.visible),
1520 scan_id: ActiveValue::set(0),
1521 is_complete: ActiveValue::set(false),
1522 }
1523 }))
1524 .exec(&*tx)
1525 .await?;
1526 }
1527
1528 project_collaborator::ActiveModel {
1529 project_id: ActiveValue::set(project.id),
1530 connection_id: ActiveValue::set(connection_id.0 as i32),
1531 connection_epoch: ActiveValue::set(self.epoch),
1532 user_id: ActiveValue::set(participant.user_id),
1533 replica_id: ActiveValue::set(ReplicaId(0)),
1534 is_host: ActiveValue::set(true),
1535 ..Default::default()
1536 }
1537 .insert(&*tx)
1538 .await?;
1539
1540 let room = self.get_room(room_id, &tx).await?;
1541 Ok((room_id, (project.id, room)))
1542 })
1543 .await
1544 }
1545
1546 pub async fn unshare_project(
1547 &self,
1548 project_id: ProjectId,
1549 connection_id: ConnectionId,
1550 ) -> Result<RoomGuard<(proto::Room, Vec<ConnectionId>)>> {
1551 self.room_transaction(|tx| async move {
1552 let guest_connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
1553
1554 let project = project::Entity::find_by_id(project_id)
1555 .one(&*tx)
1556 .await?
1557 .ok_or_else(|| anyhow!("project not found"))?;
1558 if project.host_connection_id == connection_id.0 as i32 {
1559 let room_id = project.room_id;
1560 project::Entity::delete(project.into_active_model())
1561 .exec(&*tx)
1562 .await?;
1563 let room = self.get_room(room_id, &tx).await?;
1564 Ok((room_id, (room, guest_connection_ids)))
1565 } else {
1566 Err(anyhow!("cannot unshare a project hosted by another user"))?
1567 }
1568 })
1569 .await
1570 }
1571
1572 pub async fn update_project(
1573 &self,
1574 project_id: ProjectId,
1575 connection_id: ConnectionId,
1576 worktrees: &[proto::WorktreeMetadata],
1577 ) -> Result<RoomGuard<(proto::Room, Vec<ConnectionId>)>> {
1578 self.room_transaction(|tx| async move {
1579 let project = project::Entity::find_by_id(project_id)
1580 .filter(project::Column::HostConnectionId.eq(connection_id.0))
1581 .one(&*tx)
1582 .await?
1583 .ok_or_else(|| anyhow!("no such project"))?;
1584
1585 if !worktrees.is_empty() {
1586 worktree::Entity::insert_many(worktrees.iter().map(|worktree| {
1587 worktree::ActiveModel {
1588 id: ActiveValue::set(worktree.id as i64),
1589 project_id: ActiveValue::set(project.id),
1590 abs_path: ActiveValue::set(worktree.abs_path.clone()),
1591 root_name: ActiveValue::set(worktree.root_name.clone()),
1592 visible: ActiveValue::set(worktree.visible),
1593 scan_id: ActiveValue::set(0),
1594 is_complete: ActiveValue::set(false),
1595 }
1596 }))
1597 .on_conflict(
1598 OnConflict::columns([worktree::Column::ProjectId, worktree::Column::Id])
1599 .update_column(worktree::Column::RootName)
1600 .to_owned(),
1601 )
1602 .exec(&*tx)
1603 .await?;
1604 }
1605
1606 worktree::Entity::delete_many()
1607 .filter(
1608 worktree::Column::ProjectId.eq(project.id).and(
1609 worktree::Column::Id
1610 .is_not_in(worktrees.iter().map(|worktree| worktree.id as i64)),
1611 ),
1612 )
1613 .exec(&*tx)
1614 .await?;
1615
1616 let guest_connection_ids = self.project_guest_connection_ids(project.id, &tx).await?;
1617 let room = self.get_room(project.room_id, &tx).await?;
1618 Ok((project.room_id, (room, guest_connection_ids)))
1619 })
1620 .await
1621 }
1622
1623 pub async fn update_worktree(
1624 &self,
1625 update: &proto::UpdateWorktree,
1626 connection_id: ConnectionId,
1627 ) -> Result<RoomGuard<Vec<ConnectionId>>> {
1628 self.room_transaction(|tx| async move {
1629 let project_id = ProjectId::from_proto(update.project_id);
1630 let worktree_id = update.worktree_id as i64;
1631
1632 // Ensure the update comes from the host.
1633 let project = project::Entity::find_by_id(project_id)
1634 .filter(project::Column::HostConnectionId.eq(connection_id.0))
1635 .one(&*tx)
1636 .await?
1637 .ok_or_else(|| anyhow!("no such project"))?;
1638 let room_id = project.room_id;
1639
1640 // Update metadata.
1641 worktree::Entity::update(worktree::ActiveModel {
1642 id: ActiveValue::set(worktree_id),
1643 project_id: ActiveValue::set(project_id),
1644 root_name: ActiveValue::set(update.root_name.clone()),
1645 scan_id: ActiveValue::set(update.scan_id as i64),
1646 is_complete: ActiveValue::set(update.is_last_update),
1647 abs_path: ActiveValue::set(update.abs_path.clone()),
1648 ..Default::default()
1649 })
1650 .exec(&*tx)
1651 .await?;
1652
1653 if !update.updated_entries.is_empty() {
1654 worktree_entry::Entity::insert_many(update.updated_entries.iter().map(|entry| {
1655 let mtime = entry.mtime.clone().unwrap_or_default();
1656 worktree_entry::ActiveModel {
1657 project_id: ActiveValue::set(project_id),
1658 worktree_id: ActiveValue::set(worktree_id),
1659 id: ActiveValue::set(entry.id as i64),
1660 is_dir: ActiveValue::set(entry.is_dir),
1661 path: ActiveValue::set(entry.path.clone()),
1662 inode: ActiveValue::set(entry.inode as i64),
1663 mtime_seconds: ActiveValue::set(mtime.seconds as i64),
1664 mtime_nanos: ActiveValue::set(mtime.nanos as i32),
1665 is_symlink: ActiveValue::set(entry.is_symlink),
1666 is_ignored: ActiveValue::set(entry.is_ignored),
1667 }
1668 }))
1669 .on_conflict(
1670 OnConflict::columns([
1671 worktree_entry::Column::ProjectId,
1672 worktree_entry::Column::WorktreeId,
1673 worktree_entry::Column::Id,
1674 ])
1675 .update_columns([
1676 worktree_entry::Column::IsDir,
1677 worktree_entry::Column::Path,
1678 worktree_entry::Column::Inode,
1679 worktree_entry::Column::MtimeSeconds,
1680 worktree_entry::Column::MtimeNanos,
1681 worktree_entry::Column::IsSymlink,
1682 worktree_entry::Column::IsIgnored,
1683 ])
1684 .to_owned(),
1685 )
1686 .exec(&*tx)
1687 .await?;
1688 }
1689
1690 if !update.removed_entries.is_empty() {
1691 worktree_entry::Entity::delete_many()
1692 .filter(
1693 worktree_entry::Column::ProjectId
1694 .eq(project_id)
1695 .and(worktree_entry::Column::WorktreeId.eq(worktree_id))
1696 .and(
1697 worktree_entry::Column::Id
1698 .is_in(update.removed_entries.iter().map(|id| *id as i64)),
1699 ),
1700 )
1701 .exec(&*tx)
1702 .await?;
1703 }
1704
1705 let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
1706 Ok((room_id, connection_ids))
1707 })
1708 .await
1709 }
1710
1711 pub async fn update_diagnostic_summary(
1712 &self,
1713 update: &proto::UpdateDiagnosticSummary,
1714 connection_id: ConnectionId,
1715 ) -> Result<RoomGuard<Vec<ConnectionId>>> {
1716 self.room_transaction(|tx| async move {
1717 let project_id = ProjectId::from_proto(update.project_id);
1718 let worktree_id = update.worktree_id as i64;
1719 let summary = update
1720 .summary
1721 .as_ref()
1722 .ok_or_else(|| anyhow!("invalid summary"))?;
1723
1724 // Ensure the update comes from the host.
1725 let project = project::Entity::find_by_id(project_id)
1726 .one(&*tx)
1727 .await?
1728 .ok_or_else(|| anyhow!("no such project"))?;
1729 if project.host_connection_id != connection_id.0 as i32 {
1730 return Err(anyhow!("can't update a project hosted by someone else"))?;
1731 }
1732
1733 // Update summary.
1734 worktree_diagnostic_summary::Entity::insert(worktree_diagnostic_summary::ActiveModel {
1735 project_id: ActiveValue::set(project_id),
1736 worktree_id: ActiveValue::set(worktree_id),
1737 path: ActiveValue::set(summary.path.clone()),
1738 language_server_id: ActiveValue::set(summary.language_server_id as i64),
1739 error_count: ActiveValue::set(summary.error_count as i32),
1740 warning_count: ActiveValue::set(summary.warning_count as i32),
1741 ..Default::default()
1742 })
1743 .on_conflict(
1744 OnConflict::columns([
1745 worktree_diagnostic_summary::Column::ProjectId,
1746 worktree_diagnostic_summary::Column::WorktreeId,
1747 worktree_diagnostic_summary::Column::Path,
1748 ])
1749 .update_columns([
1750 worktree_diagnostic_summary::Column::LanguageServerId,
1751 worktree_diagnostic_summary::Column::ErrorCount,
1752 worktree_diagnostic_summary::Column::WarningCount,
1753 ])
1754 .to_owned(),
1755 )
1756 .exec(&*tx)
1757 .await?;
1758
1759 let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
1760 Ok((project.room_id, connection_ids))
1761 })
1762 .await
1763 }
1764
1765 pub async fn start_language_server(
1766 &self,
1767 update: &proto::StartLanguageServer,
1768 connection_id: ConnectionId,
1769 ) -> Result<RoomGuard<Vec<ConnectionId>>> {
1770 self.room_transaction(|tx| async move {
1771 let project_id = ProjectId::from_proto(update.project_id);
1772 let server = update
1773 .server
1774 .as_ref()
1775 .ok_or_else(|| anyhow!("invalid language server"))?;
1776
1777 // Ensure the update comes from the host.
1778 let project = project::Entity::find_by_id(project_id)
1779 .one(&*tx)
1780 .await?
1781 .ok_or_else(|| anyhow!("no such project"))?;
1782 if project.host_connection_id != connection_id.0 as i32 {
1783 return Err(anyhow!("can't update a project hosted by someone else"))?;
1784 }
1785
1786 // Add the newly-started language server.
1787 language_server::Entity::insert(language_server::ActiveModel {
1788 project_id: ActiveValue::set(project_id),
1789 id: ActiveValue::set(server.id as i64),
1790 name: ActiveValue::set(server.name.clone()),
1791 ..Default::default()
1792 })
1793 .on_conflict(
1794 OnConflict::columns([
1795 language_server::Column::ProjectId,
1796 language_server::Column::Id,
1797 ])
1798 .update_column(language_server::Column::Name)
1799 .to_owned(),
1800 )
1801 .exec(&*tx)
1802 .await?;
1803
1804 let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;
1805 Ok((project.room_id, connection_ids))
1806 })
1807 .await
1808 }
1809
1810 pub async fn join_project(
1811 &self,
1812 project_id: ProjectId,
1813 connection_id: ConnectionId,
1814 ) -> Result<RoomGuard<(Project, ReplicaId)>> {
1815 self.room_transaction(|tx| async move {
1816 let participant = room_participant::Entity::find()
1817 .filter(room_participant::Column::AnsweringConnectionId.eq(connection_id.0))
1818 .one(&*tx)
1819 .await?
1820 .ok_or_else(|| anyhow!("must join a room first"))?;
1821
1822 let project = project::Entity::find_by_id(project_id)
1823 .one(&*tx)
1824 .await?
1825 .ok_or_else(|| anyhow!("no such project"))?;
1826 if project.room_id != participant.room_id {
1827 return Err(anyhow!("no such project"))?;
1828 }
1829
1830 let mut collaborators = project
1831 .find_related(project_collaborator::Entity)
1832 .all(&*tx)
1833 .await?;
1834 let replica_ids = collaborators
1835 .iter()
1836 .map(|c| c.replica_id)
1837 .collect::<HashSet<_>>();
1838 let mut replica_id = ReplicaId(1);
1839 while replica_ids.contains(&replica_id) {
1840 replica_id.0 += 1;
1841 }
1842 let new_collaborator = project_collaborator::ActiveModel {
1843 project_id: ActiveValue::set(project_id),
1844 connection_id: ActiveValue::set(connection_id.0 as i32),
1845 connection_epoch: ActiveValue::set(self.epoch),
1846 user_id: ActiveValue::set(participant.user_id),
1847 replica_id: ActiveValue::set(replica_id),
1848 is_host: ActiveValue::set(false),
1849 ..Default::default()
1850 }
1851 .insert(&*tx)
1852 .await?;
1853 collaborators.push(new_collaborator);
1854
1855 let db_worktrees = project.find_related(worktree::Entity).all(&*tx).await?;
1856 let mut worktrees = db_worktrees
1857 .into_iter()
1858 .map(|db_worktree| {
1859 (
1860 db_worktree.id as u64,
1861 Worktree {
1862 id: db_worktree.id as u64,
1863 abs_path: db_worktree.abs_path,
1864 root_name: db_worktree.root_name,
1865 visible: db_worktree.visible,
1866 entries: Default::default(),
1867 diagnostic_summaries: Default::default(),
1868 scan_id: db_worktree.scan_id as u64,
1869 is_complete: db_worktree.is_complete,
1870 },
1871 )
1872 })
1873 .collect::<BTreeMap<_, _>>();
1874
1875 // Populate worktree entries.
1876 {
1877 let mut db_entries = worktree_entry::Entity::find()
1878 .filter(worktree_entry::Column::ProjectId.eq(project_id))
1879 .stream(&*tx)
1880 .await?;
1881 while let Some(db_entry) = db_entries.next().await {
1882 let db_entry = db_entry?;
1883 if let Some(worktree) = worktrees.get_mut(&(db_entry.worktree_id as u64)) {
1884 worktree.entries.push(proto::Entry {
1885 id: db_entry.id as u64,
1886 is_dir: db_entry.is_dir,
1887 path: db_entry.path,
1888 inode: db_entry.inode as u64,
1889 mtime: Some(proto::Timestamp {
1890 seconds: db_entry.mtime_seconds as u64,
1891 nanos: db_entry.mtime_nanos as u32,
1892 }),
1893 is_symlink: db_entry.is_symlink,
1894 is_ignored: db_entry.is_ignored,
1895 });
1896 }
1897 }
1898 }
1899
1900 // Populate worktree diagnostic summaries.
1901 {
1902 let mut db_summaries = worktree_diagnostic_summary::Entity::find()
1903 .filter(worktree_diagnostic_summary::Column::ProjectId.eq(project_id))
1904 .stream(&*tx)
1905 .await?;
1906 while let Some(db_summary) = db_summaries.next().await {
1907 let db_summary = db_summary?;
1908 if let Some(worktree) = worktrees.get_mut(&(db_summary.worktree_id as u64)) {
1909 worktree
1910 .diagnostic_summaries
1911 .push(proto::DiagnosticSummary {
1912 path: db_summary.path,
1913 language_server_id: db_summary.language_server_id as u64,
1914 error_count: db_summary.error_count as u32,
1915 warning_count: db_summary.warning_count as u32,
1916 });
1917 }
1918 }
1919 }
1920
1921 // Populate language servers.
1922 let language_servers = project
1923 .find_related(language_server::Entity)
1924 .all(&*tx)
1925 .await?;
1926
1927 let room_id = project.room_id;
1928 let project = Project {
1929 collaborators,
1930 worktrees,
1931 language_servers: language_servers
1932 .into_iter()
1933 .map(|language_server| proto::LanguageServer {
1934 id: language_server.id as u64,
1935 name: language_server.name,
1936 })
1937 .collect(),
1938 };
1939 Ok((room_id, (project, replica_id as ReplicaId)))
1940 })
1941 .await
1942 }
1943
1944 pub async fn leave_project(
1945 &self,
1946 project_id: ProjectId,
1947 connection_id: ConnectionId,
1948 ) -> Result<RoomGuard<LeftProject>> {
1949 self.room_transaction(|tx| async move {
1950 let result = project_collaborator::Entity::delete_many()
1951 .filter(
1952 project_collaborator::Column::ProjectId
1953 .eq(project_id)
1954 .and(project_collaborator::Column::ConnectionId.eq(connection_id.0)),
1955 )
1956 .exec(&*tx)
1957 .await?;
1958 if result.rows_affected == 0 {
1959 Err(anyhow!("not a collaborator on this project"))?;
1960 }
1961
1962 let project = project::Entity::find_by_id(project_id)
1963 .one(&*tx)
1964 .await?
1965 .ok_or_else(|| anyhow!("no such project"))?;
1966 let collaborators = project
1967 .find_related(project_collaborator::Entity)
1968 .all(&*tx)
1969 .await?;
1970 let connection_ids = collaborators
1971 .into_iter()
1972 .map(|collaborator| ConnectionId(collaborator.connection_id as u32))
1973 .collect();
1974
1975 let left_project = LeftProject {
1976 id: project_id,
1977 host_user_id: project.host_user_id,
1978 host_connection_id: ConnectionId(project.host_connection_id as u32),
1979 connection_ids,
1980 };
1981 Ok((project.room_id, left_project))
1982 })
1983 .await
1984 }
1985
1986 pub async fn project_collaborators(
1987 &self,
1988 project_id: ProjectId,
1989 connection_id: ConnectionId,
1990 ) -> Result<RoomGuard<Vec<project_collaborator::Model>>> {
1991 self.room_transaction(|tx| async move {
1992 let project = project::Entity::find_by_id(project_id)
1993 .one(&*tx)
1994 .await?
1995 .ok_or_else(|| anyhow!("no such project"))?;
1996 let collaborators = project_collaborator::Entity::find()
1997 .filter(project_collaborator::Column::ProjectId.eq(project_id))
1998 .all(&*tx)
1999 .await?;
2000
2001 if collaborators
2002 .iter()
2003 .any(|collaborator| collaborator.connection_id == connection_id.0 as i32)
2004 {
2005 Ok((project.room_id, collaborators))
2006 } else {
2007 Err(anyhow!("no such project"))?
2008 }
2009 })
2010 .await
2011 }
2012
2013 pub async fn project_connection_ids(
2014 &self,
2015 project_id: ProjectId,
2016 connection_id: ConnectionId,
2017 ) -> Result<RoomGuard<HashSet<ConnectionId>>> {
2018 self.room_transaction(|tx| async move {
2019 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
2020 enum QueryAs {
2021 ConnectionId,
2022 }
2023
2024 let project = project::Entity::find_by_id(project_id)
2025 .one(&*tx)
2026 .await?
2027 .ok_or_else(|| anyhow!("no such project"))?;
2028 let mut db_connection_ids = project_collaborator::Entity::find()
2029 .select_only()
2030 .column_as(
2031 project_collaborator::Column::ConnectionId,
2032 QueryAs::ConnectionId,
2033 )
2034 .filter(project_collaborator::Column::ProjectId.eq(project_id))
2035 .into_values::<i32, QueryAs>()
2036 .stream(&*tx)
2037 .await?;
2038
2039 let mut connection_ids = HashSet::default();
2040 while let Some(connection_id) = db_connection_ids.next().await {
2041 connection_ids.insert(ConnectionId(connection_id? as u32));
2042 }
2043
2044 if connection_ids.contains(&connection_id) {
2045 Ok((project.room_id, connection_ids))
2046 } else {
2047 Err(anyhow!("no such project"))?
2048 }
2049 })
2050 .await
2051 }
2052
2053 async fn project_guest_connection_ids(
2054 &self,
2055 project_id: ProjectId,
2056 tx: &DatabaseTransaction,
2057 ) -> Result<Vec<ConnectionId>> {
2058 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
2059 enum QueryAs {
2060 ConnectionId,
2061 }
2062
2063 let mut db_guest_connection_ids = project_collaborator::Entity::find()
2064 .select_only()
2065 .column_as(
2066 project_collaborator::Column::ConnectionId,
2067 QueryAs::ConnectionId,
2068 )
2069 .filter(
2070 project_collaborator::Column::ProjectId
2071 .eq(project_id)
2072 .and(project_collaborator::Column::IsHost.eq(false)),
2073 )
2074 .into_values::<i32, QueryAs>()
2075 .stream(tx)
2076 .await?;
2077
2078 let mut guest_connection_ids = Vec::new();
2079 while let Some(connection_id) = db_guest_connection_ids.next().await {
2080 guest_connection_ids.push(ConnectionId(connection_id? as u32));
2081 }
2082 Ok(guest_connection_ids)
2083 }
2084
2085 // access tokens
2086
2087 pub async fn create_access_token_hash(
2088 &self,
2089 user_id: UserId,
2090 access_token_hash: &str,
2091 max_access_token_count: usize,
2092 ) -> Result<()> {
2093 self.transaction(|tx| async {
2094 let tx = tx;
2095
2096 access_token::ActiveModel {
2097 user_id: ActiveValue::set(user_id),
2098 hash: ActiveValue::set(access_token_hash.into()),
2099 ..Default::default()
2100 }
2101 .insert(&*tx)
2102 .await?;
2103
2104 access_token::Entity::delete_many()
2105 .filter(
2106 access_token::Column::Id.in_subquery(
2107 Query::select()
2108 .column(access_token::Column::Id)
2109 .from(access_token::Entity)
2110 .and_where(access_token::Column::UserId.eq(user_id))
2111 .order_by(access_token::Column::Id, sea_orm::Order::Desc)
2112 .limit(10000)
2113 .offset(max_access_token_count as u64)
2114 .to_owned(),
2115 ),
2116 )
2117 .exec(&*tx)
2118 .await?;
2119 Ok(())
2120 })
2121 .await
2122 }
2123
2124 pub async fn get_access_token_hashes(&self, user_id: UserId) -> Result<Vec<String>> {
2125 #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]
2126 enum QueryAs {
2127 Hash,
2128 }
2129
2130 self.transaction(|tx| async move {
2131 Ok(access_token::Entity::find()
2132 .select_only()
2133 .column(access_token::Column::Hash)
2134 .filter(access_token::Column::UserId.eq(user_id))
2135 .order_by_desc(access_token::Column::Id)
2136 .into_values::<_, QueryAs>()
2137 .all(&*tx)
2138 .await?)
2139 })
2140 .await
2141 }
2142
2143 async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>
2144 where
2145 F: Send + Fn(TransactionHandle) -> Fut,
2146 Fut: Send + Future<Output = Result<T>>,
2147 {
2148 let body = async {
2149 loop {
2150 let (tx, result) = self.with_transaction(&f).await?;
2151 match result {
2152 Ok(result) => {
2153 match tx.commit().await.map_err(Into::into) {
2154 Ok(()) => return Ok(result),
2155 Err(error) => {
2156 if is_serialization_error(&error) {
2157 // Retry (don't break the loop)
2158 } else {
2159 return Err(error);
2160 }
2161 }
2162 }
2163 }
2164 Err(error) => {
2165 tx.rollback().await?;
2166 if is_serialization_error(&error) {
2167 // Retry (don't break the loop)
2168 } else {
2169 return Err(error);
2170 }
2171 }
2172 }
2173 }
2174 };
2175
2176 self.run(body).await
2177 }
2178
2179 async fn room_transaction<F, Fut, T>(&self, f: F) -> Result<RoomGuard<T>>
2180 where
2181 F: Send + Fn(TransactionHandle) -> Fut,
2182 Fut: Send + Future<Output = Result<(RoomId, T)>>,
2183 {
2184 let body = async {
2185 loop {
2186 let (tx, result) = self.with_transaction(&f).await?;
2187 match result {
2188 Ok((room_id, data)) => {
2189 let lock = self.rooms.entry(room_id).or_default().clone();
2190 let _guard = lock.lock_owned().await;
2191 match tx.commit().await.map_err(Into::into) {
2192 Ok(()) => {
2193 return Ok(RoomGuard {
2194 data,
2195 _guard,
2196 _not_send: PhantomData,
2197 });
2198 }
2199 Err(error) => {
2200 if is_serialization_error(&error) {
2201 // Retry (don't break the loop)
2202 } else {
2203 return Err(error);
2204 }
2205 }
2206 }
2207 }
2208 Err(error) => {
2209 tx.rollback().await?;
2210 if is_serialization_error(&error) {
2211 // Retry (don't break the loop)
2212 } else {
2213 return Err(error);
2214 }
2215 }
2216 }
2217 }
2218 };
2219
2220 self.run(body).await
2221 }
2222
2223 async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>
2224 where
2225 F: Send + Fn(TransactionHandle) -> Fut,
2226 Fut: Send + Future<Output = Result<T>>,
2227 {
2228 let tx = self
2229 .pool
2230 .begin_with_config(Some(IsolationLevel::Serializable), None)
2231 .await?;
2232
2233 let mut tx = Arc::new(Some(tx));
2234 let result = f(TransactionHandle(tx.clone())).await;
2235 let Some(tx) = Arc::get_mut(&mut tx).and_then(|tx| tx.take()) else {
2236 return Err(anyhow!("couldn't complete transaction because it's still in use"))?;
2237 };
2238
2239 Ok((tx, result))
2240 }
2241
2242 async fn run<F, T>(&self, future: F) -> T
2243 where
2244 F: Future<Output = T>,
2245 {
2246 #[cfg(test)]
2247 {
2248 if let Some(background) = self.background.as_ref() {
2249 background.simulate_random_delay().await;
2250 }
2251
2252 self.runtime.as_ref().unwrap().block_on(future)
2253 }
2254
2255 #[cfg(not(test))]
2256 {
2257 future.await
2258 }
2259 }
2260}
2261
2262fn is_serialization_error(error: &Error) -> bool {
2263 const SERIALIZATION_FAILURE_CODE: &'static str = "40001";
2264 match error {
2265 Error::Database(
2266 DbErr::Exec(sea_orm::RuntimeErr::SqlxError(error))
2267 | DbErr::Query(sea_orm::RuntimeErr::SqlxError(error)),
2268 ) if error
2269 .as_database_error()
2270 .and_then(|error| error.code())
2271 .as_deref()
2272 == Some(SERIALIZATION_FAILURE_CODE) =>
2273 {
2274 true
2275 }
2276 _ => false,
2277 }
2278}
2279
2280struct TransactionHandle(Arc<Option<DatabaseTransaction>>);
2281
2282impl Deref for TransactionHandle {
2283 type Target = DatabaseTransaction;
2284
2285 fn deref(&self) -> &Self::Target {
2286 self.0.as_ref().as_ref().unwrap()
2287 }
2288}
2289
2290pub struct RoomGuard<T> {
2291 data: T,
2292 _guard: OwnedMutexGuard<()>,
2293 _not_send: PhantomData<Rc<()>>,
2294}
2295
2296impl<T> Deref for RoomGuard<T> {
2297 type Target = T;
2298
2299 fn deref(&self) -> &T {
2300 &self.data
2301 }
2302}
2303
2304impl<T> DerefMut for RoomGuard<T> {
2305 fn deref_mut(&mut self) -> &mut T {
2306 &mut self.data
2307 }
2308}
2309
2310#[derive(Debug, Serialize, Deserialize)]
2311pub struct NewUserParams {
2312 pub github_login: String,
2313 pub github_user_id: i32,
2314 pub invite_count: i32,
2315}
2316
2317#[derive(Debug)]
2318pub struct NewUserResult {
2319 pub user_id: UserId,
2320 pub metrics_id: String,
2321 pub inviting_user_id: Option<UserId>,
2322 pub signup_device_id: Option<String>,
2323}
2324
2325fn random_invite_code() -> String {
2326 nanoid::nanoid!(16)
2327}
2328
2329fn random_email_confirmation_code() -> String {
2330 nanoid::nanoid!(64)
2331}
2332
2333macro_rules! id_type {
2334 ($name:ident) => {
2335 #[derive(
2336 Clone,
2337 Copy,
2338 Debug,
2339 Default,
2340 PartialEq,
2341 Eq,
2342 PartialOrd,
2343 Ord,
2344 Hash,
2345 Serialize,
2346 Deserialize,
2347 )]
2348 #[serde(transparent)]
2349 pub struct $name(pub i32);
2350
2351 impl $name {
2352 #[allow(unused)]
2353 pub const MAX: Self = Self(i32::MAX);
2354
2355 #[allow(unused)]
2356 pub fn from_proto(value: u64) -> Self {
2357 Self(value as i32)
2358 }
2359
2360 #[allow(unused)]
2361 pub fn to_proto(self) -> u64 {
2362 self.0 as u64
2363 }
2364 }
2365
2366 impl std::fmt::Display for $name {
2367 fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
2368 self.0.fmt(f)
2369 }
2370 }
2371
2372 impl From<$name> for sea_query::Value {
2373 fn from(value: $name) -> Self {
2374 sea_query::Value::Int(Some(value.0))
2375 }
2376 }
2377
2378 impl sea_orm::TryGetable for $name {
2379 fn try_get(
2380 res: &sea_orm::QueryResult,
2381 pre: &str,
2382 col: &str,
2383 ) -> Result<Self, sea_orm::TryGetError> {
2384 Ok(Self(i32::try_get(res, pre, col)?))
2385 }
2386 }
2387
2388 impl sea_query::ValueType for $name {
2389 fn try_from(v: Value) -> Result<Self, sea_query::ValueTypeErr> {
2390 match v {
2391 Value::TinyInt(Some(int)) => {
2392 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2393 }
2394 Value::SmallInt(Some(int)) => {
2395 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2396 }
2397 Value::Int(Some(int)) => {
2398 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2399 }
2400 Value::BigInt(Some(int)) => {
2401 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2402 }
2403 Value::TinyUnsigned(Some(int)) => {
2404 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2405 }
2406 Value::SmallUnsigned(Some(int)) => {
2407 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2408 }
2409 Value::Unsigned(Some(int)) => {
2410 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2411 }
2412 Value::BigUnsigned(Some(int)) => {
2413 Ok(Self(int.try_into().map_err(|_| sea_query::ValueTypeErr)?))
2414 }
2415 _ => Err(sea_query::ValueTypeErr),
2416 }
2417 }
2418
2419 fn type_name() -> String {
2420 stringify!($name).into()
2421 }
2422
2423 fn array_type() -> sea_query::ArrayType {
2424 sea_query::ArrayType::Int
2425 }
2426
2427 fn column_type() -> sea_query::ColumnType {
2428 sea_query::ColumnType::Integer(None)
2429 }
2430 }
2431
2432 impl sea_orm::TryFromU64 for $name {
2433 fn try_from_u64(n: u64) -> Result<Self, DbErr> {
2434 Ok(Self(n.try_into().map_err(|_| {
2435 DbErr::ConvertFromU64(concat!(
2436 "error converting ",
2437 stringify!($name),
2438 " to u64"
2439 ))
2440 })?))
2441 }
2442 }
2443
2444 impl sea_query::Nullable for $name {
2445 fn null() -> Value {
2446 Value::Int(None)
2447 }
2448 }
2449 };
2450}
2451
2452id_type!(AccessTokenId);
2453id_type!(ContactId);
2454id_type!(RoomId);
2455id_type!(RoomParticipantId);
2456id_type!(ProjectId);
2457id_type!(ProjectCollaboratorId);
2458id_type!(ReplicaId);
2459id_type!(SignupId);
2460id_type!(UserId);
2461
2462pub struct LeftRoom {
2463 pub room: proto::Room,
2464 pub left_projects: HashMap<ProjectId, LeftProject>,
2465 pub canceled_calls_to_user_ids: Vec<UserId>,
2466}
2467
2468pub struct Project {
2469 pub collaborators: Vec<project_collaborator::Model>,
2470 pub worktrees: BTreeMap<u64, Worktree>,
2471 pub language_servers: Vec<proto::LanguageServer>,
2472}
2473
2474pub struct LeftProject {
2475 pub id: ProjectId,
2476 pub host_user_id: UserId,
2477 pub host_connection_id: ConnectionId,
2478 pub connection_ids: Vec<ConnectionId>,
2479}
2480
2481pub struct Worktree {
2482 pub id: u64,
2483 pub abs_path: String,
2484 pub root_name: String,
2485 pub visible: bool,
2486 pub entries: Vec<proto::Entry>,
2487 pub diagnostic_summaries: Vec<proto::DiagnosticSummary>,
2488 pub scan_id: u64,
2489 pub is_complete: bool,
2490}
2491
2492#[cfg(test)]
2493pub use test::*;
2494
2495#[cfg(test)]
2496mod test {
2497 use super::*;
2498 use gpui::executor::Background;
2499 use lazy_static::lazy_static;
2500 use parking_lot::Mutex;
2501 use rand::prelude::*;
2502 use sea_orm::ConnectionTrait;
2503 use sqlx::migrate::MigrateDatabase;
2504 use std::sync::Arc;
2505
2506 pub struct TestDb {
2507 pub db: Option<Arc<Database>>,
2508 pub connection: Option<sqlx::AnyConnection>,
2509 }
2510
2511 impl TestDb {
2512 pub fn sqlite(background: Arc<Background>) -> Self {
2513 let url = format!("sqlite::memory:");
2514 let runtime = tokio::runtime::Builder::new_current_thread()
2515 .enable_io()
2516 .enable_time()
2517 .build()
2518 .unwrap();
2519
2520 let mut db = runtime.block_on(async {
2521 let mut options = ConnectOptions::new(url);
2522 options.max_connections(5);
2523 let db = Database::new(options).await.unwrap();
2524 let sql = include_str!(concat!(
2525 env!("CARGO_MANIFEST_DIR"),
2526 "/migrations.sqlite/20221109000000_test_schema.sql"
2527 ));
2528 db.pool
2529 .execute(sea_orm::Statement::from_string(
2530 db.pool.get_database_backend(),
2531 sql.into(),
2532 ))
2533 .await
2534 .unwrap();
2535 db
2536 });
2537
2538 db.background = Some(background);
2539 db.runtime = Some(runtime);
2540
2541 Self {
2542 db: Some(Arc::new(db)),
2543 connection: None,
2544 }
2545 }
2546
2547 pub fn postgres(background: Arc<Background>) -> Self {
2548 lazy_static! {
2549 static ref LOCK: Mutex<()> = Mutex::new(());
2550 }
2551
2552 let _guard = LOCK.lock();
2553 let mut rng = StdRng::from_entropy();
2554 let url = format!(
2555 "postgres://postgres@localhost/zed-test-{}",
2556 rng.gen::<u128>()
2557 );
2558 let runtime = tokio::runtime::Builder::new_current_thread()
2559 .enable_io()
2560 .enable_time()
2561 .build()
2562 .unwrap();
2563
2564 let mut db = runtime.block_on(async {
2565 sqlx::Postgres::create_database(&url)
2566 .await
2567 .expect("failed to create test db");
2568 let mut options = ConnectOptions::new(url);
2569 options
2570 .max_connections(5)
2571 .idle_timeout(Duration::from_secs(0));
2572 let db = Database::new(options).await.unwrap();
2573 let migrations_path = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
2574 db.migrate(Path::new(migrations_path), false).await.unwrap();
2575 db
2576 });
2577
2578 db.background = Some(background);
2579 db.runtime = Some(runtime);
2580
2581 Self {
2582 db: Some(Arc::new(db)),
2583 connection: None,
2584 }
2585 }
2586
2587 pub fn db(&self) -> &Arc<Database> {
2588 self.db.as_ref().unwrap()
2589 }
2590 }
2591
2592 impl Drop for TestDb {
2593 fn drop(&mut self) {
2594 let db = self.db.take().unwrap();
2595 if let sea_orm::DatabaseBackend::Postgres = db.pool.get_database_backend() {
2596 db.runtime.as_ref().unwrap().block_on(async {
2597 use util::ResultExt;
2598 let query = "
2599 SELECT pg_terminate_backend(pg_stat_activity.pid)
2600 FROM pg_stat_activity
2601 WHERE
2602 pg_stat_activity.datname = current_database() AND
2603 pid <> pg_backend_pid();
2604 ";
2605 db.pool
2606 .execute(sea_orm::Statement::from_string(
2607 db.pool.get_database_backend(),
2608 query.into(),
2609 ))
2610 .await
2611 .log_err();
2612 sqlx::Postgres::drop_database(db.options.get_url())
2613 .await
2614 .log_err();
2615 })
2616 }
2617 }
2618 }
2619}