Cleanup workflow fixes

Author: evertlammerts

README.md

@@ -67,8 +67,8 @@ git config --local core.hooksPath .githooks/
 
 ### Editable installs (general)
 
-  It's good to be aware of the following when creating an editable install:
-- `uv sync` or `uv run [tool]` create editable installs by default, however, it work the way you expect. We have 
+  It's good to be aware of the following when performing an editable install:
+- `uv sync` or `uv run [tool]` perform an editable install by default. We have 
   configured the project so that scikit-build-core will use a persistent build-dir, but since the build itself 
   happens in an isolated, ephemeral environment, cmake's paths will point to non-existing directories. CMake itself 
   will be missing.

duckdb_packaging/pypi_cleanup.py

@@ -10,7 +10,6 @@
 
 import argparse
 import contextlib
-import datetime
 import heapq
 import logging
 import os
@@ -19,7 +18,7 @@
 import time
 from collections import defaultdict
 from html.parser import HTMLParser
-from typing import Dict, Optional, Set, Generator
+from typing import Optional, Set, Generator
 from urllib.parse import urlparse
 
 import pyotp
@@ -252,19 +251,20 @@ def run(self) -> int:
         logging.info(f"Max development releases to keep per unreleased version: {self._max_dev_releases}")
 
         try:
-            return self._execute_cleanup()
+            with session_with_retries() as http_session:
+                return self._execute_cleanup(http_session)
         except PyPICleanupError as e:
             logging.error(f"Cleanup failed: {e}")
             return 1
         except Exception as e:
             logging.error(f"Unexpected error: {e}", exc_info=True)
             return 1
 
-    def _execute_cleanup(self) -> int:
+    def _execute_cleanup(self, http_session: Session) -> int:
         """Execute the main cleanup logic."""
 
         # Get released versions
-        versions = self._fetch_released_versions()
+        versions = self._fetch_released_versions(http_session)
         if not versions:
             logging.info(f"No releases found for {self._package}")
             return 0
@@ -284,24 +284,23 @@ def _execute_cleanup(self) -> int:
             return 0
 
         # Perform authentication and deletion
-        self._authenticate()
-        self._delete_versions(versions_to_delete)
+        self._authenticate(http_session)
+        self._delete_versions(http_session, versions_to_delete)
 
         logging.info(f"Successfully cleaned up {len(versions_to_delete)} development versions")
         return 0
 
-    def _fetch_released_versions(self) -> Set[str]:
+    def _fetch_released_versions(self, http_session: Session) -> Set[str]:
         """Fetch package release information from PyPI API."""
         logging.debug(f"Fetching package information for '{self._package}'")
 
         try:
-            with session_with_retries() as session:
-                req = session.get(f"{self._index_url}/pypi/{self._package}/json")
-                req.raise_for_status()
-                data = req.json()
-                versions = {v for v, files in data["releases"].items() if len(files) > 0}
-                logging.debug(f"Found {len(versions)} releases with files")
-                return versions
+            req = http_session.get(f"{self._index_url}/pypi/{self._package}/json")
+            req.raise_for_status()
+            data = req.json()
+            versions = {v for v, files in data["releases"].items() if len(files) > 0}
+            logging.debug(f"Found {len(versions)} releases with files")
+            return versions
         except RequestException as e:
             raise PyPICleanupError(f"Failed to fetch package information for '{self._package}': {e}") from e
 
@@ -394,94 +393,92 @@ def _determine_versions_to_delete(self, versions: Set[str]) -> Set[str]:
 
         return versions_to_delete
 
-    def _authenticate(self) -> None:
+    def _authenticate(self, http_session: Session) -> None:
         """Authenticate with PyPI."""
         if not self._username or not self._password:
             raise AuthenticationError("Username and password are required for authentication")
 
         logging.info(f"Authenticating user '{self._username}' with PyPI")
 
         try:
-            # Get login form and CSRF token
-            csrf_token = self._get_csrf_token("/account/login/")
-            
             # Attempt login
-            login_response = self._perform_login(csrf_token)
-            
+            login_response = self._perform_login(http_session)
+
             # Handle two-factor authentication if required
             if login_response.url.startswith(f"{self._index_url}/account/two-factor/"):
                 logging.debug("Two-factor authentication required")
-                self._handle_two_factor_auth(login_response)
+                self._handle_two_factor_auth(http_session, login_response)
 
             logging.info("Authentication successful")
-            
+
         except RequestException as e:
             raise AuthenticationError(f"Network error during authentication: {e}") from e
 
-    def _get_csrf_token(self, form_action: str) -> str:
+    def _get_csrf_token(self, http_session: Session, form_action: str) -> str:
         """Extract CSRF token from a form page."""
-        with session_with_retries() as session:
-            req = session.get(f"{self._index_url}{form_action}")
-            req.raise_for_status()
-            parser = CsrfParser(form_action)
-            parser.feed(req.text)
-            if not parser.csrf:
-                raise AuthenticationError(f"No CSRF token found in {form_action}")
-            return parser.csrf
+        resp = http_session.get(f"{self._index_url}{form_action}")
+        resp.raise_for_status()
+        parser = CsrfParser(form_action)
+        parser.feed(resp.text)
+        if not parser.csrf:
+            raise AuthenticationError(f"No CSRF token found in {form_action}")
+        return parser.csrf
 
-    def _perform_login(self, csrf_token: str) -> requests.Response:
+    def _perform_login(self, http_session: Session) -> requests.Response:
         """Perform the initial login with username/password."""
+
+        # Get login form and CSRF token
+        csrf_token = self._get_csrf_token(http_session, "/account/login/")
+
         login_data = {
             "csrf_token": csrf_token,
             "username": self._username,
             "password": self._password
         }
 
-        with session_with_retries() as session:
-            response = session.post(
-                f"{self._index_url}/account/login/",
-                data=login_data,
-                headers={"referer": f"{self._index_url}/account/login/"}
-            )
-            response.raise_for_status()
+        response = http_session.post(
+            f"{self._index_url}/account/login/",
+            data=login_data,
+            headers={"referer": f"{self._index_url}/account/login/"}
+        )
+        response.raise_for_status()
 
-            # Check if login failed (redirected back to login page)
-            if response.url == f"{self._index_url}/account/login/":
-                raise AuthenticationError(f"Login failed for user '{self._username}' - check credentials")
+        # Check if login failed (redirected back to login page)
+        if response.url == f"{self._index_url}/account/login/":
+            raise AuthenticationError(f"Login failed for user '{self._username}' - check credentials")
 
-            return response
+        return response
 
-    def _handle_two_factor_auth(self, response: requests.Response) -> None:
+    def _handle_two_factor_auth(self, http_session: Session, response: requests.Response) -> None:
         """Handle two-factor authentication."""
         if not self._otp:
             raise AuthenticationError("Two-factor authentication required but no OTP secret provided")
 
         two_factor_url = response.url
         form_action = two_factor_url[len(self._index_url):]
-        csrf_token = self._get_csrf_token(form_action)
+        csrf_token = self._get_csrf_token(http_session, form_action)
 
         # Try authentication with retries
         for attempt in range(_LOGIN_RETRY_ATTEMPTS):
             try:
                 auth_code = pyotp.TOTP(self._otp).now()
                 logging.debug(f"Attempting 2FA with code (attempt {attempt + 1}/{_LOGIN_RETRY_ATTEMPTS})")
 
-                with session_with_retries() as session:
-                    auth_response = session.post(
-                        two_factor_url,
-                        data={"csrf_token": csrf_token, "method": "totp", "totp_value": auth_code},
-                        headers={"referer": two_factor_url}
-                    )
-                    auth_response.raise_for_status()
-
-                    # Check if 2FA succeeded (redirected away from 2FA page)
-                    if auth_response.url != two_factor_url:
-                        logging.debug("Two-factor authentication successful")
-                        return
-
-                    if attempt < _LOGIN_RETRY_ATTEMPTS - 1:
-                        logging.debug(f"2FA code rejected, retrying in {_LOGIN_RETRY_DELAY} seconds...")
-                        time.sleep(_LOGIN_RETRY_DELAY)
+                auth_response = http_session.post(
+                    two_factor_url,
+                    data={"csrf_token": csrf_token, "method": "totp", "totp_value": auth_code},
+                    headers={"referer": two_factor_url}
+                )
+                auth_response.raise_for_status()
+
+                # Check if 2FA succeeded (redirected away from 2FA page)
+                if auth_response.url != two_factor_url:
+                    logging.debug("Two-factor authentication successful")
+                    return
+
+                if attempt < _LOGIN_RETRY_ATTEMPTS - 1:
+                    logging.debug(f"2FA code rejected, retrying in {_LOGIN_RETRY_DELAY} seconds...")
+                    time.sleep(_LOGIN_RETRY_DELAY)
 
             except RequestException as e:
                 if attempt == _LOGIN_RETRY_ATTEMPTS - 1:
@@ -491,14 +488,14 @@ def _handle_two_factor_auth(self, response: requests.Response) -> None:
 
         raise AuthenticationError("Two-factor authentication failed after all attempts")
 
-    def _delete_versions(self, versions_to_delete: Set[str]) -> None:
+    def _delete_versions(self, http_session: Session, versions_to_delete: Set[str]) -> None:
         """Delete the specified package versions."""
         logging.info(f"Starting deletion of {len(versions_to_delete)} development versions")
 
         failed_deletions = list()
         for version in sorted(versions_to_delete):
             try:
-                self._delete_single_version(version)
+                self._delete_single_version(http_session, version)
                 logging.info(f"Successfully deleted {self._package} version {version}")
             except Exception as e:
                 # Continue with other versions rather than failing completely
@@ -510,7 +507,7 @@ def _delete_versions(self, versions_to_delete: Set[str]) -> None:
                 f"Failed to delete {len(failed_deletions)}/{len(versions_to_delete)} versions: {failed_deletions}"
             )
 
-    def _delete_single_version(self, version: str) -> None:
+    def _delete_single_version(self, http_session: Session, version: str) -> None:
         """Delete a single package version."""
         # Safety check
         if not self._is_dev_version(version) or self._is_rc_version(version):
@@ -522,19 +519,18 @@ def _delete_single_version(self, version: str) -> None:
         form_action = f"/manage/project/{self._package}/release/{version}/"
         form_url = f"{self._index_url}{form_action}"
 
-        csrf_token = self._get_csrf_token(form_action)
-
-        with session_with_retries() as session:
-            # Submit deletion request
-            delete_response = session.post(
-                form_url,
-                data={
-                    "csrf_token": csrf_token,
-                    "confirm_delete_version": version,
-                },
-                headers={"referer": form_url}
-            )
-            delete_response.raise_for_status()
+        csrf_token = self._get_csrf_token(http_session, form_action)
+
+        # Submit deletion request
+        delete_response = http_session.post(
+            form_url,
+            data={
+                "csrf_token": csrf_token,
+                "confirm_delete_version": version,
+            },
+            headers={"referer": form_url}
+        )
+        delete_response.raise_for_status()
 
 
 def main() -> int:

tests/fast/test_pypi_cleanup.py

@@ -256,7 +256,8 @@ def test_execute_cleanup_dry_run(self, mock_determine, mock_fetch, mock_delete,
         mock_fetch.return_value = {"1.0.0.dev1"}
         mock_determine.return_value = {"1.0.0.dev1"}
 
-        result = cleanup_dryrun_max_2._execute_cleanup()
+        with session_with_retries() as session:
+            result = cleanup_dryrun_max_2._execute_cleanup(session)
 
         assert result == 0
         mock_fetch.assert_called_once()
@@ -266,7 +267,8 @@ def test_execute_cleanup_dry_run(self, mock_determine, mock_fetch, mock_delete,
     @patch('duckdb_packaging.pypi_cleanup.PyPICleanup._fetch_released_versions')
     def test_execute_cleanup_no_releases(self, mock_fetch, cleanup_dryrun_max_2):
         mock_fetch.return_value = {}
-        result = cleanup_dryrun_max_2._execute_cleanup()
+        with session_with_retries() as session:
+            result = cleanup_dryrun_max_2._execute_cleanup(session)
         assert result == 0
 
     @patch('requests.Session.get')
@@ -281,7 +283,8 @@ def test_fetch_released_versions_success(self, mock_get, cleanup_dryrun_max_2):
         }
         mock_get.return_value = mock_response
 
-        releases = cleanup_dryrun_max_2._fetch_released_versions()
+        with session_with_retries() as session:
+            releases = cleanup_dryrun_max_2._fetch_released_versions(session)
 
         assert releases == {"1.0.0", "1.0.0.dev1"}
 
@@ -293,46 +296,50 @@ def test_fetch_released_versions_not_found(self, mock_get, cleanup_dryrun_max_2)
         mock_get.return_value = mock_response
 
         with pytest.raises(PyPICleanupError, match="Failed to fetch package information"):
-            cleanup_dryrun_max_2._fetch_released_versions()
+            with session_with_retries() as session:
+                cleanup_dryrun_max_2._fetch_released_versions(session)
 
     @patch('duckdb_packaging.pypi_cleanup.PyPICleanup._get_csrf_token')
-    @patch('duckdb_packaging.pypi_cleanup.PyPICleanup._perform_login')
-    def test_authenticate_success(self, mock_login, mock_csrf, cleanup_max_2):
+    @patch('requests.Session.post')
+    def test_authenticate_success(self, mock_post, mock_csrf, cleanup_max_2):
         """Test successful authentication."""
         mock_csrf.return_value = "csrf123"
         mock_response = Mock()
         mock_response.url = "https://test.pypi.org/manage/"
-        mock_login.return_value = mock_response
+        mock_post.return_value = mock_response
 
-        cleanup_max_2._authenticate()  # Should not raise
+        with session_with_retries() as session:
+            cleanup_max_2._authenticate(session)  # Should not raise
+            mock_csrf.assert_called_once_with(session, "/account/login/")
 
-        mock_csrf.assert_called_once_with("/account/login/")
-        mock_login.assert_called_once_with("csrf123")
+        mock_post.assert_called_once()
+        assert mock_post.call_args.args[0].endswith('/account/login/')
 
     @patch('duckdb_packaging.pypi_cleanup.PyPICleanup._get_csrf_token')
-    @patch('duckdb_packaging.pypi_cleanup.PyPICleanup._perform_login')
+    @patch('requests.Session.post')
     @patch('duckdb_packaging.pypi_cleanup.PyPICleanup._handle_two_factor_auth')
-    def test_authenticate_with_2fa(self, mock_2fa, mock_login, mock_csrf, cleanup_max_2):
+    def test_authenticate_with_2fa(self, mock_2fa, mock_post, mock_csrf, cleanup_max_2):
         mock_csrf.return_value = "csrf123"
         mock_response = Mock()
         mock_response.url = "https://test.pypi.org/account/two-factor/totp"
-        mock_login.return_value = mock_response
+        mock_post.return_value = mock_response
 
-        cleanup_max_2._authenticate()
-
-        mock_2fa.assert_called_once_with(mock_response)
+        with session_with_retries() as session:
+            cleanup_max_2._authenticate(session)
+            mock_2fa.assert_called_once_with(session, mock_response)
 
     def test_authenticate_missing_credentials(self, cleanup_dryrun_max_2):
         with pytest.raises(AuthenticationError, match="Username and password are required"):
-            cleanup_dryrun_max_2._authenticate()
+            cleanup_dryrun_max_2._authenticate(None)
 
     @patch('duckdb_packaging.pypi_cleanup.PyPICleanup._delete_single_version')
     def test_delete_versions_success(self, mock_delete, cleanup_max_2):
         """Test successful version deletion."""
         versions = {"1.0.0.dev1", "1.0.0.dev2"}
         mock_delete.side_effect = [None, None]  # Successful deletions
 
-        cleanup_max_2._delete_versions(versions)
+        with session_with_retries() as session:
+            cleanup_max_2._delete_versions(session, versions)
 
         assert mock_delete.call_count == 2
 
@@ -343,12 +350,12 @@ def test_delete_versions_partial_failure(self, mock_delete, cleanup_max_2):
         mock_delete.side_effect = [None, Exception("Delete failed")]
 
         with pytest.raises(PyPICleanupError, match="Failed to delete 1/2 versions"):
-            cleanup_max_2._delete_versions(versions)
+            cleanup_max_2._delete_versions(None, versions)
 
     def test_delete_single_version_safety_check(self, cleanup_max_2):
         """Test single version deletion safety check."""
         with pytest.raises(PyPICleanupError, match="Refusing to delete non-\\[dev\\|rc\\] version"):
-            cleanup_max_2._delete_single_version("1.0.0")  # Non-dev version
+            cleanup_max_2._delete_single_version(None, "1.0.0")  # Non-dev version
 
 
 class TestArgumentParser: