db.rs

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