dataset-uploader/scripts/download_geonames_robust.py

#!/usr/bin/env python3
"""Robust download script for large GeoNames RDF file with resume capability."""

import sys
import time
from pathlib import Path

import httpx
from rich.console import Console
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    TextColumn,
    TimeElapsedColumn,
    TransferSpeedColumn,
)

console = Console()


def download_with_resume(url: str, destination: Path, max_retries: int = 5) -> bool:
    """Download file with resume capability and retries.

    Args:
        url: URL to download from
        destination: Path to save file
        max_retries: Maximum number of retry attempts

    Returns:
        True if successful, False otherwise
    """
    destination.parent.mkdir(parents=True, exist_ok=True)

    for attempt in range(max_retries):
        try:
            # Check if partial file exists
            start_byte = 0
            if destination.exists():
                start_byte = destination.stat().st_size
                console.print(
                    f"[yellow]Found partial download: "
                    f"{start_byte / (1024 * 1024):.1f} MB[/yellow]"
                )

            # Get total file size
            with httpx.Client(timeout=30.0) as client:
                response = client.head(url)
                total_size = int(response.headers.get("content-length", 0))

                if total_size == 0:
                    console.print("[red]Error: Cannot determine file size[/red]")
                    return False

                console.print(
                    f"[cyan]Total file size: {total_size / (1024 * 1024):.1f} MB[/cyan]"
                )

                # Check if already complete
                if start_byte >= total_size:
                    console.print("[green]✓ File already completely downloaded[/green]")
                    return True

                # Download with resume
                headers = {}
                if start_byte > 0:
                    headers["Range"] = f"bytes={start_byte}-"
                    console.print(f"[yellow]Resuming from byte {start_byte}[/yellow]")

                with Progress(
                    TextColumn("[bold blue]{task.description}"),
                    BarColumn(),
                    DownloadColumn(),
                    TransferSpeedColumn(),
                    TimeElapsedColumn(),
                    console=console,
                ) as progress:
                    task = progress.add_task("Downloading", total=total_size)
                    progress.update(task, completed=start_byte)

                    # Use streaming with larger timeout and chunk size
                    with client.stream(
                        "GET", url, headers=headers, timeout=httpx.Timeout(60.0)
                    ) as response:
                        # Open file in append mode if resuming
                        mode = "ab" if start_byte > 0 else "wb"
                        with open(destination, mode) as f:
                            for chunk in response.iter_bytes(
                                chunk_size=1024 * 1024
                            ):  # 1MB chunks
                                f.write(chunk)
                                progress.update(task, advance=len(chunk))

                # Verify final size
                final_size = destination.stat().st_size
                if final_size == total_size:
                    console.print(
                        f"[green]✓ Download complete: "
                        f"{final_size / (1024 * 1024):.1f} MB[/green]"
                    )
                    return True
                else:
                    console.print(
                        f"[yellow]Warning: File size mismatch "
                        f"({final_size} vs {total_size} bytes)[/yellow]"
                    )
                    if attempt < max_retries - 1:
                        console.print(
                            f"[yellow]Retrying... (attempt {attempt + 2}/"
                            f"{max_retries})[/yellow]"
                        )
                        continue

        except (httpx.RequestError, httpx.HTTPStatusError) as e:
            console.print(f"[red]Download error: {e}[/red]")
            if attempt < max_retries - 1:
                wait_time = min(30, (attempt + 1) * 10)
                console.print(
                    f"[yellow]Waiting {wait_time} seconds before retry...[/yellow]"
                )
                time.sleep(wait_time)
                console.print(
                    f"[yellow]Retrying... (attempt {attempt + 2}/"
                    f"{max_retries})[/yellow]"
                )
            else:
                console.print(f"[red]Failed after {max_retries} attempts[/red]")
                return False
        except KeyboardInterrupt:
            console.print("[red]Download interrupted by user[/red]")
            return False
        except Exception as e:
            console.print(f"[red]Unexpected error: {e}[/red]")
            return False

    return False


def main():
    """Download GeoNames RDF file."""
    url = "http://download.geonames.org/all-geonames-rdf.zip"
    destination = Path(
        "dataset_processing/downloads/geonames/geonames/all-geonames-rdf.zip"
    )

    console.print("\n[bold cyan]GeoNames RDF Download[/bold cyan]")
    console.print(f"URL: {url}")
    console.print(f"Destination: {destination}\n")

    success = download_with_resume(url, destination, max_retries=5)

    if success:
        # Try to extract
        console.print("\n[yellow]Attempting extraction...[/yellow]")
        import zipfile

        try:
            with zipfile.ZipFile(destination, "r") as zf:
                # Check if it's a valid zip
                if zf.testzip() is not None:
                    console.print("[red]ZIP file has errors[/red]")
                    return 1

                # List contents
                file_list = zf.namelist()
                console.print(f"[cyan]ZIP contains {len(file_list)} files[/cyan]")

                # Extract
                extract_dir = destination.parent
                console.print(f"[yellow]Extracting to {extract_dir}...[/yellow]")
                zf.extractall(extract_dir)
                console.print("[green]✓ Extraction complete[/green]")

                # List extracted files
                for fname in file_list[:5]:  # Show first 5
                    console.print(f"  - {fname}")
                if len(file_list) > 5:
                    console.print(f"  ... and {len(file_list) - 5} more files")

                return 0
        except zipfile.BadZipFile as e:
            console.print(f"[red]Error: File is not a valid ZIP: {e}[/red]")
            return 1
    else:
        return 1


if __name__ == "__main__":
    sys.exit(main())