774 lines · 23.2 KB · Rust
RawDownload
mod ids;pub mod queries;mod tables; use crate::{Error, Result};use anyhow::{Context as _, anyhow};use cloud_api_types::{ExtensionMetadata, ExtensionProvides};use collections::{BTreeMap, BTreeSet, HashMap, HashSet};use dashmap::DashMap;use futures::StreamExt;use project_repository_statuses::StatusKind;use rpc::{    ConnectionId,    proto::{self},};use sea_orm::{    ActiveValue, Condition, ConnectionTrait, DatabaseConnection, DatabaseTransaction,    FromQueryResult, IntoActiveModel, IsolationLevel, JoinType, QueryOrder, QuerySelect, Statement,    TransactionTrait,    entity::prelude::*,    sea_query::{Alias, Expr, OnConflict},};use semver::Version;use serde::{Deserialize, Serialize};use std::ops::RangeInclusive;use std::{    future::Future,    marker::PhantomData,    ops::{Deref, DerefMut},    rc::Rc,    sync::Arc,};use time::PrimitiveDateTime;use tokio::sync::{Mutex, OwnedMutexGuard};use util::paths::PathStyle;use worktree_settings_file::LocalSettingsKind; pub use ids::*;pub use sea_orm::ConnectOptions;pub use tables::user::Model as User;pub use tables::*; #[cfg(feature = "test-support")]pub struct DatabaseTestOptions {    pub executor: gpui::BackgroundExecutor,    pub runtime: tokio::runtime::Runtime,    pub query_failure_probability: parking_lot::Mutex<f64>,} /// Database gives you a handle that lets you access the database./// It handles pooling internally.pub struct Database {    pub options: ConnectOptions,    pub pool: DatabaseConnection,    rooms: DashMap<RoomId, Arc<Mutex<()>>>,    projects: DashMap<ProjectId, Arc<Mutex<()>>>,    notification_kinds_by_id: HashMap<NotificationKindId, &'static str>,    notification_kinds_by_name: HashMap<String, NotificationKindId>,    #[cfg(feature = "test-support")]    pub test_options: Option<DatabaseTestOptions>,} // The `Database` type has so many methods that its impl blocks are split into// separate files in the `queries` folder.impl Database {    /// Connects to the database with the given options    pub async fn new(options: ConnectOptions) -> Result<Self> {        sqlx::any::install_default_drivers();        Ok(Self {            options: options.clone(),            pool: sea_orm::Database::connect(options).await?,            rooms: DashMap::with_capacity(16384),            projects: DashMap::with_capacity(16384),            notification_kinds_by_id: HashMap::default(),            notification_kinds_by_name: HashMap::default(),            #[cfg(feature = "test-support")]            test_options: None,        })    }     pub fn options(&self) -> &ConnectOptions {        &self.options    }     #[cfg(feature = "test-support")]    pub fn reset(&self) {        self.rooms.clear();        self.projects.clear();    }     pub async fn transaction<F, Fut, T>(&self, f: F) -> Result<T>    where        F: Send + Fn(TransactionHandle) -> Fut,        Fut: Send + Future<Output = Result<T>>,    {        let body = async {            let (tx, result) = self.with_transaction(&f).await?;            match result {                Ok(result) => match tx.commit().await.map_err(Into::into) {                    Ok(()) => Ok(result),                    Err(error) => Err(error),                },                Err(error) => {                    tx.rollback().await?;                    Err(error)                }            }        };         self.run(body).await    }     /// The same as room_transaction, but if you need to only optionally return a Room.    async fn optional_room_transaction<F, Fut, T>(        &self,        f: F,    ) -> Result<Option<TransactionGuard<T>>>    where        F: Send + Fn(TransactionHandle) -> Fut,        Fut: Send + Future<Output = Result<Option<(RoomId, T)>>>,    {        let body = async {            let (tx, result) = self.with_transaction(&f).await?;            match result {                Ok(Some((room_id, data))) => {                    let lock = self.rooms.entry(room_id).or_default().clone();                    let _guard = lock.lock_owned().await;                    match tx.commit().await.map_err(Into::into) {                        Ok(()) => Ok(Some(TransactionGuard {                            data,                            _guard,                            _not_send: PhantomData,                        })),                        Err(error) => Err(error),                    }                }                Ok(None) => match tx.commit().await.map_err(Into::into) {                    Ok(()) => Ok(None),                    Err(error) => Err(error),                },                Err(error) => {                    tx.rollback().await?;                    Err(error)                }            }        };         self.run(body).await    }     async fn project_transaction<F, Fut, T>(        &self,        project_id: ProjectId,        f: F,    ) -> Result<TransactionGuard<T>>    where        F: Send + Fn(TransactionHandle) -> Fut,        Fut: Send + Future<Output = Result<T>>,    {        let room_id = Database::room_id_for_project(self, project_id).await?;        let body = async {            let lock = if let Some(room_id) = room_id {                self.rooms.entry(room_id).or_default().clone()            } else {                self.projects.entry(project_id).or_default().clone()            };            let _guard = lock.lock_owned().await;            let (tx, result) = self.with_transaction(&f).await?;            match result {                Ok(data) => match tx.commit().await.map_err(Into::into) {                    Ok(()) => Ok(TransactionGuard {                        data,                        _guard,                        _not_send: PhantomData,                    }),                    Err(error) => Err(error),                },                Err(error) => {                    tx.rollback().await?;                    Err(error)                }            }        };         self.run(body).await    }     /// room_transaction runs the block in a transaction. It returns a RoomGuard, that keeps    /// the database locked until it is dropped. This ensures that updates sent to clients are    /// properly serialized with respect to database changes.    async fn room_transaction<F, Fut, T>(        &self,        room_id: RoomId,        f: F,    ) -> Result<TransactionGuard<T>>    where        F: Send + Fn(TransactionHandle) -> Fut,        Fut: Send + Future<Output = Result<T>>,    {        let body = async {            let lock = self.rooms.entry(room_id).or_default().clone();            let _guard = lock.lock_owned().await;            let (tx, result) = self.with_transaction(&f).await?;            match result {                Ok(data) => match tx.commit().await.map_err(Into::into) {                    Ok(()) => Ok(TransactionGuard {                        data,                        _guard,                        _not_send: PhantomData,                    }),                    Err(error) => Err(error),                },                Err(error) => {                    tx.rollback().await?;                    Err(error)                }            }        };         self.run(body).await    }     async fn with_transaction<F, Fut, T>(&self, f: &F) -> Result<(DatabaseTransaction, Result<T>)>    where        F: Send + Fn(TransactionHandle) -> Fut,        Fut: Send + Future<Output = Result<T>>,    {        let tx = self            .pool            .begin_with_config(Some(IsolationLevel::ReadCommitted), None)            .await?;         let mut tx = Arc::new(Some(tx));        let result = f(TransactionHandle(tx.clone())).await;        let tx = Arc::get_mut(&mut tx)            .and_then(|tx| tx.take())            .context("couldn't complete transaction because it's still in use")?;         Ok((tx, result))    }     async fn run<F, T>(&self, future: F) -> Result<T>    where        F: Future<Output = Result<T>>,    {        #[cfg(feature = "test-support")]        {            let test_options = self.test_options.as_ref().unwrap();            test_options.executor.simulate_random_delay().await;            let fail_probability = *test_options.query_failure_probability.lock();            if test_options.executor.rng().random_bool(fail_probability) {                return Err(anyhow!("simulated query failure"))?;            }             test_options.runtime.block_on(future)        }         #[cfg(not(feature = "test-support"))]        {            future.await        }    }} /// A handle to a [`DatabaseTransaction`].pub struct TransactionHandle(pub(crate) Arc<Option<DatabaseTransaction>>); impl Deref for TransactionHandle {    type Target = DatabaseTransaction;     fn deref(&self) -> &Self::Target {        self.0.as_ref().as_ref().unwrap()    }} /// [`TransactionGuard`] keeps a database transaction alive until it is dropped./// It wraps data that depends on the state of the database and prevents an additional/// transaction from starting that would invalidate that data.pub struct TransactionGuard<T> {    data: T,    _guard: OwnedMutexGuard<()>,    _not_send: PhantomData<Rc<()>>,} impl<T> Deref for TransactionGuard<T> {    type Target = T;     fn deref(&self) -> &T {        &self.data    }} impl<T> DerefMut for TransactionGuard<T> {    fn deref_mut(&mut self) -> &mut T {        &mut self.data    }} impl<T> TransactionGuard<T> {    /// Returns the inner value of the guard.    pub fn into_inner(self) -> T {        self.data    }} #[derive(Clone, Debug, PartialEq, Eq)]pub enum Contact {    Accepted { user_id: UserId, busy: bool },    Outgoing { user_id: UserId },    Incoming { user_id: UserId },} impl Contact {    pub fn user_id(&self) -> UserId {        match self {            Contact::Accepted { user_id, .. } => *user_id,            Contact::Outgoing { user_id } => *user_id,            Contact::Incoming { user_id, .. } => *user_id,        }    }} pub type NotificationBatch = Vec<(UserId, proto::Notification)>; pub struct CreatedChannelMessage {    pub message_id: MessageId,    pub participant_connection_ids: HashSet<ConnectionId>,    pub notifications: NotificationBatch,} pub struct UpdatedChannelMessage {    pub message_id: MessageId,    pub participant_connection_ids: Vec<ConnectionId>,    pub notifications: NotificationBatch,    pub reply_to_message_id: Option<MessageId>,    pub timestamp: PrimitiveDateTime,    pub deleted_mention_notification_ids: Vec<NotificationId>,    pub updated_mention_notifications: Vec<rpc::proto::Notification>,} #[derive(Clone, Debug, PartialEq, Eq, FromQueryResult, Serialize, Deserialize)]pub struct Invite {    pub email_address: String,    pub email_confirmation_code: String,} #[derive(Clone, Debug, Deserialize)]pub struct NewSignup {    pub email_address: String,    pub platform_mac: bool,    pub platform_windows: bool,    pub platform_linux: bool,    pub editor_features: Vec<String>,    pub programming_languages: Vec<String>,    pub device_id: Option<String>,    pub added_to_mailing_list: bool,    pub created_at: Option<DateTime>,} #[derive(Clone, Debug, PartialEq, Deserialize, Serialize, FromQueryResult)]pub struct WaitlistSummary {    pub count: i64,    pub linux_count: i64,    pub mac_count: i64,    pub windows_count: i64,    pub unknown_count: i64,} /// The parameters to create a new user.#[derive(Debug, Serialize, Deserialize)]pub struct NewUserParams {    pub github_login: String,    pub github_user_id: i32,} /// The result of creating a new user.#[derive(Debug)]pub struct NewUserResult {    pub user_id: UserId,} /// The result of updating a channel membership.#[derive(Debug)]pub struct MembershipUpdated {    pub channel_id: ChannelId,    pub new_channels: ChannelsForUser,    pub removed_channels: Vec<ChannelId>,} /// The result of setting a member's role.#[derive(Debug)] pub enum SetMemberRoleResult {    InviteUpdated(Channel),    MembershipUpdated(MembershipUpdated),} /// The result of inviting a member to a channel.#[derive(Debug)]pub struct InviteMemberResult {    pub channel: Channel,    pub notifications: NotificationBatch,} #[derive(Debug)]pub struct RespondToChannelInvite {    pub membership_update: Option<MembershipUpdated>,    pub notifications: NotificationBatch,} #[derive(Debug)]pub struct RemoveChannelMemberResult {    pub membership_update: MembershipUpdated,    pub notification_id: Option<NotificationId>,} #[derive(Debug, PartialEq, Eq, Hash)]pub struct Channel {    pub id: ChannelId,    pub name: String,    pub visibility: ChannelVisibility,    /// parent_path is the channel ids from the root to this one (not including this one)    pub parent_path: Vec<ChannelId>,    pub channel_order: i32,} impl Channel {    pub fn from_model(value: channel::Model) -> Self {        Channel {            id: value.id,            visibility: value.visibility,            name: value.clone().name,            parent_path: value.ancestors().collect(),            channel_order: value.channel_order,        }    }     pub fn to_proto(&self) -> proto::Channel {        proto::Channel {            id: self.id.to_proto(),            name: self.name.clone(),            visibility: self.visibility.into(),            parent_path: self.parent_path.iter().map(|c| c.to_proto()).collect(),            channel_order: self.channel_order,        }    }     pub fn root_id(&self) -> ChannelId {        self.parent_path.first().copied().unwrap_or(self.id)    }} #[derive(Debug, PartialEq, Eq, Hash)]pub struct ChannelMember {    pub role: ChannelRole,    pub user_id: UserId,    pub kind: proto::channel_member::Kind,} impl ChannelMember {    pub fn to_proto(&self) -> proto::ChannelMember {        proto::ChannelMember {            role: self.role.into(),            user_id: self.user_id.to_proto(),            kind: self.kind.into(),        }    }} #[derive(Debug, PartialEq)]pub struct ChannelsForUser {    pub channels: Vec<Channel>,    pub channel_memberships: Vec<channel_member::Model>,    pub channel_participants: HashMap<ChannelId, Vec<UserId>>,    pub invited_channels: Vec<Channel>,     pub observed_buffer_versions: Vec<proto::ChannelBufferVersion>,    pub latest_buffer_versions: Vec<proto::ChannelBufferVersion>,} #[derive(Debug)]pub struct RejoinedChannelBuffer {    pub buffer: proto::RejoinedChannelBuffer,    pub old_connection_id: ConnectionId,} #[derive(Clone)]pub struct JoinRoom {    pub room: proto::Room,    pub channel: Option<channel::Model>,} pub struct RejoinedRoom {    pub room: proto::Room,    pub rejoined_projects: Vec<RejoinedProject>,    pub reshared_projects: Vec<ResharedProject>,    pub channel: Option<channel::Model>,} pub struct ResharedProject {

Lines 1–500 of 774.