db.rs

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