JKTV-online/database/pipeline.py

import argparse
import fcntl
import os
import sqlite3
import subprocess
import sys
import time
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from database.job_store import JobStore
from database.maintenance import (
    check_managed_databases,
    create_backup_set,
    prune_backup_sets,
    restore_backup_set,
)
from database.paths import L1_DB, L2_DB, PIPELINE_LOCK


STAGES = (
    ('l1', 15, Path('database/L1/L1_Builder.py')),
    ('l2', 55, Path('database/L2/L2_Builder.py')),
    ('l3', 85, Path('database/L3/L3_Builder.py')),
)


class PipelineError(RuntimeError):
    pass


def _run_stage(store, job_id, stage, progress, script_path, replace=False):
    store.update_progress(
        job_id,
        stage,
        progress,
        f'Running {stage.upper()} builder',
    )
    command = [sys.executable, str(PROJECT_ROOT / script_path)]
    if stage == 'l1' and replace:
        command.append('--force')
    if stage == 'l3':
        command.append('--no-backup')

    started = time.monotonic()
    result = subprocess.run(
        command,
        cwd=str(PROJECT_ROOT),
        capture_output=True,
        text=True,
        timeout=1200,
    )
    duration = time.monotonic() - started
    store.append_log(
        job_id,
        (
            f'\n===== {stage.upper()} ({duration:.2f}s) =====\n'
            f'{result.stdout}\n{result.stderr}'
        ),
    )
    if result.returncode != 0:
        raise PipelineError(
            f'{stage.upper()} builder exited with code {result.returncode}'
        )


def _validate_pipeline_output(match_id=None):
    database_status = check_managed_databases()
    failures = [
        f"{name}: {status['quick_check']}"
        for name, status in database_status.items()
        if status['quick_check'] != 'ok'
    ]
    if failures:
        raise PipelineError(
            'Database quick_check failed: ' + ', '.join(failures)
        )

    if not match_id:
        return

    l1 = sqlite3.connect(str(L1_DB))
    l2 = sqlite3.connect(str(L2_DB))
    try:
        raw_count = l1.execute(
            'SELECT COUNT(*) FROM raw_iframe_network WHERE match_id = ?',
            [match_id],
        ).fetchone()[0]
        match_count = l2.execute(
            'SELECT COUNT(*) FROM fact_matches WHERE match_id = ?',
            [match_id],
        ).fetchone()[0]
        player_count = l2.execute(
            'SELECT COUNT(*) FROM fact_match_players WHERE match_id = ?',
            [match_id],
        ).fetchone()[0]
        round_count = l2.execute(
            'SELECT COUNT(*) FROM fact_rounds WHERE match_id = ?',
            [match_id],
        ).fetchone()[0]
    finally:
        l1.close()
        l2.close()

    if raw_count != 1:
        raise PipelineError(f'L1 does not contain imported match {match_id}')
    if match_count != 1:
        raise PipelineError(f'L2 does not contain imported match {match_id}')
    if player_count != 10:
        raise PipelineError(
            f'Imported match has {player_count} players; expected 10'
        )
    if round_count <= 0:
        raise PipelineError('Imported match has no round facts')


def run_pipeline(job_id, match_id=None, replace=False):
    store = JobStore()
    job = store.get_job(job_id)
    if not job:
        raise PipelineError(f'Unknown ETL job: {job_id}')

    started = time.monotonic()
    backup_dir = None
    PIPELINE_LOCK.parent.mkdir(parents=True, exist_ok=True)
    lock_file = PIPELINE_LOCK.open('w')
    try:
        try:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise PipelineError('Another database pipeline is already running') from exc

        store.start_job(job_id, 'backup', 'Creating rollback snapshot')
        backup_dir = create_backup_set(f'job-{job_id}')
        store.append_log(job_id, f'Backup created: {backup_dir}\n')

        for stage, progress, script_path in STAGES:
            _run_stage(
                store,
                job_id,
                stage,
                progress,
                script_path,
                replace=replace,
            )

        store.update_progress(
            job_id,
            'validation',
            95,
            'Validating imported data',
        )
        _validate_pipeline_output(match_id)
        removed_backups = prune_backup_sets(keep=3)
        if removed_backups:
            store.append_log(
                job_id,
                f"Pruned old backups: {', '.join(removed_backups)}\n",
            )
        duration = time.monotonic() - started
        store.finish_job(
            job_id,
            True,
            'Pipeline completed and validated',
            duration,
        )
        return True
    except Exception as exc:
        store.append_log(job_id, f'\nPIPELINE FAILED: {exc}\n')
        if backup_dir:
            try:
                restore_backup_set(backup_dir)
                store.append_log(job_id, 'Rollback snapshot restored successfully\n')
            except Exception as restore_exc:
                store.append_log(
                    job_id,
                    f'ROLLBACK FAILED: {restore_exc}\n',
                )
                exc = PipelineError(f'{exc}; rollback also failed: {restore_exc}')
        store.finish_job(
            job_id,
            False,
            str(exc),
            time.monotonic() - started,
        )
        return False
    finally:
        try:
            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
        finally:
            lock_file.close()


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--job-id', type=int, required=True)
    parser.add_argument('--match-id')
    parser.add_argument('--replace', action='store_true')
    return parser.parse_args()


if __name__ == '__main__':
    args = _parse_args()
    succeeded = run_pipeline(
        args.job_id,
        match_id=args.match_id,
        replace=args.replace,
    )
    raise SystemExit(0 if succeeded else 1)