Use QUayApiStore to manage sessions

Author: esron

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
@@ -15,28 +15,26 @@ class OrgInfo(TypedDict):
     managedRepos: bool
     mirror: OrgKey | None
     mirror_filters: dict[str, Any]
-    # Metadata for creating QuayApi instances on-demand
-    token: str
-    org_name: str
-    base_url: str
+    api: QuayApi
 
 
-QuayApiStore = dict[OrgKey, OrgInfo]
+class QuayApiStore(UserDict[OrgKey, OrgInfo]):
+    def __init__(self) -> None:
+        super().__init__(get_quay_api_store())
 
+    def cleanup(self) -> None:
+        """Close all QuayApi sessions."""
+        for org_info in self.data.values():
+            org_info["api"].cleanup()
 
-def get_quay_api_for_org(org_key: OrgKey, org_info: OrgInfo) -> QuayApi:
-    """
-    Create a QuayApi instance for a specific org on-demand.
-    The instance should be used within a context manager to ensure cleanup.
-    """
-    return QuayApi(
-        token=org_info["token"],
-        organization=org_info["org_name"],
-        base_url=org_info["base_url"],
-    )
+    def __enter__(self) -> "QuayApiStore":
+        return self
 
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.cleanup()
 
-def get_quay_api_store() -> QuayApiStore:
+
+def get_quay_api_store() -> dict[OrgKey, OrgInfo]:
     """
     Returns a dictionary with a key for each Quay organization
     managed in app-interface.
@@ -77,16 +75,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,
             "push_token": push_token,
             "teams": org_data.get("managedTeams") or [],
             "managedRepos": bool(org_data.get("managedRepos")),
             "mirror": mirror,
             "mirror_filters": mirror_filters,
-            "token": token,
-            "org_name": org_name,
-            "base_url": base_url,
+            "api": api,
         }
 
         store[org_key] = org_info

reconcile/quay_membership.py

@@ -10,11 +10,7 @@
     PermissionQuayOrgTeamV1,
     UserV1,
 )
-from reconcile.quay_base import (
-    QuayApiStore,
-    get_quay_api_for_org,
-    get_quay_api_store,
-)
+from reconcile.quay_base import QuayApiStore
 from reconcile.status import ExitCodes
 from reconcile.utils import (
     expiration,
@@ -66,30 +62,29 @@ def fetch_current_state(quay_api_store: QuayApiStore) -> AggregatedList:
         if not teams:
             continue
 
-        # Create QuayApi instance on-demand and close immediately after use
-        with get_quay_api_for_org(org_key, org_data) as quay_api:
-            for team in teams:
-                try:
-                    members = quay_api.list_team_members(team)
-                except QuayTeamNotFoundError:
-                    logging.warning(
-                        "Attempted to list members for team %s in "
-                        "org %s/%s, but it doesn't exist",
-                        team,
-                        org_key.instance,
-                        org_key.org_name,
-                    )
-                else:
-                    # 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.org_name,
-                            "team": team,
-                        },
-                        members,
-                    )
+        quay_api = org_data["api"]
+        for team in teams:
+            try:
+                members = quay_api.list_team_members(team)
+            except QuayTeamNotFoundError:
+                logging.warning(
+                    "Attempted to list members for team %s in "
+                    "org %s/%s, but it doesn't exist",
+                    team,
+                    org_key.instance,
+                    org_key.org_name,
+                )
+            else:
+                # 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.org_name,
+                        "team": team,
+                    },
+                    members,
+                )
     return state
 
 
@@ -128,19 +123,18 @@ 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
-            # Create QuayApi instance on-demand and close immediately after use
-            with get_quay_api_for_org(org, org_data) as quay_api:
-                for member in items:
-                    logging.info([label, member, org, team])
-                    user_exists = quay_api.user_exists(member)
-                    if user_exists:
-                        if not self.dry_run:
-                            quay_api.add_user_to_team(member, team)
-                    else:
-                        missing_users = True
-                        logging.error(f"quay user {member} does not exist.")
+            for member in items:
+                logging.info([label, member, org, team])
+                user_exists = quay_api.user_exists(member)
+                if user_exists:
+                    if not self.dry_run:
+                        quay_api.add_user_to_team(member, team)
+                else:
+                    missing_users = True
+                    logging.error(f"quay user {member} does not exist.")
 
             # This will be evaluated by AggregatedDiffRunner.run(). The happy
             # case is to return True: no missing users
@@ -172,9 +166,8 @@ def action(params: dict, items: list) -> bool:
             logging.info([label, org, team])
 
             if not self.dry_run:
-                # Create QuayApi instance on-demand and close immediately after use
-                with get_quay_api_for_org(org, org_data) as quay_api:
-                    quay_api.create_or_update_team(team)
+                quay_api = org_data["api"]
+                quay_api.create_or_update_team(team)
 
             return True
 
@@ -188,39 +181,38 @@ def action(params: dict, items: list) -> bool:
             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:
-                # Create QuayApi instance on-demand and close immediately after use
-                with get_quay_api_for_org(org, org_data) as quay_api:
-                    for member in items:
-                        logging.info([label, member, org, team])
-                        quay_api.remove_user_from_team(member, team)
+                for member in items:
+                    logging.info([label, member, org, team])
+                    quay_api.remove_user_from_team(member, team)
 
             return True
 
         return action
 
 
 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_for_org, 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()
@@ -112,34 +113,31 @@ def process_org_mirrors(self) -> dict[OrgKey, list[dict[str, Any]]]:
             username = push_token["user"]
             token = push_token["token"]
 
-            # Create QuayApi instances on-demand and close immediately after use
-            with get_quay_api_for_org(org_key, org_info) as quay_api:
-                org_repos = [item["name"] for item in quay_api.list_images()]
-
-            with get_quay_api_for_org(
-                upstream_org_key, upstream_org
-            ) as upstream_quay_api:
-                for repo in upstream_quay_api.list_images():
-                    if repo["name"] not in org_repos:
-                        continue
-
-                    if self.repositories and repo["name"] not in self.repositories:
-                        continue
-
-                    server_url = upstream_org["url"]
-                    url = f"{server_url}/{org_key.org_name}/{repo['name']}"
-                    data = {
-                        "name": repo["name"],
-                        "mirror": {
-                            "url": url,
-                            "username": username,
-                            "token": token,
-                        },
-                        "mirror_filters": org_info.get("mirror_filters", {}).get(
-                            repo["name"], {}
-                        ),
-                    }
-                    summary[org_key].append(data)
+            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
+
+                if self.repositories and repo["name"] not in self.repositories:
+                    continue
+
+                server_url = upstream_org["url"]
+                url = f"{server_url}/{org_key.org_name}/{repo['name']}"
+                data = {
+                    "name": repo["name"],
+                    "mirror": {
+                        "url": url,
+                        "username": username,
+                        "token": token,
+                    },
+                    "mirror_filters": org_info.get("mirror_filters", {}).get(
+                        repo["name"], {}
+                    ),
+                }
+                summary[org_key].append(data)
 
         return summary