Merge pull request #33 from middlewarehq/GROW-1224

Grow 1224: Add MTD Broker

Author: amoghjalan

apiserver/dora/service/code/__init__.py

@@ -0,0 +1,3 @@
+from sync import sync_code_repos
+from .integration import get_code_integration_service
+from .pr_filter import apply_pr_filter

apiserver/dora/service/code/pr_filter.py

@@ -1,11 +1,7 @@
-from dataclasses import dataclass
 from typing import List, Dict, Any
 
-from sqlalchemy import and_, or_
-from sqlalchemy.dialects.postgresql import Any
 from dora.service.settings.configuration_settings import get_settings_service
 from dora.service.settings.models import ExcludedPRsSetting
-
 from dora.store.models.code import PRFilter
 from dora.store.models.settings.configuration_settings import SettingType
 from dora.store.models.settings.enums import EntityType

apiserver/dora/service/code/sync/etl_handler.py

@@ -8,6 +8,10 @@
     CodeProviderETLHandler,
     CodeETLFactory,
 )
+from dora.service.merge_to_deploy_broker import (
+    get_merge_to_deploy_broker_utils_service,
+    MergeToDeployBrokerUtils,
+)
 from dora.store.models.code import OrgRepo, BookmarkType, Bookmark, PullRequest
 from dora.store.repos.code import CodeRepoService
 from dora.utils.log import LOG
@@ -18,9 +22,11 @@ def __init__(
         self,
         code_repo_service: CodeRepoService,
         etl_service: CodeProviderETLHandler,
+        mtd_broker: MergeToDeployBrokerUtils,
     ):
         self.code_repo_service = code_repo_service
         self.etl_service = etl_service
+        self.mtd_broker = mtd_broker
 
     def sync_org_repos(self, org_id: str):
         org_repos: List[OrgRepo] = self._sync_org_repos(org_id)
@@ -58,6 +64,9 @@ def _sync_repo_pull_requests_data(self, org_repo: OrgRepo) -> None:
                 pull_requests[-1].state_changed_at.astimezone(tz=pytz.UTC).isoformat()
             )
             self.code_repo_service.update_org_repo_bookmark(bookmark)
+            self.mtd_broker.pushback_merge_to_deploy_bookmark(
+                str(org_repo.id), pull_requests
+            )
             self.__sync_revert_prs_mapping(org_repo, pull_requests)
         except Exception as e:
             LOG.error(f"Error syncing pull requests for repo {org_repo.name}: {str(e)}")
@@ -91,13 +100,16 @@ def __get_org_repo_bookmark(self, org_repo: OrgRepo, default_sync_days: int = 31
 
 
 def sync_code_repos(org_id: str):
-    code_integration_service = get_code_integration_service()
-    code_repo_service = CodeRepoService()
+    code_providers: List[str] = get_code_integration_service().get_org_providers(org_id)
     etl_factory = CodeETLFactory(org_id)
-    for provider in code_integration_service.get_org_providers(org_id):
+
+    for provider in code_providers:
         try:
-            etl_handler = etl_factory(provider)
-            code_etl_handler = CodeETLHandler(code_repo_service, etl_handler)
+            code_etl_handler = CodeETLHandler(
+                CodeRepoService(),
+                etl_factory(provider),
+                get_merge_to_deploy_broker_utils_service(),
+            )
             code_etl_handler.sync_org_repos(org_id)
             LOG.info(f"Synced org repos for provider {provider}")
         except Exception as e:

apiserver/dora/service/deployments/__init__.py

@@ -0,0 +1 @@
+from .deployment_pr_mapper import DeploymentPRMapperService

apiserver/dora/service/merge_to_deploy_broker/__init__.py

@@ -0,0 +1,2 @@
+from .mtd_handler import process_merge_to_deploy_cache
+from .utils import get_merge_to_deploy_broker_utils_service, MergeToDeployBrokerUtils

apiserver/dora/service/merge_to_deploy_broker/mtd_handler.py

@@ -0,0 +1,126 @@
+from datetime import datetime
+from typing import List
+
+from dora.service.deployments import DeploymentPRMapperService
+from dora.store.models.code import (
+    PullRequest,
+    OrgRepo,
+    RepoWorkflow,
+    BookmarkMergeToDeployBroker,
+    RepoWorkflowRuns,
+    RepoWorkflowRunsStatus,
+)
+from dora.store.repos.code import CodeRepoService
+from dora.store.repos.workflows import WorkflowRepoService
+from dora.utils.lock import RedisLockService, get_redis_lock_service
+
+DEPLOYMENTS_TO_PROCESS = 500
+
+
+class MergeToDeployCacheHandler:
+    def __init__(
+        self,
+        org_id: str,
+        code_repo_service: CodeRepoService,
+        workflow_repo_service: WorkflowRepoService,
+        deployment_pr_mapper_service: DeploymentPRMapperService,
+        redis_lock_service: RedisLockService,
+    ):
+        self.org_id = org_id
+        self.code_repo_service = code_repo_service
+        self.workflow_repo_service = workflow_repo_service
+        self.deployment_pr_mapper_service = deployment_pr_mapper_service
+        self.redis_lock_service = redis_lock_service
+
+    def process_org_mtd(self):
+        org_repos: List[OrgRepo] = self.code_repo_service.get_active_org_repos(
+            self.org_id
+        )
+        for org_repo in org_repos:
+            try:
+                with self.redis_lock_service.acquire_lock(
+                    "{org_repo}:" + f"{str(org_repo.id)}:merge_to_deploy_broker"
+                ):
+                    self._process_deployments_for_merge_to_deploy_caching(
+                        str(org_repo.id)
+                    )
+            except Exception as e:
+                print(f"Error syncing workflow for repo {str(org_repo.id)}: {str(e)}")
+                continue
+
+    def _process_deployments_for_merge_to_deploy_caching(self, repo_id: str):
+        org_repo: OrgRepo = self.code_repo_service.get_repo_by_id(repo_id)
+        if not org_repo:
+            Exception(f"Repo with {repo_id} not found")
+
+        repo_workflows: List[
+            RepoWorkflow
+        ] = self.workflow_repo_service.get_repo_workflows_by_repo_id(repo_id)
+        if not repo_workflows:
+            return
+
+        broker_bookmark: BookmarkMergeToDeployBroker = (
+            self.code_repo_service.get_merge_to_deploy_broker_bookmark(repo_id)
+        )
+        if not broker_bookmark:
+            broker_bookmark = BookmarkMergeToDeployBroker(repo_id=repo_id)
+
+        bookmark_time: datetime = broker_bookmark.bookmark_date
+
+        repo_workflow_runs: List[
+            RepoWorkflowRuns
+        ] = self.workflow_repo_service.get_repo_workflow_runs_conducted_after_time(
+            repo_id, bookmark_time, DEPLOYMENTS_TO_PROCESS
+        )
+
+        if not repo_workflow_runs:
+            return
+
+        for repo_workflow_run in repo_workflow_runs:
+            try:
+                self.code_repo_service.get_merge_to_deploy_broker_bookmark(repo_id)
+                self._cache_prs_merge_to_deploy_for_repo_workflow_run(
+                    repo_id, repo_workflow_run
+                )
+                conducted_at: datetime = repo_workflow_run.conducted_at
+                broker_bookmark.bookmark = conducted_at.isoformat()
+                self.code_repo_service.update_merge_to_deploy_broker_bookmark(
+                    broker_bookmark
+                )
+            except Exception as e:
+                raise Exception(f"Error caching prs for repo {repo_id}: {str(e)}")
+
+    def _cache_prs_merge_to_deploy_for_repo_workflow_run(
+        self, repo_id: str, repo_workflow_run: RepoWorkflowRuns
+    ):
+        if repo_workflow_run.status != RepoWorkflowRunsStatus.SUCCESS:
+            return
+
+        conducted_at: datetime = repo_workflow_run.conducted_at
+        relevant_prs: List[
+            PullRequest
+        ] = self.code_repo_service.get_prs_in_repo_merged_before_given_date_with_merge_to_deploy_as_null(
+            repo_id, conducted_at
+        )
+        prs_to_update: List[
+            PullRequest
+        ] = self.deployment_pr_mapper_service.get_all_prs_deployed(
+            relevant_prs, repo_workflow_run
+        )
+
+        for pr in prs_to_update:
+            pr.merge_to_deploy = int(
+                (conducted_at - pr.state_changed_at).total_seconds()
+            )
+        self.code_repo_service.update_prs(prs_to_update)
+
+
+def process_merge_to_deploy_cache(org_id: str):
+    merge_to_deploy_cache_handler = MergeToDeployCacheHandler(
+        org_id,
+        CodeRepoService(),
+        WorkflowRepoService(),
+        DeploymentPRMapperService(),
+        get_redis_lock_service(),
+    )
+    merge_to_deploy_cache_handler.process_org_mtd()

apiserver/dora/service/merge_to_deploy_broker/utils.py

@@ -0,0 +1,49 @@
+from datetime import datetime
+from typing import List
+
+from dora.store.models.code import (
+    PullRequest,
+    PullRequestState,
+    BookmarkMergeToDeployBroker,
+)
+from dora.store.repos.code import CodeRepoService
+from dora.utils.lock import get_redis_lock_service, RedisLockService
+
+
+class MergeToDeployBrokerUtils:
+    def __init__(
+        self, code_repo_service: CodeRepoService, redis_lock_service: RedisLockService
+    ):
+        self.code_repo_service = code_repo_service
+        self.redis_lock_service = redis_lock_service
+
+    def pushback_merge_to_deploy_bookmark(self, repo_id: str, prs: List[PullRequest]):
+        with self.redis_lock_service.acquire_lock(
+            "{org_repo}:" + f"{repo_id}:merge_to_deploy_broker"
+        ):
+            self._pushback_merge_to_deploy_bookmark(repo_id, prs)
+
+    def _pushback_merge_to_deploy_bookmark(self, repo_id: str, prs: List[PullRequest]):
+        merged_prs = [pr for pr in prs if pr.state == PullRequestState.MERGED]
+        if not merged_prs:
+            return
+
+        min_merged_time: datetime = min([pr.state_changed_at for pr in merged_prs])
+
+        merge_to_deploy_broker_bookmark: BookmarkMergeToDeployBroker = (
+            self.code_repo_service.get_merge_to_deploy_broker_bookmark(repo_id)
+        )
+        if not merge_to_deploy_broker_bookmark:
+            merge_to_deploy_broker_bookmark = BookmarkMergeToDeployBroker(
+                repo_id=repo_id, bookmark=min_merged_time.isoformat()
+            )
+
+        self.code_repo_service.update_merge_to_deploy_broker_bookmark(
+            merge_to_deploy_broker_bookmark
+        )
+
+
+def get_merge_to_deploy_broker_utils_service():
+    return MergeToDeployBrokerUtils(
+        CodeRepoService(), redis_lock_service=get_redis_lock_service()
+    )

apiserver/dora/service/workflows/sync/etl_handler.py

@@ -2,7 +2,7 @@
 from typing import List, Tuple
 from uuid import uuid4
 
-from dora.service.code.integration import get_code_integration_service
+from dora.service.code import get_code_integration_service
 from dora.service.workflows.integration import get_workflows_integrations_service
 from dora.service.workflows.sync.etl_provider_handler import WorkflowProviderETLHandler
 from dora.store.models.code import (

apiserver/dora/store/repos/code.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from operator import and_
 from typing import Optional, List
 
@@ -16,6 +17,7 @@
     TeamRepos,
     PullRequestState,
     PRFilter,
+    BookmarkMergeToDeployBroker,
 )
 
 
@@ -51,6 +53,11 @@ def save_pull_requests_data(
         ]
         session.commit()
 
+    @rollback_on_exc
+    def update_prs(self, prs: List[PullRequest]):
+        [session.merge(pr) for pr in prs]
+        session.commit()
+
     @rollback_on_exc
     def save_revert_pr_mappings(
         self, revert_pr_mappings: List[PullRequestRevertPRMapping]
@@ -84,7 +91,7 @@ def get_repo_by_id(self, repo_id: str) -> Optional[OrgRepo]:
     def get_repo_pr_by_number(self, repo_id: str, pr_number) -> Optional[PullRequest]:
         return (
             session.query(PullRequest)
-            .options(defer("data"))
+            .options(defer(PullRequest.data))
             .filter(
                 and_(
                     PullRequest.repo_id == repo_id, PullRequest.number == str(pr_number)
@@ -100,7 +107,7 @@ def get_pr_events(self, pr_model: PullRequest):
 
         pr_events = (
             session.query(PullRequestEvent)
-            .options(defer("data"))
+            .options(defer(PullRequestEvent.data))
             .filter(PullRequestEvent.pull_request_id == pr_model.id)
             .all()
         )
@@ -121,7 +128,7 @@ def get_prs_by_head_branch_match_strings(
     ) -> List[PullRequest]:
         query = (
             session.query(PullRequest)
-            .options(defer("data"))
+            .options(defer(PullRequest.data))
             .filter(
                 and_(
                     PullRequest.repo_id.in_(repo_ids),
@@ -144,7 +151,7 @@ def get_reverted_prs_by_numbers(
     ) -> List[PullRequest]:
         query = (
             session.query(PullRequest)
-            .options(defer("data"))
+            .options(defer(PullRequest.data))
             .filter(
                 and_(
                     PullRequest.repo_id.in_(repo_ids),
@@ -156,13 +163,6 @@ def get_reverted_prs_by_numbers(
 
         return query.all()
 
-    @rollback_on_exc
-    def save_revert_pr_mappings(
-        self, revert_pr_mappings: List[PullRequestRevertPRMapping]
-    ):
-        [session.merge(revert_pr_map) for revert_pr_map in revert_pr_mappings]
-        session.commit()
-
     @rollback_on_exc
     def get_active_team_repos_by_team_id(self, team_id: str) -> List[TeamRepos]:
         return (
@@ -255,6 +255,39 @@ def get_team_repos(self, team_id) -> List[OrgRepo]:
         team_repo_ids = [tr.org_repo_id for tr in team_repos]
         return self.get_repos_by_ids(team_repo_ids)
 
+    @rollback_on_exc
+    def get_merge_to_deploy_broker_bookmark(
+        self, repo_id: str
+    ) -> BookmarkMergeToDeployBroker:
+        return (
+            session.query(BookmarkMergeToDeployBroker)
+            .filter(BookmarkMergeToDeployBroker.repo_id == repo_id)
+            .one_or_none()
+        )
+
+    @rollback_on_exc
+    def update_merge_to_deploy_broker_bookmark(
+        self, bookmark: BookmarkMergeToDeployBroker
+    ):
+        session.merge(bookmark)
+        session.commit()
+
+    @rollback_on_exc
+    def get_prs_in_repo_merged_before_given_date_with_merge_to_deploy_as_null(
+        self, repo_id: str, to_time: datetime
+    ):
+        return (
+            session.query(PullRequest)
+            .options(defer(PullRequest.data))
+            .filter(
+                PullRequest.repo_id == repo_id,
+                PullRequest.state == PullRequestState.MERGED,
+                PullRequest.state_changed_at <= to_time,
+                PullRequest.merge_to_deploy.is_(None),
+            )
+            .all()
+        )
+
     def _filter_prs_by_repo_ids(self, query, repo_ids: List[str]):
         return query.filter(PullRequest.repo_id.in_(repo_ids))