gitdataai/lib/service/pull_request/merge.rs

use db::sqlx;
use git::rpc::{proto as p, proto::merge_service_client::MergeServiceClient};
use serde::Deserialize;
use session::Session;

use crate::{
    AppService, error::AppError, git::rpc_err, metrics::with_op_metric,
    pull_request::types::PullRequestResponse, session_user,
};

#[derive(Debug, Clone, Deserialize, utoipa::ToSchema)]
pub struct MergePullRequest {
    pub method: Option<String>,
    pub commit_title: Option<String>,
    pub commit_message: Option<String>,
}

impl AppService {
    pub async fn pr_merge_analysis(
        &self,
        ctx: &Session,
        wk_name: &str,
        repo_name: &str,
        number: i64,
    ) -> Result<p::MergeAnalysisResponse, AppError> {
        let (repo_id, _) =
            self.pr_resolve_repo(ctx, wk_name, repo_name).await?;
        let pr = self.pr_resolve(repo_id, number).await?;

        let mut client = MergeServiceClient::new(self.git.clone());
        let resp = client
            .merge_analysis(tonic::Request::new(p::MergeAnalysisRequest {
                repo_id: repo_id.to_string(),
                oid_a: Some(p::ObjectId {
                    value: pr.source_sha.clone(),
                }),
                oid_b: Some(p::ObjectId {
                    value: pr.target_sha.clone(),
                }),
            }))
            .await
            .map_err(rpc_err)?
            .into_inner();
        Ok(resp)
    }

    pub async fn pr_merge_base(
        &self,
        ctx: &Session,
        wk_name: &str,
        repo_name: &str,
        number: i64,
    ) -> Result<p::MergeBaseResponse, AppError> {
        let (repo_id, _) =
            self.pr_resolve_repo(ctx, wk_name, repo_name).await?;
        let pr = self.pr_resolve(repo_id, number).await?;

        let mut client = MergeServiceClient::new(self.git.clone());
        let resp = client
            .merge_base(tonic::Request::new(p::MergeBaseRequest {
                repo_id: repo_id.to_string(),
                oid_a: Some(p::ObjectId {
                    value: pr.source_sha.clone(),
                }),
                oid_b: Some(p::ObjectId {
                    value: pr.target_sha.clone(),
                }),
            }))
            .await
            .map_err(rpc_err)?
            .into_inner();
        Ok(resp)
    }

    #[tracing::instrument(skip(self, ctx, params), fields(workspace = %wk_name, repo = %repo_name, pr = %number))]
    pub async fn pr_merge(
        &self,
        ctx: &Session,
        wk_name: &str,
        repo_name: &str,
        number: i64,
        params: MergePullRequest,
    ) -> Result<PullRequestResponse, AppError> {
        let method = params.method.unwrap_or_else(|| "merge".to_string());
        with_op_metric(&self.metrics.pr_merge_total, &[&method], async {
        let user_uid = session_user(ctx)?;
        let (repo_id, _) =
            self.pr_resolve_repo_admin(ctx, wk_name, repo_name).await?;
        let pr = self.pr_resolve(repo_id, number).await?;

        if pr.state != "open" {
            return Err(AppError::BadRequest(
                "pull request is not open".to_string(),
            ));
        }
        if pr.draft {
            return Err(AppError::BadRequest(
                "draft pull request cannot be merged".to_string(),
            ));
        }
        let now = chrono::Utc::now();

        let merge_result_sha = match method.as_str() {
            "merge" => {
                let mut client = MergeServiceClient::new(self.git.clone());
                let resp = client
                    .merge_commit(tonic::Request::new(p::MergeCommitRequest {
                        repo_id: repo_id.to_string(),
                        params: Some(p::MergeCommitParams {
                            their_commit: Some(p::ObjectId {
                                value: pr.source_sha.clone(),
                            }),
                            author: Some(p::CommitSignature {
                                name: format!("merge: {}", pr.title),
                                email: "noreply@gitdata.ai".to_string(),
                                time_secs: now.timestamp(),
                                offset_minutes: 0,
                            }),
                            committer: Some(p::CommitSignature {
                                name: "gitdata".to_string(),
                                email: "noreply@gitdata.ai".to_string(),
                                time_secs: now.timestamp(),
                                offset_minutes: 0,
                            }),
                            message: params.commit_message.unwrap_or_else(
                                || {
                                    format!(
                                        "Merge pull request #{}: {}",
                                        pr.number, pr.title
                                    )
                                },
                            ),
                            update_ref: Some(format!(
                                "refs/heads/{}",
                                pr.target_branch
                            )),
                            options: None,
                        }),
                    }))
                    .await
                    .map_err(rpc_err)?
                    .into_inner();
                resp.oid.map(|oid| oid.value).unwrap_or_default()
            }
            "squash" => {
                let mut client = MergeServiceClient::new(self.git.clone());
                let resp = client
                    .squash_commit(tonic::Request::new(
                        p::SquashCommitRequest {
                            repo_id: repo_id.to_string(),
                            params: Some(p::SquashCommitParams {
                                their_commit: Some(p::ObjectId {
                                    value: pr.source_sha.clone(),
                                }),
                                options: None,
                            }),
                        },
                    ))
                    .await
                    .map_err(rpc_err)?
                    .into_inner();
                resp.oid.map(|oid| oid.value).unwrap_or_default()
            }
            _ => {
                return Err(AppError::BadRequest(
                    "merge method must be 'merge' or 'squash'".to_string(),
                ));
            }
        };

        sqlx::query(
            "UPDATE pull_request SET state = 'merged', merged_by = $1, merged_at = $2, \
             target_sha = $3, updated_at = $2 WHERE id = $4",
        )
        .bind(user_uid)
        .bind(now)
        .bind(&merge_result_sha)
        .bind(pr.id)
        .execute(self.db.writer())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?;

        let pr = self.pr_resolve(repo_id, number).await?;
        self.pr_build_response(pr).await
        }).await
    }

    pub async fn pr_merge_abort(
        &self,
        ctx: &Session,
        wk_name: &str,
        repo_name: &str,
    ) -> Result<(), AppError> {
        let (repo_id, _) =
            self.pr_resolve_repo_admin(ctx, wk_name, repo_name).await?;

        let mut client = MergeServiceClient::new(self.git.clone());
        client
            .merge_abort(tonic::Request::new(p::MergeAbortRequest {
                repo_id: repo_id.to_string(),
            }))
            .await
            .map_err(rpc_err)?;

        Ok(())
    }

    pub async fn pr_update_branch(
        &self,
        ctx: &Session,
        wk_name: &str,
        repo_name: &str,
        number: i64,
    ) -> Result<PullRequestResponse, AppError> {
        let (repo_id, _) =
            self.pr_resolve_repo(ctx, wk_name, repo_name).await?;
        let pr = self.pr_resolve(repo_id, number).await?;

        if pr.state != "open" {
            return Err(AppError::BadRequest(
                "pull request is not open".to_string(),
            ));
        }

        let source_sha = self
            .branch_head_sha(pr.source_repo, &pr.source_branch)
            .await?;
        let target_sha =
            self.branch_head_sha(repo_id, &pr.target_branch).await?;
        let now = chrono::Utc::now();

        sqlx::query(
            "UPDATE pull_request SET source_sha = $1, target_sha = $2, updated_at = $3 WHERE id = $4",
        )
        .bind(&source_sha)
        .bind(&target_sha)
        .bind(now)
        .bind(pr.id)
        .execute(self.db.writer())
        .await
        .map_err(|e| AppError::DatabaseError(e.to_string()))?;

        let pr = self.pr_resolve(repo_id, number).await?;
        self.pr_build_response(pr).await
    }
}