feat: support gh refresh tokens

conftest.py

@@ -271,12 +271,10 @@ def mock_submit_fn(metric, start, end):
     mock_submit = mocker.Mock()
     mock_submit.side_effect = mock_submit_fn
 
-    import helpers
+    from helpers.checkpoint_logger import CheckpointLogger
 
     return mocker.patch.object(
-        helpers.checkpoint_logger.CheckpointLogger,
+        CheckpointLogger,
         "submit_subflow",
         mock_submit,
     )
-
-    return mock_submit

helpers/token_refresh.py

@@ -0,0 +1,35 @@
+import logging
+from typing import Callable, Dict
+
+from shared.encryption.token import encode_token
+
+from database.models.core import Owner
+from services.encryption import encryptor
+
+log = logging.getLogger(__name__)
+
+
+def get_token_refresh_callback(owner: Owner) -> Callable[[Dict], None]:
+    """
+    Produces a callback function that will encode and update the oauth token of a user.
+    This callback is passed to the TorngitAdapter for the service.
+    """
+    # Some tokens don't have to be refreshed (GH integration, default bots)
+    # They don't belong to any owners.
+    if owner is None:
+        return None
+
+    service = owner.service
+    if service == "bitbucket" or service == "bitbucket_server":
+        return None
+
+    async def callback(new_token: Dict) -> None:
+        log.info(
+            "Saving new token after refresh",
+            extra=dict(owner=owner.username, ownerid=owner.ownerid),
+        )
+        string_to_save = encode_token(new_token)
+        oauth_token = encryptor.encode(string_to_save).decode()
+        owner.oauth_token = oauth_token
+
+    return callback

requirements.in

@@ -1,4 +1,4 @@
-git+ssh://[email protected]/codecov/shared.git@6a8a33248804a9c101d34f417efda7c11e4bbe63#egg=shared
+git+ssh://[email protected]/codecov/shared.git@f0174635ccaebc3463a5641b9ad640a46b5fd472#egg=shared
 git+ssh://[email protected]/codecov/[email protected]#egg=codecovopentelem
 boto3
 celery

requirements.txt

@@ -251,7 +251,7 @@ s3transfer==0.3.4
     # via boto3
 sentry-sdk==1.19.1
     # via -r requirements.in
-shared @ git+ssh://[email protected]/codecov/shared.git@4f65b0040ab6bb3dba2190ce544ade9a0ea3e354
+shared @ git+ssh://[email protected]/codecov/shared.git@f0174635ccaebc3463a5641b9ad640a46b5fd472
     # via -r requirements.in
 six==1.15.0
     # via

services/owner.py

@@ -3,7 +3,7 @@
 import shared.torngit as torngit
 from shared.config import get_config, get_verify_ssl
 
-from database.models import Owner
+from helpers.token_refresh import get_token_refresh_callback
 from services.bots import get_owner_appropriate_bot_token
 
 log = logging.getLogger(__name__)
@@ -27,6 +27,11 @@ def get_owner_provider_service(owner, using_integration=False):
             key=get_config(service, "client_id"),
             secret=get_config(service, "client_secret"),
         ),
+        # if using integration we will use the integration token
+        # not the owner's token
+        on_token_refresh=(
+            get_token_refresh_callback(owner) if not using_integration else None
+        ),
     )
     return _get_owner_provider_service_instance(service, **adapter_params)
 

services/repository.py

@@ -2,56 +2,28 @@
 import re
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Callable, Dict, Mapping, Optional, Tuple
+from typing import Any, Mapping, Optional, Tuple
 
 import shared.torngit as torngit
 from shared.config import get_config, get_verify_ssl
-from shared.encryption.token import encode_token
 from shared.torngit.exceptions import (
     TorngitClientError,
     TorngitError,
     TorngitObjectNotFoundError,
 )
 from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.exc import IntegrityError
-from sqlalchemy.orm import Session
 
 from database.models import Commit, Owner, Pull, Repository
+from helpers.token_refresh import get_token_refresh_callback
 from services.bots import get_repo_appropriate_bot_token, get_token_type_mapping
-from services.encryption import encryptor
 from services.yaml import read_yaml_field
 
 log = logging.getLogger(__name__)
 
 merged_pull = re.compile(r".*Merged in [^\s]+ \(pull request \#(\d+)\).*").match
 
 
-def get_token_refresh_callback(owner: Owner) -> Callable[[Dict], None]:
-    """
-    Produces a callback function that will encode and update the oauth token of a user.
-    This callback is passed to the TorngitAdapter for the service.
-    """
-    # Some tokens don't have to be refreshed (GH integration, default bots)
-    # They don't belong to any owners.
-    if owner is None:
-        return None
-
-    service = owner.service
-    if service != "gitlab" and service != "gitlab_enterprise":
-        return None
-
-    async def callback(new_token: Dict) -> None:
-        log.info(
-            "Saving new token after refresh",
-            extra=dict(owner=owner.username, ownerid=owner.ownerid),
-        )
-        string_to_save = encode_token(new_token)
-        oauth_token = encryptor.encode(string_to_save).decode()
-        owner.oauth_token = oauth_token
-
-    return callback
-
-
 def get_repo_provider_service(
     repository, commit=None
 ) -> torngit.base.TorngitBaseAdapter:

services/tests/test_repository_service.py

@@ -31,7 +31,7 @@
 
 
 class TestRepositoryServiceTestCase(object):
-    def test_get_repo_provider_service(self, dbsession):
+    def test_get_repo_provider_service_github(self, dbsession):
         repo = RepositoryFactory.create(
             owner__unencrypted_oauth_token="testyftq3ovzkb3zmt823u3t04lkrt9w",
             owner__service="github",
@@ -55,7 +55,39 @@ def test_get_repo_provider_service(self, dbsession):
         }
         assert res.data == expected_data
         assert repo.owner.service == "github"
-        assert res._on_token_refresh is None  # GH doesn't have callback implemented
+        assert res._on_token_refresh is not None
+        assert inspect.isawaitable(res._on_token_refresh(None))
+        assert res.token == {
+            "username": repo.owner.username,
+            "key": "testyftq3ovzkb3zmt823u3t04lkrt9w",
+            "secret": None,
+        }
+
+    def test_get_repo_provider_service_bitbucket(self, dbsession):
+        repo = RepositoryFactory.create(
+            owner__unencrypted_oauth_token="testyftq3ovzkb3zmt823u3t04lkrt9w",
+            owner__service="bitbucket",
+            name="example-python",
+        )
+        dbsession.add(repo)
+        dbsession.flush()
+        res = get_repo_provider_service(repo)
+        expected_data = {
+            "owner": {
+                "ownerid": repo.owner.ownerid,
+                "service_id": repo.owner.service_id,
+                "username": repo.owner.username,
+            },
+            "repo": {
+                "name": "example-python",
+                "using_integration": False,
+                "service_id": repo.service_id,
+                "repoid": repo.repoid,
+            },
+        }
+        assert res.data == expected_data
+        assert repo.owner.service == "bitbucket"
+        assert res._on_token_refresh is None
         assert res.token == {
             "username": repo.owner.username,
             "key": "testyftq3ovzkb3zmt823u3t04lkrt9w",