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