Merge "Allow to reuse the code handling the branch cache"

zuul/connection/__init__.py

@@ -13,8 +13,12 @@
 # under the License.
 
 import abc
+import logging
+
+from typing import Dict, List, Optional
 
 from zuul.lib.logutil import get_annotated_logger
+from zuul.model import Project, Tenant
 
 
 class BaseConnection(object, metaclass=abc.ABCMeta):
@@ -33,6 +37,8 @@ class BaseConnection(object, metaclass=abc.ABCMeta):
     into. For example, a trigger will likely require some kind of query method
     while a reporter may need a review method."""
 
+    log = logging.getLogger('zuul.BaseConnection')
+
     def __init__(self, driver, connection_name, connection_config):
         # connection_name is the name given to this connection in zuul.ini
         # connection_config is a dictionary of config_section from zuul.ini for
@@ -70,11 +76,11 @@ class BaseConnection(object, metaclass=abc.ABCMeta):
     def registerScheduler(self, sched):
         self.sched = sched
 
-    def clearBranchCache(self):
+    def clearCache(self):
-        """Clear the branch cache for this connection.
+        """Clear the cache for this connection.
 
         This is called immediately prior to performing a full
-        reconfiguration.  The branch cache should be cleared so that a
+        reconfiguration. The cache should be cleared so that a
         full reconfiguration can be used to correct any errors in
         cached data.
 
@@ -118,3 +124,170 @@ class BaseConnection(object, metaclass=abc.ABCMeta):
             "name": self.connection_name,
             "driver": self.driver.name,
         }
+
+
+class CachedBranchConnection(BaseConnection):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._project_branch_cache_exclude_unprotected = {}
+        self._project_branch_cache_include_unprotected = {}
+
+    @abc.abstractmethod
+    def isBranchProtected(self, project_name: str, branch_name: str,
+                          zuul_event_id) -> Optional[bool]:
+        """Return if the branch is protected or None if the branch is unknown.
+
+        :param str project_name:
+            The name of the project.
+        :param str branch_name:
+            The name of the branch.
+        :param zuul_event_id:
+            Unique id associated to the handled event.
+        """
+        pass
+
+    @abc.abstractmethod
+    def _fetchProjectBranches(self, project: Project,
+                              exclude_unprotected: bool) -> List[str]:
+        pass
+
+    def clearConnectionCacheOnBranchEvent(self, event):
+        """Update event and clear connection cache if needed.
+
+        This checks whether the event created or deleted a branch so
+        that Zuul may know to perform a reconfiguration on the
+        project. Drivers must call this method when a branch event is
+        received.
+
+        :param event:
+            The event, inherit from `zuul.model.TriggerEvent` class.
+        """
+        if event.oldrev == '0' * 40:
+            event.branch_created = True
+        elif event.newrev == '0' * 40:
+            event.branch_deleted = True
+        else:
+            event.branch_updated = True
+
+        project = self.source.getProject(event.project_name)
+        if event.branch:
+            if event.branch_deleted:
+                # We currently cannot determine if a deleted branch was
+                # protected so we need to assume it was. GitHub/GitLab don't
+                # allow deletion of protected branches but we don't get a
+                # notification about branch protection settings. Thus we don't
+                # know if branch protection has been disabled before deletion
+                # of the branch.
+                # FIXME(tobiash): Find a way to handle that case
+                self.clearProjectCache(project)
+            elif event.branch_created:
+                # A new branch never can be protected because that needs to be
+                # configured after it has been created.
+                self.clearProjectCache(project)
+        return event
+
+    def getCachedBranches(self, exclude_unprotected) -> Dict[str, List[str]]:
+        """Get the connection cache: the branches foreach project.
+
+        :param bool exclude_unprotected:
+            Specify whether the cache excludes or contains unprotected
+            branches.
+        :returns: A dictionary where keys are project names and values are list
+            of branch names.
+        """
+        if exclude_unprotected:
+            cache = self._project_branch_cache_exclude_unprotected
+        else:
+            cache = self._project_branch_cache_include_unprotected
+
+        return cache
+
+    def getProjectBranches(self, project: Project,
+                           tenant: Tenant) -> List[str]:
+        """Get the branch names for the given project.
+
+        :param zuul.model.Project project:
+            The project for which the branches are returned.
+        :param zuul.model.Tenant tenant:
+            The related tenant.
+
+        :returns: The list of branch names.
+        """
+        exclude_unprotected = tenant.getExcludeUnprotectedBranches(project)
+        cache = self.getCachedBranches(exclude_unprotected)
+        branches = cache.get(project.name)
+
+        if branches is not None:
+            return branches
+
+        branches = self._fetchProjectBranches(project, exclude_unprotected)
+        self.log.info("Got branches for %s" % project.name)
+
+        cache[project.name] = branches
+        return branches
+
+    def clearCache(self) -> None:
+        """Clear the connection cache for all projects.
+
+        This method is called by the scheduler in order to perform a full
+        reconfiguration.
+        """
+        self.log.debug("Clearing branch cache for all branches: %s",
+                       self.connection_name)
+        self._project_branch_cache_exclude_unprotected = {}
+        self._project_branch_cache_include_unprotected = {}
+
+    def clearProjectCache(self, project: Project) -> None:
+        """Clear the connection cache for this project.
+        """
+
+        self.log.debug("Clearing cache for %s:%s", self.connection_name,
+                       project.name)
+        self._project_branch_cache_exclude_unprotected.pop(project.name, None)
+        self._project_branch_cache_include_unprotected.pop(project.name, None)
+
+    def checkBranchCache(self, project_name: str, event) -> None:
+        """Clear the cache for a project when a branch event is processed
+
+        This method must be called when a branch event is processed: if the
+        event references a branch and the unprotected branches are excluded,
+        the branch protection status could have been changed.
+
+        :params str project_name:
+            The project name.
+        :params event:
+            The event, inherit from `zuul.model.TriggerEvent` class.
+        """
+        protected = self.isBranchProtected(project_name, event.branch,
+                                           zuul_event_id=event)
+        if protected is not None:
+
+            # If the branch appears in the exclude_unprotected cache but
+            # is unprotected, clear the exclude cache.
+
+            # If the branch does not appear in the exclude_unprotected
+            # cache but is protected, clear the exclude cache.
+
+            # All branches should always appear in the include_unprotected
+            # cache, so we never clear it.
+
+            cache = self._project_branch_cache_exclude_unprotected
+            branches = cache.get(project_name, [])
+            if (event.branch in branches) and (not protected):
+                self.log.debug("Clearing protected branch cache for %s",
+                               project_name)
+                cache.pop(project_name, None)
+            if (event.branch not in branches) and (protected):
+                self.log.debug("Clearing protected branch cache for %s",
+                               project_name)
+                cache.pop(project_name, None)
+
+            event.branch_protected = protected
+        else:
+            # This can happen if the branch was deleted in GitHub/GitLab.
+            # In this case we assume that the branch COULD have
+            # been protected before. The cache update is handled by
+            # the push event, so we don't touch the cache here
+            # again.
+            event.branch_protected = True

zuul/driver/gerrit/gerritconnection.py

@@ -669,7 +669,7 @@ class GerritConnection(BaseConnection):
     def addProject(self, project: Project) -> None:
         self.projects[project.name] = project
 
-    def clearBranchCache(self):
+    def clearCache(self):
         self._project_branch_cache = {}
 
     def _clearBranchCache(self, project):

zuul/driver/github/githubconnection.py

@@ -24,6 +24,7 @@ import time
 import json
 from collections import OrderedDict
 from json.decoder import JSONDecodeError
+from typing import List, Optional
 
 import cherrypy
 import cachecontrol
@@ -38,7 +39,7 @@ import github3.exceptions
 import github3.pulls
 from github3.session import AppInstallationTokenAuth
 
-from zuul.connection import BaseConnection
+from zuul.connection import CachedBranchConnection
 from zuul.driver.github.graphql import GraphQLClient
 from zuul.lib.gearworker import ZuulGearWorker
 from zuul.web.handler import BaseWebController
@@ -447,20 +448,7 @@ class GithubEventProcessor(object):
                 # unprotected branches, we might need to check whether the
                 # branch is now protected.
                 if hasattr(event, "branch") and event.branch:
-                    b = self.connection.getBranch(
+                    self.connection.checkBranchCache(project.name, event)
-                        project.name, event.branch, zuul_event_id=event)
-                    if b is not None:
-                        branch_protected = b.get('protected')
-                        self.connection.checkBranchCache(
-                            project, event.branch, branch_protected, self.log)
-                        event.branch_protected = branch_protected
-                    else:
-                        # This can happen if the branch was deleted in GitHub.
-                        # In this case we assume that the branch COULD have
-                        # been protected before. The cache update is handled by
-                        # the push event, so we don't touch the cache here
-                        # again.
-                        event.branch_protected = True
 
             event.project_hostname = self.connection.canonical_hostname
             self.event = event
@@ -484,31 +472,7 @@ class GithubEventProcessor(object):
             # necessary for the scheduler to match against particular branches
             event.branch = ref_parts[2]
 
-        # This checks whether the event created or deleted a branch so
+        self.connection.clearConnectionCacheOnBranchEvent(event)
-        # that Zuul may know to perform a reconfiguration on the
-        # project.
-        if event.oldrev == '0' * 40:
-            event.branch_created = True
-        elif event.newrev == '0' * 40:
-            event.branch_deleted = True
-        else:
-            event.branch_updated = True
-
-        if event.branch:
-            project = self.connection.source.getProject(event.project_name)
-            if event.branch_deleted:
-                # We currently cannot determine if a deleted branch was
-                # protected so we need to assume it was. GitHub doesn't allow
-                # deletion of protected branches but we don't get a
-                # notification about branch protection settings. Thus we don't
-                # know if branch protection has been disabled before deletion
-                # of the branch.
-                # FIXME(tobiash): Find a way to handle that case
-                self.connection._clearBranchCache(project, self.log)
-            elif event.branch_created:
-                # A new branch never can be protected because that needs to be
-                # configured after it has been created.
-                self.connection._clearBranchCache(project, self.log)
 
         return event
 
@@ -1175,19 +1139,17 @@ class GithubClientManager:
         return github
 
 
-class GithubConnection(BaseConnection):
+class GithubConnection(CachedBranchConnection):
     driver_name = 'github'
     log = logging.getLogger("zuul.GithubConnection")
     payload_path = 'payload'
     client_manager_class = GithubClientManager
 
     def __init__(self, driver, connection_name, connection_config):
-        super(GithubConnection, self).__init__(
+        super(GithubConnection, self).__init__(driver, connection_name,
-            driver, connection_name, connection_config)
+                                               connection_config)
         self._change_cache = {}
         self._change_update_lock = {}
-        self._project_branch_cache_include_unprotected = {}
-        self._project_branch_cache_exclude_unprotected = {}
         self.projects = {}
         self.git_ssh_key = self.connection_config.get('sshkey')
         self.server = self.connection_config.get('server', 'github.com')
@@ -1560,21 +1522,8 @@ class GithubConnection(BaseConnection):
     def addProject(self, project):
         self.projects[project.name] = project
 
-    def clearBranchCache(self):
+    def _fetchProjectBranches(self, project: Project,
-        self._project_branch_cache_exclude_unprotected = {}
+                              exclude_unprotected: bool) -> List[str]:
-        self._project_branch_cache_include_unprotected = {}
-
-    def getProjectBranches(self, project, tenant):
-        exclude_unprotected = tenant.getExcludeUnprotectedBranches(project)
-        if exclude_unprotected:
-            cache = self._project_branch_cache_exclude_unprotected
-        else:
-            cache = self._project_branch_cache_include_unprotected
-
-        branches = cache.get(project.name)
-        if branches is not None:
-            return branches
-
         github = self.getGithubClient(project.name)
         url = github.session.build_url('repos', project.name,
                                        'branches')
@@ -1605,25 +1554,25 @@ class GithubConnection(BaseConnection):
 
             branches.extend([x['name'] for x in resp.json()])
 
-        cache[project.name] = branches
         return branches
 
-    def getBranch(self, project_name, branch, zuul_event_id=None):
+    def isBranchProtected(self, project_name: str, branch_name: str,
+                          zuul_event_id=None) -> Optional[bool]:
         github = self.getGithubClient(
             project_name, zuul_event_id=zuul_event_id)
 
         # Note that we directly use a web request here because if we use the
         # github3.py api directly we need a repository object which needs
         # an unneeded web request during creation.
-        url = github.session.build_url('repos', project_name,
+        url = github.session.build_url('repos', project_name, 'branches',
-                                       'branches', branch)
+                                       branch_name)
 
         resp = github.session.get(url)
 
         if resp.status_code == 404:
             return None
 
-        return resp.json()
+        return resp.json().get('protected')
 
     def getPullUrl(self, project, number):
         return '%s/pull/%s' % (self.getGitwebUrl(project), number)
@@ -2254,46 +2203,6 @@ class GithubConnection(BaseConnection):
                 self.connection_name)
         return True
 
-    def _clearBranchCache(self, project, log):
-        log.debug("Clearing branch cache for %s", project.name)
-        for cache in [
-                self._project_branch_cache_exclude_unprotected,
-                self._project_branch_cache_include_unprotected,
-        ]:
-            try:
-                del cache[project.name]
-            except KeyError:
-                pass
-
-    def checkBranchCache(self, project, branch, protected, log):
-        # If the branch appears in the exclude_unprotected cache but
-        # is unprotected, clear the exclude cache.
-
-        # If the branch does not appear in the exclude_unprotected
-        # cache but is protected, clear the exclude cache.
-
-        # All branches should always appear in the include_unprotected
-        # cache, so we never clear it.
-
-        cache = self._project_branch_cache_exclude_unprotected
-        branches = cache.get(project.name, [])
-        if (branch in branches) and (not protected):
-            log.debug("Clearing protected branch cache for %s",
-                      project.name)
-            try:
-                del cache[project.name]
-            except KeyError:
-                pass
-            return
-        if (branch not in branches) and (protected):
-            log.debug("Clearing protected branch cache for %s",
-                      project.name)
-            try:
-                del cache[project.name]
-            except KeyError:
-                pass
-            return
-
 
 class GithubWebController(BaseWebController):
 

zuul/driver/gitlab/gitlabconnection.py

@@ -444,7 +444,7 @@ class GitlabConnection(BaseConnection):
         self.log.info("Got branches for %s" % project.name)
         return branches
 
-    def clearBranchCache(self):
+    def clearCache(self):
         self.project_branch_cache = {}
 
     def getGitwebUrl(self, project, sha=None):

zuul/driver/pagure/pagureconnection.py

@@ -610,7 +610,7 @@ class PagureConnection(BaseConnection):
         for key in remove:
             del self._change_cache[key]
 
-    def clearBranchCache(self):
+    def clearCache(self):
         self.project_branch_cache = {}
 
     def getWebController(self, zuul_web):

zuul/scheduler.py

@@ -857,8 +857,8 @@ class Scheduler(threading.Thread):
                 default_version=default_ansible_version)
 
             for connection in self.connections.connections.values():
-                self.log.debug("Clear branch cache for: %s" % connection)
+                self.log.debug("Clear cache for: %s" % connection)
-                connection.clearBranchCache()
+                connection.clearCache()
 
             loader = configloader.ConfigLoader(
                 self.connections, self, self.merger,