Add PyPI fallback for APIStub wheel downloads

eng/tools/azure-sdk-tools/azpysdk/apistub.py

@@ -13,6 +13,8 @@
 from ci_tools.parsing import ParsedSetup
 
 REPO_ROOT = discover_repo_root()
+AZURE_SDK_INDEX_URL = "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/"
+PYPI_INDEX_URL = "https://pypi.org/simple/"
 
 
 def get_package_wheel_path(pkg_root: str) -> str:
@@ -103,22 +105,35 @@ def ensure_apistub_dependencies(self, executable: str, package_dir: str, staging
 
     def download_pypi_wheel(self, executable: str, package_name: str, version: str, staging_directory: str) -> str:
         """Download a released wheel from PyPI into the staging directory and return its path."""
-        logger.info(f"Downloading {package_name}=={version} from PyPI.")
-        self.run_venv_command(
-            executable,
-            [
-                "-m",
-                "pip",
-                "download",
-                f"{package_name}=={version}",
-                "--no-deps",
-                "--only-binary=:all:",
-                "-d",
-                staging_directory,
-            ],
-            cwd=staging_directory,
-            check=True,
-        )
+        for index_url in (AZURE_SDK_INDEX_URL, PYPI_INDEX_URL):
+            logger.info(f"Downloading {package_name}=={version} from {index_url}.")
+            try:
+                self.run_venv_command(
+                    executable,
+                    [
+                        "-m",
+                        "pip",
+                        "download",
+                        f"{package_name}=={version}",
+                        "--no-deps",
+                        "--only-binary=:all:",
+                        f"--index-url={index_url}",
+                        "-d",
+                        staging_directory,
+                    ],
+                    cwd=staging_directory,
+                    check=True,
+                    additional_environment_settings={"PIP_EXTRA_INDEX_URL": ""},
+                )
+                break
+            except CalledProcessError as error:
+                if index_url == PYPI_INDEX_URL:
+                    error_details = error.stderr or error.stdout or str(error)
+                    logger.error(
+                        f"Failed to download {package_name}=={version} from both package indexes: {error_details}"
+                    )
+                    raise
+                logger.warning(f"Failed to download from the Azure SDK feed: {error}. Retrying from public PyPI.")
         found_whl = find_whl(staging_directory, package_name, version)
         if not found_whl:
             raise FileNotFoundError(

eng/tools/azure-sdk-tools/tests/test_apistub.py

@@ -496,8 +496,43 @@ def test_download_pypi_wheel_runs_pip_download(self, _find_whl, tmp_path):
         cmds = run_venv_command.call_args.args[1]
         assert cmds[0:4] == ["-m", "pip", "download", "azure-core==1.0.0"]
         assert "--no-deps" in cmds
+        assert (
+            "--index-url=https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/"
+            in cmds
+        )
+        assert run_venv_command.call_args.kwargs["additional_environment_settings"] == {"PIP_EXTRA_INDEX_URL": ""}
         assert result == os.path.join(staging, "azure_core-1.0.0-py3-none-any.whl")
 
+    @patch("azpysdk.apistub.find_whl", return_value="azure_core-1.0.0-py3-none-any.whl")
+    def test_download_pypi_wheel_falls_back_to_public_pypi(self, _find_whl, tmp_path):
+        """download_pypi_wheel should retry public PyPI when the Azure SDK feed fails."""
+        stub = apistub()
+
+        with patch.object(stub, "run_venv_command", side_effect=[CalledProcessError(1, "pip"), None]) as run:
+            result = stub.download_pypi_wheel(sys.executable, "azure-core", "1.0.0", str(tmp_path))
+
+        assert run.call_count == 2
+        assert "--index-url=https://pypi.org/simple/" in run.call_args_list[1].args[1]
+        assert all(
+            call.kwargs["additional_environment_settings"] == {"PIP_EXTRA_INDEX_URL": ""} for call in run.call_args_list
+        )
+        assert result == os.path.join(str(tmp_path), "azure_core-1.0.0-py3-none-any.whl")
+
+    def test_download_pypi_wheel_reports_both_index_failures(self, tmp_path, caplog):
+        """download_pypi_wheel should propagate the public PyPI failure after both indexes fail."""
+        stub = apistub()
+        public_pypi_error = CalledProcessError(2, "pip", stderr="public PyPI unavailable")
+
+        with patch.object(
+            stub, "run_venv_command", side_effect=[CalledProcessError(1, "pip"), public_pypi_error]
+        ) as run:
+            with pytest.raises(CalledProcessError) as exc_info:
+                stub.download_pypi_wheel(sys.executable, "azure-core", "1.0.0", str(tmp_path))
+
+        assert run.call_count == 2
+        assert exc_info.value is public_pypi_error
+        assert "public PyPI unavailable" in caplog.text
+
     @patch("azpysdk.apistub.find_whl", return_value=None)
     def test_download_pypi_wheel_raises_when_no_wheel(self, _find_whl, tmp_path):
         """download_pypi_wheel should raise FileNotFoundError when no wheel is downloaded."""