1141 lines · 44.9 KB · Rust
RawDownload
use anyhow::Context as _;use util::ResultExt; use super::*; impl Database {    /// Returns the count of all projects, excluding ones marked as admin.    pub async fn project_count_excluding_admins(&self) -> Result<usize> {        #[derive(Copy, Clone, Debug, EnumIter, DeriveColumn)]        enum QueryAs {            Count,        }         self.transaction(|tx| async move {            Ok(project::Entity::find()                .select_only()                .column_as(project::Column::Id.count(), QueryAs::Count)                .inner_join(user::Entity)                .filter(user::Column::Admin.eq(false))                .into_values::<_, QueryAs>()                .one(&*tx)                .await?                .unwrap_or(0i64) as usize)        })        .await    }     /// Shares a project with the given room.    pub async fn share_project(        &self,        room_id: RoomId,        connection: ConnectionId,        worktrees: &[proto::WorktreeMetadata],        is_ssh_project: bool,    ) -> Result<TransactionGuard<(ProjectId, proto::Room)>> {        self.room_transaction(room_id, |tx| async move {            let participant = room_participant::Entity::find()                .filter(                    Condition::all()                        .add(                            room_participant::Column::AnsweringConnectionId                                .eq(connection.id as i32),                        )                        .add(                            room_participant::Column::AnsweringConnectionServerId                                .eq(connection.owner_id as i32),                        ),                )                .one(&*tx)                .await?                .ok_or_else(|| anyhow!("could not find participant"))?;            if participant.room_id != room_id {                return Err(anyhow!("shared project on unexpected room"))?;            }            if !participant                .role                .unwrap_or(ChannelRole::Member)                .can_edit_projects()            {                return Err(anyhow!("guests cannot share projects"))?;            }             let project = project::ActiveModel {                room_id: ActiveValue::set(Some(participant.room_id)),                host_user_id: ActiveValue::set(Some(participant.user_id)),                host_connection_id: ActiveValue::set(Some(connection.id as i32)),                host_connection_server_id: ActiveValue::set(Some(ServerId(                    connection.owner_id as i32,                ))),                id: ActiveValue::NotSet,            }            .insert(&*tx)            .await?;             if !worktrees.is_empty() {                worktree::Entity::insert_many(worktrees.iter().map(|worktree| {                    worktree::ActiveModel {                        id: ActiveValue::set(worktree.id as i64),                        project_id: ActiveValue::set(project.id),                        abs_path: ActiveValue::set(worktree.abs_path.clone()),                        root_name: ActiveValue::set(worktree.root_name.clone()),                        visible: ActiveValue::set(worktree.visible),                        scan_id: ActiveValue::set(0),                        completed_scan_id: ActiveValue::set(0),                    }                }))                .exec(&*tx)                .await?;            }             let replica_id = if is_ssh_project { 1 } else { 0 };             project_collaborator::ActiveModel {                project_id: ActiveValue::set(project.id),                connection_id: ActiveValue::set(connection.id as i32),                connection_server_id: ActiveValue::set(ServerId(connection.owner_id as i32)),                user_id: ActiveValue::set(participant.user_id),                replica_id: ActiveValue::set(ReplicaId(replica_id)),                is_host: ActiveValue::set(true),                ..Default::default()            }            .insert(&*tx)            .await?;             let room = self.get_room(room_id, &tx).await?;            Ok((project.id, room))        })        .await    }     pub async fn delete_project(&self, project_id: ProjectId) -> Result<()> {        self.weak_transaction(|tx| async move {            project::Entity::delete_by_id(project_id).exec(&*tx).await?;            Ok(())        })        .await    }     /// Unshares the given project.    pub async fn unshare_project(        &self,        project_id: ProjectId,        connection: ConnectionId,    ) -> Result<TransactionGuard<(bool, Option<proto::Room>, Vec<ConnectionId>)>> {        self.project_transaction(project_id, |tx| async move {            let guest_connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;            let project = project::Entity::find_by_id(project_id)                .one(&*tx)                .await?                .ok_or_else(|| anyhow!("project not found"))?;            let room = if let Some(room_id) = project.room_id {                Some(self.get_room(room_id, &tx).await?)            } else {                None            };            if project.host_connection()? == connection {                return Ok((true, room, guest_connection_ids));            }            Err(anyhow!("cannot unshare a project hosted by another user"))?        })        .await    }     /// Updates the worktrees associated with the given project.    pub async fn update_project(        &self,        project_id: ProjectId,        connection: ConnectionId,        worktrees: &[proto::WorktreeMetadata],    ) -> Result<TransactionGuard<(Option<proto::Room>, Vec<ConnectionId>)>> {        self.project_transaction(project_id, |tx| async move {            let project = project::Entity::find_by_id(project_id)                .filter(                    Condition::all()                        .add(project::Column::HostConnectionId.eq(connection.id as i32))                        .add(                            project::Column::HostConnectionServerId.eq(connection.owner_id as i32),                        ),                )                .one(&*tx)                .await?                .ok_or_else(|| anyhow!("no such project"))?;             self.update_project_worktrees(project.id, worktrees, &tx)                .await?;             let guest_connection_ids = self.project_guest_connection_ids(project.id, &tx).await?;             let room = if let Some(room_id) = project.room_id {                Some(self.get_room(room_id, &tx).await?)            } else {                None            };             Ok((room, guest_connection_ids))        })        .await    }     pub(in crate::db) async fn update_project_worktrees(        &self,        project_id: ProjectId,        worktrees: &[proto::WorktreeMetadata],        tx: &DatabaseTransaction,    ) -> Result<()> {        if !worktrees.is_empty() {            worktree::Entity::insert_many(worktrees.iter().map(|worktree| worktree::ActiveModel {                id: ActiveValue::set(worktree.id as i64),                project_id: ActiveValue::set(project_id),                abs_path: ActiveValue::set(worktree.abs_path.clone()),                root_name: ActiveValue::set(worktree.root_name.clone()),                visible: ActiveValue::set(worktree.visible),                scan_id: ActiveValue::set(0),                completed_scan_id: ActiveValue::set(0),            }))            .on_conflict(                OnConflict::columns([worktree::Column::ProjectId, worktree::Column::Id])                    .update_column(worktree::Column::RootName)                    .to_owned(),            )            .exec(tx)            .await?;        }         worktree::Entity::delete_many()            .filter(worktree::Column::ProjectId.eq(project_id).and(                worktree::Column::Id.is_not_in(worktrees.iter().map(|worktree| worktree.id as i64)),            ))            .exec(tx)            .await?;         Ok(())    }     pub async fn update_worktree(        &self,        update: &proto::UpdateWorktree,        connection: ConnectionId,    ) -> Result<TransactionGuard<Vec<ConnectionId>>> {        if update.removed_entries.len() > proto::MAX_WORKTREE_UPDATE_MAX_CHUNK_SIZE            || update.updated_entries.len() > proto::MAX_WORKTREE_UPDATE_MAX_CHUNK_SIZE        {            return Err(anyhow!(                "invalid worktree update. removed entries: {}, updated entries: {}",                update.removed_entries.len(),                update.updated_entries.len()            ))?;        }         let project_id = ProjectId::from_proto(update.project_id);        let worktree_id = update.worktree_id as i64;        self.project_transaction(project_id, |tx| async move {            // Ensure the update comes from the host.            let _project = project::Entity::find_by_id(project_id)                .filter(                    Condition::all()                        .add(project::Column::HostConnectionId.eq(connection.id as i32))                        .add(                            project::Column::HostConnectionServerId.eq(connection.owner_id as i32),                        ),                )                .one(&*tx)                .await?                .ok_or_else(|| anyhow!("no such project: {project_id}"))?;             // Update metadata.            worktree::Entity::update(worktree::ActiveModel {                id: ActiveValue::set(worktree_id),                project_id: ActiveValue::set(project_id),                root_name: ActiveValue::set(update.root_name.clone()),                scan_id: ActiveValue::set(update.scan_id as i64),                completed_scan_id: if update.is_last_update {                    ActiveValue::set(update.scan_id as i64)                } else {                    ActiveValue::default()                },                abs_path: ActiveValue::set(update.abs_path.clone()),                ..Default::default()            })            .exec(&*tx)            .await?;             if !update.updated_entries.is_empty() {                worktree_entry::Entity::insert_many(update.updated_entries.iter().map(|entry| {                    let mtime = entry.mtime.clone().unwrap_or_default();                    worktree_entry::ActiveModel {                        project_id: ActiveValue::set(project_id),                        worktree_id: ActiveValue::set(worktree_id),                        id: ActiveValue::set(entry.id as i64),                        is_dir: ActiveValue::set(entry.is_dir),                        path: ActiveValue::set(entry.path.clone()),                        inode: ActiveValue::set(entry.inode as i64),                        mtime_seconds: ActiveValue::set(mtime.seconds as i64),                        mtime_nanos: ActiveValue::set(mtime.nanos as i32),                        canonical_path: ActiveValue::set(entry.canonical_path.clone()),                        is_ignored: ActiveValue::set(entry.is_ignored),                        is_external: ActiveValue::set(entry.is_external),                        git_status: ActiveValue::set(entry.git_status.map(|status| status as i64)),                        is_deleted: ActiveValue::set(false),                        scan_id: ActiveValue::set(update.scan_id as i64),                        is_fifo: ActiveValue::set(entry.is_fifo),                    }                }))                .on_conflict(                    OnConflict::columns([                        worktree_entry::Column::ProjectId,                        worktree_entry::Column::WorktreeId,                        worktree_entry::Column::Id,                    ])                    .update_columns([                        worktree_entry::Column::IsDir,                        worktree_entry::Column::Path,                        worktree_entry::Column::Inode,                        worktree_entry::Column::MtimeSeconds,                        worktree_entry::Column::MtimeNanos,                        worktree_entry::Column::CanonicalPath,                        worktree_entry::Column::IsIgnored,                        worktree_entry::Column::GitStatus,                        worktree_entry::Column::ScanId,                    ])                    .to_owned(),                )                .exec(&*tx)                .await?;            }             if !update.removed_entries.is_empty() {                worktree_entry::Entity::update_many()                    .filter(                        worktree_entry::Column::ProjectId                            .eq(project_id)                            .and(worktree_entry::Column::WorktreeId.eq(worktree_id))                            .and(                                worktree_entry::Column::Id                                    .is_in(update.removed_entries.iter().map(|id| *id as i64)),                            ),                    )                    .set(worktree_entry::ActiveModel {                        is_deleted: ActiveValue::Set(true),                        scan_id: ActiveValue::Set(update.scan_id as i64),                        ..Default::default()                    })                    .exec(&*tx)                    .await?;            }             if !update.updated_repositories.is_empty() {                worktree_repository::Entity::insert_many(update.updated_repositories.iter().map(                    |repository| worktree_repository::ActiveModel {                        project_id: ActiveValue::set(project_id),                        worktree_id: ActiveValue::set(worktree_id),                        work_directory_id: ActiveValue::set(repository.work_directory_id as i64),                        scan_id: ActiveValue::set(update.scan_id as i64),                        branch: ActiveValue::set(repository.branch.clone()),                        is_deleted: ActiveValue::set(false),                    },                ))                .on_conflict(                    OnConflict::columns([                        worktree_repository::Column::ProjectId,                        worktree_repository::Column::WorktreeId,                        worktree_repository::Column::WorkDirectoryId,                    ])                    .update_columns([                        worktree_repository::Column::ScanId,                        worktree_repository::Column::Branch,                    ])                    .to_owned(),                )                .exec(&*tx)                .await?;            }             if !update.removed_repositories.is_empty() {                worktree_repository::Entity::update_many()                    .filter(                        worktree_repository::Column::ProjectId                            .eq(project_id)                            .and(worktree_repository::Column::WorktreeId.eq(worktree_id))                            .and(                                worktree_repository::Column::WorkDirectoryId                                    .is_in(update.removed_repositories.iter().map(|id| *id as i64)),                            ),                    )                    .set(worktree_repository::ActiveModel {                        is_deleted: ActiveValue::Set(true),                        scan_id: ActiveValue::Set(update.scan_id as i64),                        ..Default::default()                    })                    .exec(&*tx)                    .await?;            }             let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;            Ok(connection_ids)        })        .await    }     /// Updates the diagnostic summary for the given connection.    pub async fn update_diagnostic_summary(        &self,        update: &proto::UpdateDiagnosticSummary,        connection: ConnectionId,    ) -> Result<TransactionGuard<Vec<ConnectionId>>> {        let project_id = ProjectId::from_proto(update.project_id);        let worktree_id = update.worktree_id as i64;        self.project_transaction(project_id, |tx| async move {            let summary = update                .summary                .as_ref()                .ok_or_else(|| anyhow!("invalid summary"))?;             // Ensure the update comes from the host.            let project = project::Entity::find_by_id(project_id)                .one(&*tx)                .await?                .ok_or_else(|| anyhow!("no such project"))?;            if project.host_connection()? != connection {                return Err(anyhow!("can't update a project hosted by someone else"))?;            }             // Update summary.            worktree_diagnostic_summary::Entity::insert(worktree_diagnostic_summary::ActiveModel {                project_id: ActiveValue::set(project_id),                worktree_id: ActiveValue::set(worktree_id),                path: ActiveValue::set(summary.path.clone()),                language_server_id: ActiveValue::set(summary.language_server_id as i64),                error_count: ActiveValue::set(summary.error_count as i32),                warning_count: ActiveValue::set(summary.warning_count as i32),            })            .on_conflict(                OnConflict::columns([                    worktree_diagnostic_summary::Column::ProjectId,                    worktree_diagnostic_summary::Column::WorktreeId,                    worktree_diagnostic_summary::Column::Path,                ])                .update_columns([                    worktree_diagnostic_summary::Column::LanguageServerId,                    worktree_diagnostic_summary::Column::ErrorCount,                    worktree_diagnostic_summary::Column::WarningCount,                ])                .to_owned(),            )            .exec(&*tx)            .await?;             let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;            Ok(connection_ids)        })        .await    }     /// Starts the language server for the given connection.    pub async fn start_language_server(        &self,        update: &proto::StartLanguageServer,        connection: ConnectionId,    ) -> Result<TransactionGuard<Vec<ConnectionId>>> {        let project_id = ProjectId::from_proto(update.project_id);        self.project_transaction(project_id, |tx| async move {            let server = update                .server                .as_ref()                .ok_or_else(|| anyhow!("invalid language server"))?;             // Ensure the update comes from the host.            let project = project::Entity::find_by_id(project_id)                .one(&*tx)                .await?                .ok_or_else(|| anyhow!("no such project"))?;            if project.host_connection()? != connection {                return Err(anyhow!("can't update a project hosted by someone else"))?;            }             // Add the newly-started language server.            language_server::Entity::insert(language_server::ActiveModel {                project_id: ActiveValue::set(project_id),                id: ActiveValue::set(server.id as i64),                name: ActiveValue::set(server.name.clone()),            })            .on_conflict(                OnConflict::columns([                    language_server::Column::ProjectId,                    language_server::Column::Id,                ])                .update_column(language_server::Column::Name)                .to_owned(),            )            .exec(&*tx)            .await?;             let connection_ids = self.project_guest_connection_ids(project_id, &tx).await?;            Ok(connection_ids)        })        .await    }     /// Updates the worktree settings for the given connection.    pub async fn update_worktree_settings(        &self,        update: &proto::UpdateWorktreeSettings,        connection: ConnectionId,    ) -> Result<TransactionGuard<Vec<ConnectionId>>> {        let project_id = ProjectId::from_proto(update.project_id);        let kind = match update.kind {            Some(kind) => proto::LocalSettingsKind::from_i32(kind)                .with_context(|| format!("unknown worktree settings kind: {kind}"))?,            None => proto::LocalSettingsKind::Settings,        };        let kind = LocalSettingsKind::from_proto(kind);        self.project_transaction(project_id, |tx| async move {            // Ensure the update comes from the host.            let project = project::Entity::find_by_id(project_id)                .one(&*tx)                .await?                .ok_or_else(|| anyhow!("no such project"))?;            if project.host_connection()? != connection {                return Err(anyhow!("can't update a project hosted by someone else"))?;            }

Lines 1–500 of 1141.