Move common change cache related methods to mixin

Remove the duplicate methods in the connections that use a
Zookeeper-backed change cache by moving them to a common mixin-class.

Change-Id: I527522e79444dfb69b41c70e7698a281166b2128

6 files changed, 32 insertions, 83 deletions

diff --git a/zuul/connection/__init__.py b/zuul/connection/__init__.py
index 2a8d4f36e..aa4c16967 100644
--- a/zuul/connection/__init__.py
+++ b/zuul/connection/__init__.py
@@ -291,3 +291,25 @@ class CachedBranchConnection(BaseConnection):
             # the push event, so we don't touch the cache here
             # again.
             event.branch_protected = True
+
+
+class ZKChangeCacheMixin:
+    # Expected to be defined by the connection and to be an instance
+    # that implements the AbstractChangeCache API.
+    _change_cache = None
+
+    def cleanupCache(self):
+        self._change_cache.cleanup()
+
+    def maintainCache(self, relevant, max_age):
+        self._change_cache.prune(relevant, max_age)
+
+    def updateChangeAttributes(self, change, **attrs):
+        def _update_attrs(c):
+            for name, value in attrs.items():
+                setattr(c, name, value)
+        self._change_cache.updateChangeWithRetry(change.cache_stat.key,
+                                                 change, _update_attrs)
+
+    def getChangeByKey(self, key):
+        return self._change_cache.get(key)
diff --git a/zuul/driver/gerrit/gerritconnection.py b/zuul/driver/gerrit/gerritconnection.py
index ea884a55a..dbb2017d1 100644
--- a/zuul/driver/gerrit/gerritconnection.py
+++ b/zuul/driver/gerrit/gerritconnection.py
@@ -35,7 +35,7 @@ from typing import Dict, List
 from uuid import uuid4
 
 from zuul import version as zuul_version
-from zuul.connection import BaseConnection
+from zuul.connection import BaseConnection, ZKChangeCacheMixin
 from zuul.driver.gerrit.auth import FormAuth
 from zuul.driver.gerrit.gcloudauth import GCloudAuth
 from zuul.driver.gerrit.gerritmodel import GerritChange, GerritTriggerEvent
@@ -547,7 +547,7 @@ class GerritPoller(threading.Thread):
         self.poller_election.cancel()
 
 
-class GerritConnection(BaseConnection):
+class GerritConnection(ZKChangeCacheMixin, BaseConnection):
     driver_name = 'gerrit'
     log = logging.getLogger("zuul.GerritConnection")
     iolog = logging.getLogger("zuul.GerritConnection.io")
@@ -745,19 +745,6 @@ class GerritConnection(BaseConnection):
         except KeyError:
             pass
 
-    def cleanupCache(self):
-        self._change_cache.cleanup()
-
-    def maintainCache(self, relevant, max_age):
-        self._change_cache.prune(relevant, max_age)
-
-    def updateChangeAttributes(self, change, **attrs):
-        def _update_attrs(c):
-            for name, value in attrs.items():
-                setattr(c, name, value)
-        self._change_cache.updateChangeWithRetry(change.cache_stat.key,
-                                                 change, _update_attrs)
-
     def getChange(self, event, refresh=False):
         if event.change_number:
             change = self._getChange(event.change_number, event.patch_number,
@@ -1011,9 +998,6 @@ class GerritConnection(BaseConnection):
             git_needed_by_changes=git_needed_by_changes,
             compat_needed_by_changes=compat_needed_by_changes)
 
-    def getChangeByKey(self, key):
-        return self._change_cache.get(key)
-
     def isMerged(self, change, head=None):
         self.log.debug("Checking if change %s is merged" % change)
         if not change.number:
diff --git a/zuul/driver/git/gitconnection.py b/zuul/driver/git/gitconnection.py
index b2cd008ad..d63ef89fb 100644
--- a/zuul/driver/git/gitconnection.py
+++ b/zuul/driver/git/gitconnection.py
@@ -19,7 +19,7 @@ import time
 import logging
 import urllib
 
-from zuul.connection import BaseConnection
+from zuul.connection import BaseConnection, ZKChangeCacheMixin
 from zuul.driver.git.gitmodel import GitTriggerEvent
 from zuul.driver.git.gitwatcher import GitWatcher
 from zuul.model import Ref, Branch
@@ -41,7 +41,7 @@ class GitChangeCache(AbstractChangeCache):
         return type(change).__name__
 
 
-class GitConnection(BaseConnection):
+class GitConnection(ZKChangeCacheMixin, BaseConnection):
     driver_name = 'git'
     log = logging.getLogger("zuul.connection.git")
 
@@ -103,12 +103,6 @@ class GitConnection(BaseConnection):
                 refs[ref] = sha
         return refs
 
-    def cleanupCache(self):
-        self._change_cache.cleanup()
-
-    def maintainCache(self, relevant, max_age):
-        self._change_cache.prune(relevant, max_age)
-
     def getChange(self, event, refresh=False):
         key = str((event.project_name, event.ref, event.newrev))
         change = self._change_cache.get(key)
@@ -150,9 +144,6 @@ class GitConnection(BaseConnection):
                     refs if ref.startswith('refs/heads/')]
         return branches
 
-    def getChangeByKey(self, key):
-        return self._change_cache.get(key)
-
     def getGitUrl(self, project):
         return os.path.join(self.baseurl, project.name)
 
diff --git a/zuul/driver/github/githubconnection.py b/zuul/driver/github/githubconnection.py
index f08cd8de3..6c0b3959e 100644
--- a/zuul/driver/github/githubconnection.py
+++ b/zuul/driver/github/githubconnection.py
@@ -40,7 +40,7 @@ import github3.exceptions
 import github3.pulls
 from github3.session import AppInstallationTokenAuth
 
-from zuul.connection import CachedBranchConnection
+from zuul.connection import CachedBranchConnection, ZKChangeCacheMixin
 from zuul.driver.github.graphql import GraphQLClient
 from zuul.web.handler import BaseWebController
 from zuul.lib.logutil import get_annotated_logger
@@ -1175,7 +1175,7 @@ class GithubClientManager:
         return github
 
 
-class GithubConnection(CachedBranchConnection):
+class GithubConnection(ZKChangeCacheMixin, CachedBranchConnection):
     driver_name = 'github'
     log = logging.getLogger("zuul.GithubConnection")
     payload_path = 'payload'
@@ -1262,19 +1262,6 @@ class GithubConnection(CachedBranchConnection):
         return self._github_client_manager.getGithubClient(
             project_name=project_name, zuul_event_id=zuul_event_id)
 
-    def cleanupCache(self):
-        self._change_cache.cleanup()
-
-    def maintainCache(self, relevant, max_age):
-        self._change_cache.prune(relevant, max_age)
-
-    def updateChangeAttributes(self, change, **attrs):
-        def _update_attrs(c):
-            for name, value in attrs.items():
-                setattr(c, name, value)
-        self._change_cache.updateChangeWithRetry(change.cache_stat.key,
-                                                 change, _update_attrs)
-
     def getChange(self, event, refresh=False):
         """Get the change representing an event."""
 
@@ -1407,9 +1394,6 @@ class GithubConnection(CachedBranchConnection):
             change = self._change_cache.get(key)
         return change
 
-    def getChangeByKey(self, key):
-        return self._change_cache.get(key)
-
     def getChangesDependingOn(self, change, projects, tenant):
         changes = []
         if not change.uris:
diff --git a/zuul/driver/gitlab/gitlabconnection.py b/zuul/driver/gitlab/gitlabconnection.py
index 039e69409..8e69140d6 100644
--- a/zuul/driver/gitlab/gitlabconnection.py
+++ b/zuul/driver/gitlab/gitlabconnection.py
@@ -28,7 +28,7 @@ import dateutil.parser
 from urllib.parse import quote_plus
 from typing import List, Optional
 
-from zuul.connection import CachedBranchConnection
+from zuul.connection import CachedBranchConnection, ZKChangeCacheMixin
 from zuul.web.handler import BaseWebController
 from zuul.lib.logutil import get_annotated_logger
 from zuul.exceptions import MergeFailure
@@ -409,7 +409,7 @@ class GitlabAPIClient():
         return resp[0]
 
 
-class GitlabConnection(CachedBranchConnection):
+class GitlabConnection(ZKChangeCacheMixin, CachedBranchConnection):
     driver_name = 'gitlab'
     log = logging.getLogger("zuul.GitlabConnection")
     payload_path = 'payload'
@@ -459,19 +459,6 @@ class GitlabConnection(CachedBranchConnection):
         if hasattr(self, 'gitlab_event_connector'):
             self._stop_event_connector()
 
-    def cleanupCache(self):
-        self._change_cache.cleanup()
-
-    def maintainCache(self, relevant, max_age):
-        self._change_cache.prune(relevant, max_age)
-
-    def updateChangeAttributes(self, change, **attrs):
-        def _update_attrs(c):
-            for name, value in attrs.items():
-                setattr(c, name, value)
-        self._change_cache.updateChangeWithRetry(change.cache_stat.key,
-                                                 change, _update_attrs)
-
     def getWebController(self, zuul_web):
         return GitlabWebController(zuul_web, self)
 
@@ -658,9 +645,6 @@ class GitlabConnection(CachedBranchConnection):
                 "Set approval: %s on MR %s#%s (%s)", approve,
                 project_name, number, patchset)
 
-    def getChangeByKey(self, key):
-        return self._change_cache.get(key)
-
     def getChangesDependingOn(self, change, projects, tenant):
         """ Reverse lookup of MR depending on this one
         """
diff --git a/zuul/driver/pagure/pagureconnection.py b/zuul/driver/pagure/pagureconnection.py
index 1f5331662..b99598153 100644
--- a/zuul/driver/pagure/pagureconnection.py
+++ b/zuul/driver/pagure/pagureconnection.py
@@ -23,7 +23,7 @@ import requests
 import cherrypy
 import voluptuous as v
 
-from zuul.connection import BaseConnection
+from zuul.connection import BaseConnection, ZKChangeCacheMixin
 from zuul.lib.logutil import get_annotated_logger
 from zuul.web.handler import BaseWebController
 from zuul.model import Ref, Branch, Tag
@@ -466,7 +466,7 @@ class PagureAPIClient():
         return resp[0]['webhook']['token']
 
 
-class PagureConnection(BaseConnection):
+class PagureConnection(ZKChangeCacheMixin, BaseConnection):
     driver_name = 'pagure'
     log = logging.getLogger("zuul.PagureConnection")
 
@@ -545,19 +545,6 @@ class PagureConnection(BaseConnection):
                 "Fetching project %s webhook token from API" % project)
             return token
 
-    def cleanupCache(self):
-        self._change_cache.cleanup()
-
-    def maintainCache(self, relevant, max_age):
-        self._change_cache.prune(relevant, max_age)
-
-    def updateChangeAttributes(self, change, **attrs):
-        def _update_attrs(c):
-            for name, value in attrs.items():
-                setattr(c, name, value)
-        self._change_cache.updateChangeWithRetry(change.cache_stat.key,
-                                                 change, _update_attrs)
-
     def getWebController(self, zuul_web):
         return PagureWebController(zuul_web, self)
 
@@ -797,9 +784,6 @@ class PagureConnection(BaseConnection):
                 number, project, flag.get('status')))
         return flag.get('status')
 
-    def getChangeByKey(self, key):
-        return self._change_cache.get(key)
-
     def getChangesDependingOn(self, change, projects, tenant):
         """ Reverse lookup of PR depending on this one
         """