286 lines
10 KiB
Python
286 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from rich.console import Console # noqa: E402
|
|
|
|
from scripts.dataset_registry import DatasetInfo, DATASET_REGISTRY
|
|
|
|
console = Console()
|
|
|
|
def download(
|
|
dataset_id: str,
|
|
base_dir: Path,
|
|
skip_download: bool = False,
|
|
dry_run: bool = False,
|
|
) -> tuple[bool, bool, Path | None]:
|
|
"""
|
|
This function downloads the dataset from the source URL.
|
|
|
|
It is responsible only for downloading the dataset and checking that the
|
|
dataset looks reasonable.
|
|
|
|
At the end of the function, the (possibly compressed) dataset should be available
|
|
in the downloads directory.
|
|
|
|
Returns:
|
|
Whether the download was successful.
|
|
Whether there should be another attempt to download the dataset.
|
|
The name of the downloaded file, if successful.
|
|
"""
|
|
dataset_info = get_dataset_info(dataset_id)
|
|
rdf_file = None
|
|
|
|
if not dataset_info:
|
|
return (False, False, None)
|
|
|
|
# Check if dataset is available
|
|
if not dataset_info.available:
|
|
console.print(f"\n[bold yellow]⊘ Skipping '{dataset_info.name}'[/bold yellow]")
|
|
console.print("[yellow]Status:[/yellow] Dataset currently unavailable")
|
|
console.print(f"[yellow]Reason:[/yellow] {dataset_info.notes}")
|
|
console.print(f"[dim]URL: {dataset_info.url}[/dim]")
|
|
return (False, False, None)
|
|
|
|
# Create directories if necessary
|
|
console.print(
|
|
f"\n[bold cyan]{'[DRY RUN] ' if dry_run else ''}"
|
|
f"Processing: {dataset_info.name}[/bold cyan]"
|
|
)
|
|
|
|
# Setup directories
|
|
download_dir = base_dir / "downloads" / dataset_id
|
|
try:
|
|
download_dir.mkdir(parents=True, exist_ok=True)
|
|
except OSError as e:
|
|
console.print(f"[red]Error creating directories: {e}[/red]")
|
|
console.print(f"[yellow]Download dir: {download_dir}[/yellow]")
|
|
return (False, False, None)
|
|
|
|
# Download the dataset
|
|
if not skip_download:
|
|
path = str(Path(__file__).parent / "rdf_dataset_downloader.py")
|
|
console.print("\n[yellow]Step 1: Downloading raw RDF dataset...[/yellow]")
|
|
if dry_run:
|
|
console.print(
|
|
f"[dim]Would run: {path} "
|
|
f"{dataset_id} -o {download_dir}[/dim]"
|
|
)
|
|
return (True, False, None)
|
|
else:
|
|
console.print(
|
|
f"[dim]Running: {path} {dataset_id} -o {download_dir}[/dim]"
|
|
)
|
|
try:
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
path,
|
|
dataset_id,
|
|
"-o",
|
|
str(download_dir),
|
|
],
|
|
capture_output=False,
|
|
timeout=3600 * 24, # 24 hour timeout
|
|
)
|
|
if result.returncode != 0:
|
|
console.print(
|
|
f"[red]Download failed for {dataset_id} "
|
|
f"(exit code: {result.returncode})[/red]"
|
|
)
|
|
console.print(
|
|
"[yellow]Try running manually: "
|
|
"python scripts/rdf_dataset_downloader.py "
|
|
f"{dataset_id} -o {download_dir}[/yellow]"
|
|
)
|
|
return (False, False, None)
|
|
except subprocess.TimeoutExpired:
|
|
console.print(
|
|
f"[red]Download timeout for {dataset_id} (exceeded 24 hours)[/red]"
|
|
)
|
|
console.print(
|
|
"[yellow]Large datasets may need more time. "
|
|
"Try running manually using rdf_dataset_downloader.py "
|
|
"with --force[/yellow]"
|
|
)
|
|
return (False, True, None)
|
|
except FileNotFoundError:
|
|
console.print(
|
|
"[red]Download script not found: rdf_dataset_downloader.py[/red]"
|
|
)
|
|
console.print(
|
|
"[yellow]Make sure you're running "
|
|
"from the correct directory[/yellow]"
|
|
)
|
|
return (False, False, None)
|
|
except Exception as e:
|
|
console.print(f"[red]Unexpected error during download: {e}[/red]")
|
|
return (False, False, None)
|
|
|
|
# Check the dataset.
|
|
# Find the downloaded file (only if we need it for conversion)
|
|
|
|
# This code assumes that the most recent file is the one we want.
|
|
rdf_file = get_most_recent_file(download_dir) if not skip_download else None
|
|
|
|
if rdf_file is None:
|
|
console.print(f"[red]No RDF file found in {download_dir}[/red]")
|
|
return (False, False, None)
|
|
|
|
if str(rdf_file).endswith(".bz2"):
|
|
# Verify that the file is not corrupted
|
|
console.print(
|
|
"[yellow]Found a BZ2 file, checking integrity...[/yellow]"
|
|
)
|
|
|
|
# TEST 1: Check the size of the file
|
|
file_size_gb = rdf_file.stat().st_size / (1024**3)
|
|
if dataset_info.compressed_size_gb:
|
|
expected_gb = dataset_info.compressed_size_gb
|
|
if file_size_gb < expected_gb * 0.9:
|
|
console.print(
|
|
f"[red]Incomplete BZ2 file: {rdf_file.name}[/red]"
|
|
)
|
|
console.print(
|
|
f"[yellow]File size: {file_size_gb:.3f} GB, "
|
|
f"Expected: ~{expected_gb:.2f} GB[/yellow]"
|
|
)
|
|
console.print(
|
|
"[red]File is only "
|
|
f"{(file_size_gb / expected_gb * 100):.1f}% "
|
|
"of expected size[/red]"
|
|
)
|
|
console.print(
|
|
"[yellow]Removing incomplete file "
|
|
"and re-downloading...[/yellow]"
|
|
)
|
|
rdf_file.unlink()
|
|
return (False, True, None) # Re-download.
|
|
|
|
# TEST 2: Check the integrity of the file
|
|
try:
|
|
# Test the bz2 file integrity
|
|
test_result = subprocess.run(
|
|
["bunzip2", "-t", str(rdf_file)], capture_output=True, text=True
|
|
)
|
|
if test_result.returncode != 0:
|
|
console.print(f"[red]Corrupted BZ2 file: {rdf_file.name}[/red]")
|
|
console.print(
|
|
"[yellow]File size: "
|
|
f"{rdf_file.stat().st_size / (1024**2):.2f} MB[/yellow]"
|
|
)
|
|
console.print(
|
|
"[yellow]Removing corrupted file "
|
|
"and attempting re-download...[/yellow]"
|
|
)
|
|
rdf_file.unlink()
|
|
return (False, True, None) # Re-download.
|
|
except FileNotFoundError:
|
|
console.print("[red]bunzip2 not found: please install it[/red]")
|
|
return (False, False, None)
|
|
except Exception as e:
|
|
if (
|
|
"end-of-stream marker" in str(e).lower()
|
|
or "compressed file ended" in str(e).lower()
|
|
):
|
|
console.print(
|
|
"[yellow]The file is incomplete. "
|
|
"Attempting to re-download...[/yellow]"
|
|
)
|
|
rdf_file.unlink()
|
|
return (False, True, None)
|
|
|
|
console.print(f"[red]Unexpected error checking BZ2 file: {e}[/red]")
|
|
return (False, False, None)
|
|
else:
|
|
console.print("[dim]Skipping download[/dim]")
|
|
|
|
return (True, False, rdf_file)
|
|
|
|
|
|
def decompress(rdf_file: Path) -> bool:
|
|
"""
|
|
This function decompresses the dataset if necessary.
|
|
|
|
At the end of the function, the uncompressid dataset should be available
|
|
in the downloads directory.
|
|
|
|
Returns:
|
|
Whether the decompression was successful.
|
|
"""
|
|
decompress_result = None
|
|
|
|
if str(rdf_file).endswith(".bz2"):
|
|
decompress_result = subprocess.run(
|
|
["bunzip2", "-d", "-f", str(rdf_file)],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
if str(rdf_file).endswith(".gz"):
|
|
decompress_result = subprocess.run(
|
|
["gunzip", "-f", str(rdf_file)], capture_output=True, text=True
|
|
)
|
|
|
|
if str(rdf_file).endswith(".zip"):
|
|
decompress_result = subprocess.run(
|
|
["unzip", "-o", "-f", str(rdf_file)], capture_output=True, text=True
|
|
)
|
|
|
|
if decompress_result is not None and decompress_result.returncode != 0:
|
|
console.print(f"[red]Error decompressing {rdf_file}[/red]")
|
|
console.print("[yellow]Output:[/yellow]")
|
|
console.print(decompress_result.stdout)
|
|
console.print("[yellow]Error output:[/yellow]")
|
|
console.print(decompress_result.stderr)
|
|
return False
|
|
return True
|
|
|
|
|
|
def get_dataset_info(dataset_id: str) -> DatasetInfo | None:
|
|
"""
|
|
Gets the dataset info from the dataset registry.
|
|
Args:
|
|
dataset_id: An identifier for the dataset
|
|
|
|
Returns:
|
|
The dataset info, or None if not found.
|
|
"""
|
|
try:
|
|
dataset_info = DATASET_REGISTRY.get(dataset_id)
|
|
if not dataset_info:
|
|
console.print(f"[red]Dataset '{dataset_id}' not found in registry[/red]")
|
|
console.print(
|
|
"[yellow]Available datasets: "
|
|
f"{', '.join(list(DATASET_REGISTRY.keys())[:5])}...[/yellow]"
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
console.print(f"[red]Error accessing dataset registry: {e}[/red]")
|
|
return None
|
|
|
|
return dataset_info
|
|
|
|
|
|
def get_most_recent_file(download_dir: Path) -> Path | None:
|
|
"""
|
|
Get the file with the most recent creation date in a directory.
|
|
|
|
Args:
|
|
download_dir: Path object pointing to the directory to search
|
|
|
|
Returns:
|
|
Path to the most recent file, or None if no files found
|
|
"""
|
|
# Get all files (not directories) in the directory and subdirectories
|
|
files = [f for f in download_dir.rglob("*") if f.is_file()]
|
|
|
|
if not files:
|
|
return None
|
|
|
|
# Find the file with the most recent creation time
|
|
most_recent_file = max(files, key=lambda f: f.stat().st_ctime)
|
|
|
|
return most_recent_file
|