support chunks in git export; rebase before merging to avoid merging manually

Closes #32 (closed)

arkindex_worker/git.py

@@ -6,51 +6,130 @@ from datetime import datetime
 from pathlib import Path
 
 import gitlab
+import requests
 import sh
 
 from arkindex_worker import logger
 
 NOTHING_TO_COMMIT_MSG = "nothing to commit, working tree clean"
+MR_HAS_CONFLICTS_ERROR_CODE = 406
 
 
 class GitlabHelper:
     """Helper class to save files to GitLab repository"""
 
-    def __init__(self, project_id, gitlab_url, gitlab_token, branch):
+    def __init__(
+        self,
+        project_id,
+        gitlab_url,
+        gitlab_token,
+        branch,
+        rebase_wait_period=1,
+        delete_source_branch=True,
+        max_rebase_tries=10,
+    ):
         """
-
         :param project_id: the id of the gitlab project
         :param gitlab_url: gitlab server url
         :param gitlab_token: gitlab private token of user with permission to accept merge requests
         :param branch: name of the branch to where the exported branch will be merged
+        :param rebase_wait_period: seconds to wait between each poll to check whether rebase has finished
+        :param delete_source_branch: should delete the source branch after merging?
+        :param max_rebase_tries: max number of tries to rebase when merging before giving up
         """
         self.project_id = project_id
         self.gitlab_url = gitlab_url
         self.gitlab_token = str(gitlab_token).strip()
         self.branch = branch
+        self.rebase_wait_period = rebase_wait_period
+        self.delete_source_branch = delete_source_branch
+        self.max_rebase_tries = max_rebase_tries
 
         logger.info("Creating a Gitlab client")
         self._api = gitlab.Gitlab(self.gitlab_url, private_token=self.gitlab_token)
         self.project = self._api.projects.get(self.project_id)
+        self.is_rebase_finished = False
 
-    def merge(self, branch_name, title):
-        """Create a merge request and try to merge"""
+    def merge(self, branch_name, title) -> bool:
+        """
+        Create a merge request and try to merge.
+        Always rebase first to avoid conflicts from MRs made in parallel
+        :param branch_name: source branch name
+        :param title: title of the merge request
+        :return: was the branch successfully merged?
+        """
+        mr = None
+        # always rebase first, because other workers might have merged already
+        for i in range(self.max_rebase_tries):
+            logger.info(f"Trying to merge, try nr: {i}")
+            try:
+                if mr is None:
+                    mr = self._create_merge_request(branch_name, title)
+
+                mr.rebase()
+                rebase_success = self._wait_for_rebase_to_finish(mr.iid)
+                if not rebase_success:
+                    logger.error("Rebase failed, won't be able to merge!")
+                    return False
+
+                mr.merge(should_remove_source_branch=self.delete_source_branch)
+                logger.info("Merge successful")
+                return True
+            except gitlab.GitlabMRClosedError as e:
+                if e.response_code == MR_HAS_CONFLICTS_ERROR_CODE:
+                    logger.info("Merge failed, trying to rebase and merge again.")
+                    continue
+                else:
+                    logger.error(f"Merge was not successful: {e}")
+                    return False
+            except gitlab.GitlabError as e:
+                logger.error(f"Gitlab error: {e}")
+                if 400 <= e.response_code < 500:
+                    # 4XX errors shouldn't be fixed by retrying
+                    raise e
+            except requests.exceptions.ConnectionError as e:
+                logger.error(f"Server connection error, will wait and retry: {e}")
+                time.sleep(self.rebase_wait_period)
+
+        return False
+
+    def _create_merge_request(self, branch_name, title):
         logger.info(f"Creating a merge request for {branch_name}")
+        # retry_transient_error will retry the request on 50X errors
+        # https://github.com/python-gitlab/python-gitlab/blob/265dbbdd37af88395574564aeb3fd0350288a18c/gitlab/__init__.py#L539
         mr = self.project.mergerequests.create(
             {
                 "source_branch": branch_name,
                 "target_branch": self.branch,
                 "title": title,
-            }
+            },
         )
-        logger.info("Attempting to merge")
-        try:
-            mr.merge()
-            logger.info("Merge successful")
+        return mr
+
+    def _get_merge_request(self, merge_request_id, include_rebase_in_progress=True):
+        return self.project.mergerequests.get(
+            merge_request_id, include_rebase_in_progress=include_rebase_in_progress
+        )
+
+    def _wait_for_rebase_to_finish(self, merge_request_id) -> bool:
+        """
+        Poll the merge request until it has finished rebasing
+        :param merge_request_id:
+        :return: rebase finished successfully?
+        """
+
+        logger.info("Checking if rebase has finished..")
+        self.is_rebase_finished = False
+        while not self.is_rebase_finished:
+            time.sleep(self.rebase_wait_period)
+            mr = self._get_merge_request(merge_request_id)
+            self.is_rebase_finished = not mr.rebase_in_progress
+        if mr.merge_error is None:
+            logger.info("Rebase has finished")
             return True
-        except gitlab.GitlabMRClosedError as e:
-            logger.error(f"Merge was not successful: {e}")
-            return False
+
+        logger.error(f"Rebase failed: {mr.merge_error}")
+        return False
 
 
 def make_backup(path):
@@ -220,12 +299,19 @@ class GitHelper:
         # move exported files to git directory
         file_count = self._move_files_to_git(export_out_dir)
 
+        # use timestamp to avoid branch name conflicts with multiple chunks
+        current_timestamp = datetime.isoformat(datetime.now())
+        # ":" is not allowed in a branch name
+        branch_timestamp = current_timestamp.replace(":", ".")
         # add files to a new branch
-        branch_name = f"workflow_{self.workflow_id}"
+        branch_name = f"workflow_{self.workflow_id}_{branch_timestamp}"
         self._git.checkout("-b", branch_name)
         self._git.add("-A")
         try:
-            self._git.commit("-m", f"Exported files from workflow: {self.workflow_id}")
+            self._git.commit(
+                "-m",
+                f"Exported files from workflow: {self.workflow_id} at {current_timestamp}",
+            )
         except sh.ErrorReturnCode as e:
             if NOTHING_TO_COMMIT_MSG in str(e.stdout):
                 logger.warning("Nothing to commit (no changes)")
@@ -249,7 +335,11 @@ class GitHelper:
         self._git.push("-u", "origin", "HEAD")
 
         if self.gitlab_helper:
-            self.gitlab_helper.merge(branch_name, f"Merge {branch_name}")
+            try:
+                self.gitlab_helper.merge(branch_name, f"Merge {branch_name}")
+            except Exception as e:
+                logger.error(f"Merge failed: {e}")
+                raise e
         else:
             logger.info(
                 "No gitlab_helper defined, not trying to merge the pushed branch"