Files
dataset-uploader/scripts/handle_partial_downloads.py

81 lines
2.7 KiB
Python

#!/usr/bin/env python3
"""Handle partial or corrupted RDF downloads by extracting usable data."""
import gzip
import sys
from pathlib import Path
from rich.console import Console
console = Console()
def extract_partial_gz(gz_file: Path, output_file: Path, max_errors: int = 10):
"""Extract as much as possible from a potentially corrupted gzip file.
Args:
gz_file: Path to the gzip file
output_file: Path to output the extracted data
max_errors: Maximum number of consecutive errors before stopping
Returns:
Number of lines extracted
"""
lines_extracted = 0
consecutive_errors = 0
try:
with gzip.open(gz_file, "rt", encoding="utf-8", errors="ignore") as f_in:
with open(output_file, "w", encoding="utf-8") as f_out:
while consecutive_errors < max_errors:
try:
line = f_in.readline()
if not line:
break
f_out.write(line)
lines_extracted += 1
consecutive_errors = 0 # Reset error counter on success
# Show progress every 100k lines
if lines_extracted % 100000 == 0:
console.print(f"[dim]Extracted {lines_extracted:,} lines...[/dim]")
except Exception as e:
consecutive_errors += 1
if consecutive_errors >= max_errors:
console.print(f"[yellow]Stopping after {consecutive_errors} consecutive errors[/yellow]")
break
except EOFError:
console.print(f"[yellow]Reached end of compressed data[/yellow]")
except Exception as e:
console.print(f"[red]Error opening file: {e}[/red]")
return lines_extracted
def main():
if len(sys.argv) != 3:
console.print("Usage: python handle_partial_downloads.py <input.gz> <output>")
sys.exit(1)
input_file = Path(sys.argv[1])
output_file = Path(sys.argv[2])
if not input_file.exists():
console.print(f"[red]Input file not found: {input_file}[/red]")
sys.exit(1)
console.print(f"[cyan]Extracting partial data from: {input_file}[/cyan]")
console.print(f"[cyan]Output: {output_file}[/cyan]")
lines = extract_partial_gz(input_file, output_file)
if lines > 0:
output_size_mb = output_file.stat().st_size / (1024 * 1024)
console.print(f"[green]✓ Successfully extracted {lines:,} lines ({output_size_mb:.2f} MB)[/green]")
else:
console.print(f"[red]No data could be extracted[/red]")
sys.exit(1)
if __name__ == "__main__":
main()