refactor: extract aria2 URL preflight
This commit is contained in:
@@ -152,6 +152,42 @@ def aria2_action_error_is_ok(status: str, message: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def resolve_download_url_for_aria2(
|
||||
url: str,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
) -> tuple[str, Dict[str, str]]:
|
||||
"""Preflight an aria2 URL and validate every redirect before RPC handoff."""
|
||||
facade = _downloader_module()
|
||||
request_headers = facade.build_download_headers(url, headers)
|
||||
source_host = facade.urlparse(str(url or "")).hostname
|
||||
is_huggingface_source = facade.host_matches_domain(
|
||||
source_host,
|
||||
"huggingface.co",
|
||||
)
|
||||
response = None
|
||||
try:
|
||||
response, resolved_url, resolved_headers = facade.request_public_url(
|
||||
"GET",
|
||||
url,
|
||||
headers=request_headers,
|
||||
timeout=20,
|
||||
stream=True,
|
||||
trusted_sensitive_redirect_hosts=(
|
||||
facade.HF_XET_ARIA2_AUTH_HOSTS
|
||||
if is_huggingface_source
|
||||
else None
|
||||
),
|
||||
trusted_sensitive_redirect_headers=(
|
||||
{"authorization"} if is_huggingface_source else None
|
||||
),
|
||||
)
|
||||
response.raise_for_status()
|
||||
return resolved_url, resolved_headers
|
||||
finally:
|
||||
if response is not None:
|
||||
response.close()
|
||||
|
||||
|
||||
def read_aria2_version(executable: str) -> str:
|
||||
"""Read the installed aria2 version without raising process errors."""
|
||||
facade = _downloader_module()
|
||||
|
||||
+7
-34
@@ -14,7 +14,7 @@ import time # noqa: F401
|
||||
from collections import deque # noqa: F401
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Optional # noqa: F401
|
||||
from urllib.parse import urlparse
|
||||
from urllib.parse import urlparse # noqa: F401
|
||||
|
||||
import requests # noqa: F401
|
||||
|
||||
@@ -78,6 +78,9 @@ from .download.aria2_backend import (
|
||||
from .download.aria2_backend import (
|
||||
resolve_aria2c_executable as _resolve_aria2c_executable, # noqa: F401
|
||||
)
|
||||
from .download.aria2_backend import (
|
||||
resolve_download_url_for_aria2 as _resolve_download_url_for_aria2, # noqa: F401
|
||||
)
|
||||
from .download.aria2_backend import (
|
||||
run_aria2_desired_state_worker as _run_aria2_desired_state_worker, # noqa: F401
|
||||
)
|
||||
@@ -127,8 +130,8 @@ from .download.previews import (
|
||||
get_existing_model_preview_path, # noqa: F401
|
||||
)
|
||||
from .network_utils import (
|
||||
host_matches_domain,
|
||||
request_public_url,
|
||||
host_matches_domain, # noqa: F401
|
||||
request_public_url, # noqa: F401
|
||||
validate_public_http_url, # noqa: F401
|
||||
)
|
||||
from .path_utils import (
|
||||
@@ -204,7 +207,7 @@ from .download.validation import (
|
||||
_is_sensitive_metadata_key, # noqa: F401
|
||||
_sanitize_download_error, # noqa: F401
|
||||
_strip_sensitive_url_params, # noqa: F401
|
||||
build_download_headers,
|
||||
build_download_headers, # noqa: F401
|
||||
is_allowed_model_download_filename, # noqa: F401
|
||||
sanitize_download_filename, # noqa: F401
|
||||
)
|
||||
@@ -215,36 +218,6 @@ from .settings import (
|
||||
)
|
||||
|
||||
|
||||
def _resolve_download_url_for_aria2(
|
||||
url: str,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
) -> tuple[str, Dict[str, str]]:
|
||||
"""Preflight an aria2 URL and validate every redirect before RPC handoff."""
|
||||
request_headers = build_download_headers(url, headers)
|
||||
source_host = urlparse(str(url or "")).hostname
|
||||
is_huggingface_source = host_matches_domain(source_host, "huggingface.co")
|
||||
response = None
|
||||
try:
|
||||
response, resolved_url, resolved_headers = request_public_url(
|
||||
"GET",
|
||||
url,
|
||||
headers=request_headers,
|
||||
timeout=20,
|
||||
stream=True,
|
||||
trusted_sensitive_redirect_hosts=(
|
||||
HF_XET_ARIA2_AUTH_HOSTS if is_huggingface_source else None
|
||||
),
|
||||
trusted_sensitive_redirect_headers=(
|
||||
{"authorization"} if is_huggingface_source else None
|
||||
),
|
||||
)
|
||||
response.raise_for_status()
|
||||
return resolved_url, resolved_headers
|
||||
finally:
|
||||
if response is not None:
|
||||
response.close()
|
||||
|
||||
|
||||
def write_model_resolver_metadata(
|
||||
dest_path: str,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
|
||||
Reference in New Issue
Block a user