1274 lines · 51.7 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),                        git_status: ActiveValue::set(None),                        is_external: ActiveValue::set(entry.is_external),                        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::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),                            branch_summary: ActiveValue::Set(                                repository                                    .branch_summary                                    .as_ref()                                    .map(|summary| serde_json::to_string(summary).unwrap()),                            ),                            current_merge_conflicts: ActiveValue::Set(Some(                                serde_json::to_string(&repository.current_merge_conflicts).unwrap(),                            )),                        }                    },                ))                .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,                        worktree_repository::Column::BranchSummary,                        worktree_repository::Column::CurrentMergeConflicts,                    ])                    .to_owned(),                )                .exec(&*tx)                .await?;                 let has_any_statuses = update                    .updated_repositories                    .iter()                    .any(|repository| !repository.updated_statuses.is_empty());                 if has_any_statuses {                    worktree_repository_statuses::Entity::insert_many(                        update.updated_repositories.iter().flat_map(                            |repository: &proto::RepositoryEntry| {                                repository.updated_statuses.iter().map(|status_entry| {                                    let (repo_path, status_kind, first_status, second_status) =                                        proto_status_to_db(status_entry.clone());                                    worktree_repository_statuses::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),                                        is_deleted: ActiveValue::set(false),                                        repo_path: ActiveValue::set(repo_path),                                        status: ActiveValue::set(0),                                        status_kind: ActiveValue::set(status_kind),                                        first_status: ActiveValue::set(first_status),                                        second_status: ActiveValue::set(second_status),                                    }                                })                            },                        ),                    )                    .on_conflict(                        OnConflict::columns([                            worktree_repository_statuses::Column::ProjectId,                            worktree_repository_statuses::Column::WorktreeId,                            worktree_repository_statuses::Column::WorkDirectoryId,                            worktree_repository_statuses::Column::RepoPath,                        ])                        .update_columns([                            worktree_repository_statuses::Column::ScanId,                            worktree_repository_statuses::Column::StatusKind,                            worktree_repository_statuses::Column::FirstStatus,                            worktree_repository_statuses::Column::SecondStatus,                        ])                        .to_owned(),                    )                    .exec(&*tx)                    .await?;                }                 let has_any_removed_statuses = update                    .updated_repositories                    .iter()                    .any(|repository| !repository.removed_statuses.is_empty());                 if has_any_removed_statuses {                    worktree_repository_statuses::Entity::update_many()                        .filter(                            worktree_repository_statuses::Column::ProjectId                                .eq(project_id)                                .and(                                    worktree_repository_statuses::Column::WorktreeId                                        .eq(worktree_id),                                )                                .and(                                    worktree_repository_statuses::Column::RepoPath.is_in(                                        update.updated_repositories.iter().flat_map(|repository| {                                            repository.removed_statuses.iter()                                        }),                                    ),                                ),                        )                        .set(worktree_repository_statuses::ActiveModel {                            is_deleted: ActiveValue::Set(true),                            scan_id: ActiveValue::Set(update.scan_id as i64),                            ..Default::default()                        })                        .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),

Lines 1–500 of 1274.