Add processing for mr messages

Signed-off-by: Andrew Brain <[email protected]>

augur/application/db/data_parse.py

@@ -869,3 +869,19 @@ def extract_needed_gitlab_message_data(comment: dict, platform_id: int, repo_id:
     }
 
     return comment_dict
+
+# retrieve only the needed data for pr labels from the api response
+def extract_needed_gitlab_mr_message_ref_data(comment: dict, pull_request_id: int, repo_id: int, tool_source: str, tool_version: str, data_source: str) -> List[dict]:
+
+    pr_msg_ref = {
+        'pull_request_id': pull_request_id,
+        'pr_message_ref_src_comment_id': comment['id'],
+        'repo_id': repo_id,
+        'pr_message_ref_src_node_id': None,
+        'tool_source': tool_source,
+        'tool_version': tool_version,
+        'data_source': data_source
+    }
+
+    return pr_msg_ref 
+

augur/tasks/gitlab/merge_request_task.py

@@ -4,11 +4,12 @@
 from augur.tasks.init.celery_app import AugurCoreRepoCollectionTask
 from augur.tasks.gitlab.gitlab_api_handler import GitlabApiHandler
 from augur.tasks.gitlab.gitlab_task_session import GitlabTaskManifest
-from augur.application.db.data_parse import extract_needed_pr_data_from_gitlab_merge_request, extract_needed_merge_request_assignee_data, extract_needed_mr_label_data, extract_needed_pr_reviewer_data, extract_needed_mr_commit_data, extract_needed_mr_file_data, extract_needed_mr_metadata
+from augur.application.db.data_parse import extract_needed_pr_data_from_gitlab_merge_request, extract_needed_merge_request_assignee_data, extract_needed_mr_label_data, extract_needed_pr_reviewer_data, extract_needed_mr_commit_data, extract_needed_mr_file_data, extract_needed_mr_metadata, extract_needed_gitlab_mr_message_ref_data, extract_needed_gitlab_message_data
 from augur.tasks.github.util.util import get_owner_repo, add_key_value_pair_to_dicts
-from augur.application.db.models import PullRequest, PullRequestAssignee, PullRequestLabel, PullRequestReviewer, PullRequestMeta, PullRequestCommit, PullRequestFile, Repo
+from augur.application.db.models import PullRequest, PullRequestAssignee, PullRequestLabel, PullRequestReviewer, PullRequestMeta, PullRequestCommit, PullRequestFile, PullRequestMessageRef, Repo, Message
 from augur.application.db.util import execute_session_query
 
+platform_id = 2
 
 @celery.task(base=AugurCoreRepoCollectionTask)
 def collect_gitlab_merge_requests(repo_git: str) -> int:
@@ -134,18 +135,81 @@ def collect_merge_request_comments(mr_ids, repo_git) -> int:
     logger = logging.getLogger(collect_merge_request_comments.__name__) 
     with GitlabTaskManifest(logger) as manifest:
 
+        augur_db = manifest.augur_db
+
+        query = augur_db.session.query(Repo).filter(Repo.repo_git == repo_git)
+        repo_obj = execute_session_query(query, 'one')
+        repo_id = repo_obj.repo_id
+
         url = "https://gitlab.com/api/v4/projects/{owner}%2f{repo}/merge_requests/{id}/notes".format(owner=owner, repo=repo, id="{id}")
         comments = retrieve_merge_request_data(mr_ids, url, "comments", owner, repo, manifest.key_auth, logger, response_type="list")
 
         if comments:
             logger.info(f"Length of merge request comments: {len(comments)}")
-            logger.info(f"Mr comment: {comments[0]}")
-            #issue_ids = process_issues(issue_data, f"{owner}/{repo}: Gitlab Issue task", repo_id, logger, augur_db)
+            process_gitlab_mr_messages(comments, f"{owner}/{repo}: Gitlab mr messages task", repo_id, logger, augur_db)
         else:
             logger.info(f"{owner}/{repo} has no gitlab merge request comments")
 
 
+def process_gitlab_mr_messages(data, task_name, repo_id, logger, augur_db):
+
+    tool_source = "Gitlab mr comments"
+    tool_version = "2.0"
+    data_source = "Gitlab API"
+
+    # create mapping from mr number to pull request id of current mrs
+    mr_number_to_id_map = {}
+    mrs = augur_db.session.query(PullRequest).filter(PullRequest.repo_id == repo_id).all()
+    for mr in mrs:
+        mr_number_to_id_map[mr.pr_src_number] = mr.pull_request_id
+
+    message_dicts = []
+    message_ref_mapping_data = {}
+    for id, messages in data.items():
+
+        try:
+                pull_request_id = mr_number_to_id_map[id]
+        except KeyError:
+            logger.info(f"{task_name}: Could not find related mr")
+            logger.info(f"{task_name}: We were searching for mr number {id} in repo {repo_id}")
+            logger.info(f"{task_name}: Skipping")
+            continue
+
+        for message in messages:
+
+            mr_message_ref_data = extract_needed_gitlab_mr_message_ref_data(message, pull_request_id, repo_id, tool_source, tool_version, data_source)
+
+            message_ref_mapping_data[message["id"]] = {
+                "msg_ref_data": mr_message_ref_data
+            }
+
+            message_dicts.append(
+                extract_needed_gitlab_message_data(message, platform_id, repo_id, tool_source, tool_version, data_source)
+            )
+
+
+    logger.info(f"{task_name}: Inserting {len(message_dicts)} messages")
+    message_natural_keys = ["platform_msg_id"]
+    message_return_columns = ["msg_id", "platform_msg_id"]
+    message_string_fields = ["msg_text"]
+    message_return_data = augur_db.insert_data(message_dicts, Message, message_natural_keys, 
+                                                return_columns=message_return_columns, string_fields=message_string_fields)
+
+    mr_message_ref_dicts = []
+    for data in message_return_data:
+
+        augur_msg_id = data["msg_id"]
+        platform_message_id = data["platform_msg_id"]
+
+        ref = message_ref_mapping_data[platform_message_id]
+        message_ref_data = ref["msg_ref_data"]
+        message_ref_data["msg_id"] = augur_msg_id
+
+        mr_message_ref_dicts.append(message_ref_data)
 
+    logger.info(f"{task_name}: Inserting {len(mr_message_ref_dicts)} mr messages ref rows")
+    mr_message_ref_natural_keys = ["pull_request_id", "pr_message_ref_src_comment_id"]
+    augur_db.insert_data(mr_message_ref_dicts, PullRequestMessageRef, mr_message_ref_natural_keys)
 
 
 @celery.task(base=AugurCoreRepoCollectionTask)
@@ -347,7 +411,7 @@ def process_mr_files(data, task_name, repo_id, logger, augur_db):
 def retrieve_merge_request_data(ids, url, name, owner, repo, key_auth, logger, response_type):
 
     all_data = {}
-    issue_count = len(ids)
+    mr_count = len(ids)
     index = 1
 
     api_handler = GitlabApiHandler(key_auth, logger)
@@ -356,7 +420,7 @@ def retrieve_merge_request_data(ids, url, name, owner, repo, key_auth, logger, r
         if len(all_data) > 10:
             return all_data
 
-        print(f"Collecting {owner}/{repo} gitlab merge request {name} for merge request {index} of {issue_count}")
+        print(f"Collecting {owner}/{repo} gitlab merge request {name} for merge request {index} of {mr_count}")
         formatted_url = url.format(id=id)
 
         if response_type == "dict":