Refactor QuayApi to extend ApiBase

reconcile/quay_base.py

@@ -1,4 +1,4 @@
-from collections import namedtuple
+from collections import UserDict, namedtuple
 from typing import Any, TypedDict
 
 from reconcile import queries
@@ -10,22 +10,34 @@
 
 class OrgInfo(TypedDict):
     url: str
-    api: QuayApi
     push_token: dict[str, str] | None
     teams: list[str]
     managedRepos: bool
     mirror: OrgKey | None
     mirror_filters: dict[str, Any]
+    api: QuayApi
+
 
+class QuayApiStore(UserDict[OrgKey, OrgInfo]):
+    def __init__(self) -> None:
+        super().__init__(get_quay_api_store())
 
-QuayApiStore = dict[OrgKey, OrgInfo]
+    def cleanup(self) -> None:
+        """Close all QuayApi sessions."""
+        for org_info in self.data.values():
+            org_info["api"].cleanup()
 
+    def __enter__(self) -> "QuayApiStore":
+        return self
 
-def get_quay_api_store() -> QuayApiStore:
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.cleanup()
+
+
+def get_quay_api_store() -> dict[OrgKey, OrgInfo]:
     """
     Returns a dictionary with a key for each Quay organization
     managed in app-interface.
-    Each key contains an initiated QuayApi instance.
     """
     quay_orgs = queries.get_quay_orgs()
     settings = queries.get_app_interface_settings()
@@ -61,14 +73,21 @@ def get_quay_api_store() -> QuayApiStore:
         else:
             push_token = None
 
+        # Create QuayApi instance for this org
+        api = QuayApi(
+            token=token,
+            organization=org_name,
+            base_url=base_url,
+        )
+
         org_info: OrgInfo = {
             "url": base_url,
-            "api": QuayApi(token, org_name, base_url=base_url),
             "push_token": push_token,
             "teams": org_data.get("managedTeams") or [],
             "managedRepos": bool(org_data.get("managedRepos")),
             "mirror": mirror,
             "mirror_filters": mirror_filters,
+            "api": api,
         }
 
         store[org_key] = org_info

reconcile/quay_membership.py

@@ -10,7 +10,7 @@
     PermissionQuayOrgTeamV1,
     UserV1,
 )
-from reconcile.quay_base import QuayApiStore, get_quay_api_store
+from reconcile.quay_base import QuayApiStore
 from reconcile.status import ExitCodes
 from reconcile.utils import (
     expiration,
@@ -58,10 +58,11 @@ def fetch_current_state(quay_api_store: QuayApiStore) -> AggregatedList:
     state = AggregatedList()
 
     for org_key, org_data in quay_api_store.items():
-        quay_api = org_data["api"]
         teams = org_data["teams"]
         if not teams:
             continue
+
+        quay_api = org_data["api"]
         for team in teams:
             try:
                 members = quay_api.list_team_members(team)
@@ -77,7 +78,11 @@ def fetch_current_state(quay_api_store: QuayApiStore) -> AggregatedList:
                 # Teams are only added to the state if they exist so that
                 # there is a proper diff between the desired and current state.
                 state.add(
-                    {"service": "quay-membership", "org": org_key, "team": team},
+                    {
+                        "service": "quay-membership",
+                        "org": org_key.org_name,
+                        "team": team,
+                    },
                     members,
                 )
     return state
@@ -117,9 +122,10 @@ def add_to_team(self) -> Action:
         def action(params: dict, items: list) -> bool:
             org = params["org"]
             team = params["team"]
+            org_data = self.quay_api_store[org]
+            quay_api = org_data["api"]
 
             missing_users = False
-            quay_api = self.quay_api_store[org]["api"]
             for member in items:
                 logging.info([label, member, org, team])
                 user_exists = quay_api.user_exists(member)
@@ -147,10 +153,11 @@ def create_team(self) -> Action:
         def action(params: dict, items: list) -> bool:
             org = params["org"]
             team = params["team"]
+            org_data = self.quay_api_store[org]
 
             # Ensure all quay org/teams are declared as dependencies in a
             # `/dependencies/quay-org-1.yml` datafile.
-            if team not in self.quay_api_store[org]["teams"]:
+            if team not in org_data["teams"]:
                 raise RunnerError(
                     f"Quay team {team} is not defined as a "
                     f"managedTeam in the {org} org."
@@ -159,7 +166,7 @@ def action(params: dict, items: list) -> bool:
             logging.info([label, org, team])
 
             if not self.dry_run:
-                quay_api = self.quay_api_store[org]["api"]
+                quay_api = org_data["api"]
                 quay_api.create_or_update_team(team)
 
             return True
@@ -172,12 +179,13 @@ def del_from_team(self) -> Action:
         def action(params: dict, items: list) -> bool:
             org = params["org"]
             team = params["team"]
+            org_data = self.quay_api_store[org]
 
+            quay_api = org_data["api"]
             if self.dry_run:
                 for member in items:
                     logging.info([label, member, org, team])
             else:
-                quay_api = self.quay_api_store[org]["api"]
                 for member in items:
                     logging.info([label, member, org, team])
                     quay_api.remove_user_from_team(member, team)
@@ -188,24 +196,23 @@ def action(params: dict, items: list) -> bool:
 
 
 def run(dry_run: bool) -> None:
-    quay_api_store = get_quay_api_store()
-
-    current_state = fetch_current_state(quay_api_store)
-    desired_state = fetch_desired_state()
-
-    # calculate diff
-    diff = current_state.diff(desired_state)
-    logging.debug("State diff: %s", diff)
-
-    # Run actions
-    runner_action = RunnerAction(dry_run, quay_api_store)
-    runner = AggregatedDiffRunner(diff)
-
-    runner.register("insert", runner_action.create_team())
-    runner.register("update-insert", runner_action.add_to_team())
-    runner.register("update-delete", runner_action.del_from_team())
-    runner.register("delete", runner_action.del_from_team())
-
-    status = runner.run()
-    if not status:
-        sys.exit(ExitCodes.ERROR)
+    with QuayApiStore() as quay_api_store:
+        current_state = fetch_current_state(quay_api_store)
+        desired_state = fetch_desired_state()
+
+        # calculate diff
+        diff = current_state.diff(desired_state)
+        logging.debug("State diff: %s", diff)
+
+        # Run actions
+        runner_action = RunnerAction(dry_run, quay_api_store)
+        runner = AggregatedDiffRunner(diff)
+
+        runner.register("insert", runner_action.create_team())
+        runner.register("update-insert", runner_action.add_to_team())
+        runner.register("update-delete", runner_action.del_from_team())
+        runner.register("delete", runner_action.del_from_team())
+
+        status = runner.run()
+        if not status:
+            sys.exit(ExitCodes.ERROR)

reconcile/quay_mirror_org.py

@@ -17,7 +17,7 @@
 )
 from sretoolbox.container.skopeo import SkopeoCmdError
 
-from reconcile.quay_base import get_quay_api_store
+from reconcile.quay_base import QuayApiStore
 from reconcile.quay_mirror import QuayMirror
 from reconcile.utils.quay_mirror import record_timestamp, sync_tag
 
@@ -45,7 +45,7 @@ def __init__(
     ) -> None:
         self.dry_run = dry_run
         self.skopeo_cli = Skopeo(dry_run)
-        self.quay_api_store = get_quay_api_store()
+        self.quay_api_store = QuayApiStore()
         self.compare_tags = compare_tags
         self.compare_tags_interval = compare_tags_interval
         self.orgs = orgs
@@ -71,6 +71,7 @@ def __enter__(self) -> Self:
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         self.session.close()
+        self.quay_api_store.cleanup()
 
     def run(self) -> None:
         sync_tasks = self.process_sync_tasks()
@@ -101,11 +102,9 @@ def process_org_mirrors(self) -> dict[OrgKey, list[dict[str, Any]]]:
             if self.orgs and org_key.org_name not in self.orgs:
                 continue
 
-            quay_api = org_info["api"]
             upstream_org_key = org_info["mirror"]
             assert upstream_org_key is not None
             upstream_org = self.quay_api_store[upstream_org_key]
-            upstream_quay_api = upstream_org["api"]
 
             push_token = upstream_org["push_token"]
 
@@ -114,7 +113,10 @@ def process_org_mirrors(self) -> dict[OrgKey, list[dict[str, Any]]]:
             username = push_token["user"]
             token = push_token["token"]
 
+            quay_api = org_info["api"]
             org_repos = [item["name"] for item in quay_api.list_images()]
+
+            upstream_quay_api = upstream_org["api"]
             for repo in upstream_quay_api.list_images():
                 if repo["name"] not in org_repos:
                     continue