getDomain/domain-api/app/services/export_service.py

from __future__ import annotations

import csv
from datetime import datetime

from openpyxl import Workbook

from app.core.db import get_db
from app.core.files import exports_root, load_export_records, save_export_record, timestamp_filename
from app.services.domains_service import (
    BEIAN_STATUS_LABELS,
    DETECT_STATUS_LABELS,
    REGISTER_STATUS_LABELS,
    REVIEW_STATUS_LABELS,
    USE_STATUS_LABELS,
    _build_domain_query_parts,
)


EXPORT_HEADERS = [
    ("domain", "域名"),
    ("register_status", "注册状态"),
    ("use_status", "使用状态"),
    ("detect_status", "检测状态"),
    ("review_status", "复核状态"),
    ("has_beian", "备案状态"),
    ("website_url", "首页网址"),
    ("beian_year", "备案年份"),
    ("snapshot_years", "快照年份"),
    ("backlink_count", "友链数"),
    ("backlink_gt_10", "友链>10"),
    ("detect_time", "检测时间"),
]


def _normalize_payload(payload: dict) -> dict:
    data = dict(payload or {})
    data["page"] = int(data.get("page", 1) or 1)
    data["page_size"] = int(data.get("page_size", 100) or 100)
    data["page_count"] = int(data.get("page_count", 1) or 1)
    data["scope"] = data.get("scope", "page")
    data["type"] = data.get("type", "txt")
    return data


def _query_export_rows(payload: dict) -> list[dict]:
    data = _normalize_payload(payload)
    from_clause, where_clause, params = _build_domain_query_parts(data)

    limit_offset = ""
    if data["scope"] == "page":
        offset = (data["page"] - 1) * data["page_size"]
        limit_offset = " limit %s offset %s"
        params.extend([data["page_size"], offset])
    elif data["scope"] == "pages":
        offset = (data["page"] - 1) * data["page_size"]
        limit_value = data["page_size"] * max(data["page_count"], 1)
        limit_offset = " limit %s offset %s"
        params.extend([limit_value, offset])

    with get_db() as conn:
        with conn.cursor() as cur:
            cur.execute(
                f"""
                select
                    d.domain,
                    d.register_status,
                    d.use_status,
                    d.detect_status,
                    d.review_status,
                    d.has_beian,
                    d.website_url,
                    d.beian_year,
                    d.snapshot_years,
                    d.backlink_count,
                    coalesce(dd.backlink_count_gt_10, false) as backlink_gt_10,
                    d.detect_time
                {from_clause}
                {where_clause}
                order by d.id desc
                {limit_offset}
                """,
                tuple(params),
            )
            rows = cur.fetchall()

    result: list[dict] = []
    for row in rows:
        result.append(
            {
                "domain": row[0],
                "register_status": REGISTER_STATUS_LABELS.get(row[1], str(row[1])),
                "use_status": USE_STATUS_LABELS.get(row[2], str(row[2])),
                "detect_status": DETECT_STATUS_LABELS.get(row[3], str(row[3])),
                "review_status": REVIEW_STATUS_LABELS.get(row[4], str(row[4])),
                "has_beian": BEIAN_STATUS_LABELS.get(row[5], str(row[5])),
                "website_url": row[6] or "",
                "beian_year": row[7] or "",
                "snapshot_years": row[8] or "",
                "backlink_count": row[9] or 0,
                "backlink_gt_10": "是" if row[10] else "否",
                "detect_time": row[11].isoformat(sep=" ", timespec="seconds") if row[11] else "",
            }
        )
    return result


def _write_txt(path, rows: list[dict]) -> None:
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(f"{row['domain']}\n")


def _write_csv(path, rows: list[dict]) -> None:
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow([label for _, label in EXPORT_HEADERS])
        for row in rows:
            writer.writerow([row[key] for key, _ in EXPORT_HEADERS])


def _write_xlsx(path, rows: list[dict]) -> None:
    workbook = Workbook()
    sheet = workbook.active
    sheet.title = "domains"
    sheet.append([label for _, label in EXPORT_HEADERS])
    for row in rows:
        sheet.append([row[key] for key, _ in EXPORT_HEADERS])
    workbook.save(path)


def create_export_file(payload: dict) -> dict:
    data = _normalize_payload(payload)
    rows = _query_export_rows(data)

    ext = data["type"] if data["type"] in {"txt", "csv", "xlsx"} else "txt"
    filename = timestamp_filename("domain_export", ext)
    output_path = exports_root() / filename

    if ext == "txt":
        _write_txt(output_path, rows)
    elif ext == "csv":
        _write_csv(output_path, rows)
    else:
        _write_xlsx(output_path, rows)

    created_at = datetime.fromtimestamp(output_path.stat().st_mtime)
    record = {
        "filename": filename,
        "type": ext,
        "scope": data["scope"],
        "page": data["page"],
        "page_size": data["page_size"],
        "page_count": data["page_count"],
        "count": len(rows),
        "created_at": created_at.isoformat(sep=" ", timespec="seconds"),
        "download_path": f"/api/v1/exports/download/{filename}",
    }
    save_export_record(record)
    return record


def list_exports() -> list[dict]:
    records = load_export_records()
    normalized: list[dict] = []
    for record in records:
        item = dict(record)
        created_at = item.get("created_at")
        if isinstance(created_at, (int, float)):
            item["created_at"] = datetime.fromtimestamp(created_at).isoformat(sep=" ", timespec="seconds")
        normalized.append(item)
    return normalized