mtgonline/backend/scripts/download_mtgjson_data.py

#!/usr/bin/env python3
"""
MTGJSON Data Downloader

Downloads MTGJSON data files into the container for loading.
Handles decompression of gzip files and unzipping of zip files.

Usage:
    python download_mtgjson_data.py
"""

import asyncio
import json
import gzip
import logging
import os
import sys
import zipfile
from pathlib import Path
from urllib.request import urlretrieve, urlopen

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# MTGJSON download URLs
MTGJSON_BASE_URL = "https://mtgjson.com/api/5x"
MTGJSON_FILES = {
    "AllPrintings.psql.gz": "AllPrintings.psql.gz",
    "AllSetFiles.zip": "AllSetFiles.zip",
    "AllDeckFiles.zip": "AllDeckFiles.zip",
    "AllIdentifiers.json.gz": "AllIdentifiers.json.gz",
    "CardTypes.json.gz": "CardTypes.json.gz",
    "DeckList.json.gz": "DeckList.json.gz",
    "Keywords.json.gz": "Keywords.json.gz",
    "SetList.json.gz": "SetList.json.gz",
}


def download_file(url, dest_dir, filename):
    """Download a file from URL to destination directory"""
    dest_path = dest_dir / filename
    if dest_path.exists():
        logger.info(f"{filename} already exists, skipping download")
        return True
    
    logger.info(f"Downloading {filename}...")
    try:
        urlretrieve(url, dest_path)
        logger.info(f"Downloaded {filename} to {dest_path}")
        return True
    except Exception as e:
        logger.error(f"Failed to download {filename}: {e}")
        return False


def extract_zip(zip_path, dest_dir):
    """Extract a zip file to destination directory"""
    logger.info(f"Extracting {zip_path.name} to {dest_dir}...")
    try:
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(dest_dir)
        logger.info(f"Extracted {zip_path.name} successfully")
        # Remove the zip file after extraction
        zip_path.unlink()
        return True
    except Exception as e:
        logger.error(f"Failed to extract {zip_path}: {e}")
        return False


def extract_gzip(gz_path, dest_path=None):
    """Extract a gzip file"""
    if not dest_path:
        dest_path = gz_path.with_suffix('')
    
    logger.info(f"Extracting {gz_path.name} to {dest_path}...")
    try:
        with gzip.open(gz_path, 'rt', encoding='utf-8') as f_in:
            with open(dest_path, 'w', encoding='utf-8') as f_out:
                f_out.write(f_in.read())
        logger.info(f"Extracted {gz_path.name} to {dest_path}")
        # Remove the gz file after extraction
        gz_path.unlink()
        return True
    except Exception as e:
        logger.error(f"Failed to extract {gz_path}: {e}")
        return False


async def download_all():
    """Download and extract all MTGJSON data files"""
    logger.info("Starting MTGJSON data download...")
    
    # Create data directory
    data_dir = Path("/app/data/mtgjson")
    data_dir.mkdir(parents=True, exist_ok=True)
    
    # Download AllPrintings.psql.gz
    if download_file(MTGJSON_BASE_URL + "/AllPrintings.psql.gz", data_dir, "AllPrintings.psql.gz"):
        # Extract gzip
        gz_path = data_dir / "AllPrintings.psql.gz"
        if extract_gzip(gz_path, data_dir / "AllPrintings.psql"):
            logger.info("AllPrintings.psql extracted successfully")
    
    # Download AllSetFiles.zip
    if download_file(MTGJSON_BASE_URL + "/AllSetFiles.zip", data_dir, "AllSetFiles.zip"):
        zip_path = data_dir / "AllSetFiles.zip"
        allsetfiles_dir = data_dir / "allsetfiles"
        allsetfiles_dir.mkdir(exist_ok=True)
        if extract_zip(zip_path, allsetfiles_dir):
            logger.info("AllSetFiles extracted successfully")
    
    # Download AllDeckFiles.zip
    if download_file(MTGJSON_BASE_URL + "/AllDeckFiles.zip", data_dir, "AllDeckFiles.zip"):
        zip_path = data_dir / "AllDeckFiles.zip"
        alldeckfiles_dir = data_dir / "alldeckfiles"
        alldeckfiles_dir.mkdir(exist_ok=True)
        if extract_zip(zip_path, alldeckfiles_dir):
            logger.info("AllDeckFiles extracted successfully")
    
    # Download AllIdentifiers.json.gz
    if download_file(MTGJSON_BASE_URL + "/AllIdentifiers.json.gz", data_dir, "AllIdentifiers.json.gz"):
        gz_path = data_dir / "AllIdentifiers.json.gz"
        if extract_gzip(gz_path, data_dir / "AllIdentifiers.json"):
            logger.info("AllIdentifiers.json extracted successfully")
    
    # Download CardTypes.json.gz
    if download_file(MTGJSON_BASE_URL + "/CardTypes.json.gz", data_dir, "CardTypes.json.gz"):
        gz_path = data_dir / "CardTypes.json.gz"
        if extract_gzip(gz_path, data_dir / "CardTypes.json"):
            logger.info("CardTypes.json extracted successfully")
    
    # Download DeckList.json.gz
    if download_file(MTGJSON_BASE_URL + "/DeckList.json.gz", data_dir, "DeckList.json.gz"):
        gz_path = data_dir / "DeckList.json.gz"
        if extract_gzip(gz_path, data_dir / "DeckList.json"):
            logger.info("DeckList.json extracted successfully")
    
    # Download Keywords.json.gz
    if download_file(MTGJSON_BASE_URL + "/Keywords.json.gz", data_dir, "Keywords.json.gz"):
        gz_path = data_dir / "Keywords.json.gz"
        if extract_gzip(gz_path, data_dir / "Keywords.json"):
            logger.info("Keywords.json extracted successfully")
    
    # Download SetList.json.gz
    if download_file(MTGJSON_BASE_URL + "/SetList.json.gz", data_dir, "SetList.json.gz"):
        gz_path = data_dir / "SetList.json.gz"
        if extract_gzip(gz_path, data_dir / "SetList.json"):
            logger.info("SetList.json extracted successfully")
    
    logger.info("MTGJSON data download and extraction complete!")


if __name__ == "__main__":
    asyncio.run(download_all())