fix: add on_token_refresh in owner.py

* Updates shared to a newer version
* Adds a on_token_refresh callback if the torngit adapter comes from owner.py

context codecov/engineering-team#162

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@680951c4849074db131353b251581ee03029a72a#egg=shared
+git+ssh://[email protected]/codecov/shared.git@96376263de4d4684d7b7d06c2d612f340eab5611#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@680951c4849074db131353b251581ee03029a72a
+shared @ git+ssh://[email protected]/codecov/shared.git@96376263de4d4684d7b7d06c2d612f340eab5611
     # 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 == "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
-
-
 def get_repo_provider_service(
     repository, commit=None
 ) -> torngit.base.TorngitBaseAdapter: