scripts/update_db.py - master - ave6990/rst_orders

scripts/update_db.py
336 строк13 KB
dev
admin: панель обновления БД (логин, фоновые задачи, лог, отмена/удаление); ГРСИ извлекается только для записей без данных (registry_processed)
05 авг 2026, 13:34
05 авг 2026, 13:34f9e771c
#!/usr/bin/env python3
"""
Полное обновление базы данных:

  1. Загрузка приказов РСТ с сайта в MariaDB (fetch_rst_orders).
  2. Скачивание отсутствующих PDF-файлов (download_pdfs).
  3. Извлечение номеров ГРСИ из PDF и запись в orders.registry_numbers.

Использование:
    python3 update_db.py                 # обновить всё по умолчанию
    python3 update_db.py --skip-fetch    # только скачивание PDF + извлечение ГРСИ
    python3 update_db.py --grs-only      # извлекать номера только у приказов ГРСИ
"""

import argparse
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from scripts.fetch_rst_orders import (
    DB_NAME,
    Order,
    OrderFile,
    create_session,
    ensure_database,
    fetch_all_orders,
    format_order,
    init_db,
    save_orders_to_db,
)
from scripts.download_pdfs import download_pdf
from scripts.extract_registry_numbers import (
    GRSII_FILTERS,
    extract_registry_numbers,
    extract_text_from_pdf,
)

from sqlalchemy import text


def download_missing_pdfs(db_factory, verbose=False, dry_run=False, cancel_check=None):
    """Скачивает PDF-файлы, у которых нет file_data."""
    engine, db = None, None
    db = db_factory()
    try:
        rows = db.execute(
            text("""
                SELECT f.id, f.order_id, f.local_name, f.uid_file, o.number
                FROM order_files f
                JOIN orders o ON o.order_id = f.order_id
                WHERE f.download_url IS NOT NULL AND f.file_data IS NULL
                ORDER BY f.id
            """)
        ).fetchall()

        if not rows:
            print("  PDF для скачивания не найдены (все файлы уже в БД).")
            return 0, 0

        total = len(rows)
        print(f"  Найдено файлов для скачивания: {total}")

        if dry_run:
            return 0, 0

        http_session = create_session()
        ok = 0
        fail = 0
        for i, r in enumerate(rows, 1):
            if cancel_check and cancel_check():
                print("  Прервано пользователем.")
                break
            if not r.uid_file:
                if verbose:
                    print(f"  [{i}/{total}] SKIP {r.local_name} — нет uid")
                fail += 1
                continue

            if verbose:
                print(f"  [{i}/{total}] {r.local_name}...", end=" ", flush=True)
            data, used_url = download_pdf(http_session, r.uid_file)

            if data:
                db.execute(
                    text(
                        "UPDATE order_files SET file_data = :data, download_url = :url "
                        "WHERE id = :id"
                    ),
                    {"data": data, "url": used_url, "id": r.id},
                )
                db.execute(
                    text(
                        "UPDATE orders SET registry_processed = 0 "
                        "WHERE order_id = :oid AND registry_processed = 1"
                    ),
                    {"oid": r.order_id},
                )
                db.commit()
                if verbose:
                    print(f"OK ({len(data) / 1024:.0f} КБ)")
                ok += 1
            else:
                if verbose:
                    print("ОШИБКА")
                fail += 1

            time.sleep(0.3)
        return ok, fail
    finally:
        db.close()


def extract_worker(order_id, db_factory, verbose):
    db = db_factory()
    try:
        return order_id, process_order_numbers(order_id, db, verbose)
    finally:
        db.close()


def process_order_numbers(order_id, db, verbose=False):
    """Находит номера ГРСИ во всех PDF-файлах приказа."""
    files = (
        db.query(OrderFile)
        .filter(OrderFile.order_id == order_id, OrderFile.file_data != None)
        .all()
    )
    if not files:
        return []

    found = set()
    for f in files:
        if (f.file_ext or "").lower() != "pdf":
            continue
        try:
            text = extract_text_from_pdf(f.file_data)
            found.update(extract_registry_numbers(text))
        except Exception as e:
            if verbose:
                print(f"  [warn] file {f.id}: {e}", file=sys.stderr)
    return sorted(found)


def extract_registry_numbers_all(db_factory, grs_only=False, workers=4,
                                 verbose=False, dry_run=False, force=False,
                                 cancel_check=None):
    """Извлекает номера ГРСИ из PDF и записывает в orders.registry_numbers.

    Обрабатывает только приказы, у которых ГРСИ ещё не извлечены
    (registry_processed IS NULL/0). Пропускает уже обработанные, даже если
    номера не найдены. При force=True обрабатывает все подходящие приказы.
    """
    db = db_factory()
    try:
        q = db.query(Order)
        if grs_only:
            filters = [Order.title.ilike(f"%{f}%") for f in GRSII_FILTERS]
            cond = filters[0]
            for flt in filters[1:]:
                cond = cond | flt
            q = q.filter(cond)
        if not force:
            q = q.filter(
                (Order.registry_processed.is_(None)) | (Order.registry_processed == 0)
            )
        orders = q.all()
    finally:
        db.close()

    if not orders:
        print("  Приказы для извлечения не найдены (все уже обработаны).")
        return

    total = len(orders)
    print(f"  Обработка {total} приказов в {workers} потоков...")

    results = []
    cancelled = False
    done = 0
    batch_size = max(workers * 4, 8)
    with ThreadPoolExecutor(max_workers=workers) as pool:
        for start in range(0, total, batch_size):
            if cancel_check and cancel_check():
                print("  Прервано пользователем.")
                cancelled = True
                break
            batch = orders[start:start + batch_size]
            futures = {
                pool.submit(extract_worker, o.order_id, db_factory, verbose): o
                for o in batch
            }
            for f in as_completed(futures):
                o = futures[f]
                done += 1
                try:
                    oid, numbers = f.result()
                    results.append((oid, numbers))
                    if verbose or dry_run:
                        nums_str = ", ".join(numbers) if numbers else "—"
                        print(f"  [{done}/{total}] #{o.order_id} {o.number or '-'}: "
                              f"{len(numbers)} — {nums_str}")
                except Exception as e:
                    print(f"  [{done}/{total}] #{o.order_id} ERROR: {e}", file=sys.stderr)
                    results.append((o.order_id, []))

    if dry_run:
        print("  Dry-run: запись в БД не производилась.")
        return

    if cancelled:
        print("  Прервано пользователем: частичные результаты не записаны.")
        return

    db = db_factory()
    try:
        updated = 0
        for oid, numbers in results:
            db.query(Order).filter(Order.order_id == oid).update(
                {
                    "registry_numbers": ", ".join(numbers),
                    "registry_processed": 1,
                }
            )
            if numbers:
                updated += 1
        db.commit()
        print(f"  Готово. Обновлено приказов с номерами ГРСИ: {updated} из {total}")
    finally:
        db.close()


def main():
    parser = argparse.ArgumentParser(
        description="Обновление БД: приказы с сайта + PDF + номера ГРСИ"
    )
    parser.add_argument("--search", default="", help="Поисковый запрос")
    parser.add_argument("--source", choices=["all", "site", "sedo"], default="all",
                        help="Источник: all/site/sedo")
    parser.add_argument("--date-from", default=None, help="Дата начала (ДД.ММ.ГГГГ)")
    parser.add_argument("--date-to", default=None, help="Дата окончания (ДД.ММ.ГГГГ)")
    parser.add_argument("--year", type=int, default=None, help="Год (для архива)")
    parser.add_argument("--archive", action="store_true", help="Искать в архиве")
    parser.add_argument("--page-size", type=int, default=100, help="Размер страницы API")
    parser.add_argument("--workers", type=int, default=min(os.cpu_count() or 4, 8),
                        help="Потоков для извлечения ГРСИ (по умолчанию min(CPU, 8))")
    parser.add_argument("--grs-only", action="store_true",
                        help="Извлекать номера ГРСИ только у приказов с ГРСИ-заголовками")
    parser.add_argument("--force-extract", action="store_true",
                        help="Извлечь номера ГРСИ заново у всех приказов (игнорировать флаг обработки)")
    parser.add_argument("--skip-fetch", action="store_true",
                        help="Не загружать приказы с сайта")
    parser.add_argument("--skip-download", action="store_true",
                        help="Не скачивать отсутствующие PDF")
    parser.add_argument("--skip-extract", action="store_true",
                        help="Не извлекать номера ГРСИ")
    parser.add_argument("--dry-run", action="store_true",
                        help="Показать что будет сделано, без записи")
    parser.add_argument("--verbose", "-v", action="store_true", help="Подробный вывод")
    args = parser.parse_args()

    print("=" * 60)
    print("  Обновление базы данных РСТ")
    print("=" * 60)

    print("\n[1/4] Подключение к БД...")
    ensure_database()
    engine, db_factory = init_db()
    print(f"  БД '{DB_NAME}': таблицы готовы")

    if not args.skip_fetch:
        print("\n[2/4] Загрузка приказов с сайта РСТ...")
        http_session = create_session()
        orders, total = fetch_all_orders(
            http_session,
            search=args.search,
            source=args.source,
            date_from=args.date_from,
            date_to=args.date_to,
            archive=args.archive,
            year=args.year,
            page_size=args.page_size,
        )
        print(f"  Загружено из API: {len(orders)} из {total}")

        if orders:
            formatted = [format_order(o) for o in orders]
            db_session = db_factory()
            try:
                stats = save_orders_to_db(
                    db_session, formatted, http_session, download_pdfs=True
                )
                print(f"  Новых приказов:    {stats['new_orders']}")
                print(f"  Обновлено:         {stats['updated_orders']}")
                print(f"  Новых файлов:      {stats['new_files']}")
                print(f"  Скачано файлов:    {stats['downloaded_files']}")
                print(f"  Не удалось скачать: {stats['skipped_files']}")
            except Exception as e:
                db_session.rollback()
                print(f"  ОШИБКА при сохранении: {e}", file=sys.stderr)
                raise
            finally:
                db_session.close()
        else:
            print("  Приказы не найдены.")

    if not args.skip_download:
        print("\n[3/4] Скачивание отсутствующих PDF...")
        ok, fail = download_missing_pdfs(
            db_factory, verbose=args.verbose, dry_run=args.dry_run
        )
        if not args.dry_run:
            print(f"  Скачано: {ok}, ошибок: {fail}")

    if not args.skip_extract:
        print("\n[4/4] Извлечение номеров ГРСИ из PDF...")
        if args.dry_run:
            print("  (dry-run — без записи в БД)")
        extract_registry_numbers_all(
            db_factory,
            grs_only=args.grs_only,
            workers=args.workers,
            verbose=args.verbose,
            dry_run=args.dry_run,
            force=args.force_extract,
        )

    engine.dispose()
    print("\nГотово.")


if __name__ == "__main__":
    main()