chore(types): Fix few type errors around bots/ (#463)

Fixes some type errors flagged by mypy in the `bots/` folder (mostly).

Pydantic is being added to create the `DedicatedApp` model, that validates config
pulled from the install YAML. Might be overkill, but it's a direction the team wants
to go in, I believe.

Author: giovanni-guidini

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "sentry-sdk>=2.13.0",
     "sqlalchemy<2",
     "zstandard>=0.23.0",
+    "pydantic>=2.10.4",
 ]
 
 [build-system]
@@ -61,4 +62,5 @@ dev-dependencies = [
     # `oauth2` / `minio` deal with requests forces us to downgrade `urllib3`:
     "urllib3==1.26.19",
     "vcrpy>=4.1.1",
+    "types-requests>=2.31.0.6",
 ]

shared/bots/__init__.py

@@ -1,13 +1,11 @@
 import logging
-from typing import Any, List, Optional
+from typing import Any, List
 
 from shared.bots.github_apps import get_github_app_info_for_owner
 from shared.bots.owner_bots import get_owner_appropriate_bot_token
 from shared.bots.public_bots import get_token_type_mapping
 from shared.bots.repo_bots import get_repo_appropriate_bot_token
-from shared.bots.types import (
-    AdapterAuthInformation,
-)
+from shared.bots.types import AdapterAuthInformation
 from shared.django_apps.codecov_auth.models import (
     GITHUB_APP_INSTALLATION_DEFAULT_NAME,
     Owner,
@@ -18,12 +16,13 @@
 
 log = logging.getLogger(__name__)
 
+SQLAlchemyOwner = Any
+SQLAlchemyRepository = Any
+
 
 def get_adapter_auth_information(
-    # Owner type can be a Django Owner | SQLAlchemy Owner - added Any to help with IDE typing
-    owner: Owner | Any,
-    # Owner type can be a Django Repository | SQLAlchemy Repository - added Any to help with IDE typing
-    repository: Optional[Repository] | Any = None,
+    owner: Owner | SQLAlchemyOwner,
+    repository: Repository | SQLAlchemyRepository | None = None,
     *,
     ignore_installations: bool = False,
     installation_name_to_use: str = GITHUB_APP_INSTALLATION_DEFAULT_NAME,

shared/bots/github_apps.py

@@ -1,7 +1,7 @@
 import logging
 import random
 from datetime import datetime, timezone
-from typing import Dict, List, Optional
+from typing import Dict, List
 
 from shared.bots.exceptions import NoConfiguredAppsAvailable, RequestedGithubAppNotFound
 from shared.bots.types import TokenWithOwner
@@ -12,10 +12,7 @@
     Service,
 )
 from shared.django_apps.core.models import Repository
-from shared.github import (
-    InvalidInstallationError,
-    get_github_integration_token,
-)
+from shared.github import InvalidInstallationError, get_github_integration_token
 from shared.helpers.redis import get_redis_connection
 from shared.orms.owner_helper import DjangoSQLAlchemyOwnerWrapper
 from shared.rate_limits import determine_if_entity_is_rate_limited, gh_app_key_name
@@ -42,7 +39,7 @@ def _get_installation_weight(installation: GithubAppInstallation) -> int:
 
 
 def _can_use_this_app(
-    app: GithubAppInstallation, installation_name: str, repository: Optional[Repository]
+    app: GithubAppInstallation, installation_name: str, repository: Repository | None
 ) -> bool:
     return (
         app.name == installation_name
@@ -54,7 +51,7 @@ def _can_use_this_app(
 
 
 def _get_apps_from_weighted_selection(
-    owner: Owner, installation_name: str, repository: Optional[Repository]
+    owner: Owner, installation_name: str, repository: Repository | None
 ) -> List[GithubAppInstallation]:
     """This function returns an ordered list of GithubAppInstallations that can be used to communicate with GitHub
     in behalf of the owner. The list is ordered in such a way that the 1st element is the app to be used in Torngit,
@@ -162,12 +159,16 @@ def get_github_app_token(
         github_token = get_github_integration_token(
             service.value,
             installation_id,
-            app_id=installation_info.get("app_id", None),
+            app_id=str(app_id) if app_id else None,
             pem_path=installation_info.get("pem_path", None),
         )
         installation_token = Token(
             key=github_token,
-            entity_name=gh_app_key_name(installation_id=installation_id, app_id=app_id),
+            username=f"installation_{installation_id}",
+            entity_name=gh_app_key_name(
+                installation_id=installation_id,
+                app_id=app_id,
+            ),
         )
         # The token is not owned by an Owner object, so 2nd arg is None
         return installation_token, None
@@ -191,7 +192,7 @@ def get_specific_github_app_details(
             The assumption is that we need this specific app for a reason, and if we can't find the app
             it's better to just fail
     """
-    app: GithubAppInstallation = next(
+    app: GithubAppInstallation | None = next(
         (
             obj
             for obj in DjangoSQLAlchemyOwnerWrapper.get_github_app_installations(owner)
@@ -242,7 +243,7 @@ def _filter_suspended_apps(
 def get_github_app_info_for_owner(
     owner: Owner,
     *,
-    repository: Optional[Repository] = None,
+    repository: Repository | None = None,
     installation_name: str = GITHUB_APP_INSTALLATION_DEFAULT_NAME,
 ) -> List[GithubInstallationInfo]:
     """Gets the GitHub app info needed to communicate with GitHub using an app for this owner.

shared/bots/helpers.py

@@ -1,5 +1,7 @@
 import logging
-from typing import Any, Dict, Optional
+from typing import Self
+
+from pydantic import BaseModel, ValidationError
 
 from shared.bots.exceptions import RepositoryWithoutValidBotError
 from shared.config import get_config
@@ -21,38 +23,59 @@
 @cache.cache_function(ttl=1800)
 def get_github_integration_token(
     service: str,
-    installation_id: int = None,
-    app_id: Optional[str] = None,
-    pem_path: Optional[str] = None,
+    installation_id: int | None = None,
+    app_id: str | None = None,
+    pem_path: str | None = None,
 ):
     try:
         return _get_github_integration_token(
-            service, integration_id=installation_id, app_id=app_id, pem_path=pem_path
+            service,
+            integration_id=installation_id,
+            app_id=app_id,
+            pem_path=pem_path,
         )
     except InvalidInstallationError:
         log.warning("Failed to get installation token")
         raise RepositoryWithoutValidBotError()
 
 
+class DedicatedApp(BaseModel):
+    id: str | int
+    installation_id: int
+    pem: str
+
+    @classmethod
+    def validate_or_none(cls, value: object) -> Self | None:
+        try:
+            return cls.model_validate(value)
+        except ValidationError:
+            return None
+
+    def pem_path(self, service: str, token_type: TokenType) -> str:
+        return f"yaml+file://{service}.dedicated_apps.{token_type.value}.pem"
+
+
 def get_dedicated_app_token_from_config(
     service: str, token_type: TokenType
 ) -> Token | None:
     # GitHub can have 'dedicated_apps', and those are preferred
-    dedicated_app: Dict[str, Any] = get_config(
-        service, "dedicated_apps", token_type.value, default={}
+    dedicated_app = DedicatedApp.validate_or_none(
+        get_config(service, "dedicated_apps", token_type.value, default={})
+    )
+    if dedicated_app is None:
+        return None
+
+    actual_token = get_github_integration_token(
+        service,
+        app_id=str(dedicated_app.id),
+        installation_id=dedicated_app.installation_id,
+        pem_path=dedicated_app.pem_path(service, token_type),
+    )
+    return Token(
+        key=actual_token,
+        username=f"{token_type.value}_dedicated_app",
+        entity_name=token_type.value,
     )
-    app_id = dedicated_app.get("id")
-    installation_id = dedicated_app.get("installation_id")
-    pem_exists = "pem" in dedicated_app
-    if app_id and installation_id and pem_exists:
-        actual_token = get_github_integration_token(
-            service,
-            app_id=app_id,
-            installation_id=installation_id,
-            pem_path=f"yaml+file://{service}.dedicated_apps.{token_type.value}.pem",
-        )
-        return Token(key=actual_token, username=f"{token_type.value}_dedicated_app")
-    return None
 
 
 def get_token_type_from_config(service: str, token_type: TokenType) -> Token | None:
@@ -65,4 +88,7 @@ def get_token_type_from_config(service: str, token_type: TokenType) -> Token | N
         if dedicated_app_config:
             return dedicated_app_config
 
-    return get_config(service, "bots", token_type.value)
+    token_from_config = get_config(service, "bots", token_type.value)
+    if token_from_config:
+        token_from_config["entity_name"] = token_type.value
+    return token_from_config

shared/github/__init__.py

@@ -1,6 +1,6 @@
 import logging
 from time import time
-from typing import Literal, Optional
+from typing import Literal
 from urllib.parse import urlparse
 
 import jwt
@@ -29,7 +29,7 @@ def load_pem_from_path(pem_path: str) -> str:
     raise Exception("Unknown schema to load PEM")
 
 
-def get_pem(*, pem_name: Optional[str] = None, pem_path: Optional[str] = None) -> str:
+def get_pem(pem_name: str | None = None, pem_path: str | None = None) -> str:
     if pem_path:
         return load_pem_from_path(pem_path)
     if pem_name:
@@ -38,14 +38,15 @@ def get_pem(*, pem_name: Optional[str] = None, pem_path: Optional[str] = None) -
     raise Exception("No PEM provided to get installation token")
 
 
-InstallationErrorCause = (
-    Literal["installation_not_found"]
-    | Literal["permission_error"]
-    | Literal["requires_authentication"]
-    | Literal["installation_suspended"]
-    | Literal["validation_failed_or_spammed"]
-    | Literal["rate_limit"]
-)
+InstallationErrorCause = Literal[
+    "installation_not_found",
+    "permission_error",
+    "requires_authentication",
+    "installation_suspended",
+    "validation_failed_or_spammed",
+    "rate_limit",
+    "unknown_error",
+]
 
 
 class InvalidInstallationError(Exception):
@@ -78,11 +79,13 @@ def decide_installation_error_cause(
                 return "rate_limit"
             else:
                 return "permission_error"
+        case _:
+            return "unknown_error"
 
 
 def get_github_jwt_token(
-    service: str, app_id: Optional[str] = None, pem_path: Optional[str] = None
-) -> Optional[str]:
+    service: str, app_id: str | None = None, pem_path: str | None = None
+) -> str:
     # https://developer.github.com/apps/building-github-apps/authenticating-with-github-apps/
     now = int(time())
     payload = {
@@ -100,9 +103,9 @@ def get_github_jwt_token(
 def get_github_integration_token(
     service,
     integration_id=None,
-    app_id: Optional[str] = None,
-    pem_path: Optional[str] = None,
-) -> Optional[str]:
+    app_id: str | None = None,
+    pem_path: str | None = None,
+) -> str:
     # https://developer.github.com/apps/building-github-apps/authenticating-with-github-apps/
     token = get_github_jwt_token(service, app_id, pem_path)
     if integration_id:

shared/rate_limits/__init__.py

@@ -11,8 +11,9 @@
 RATE_LIMIT_REDIS_KEY_PREFIX = "rate_limited_entity_"
 
 
-def gh_app_key_name(installation_id: int, app_id: Optional[int] = None) -> str:
-    app_id = app_id or "default_app"
+def gh_app_key_name(installation_id: int, app_id: str | int | None = None) -> str:
+    if app_id is None:
+        return f"default_app_{installation_id}"
     return f"{app_id}_{installation_id}"
 
 

shared/typings/torngit.py

@@ -18,14 +18,19 @@ class RepoInfo(TypedDict):
 
 
 class GithubInstallationInfo(TypedDict):
-    """Required info to get a token from Github for a given installation"""
-
-    id: int
+    """
+    Information about a Github installation.
+    `id` - The id of the GithubAppInstallation object in the database
+           If using the deprecated owner.integration_id it doesn't exist.
+    `installation_id` - Required info to get a token from Github for a given installation.
+    """
+
+    id: NotRequired[int]
     installation_id: int
-    # The default app (configured via yaml) doesn't need this info.
-    # All other apps need app_id and pem_path
-    app_id: Optional[int]
-    pem_path: Optional[str]
+    # The default app (configured via yaml) doesn't need `app_id` and `pem_path`.
+    # All other apps need `app_id` and `pem_path`.
+    app_id: NotRequired[int | None]
+    pem_path: NotRequired[str | None]
 
 
 class AdditionalData(TypedDict):