From 2dcfba2a9020846377e88a8d6813979f00caa0bc Mon Sep 17 00:00:00 2001
From: joseph-sentry <joseph.sawaya@sentry.io>
Date: Tue, 19 Nov 2024 17:08:35 -0500
Subject: [PATCH] fix: upsert branch and pull when updating commit

when we update a commit, we want to ensure the commit's corresponding
pull and branch are also upserted to reflect whatever changes were made
to the commit
---
 tasks/commit_update.py                 | 65 +++++++++++++++++++++++++-
 tasks/tests/unit/test_commit_update.py | 19 ++++++++
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/tasks/commit_update.py b/tasks/commit_update.py
index 837963473..f1cd7ea74 100644
--- a/tasks/commit_update.py
+++ b/tasks/commit_update.py
@@ -4,7 +4,7 @@
 from shared.torngit.exceptions import TorngitClientError, TorngitRepoNotFoundError
 
 from app import celery_app
-from database.models import Commit
+from database.models import Branch, Commit, Pull
 from helpers.exceptions import RepositoryWithoutValidBotError
 from helpers.github_installation import get_installation_name_for_owner_for_task
 from services.repository import (
@@ -43,6 +43,69 @@ def run_impl(
             was_updated = possibly_update_commit_from_provider_info(
                 commit, repository_service
             )
+
+            # upsert pull
+            pull = (
+                db_session.query(Pull)
+                .filter(Pull.repoid == repoid, Pull.pullid == commit.pullid)
+                .first()
+            )
+
+            if pull is None:
+                pull = Pull(
+                    repoid=repoid,
+                    pullid=commit.pullid,
+                    author_id=commit.author_id,
+                    head=commit.commitid,
+                )
+                db_session.add(pull)
+            else:
+                previous_pull_head = (
+                    db_session.query(Commit)
+                    .filter(Commit.repoid == repoid, Commit.commitid == pull.head)
+                    .first()
+                )
+                if (
+                    previous_pull_head is None
+                    or previous_pull_head.deleted == True
+                    or previous_pull_head.timestamp < commit.timestamp
+                ):
+                    pull.head = commit.commitid
+
+            # upsert branch
+            branch = (
+                db_session.query(Branch)
+                .filter(Branch.repoid == repoid, Branch.branch == commit.branch)
+                .first()
+            )
+
+            if branch is None:
+                branch = Branch(
+                    repoid=repoid,
+                    branch=commit.branch,
+                    head=commit.commitid,
+                    authors=[commit.author_id],
+                )
+                db_session.add(branch)
+            else:
+                if commit.author_id not in branch.authors:
+                    branch.authors.append(commit.author_id)
+
+                previous_branch_head = (
+                    db_session.query(Commit)
+                    .filter(Commit.repoid == repoid, Commit.commitid == branch.head)
+                    .first()
+                )
+
+                if (
+                    previous_branch_head is None
+                    or previous_branch_head.deleted == True
+                    or previous_branch_head.timestamp < commit.timestamp
+                ):
+                    branch.head = commit.commitid
+
+            db_session.flush()
+
         except RepositoryWithoutValidBotError:
             log.warning(
                 "Unable to reach git provider because repo doesn't have a valid bot",
diff --git a/tasks/tests/unit/test_commit_update.py b/tasks/tests/unit/test_commit_update.py
index e51fb41df..44ab7c97e 100644
--- a/tasks/tests/unit/test_commit_update.py
+++ b/tasks/tests/unit/test_commit_update.py
@@ -5,6 +5,7 @@
     TorngitRepoNotFoundError,
 )
 
+from database.models import Branch, Pull
 from database.tests.factories import CommitFactory
 from helpers.exceptions import RepositoryWithoutValidBotError
 from tasks.commit_update import CommitUpdateTask
@@ -184,6 +185,8 @@ def test_update_commit_already_populated(
             repository__owner__username="test-acc9",
             repository__yaml={"codecov": {"max_report_age": "764y ago"}},
             repository__name="test_example",
+            branch="main",
+            pullid=1,
         )
         dbsession.add(commit)
         dbsession.flush()
@@ -193,3 +196,19 @@ def test_update_commit_already_populated(
         assert expected_result == result
         assert commit.message == "commit_msg"
         assert commit.parent_commit_id is None
+
+        branch = (
+            dbsession.query(Branch)
+            .filter(Branch.repoid == commit.repoid, Branch.branch == commit.branch)
+            .first()
+        )
+        assert branch.head == commit.commitid
+        assert branch.authors == [commit.author_id]
+
+        pull = (
+            dbsession.query(Pull)
+            .filter(Pull.repoid == commit.repoid, Pull.pullid == commit.pullid)
+            .first()
+        )
+        assert pull.head == commit.commitid
+        assert pull.author_id == commit.author_id