boursecu_server/database/sql/generate_cities_sql.py

#!/usr/bin/env python3
"""
Generate iran_cities_seed.sql from goldAli cities API.

Usage:
  python generate_cities_sql.py
  python generate_cities_sql.py --source /path/to/cities.json
"""

from __future__ import annotations

import json
import sys
import urllib.request
from pathlib import Path

SQL_DIR = Path(__file__).resolve().parent
DEFAULT_SOURCE = "https://goldaliii.com/server/api/city/all"
OUTPUT = SQL_DIR / "iran_cities_seed.sql"

ARABIC_TO_PERSIAN = str.maketrans(
    {"ي": "ی", "ك": "ک", "ة": "ه", "ۀ": "ه", "ؤ": "و", "إ": "ا", "أ": "ا"}
)

TITLE_FIXES = {
    "چهارمحال وبختیاری": "چهارمحال و بختیاری",
    "سیستان وبلوچستان": "سیستان و بلوچستان",
    "کهگیلویه وبویراحمد": "کهگیلویه و بویراحمد",
}


def normalize_title(title: str) -> str:
    title = title.translate(ARABIC_TO_PERSIAN)
    return TITLE_FIXES.get(title, title)


def sql_str(value: str | None) -> str:
    if value is None:
        return "NULL"
    return "'" + str(value).replace("\\", "\\\\").replace("'", "''") + "'"


def sql_bool(value: bool | None) -> str:
    return "1" if value else "0"


def load_rows(source: str) -> list[dict]:
    if source.startswith("http://") or source.startswith("https://"):
        with urllib.request.urlopen(source, timeout=60) as resp:
            data = json.load(resp)
    else:
        data = json.loads(Path(source).read_text(encoding="utf-8"))

    return [
        row
        for row in data
        if row.get("id") != 1600 and row.get("title") != "تست"
    ]


def generate_sql(rows: list[dict]) -> str:
    provinces = sorted(
        [row for row in rows if row.get("parent_id") is None], key=lambda r: r["id"]
    )
    others = sorted(
        [row for row in rows if row.get("parent_id") is not None], key=lambda r: r["id"]
    )
    ordered = provinces + others

    lines = [
        "-- =============================================================================",
        "-- استان‌ها + شهرها + زیرمجموعه‌ها (جدول cities)",
        "--",
        "-- ساختار مطابق City model:",
        "--   level=1 → استان (parent_id = NULL, has_children = 1)",
        "--   level=2 → شهر/شهرستان (parent_id = id استان)",
        "--   level=3 → محله/بخش (parent_id = id شهر level=2)",
        "--",
        f"-- تعداد: {len(ordered)} رکورد (31 استان + {len(ordered) - 31} زیرمجموعه)",
        "-- =============================================================================",
        "",
        "START TRANSACTION;",
        "",
        "SET NAMES utf8mb4 COLLATE utf8mb4_unicode_ci;",
        "",
        "CREATE TABLE IF NOT EXISTS `cities` (",
        "  `id` bigint unsigned NOT NULL AUTO_INCREMENT,",
        "  `title` varchar(255) NOT NULL,",
        "  `code` varchar(50) DEFAULT NULL,",
        "  `parent_id` bigint unsigned DEFAULT NULL,",
        "  `has_children` tinyint(1) NOT NULL DEFAULT 0,",
        "  `level` tinyint unsigned DEFAULT NULL,",
        "  PRIMARY KEY (`id`),",
        "  KEY `cities_parent_id_index` (`parent_id`)",
        ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci;",
        "",
        "ALTER TABLE `cities` MODIFY COLUMN `code` varchar(50) DEFAULT NULL;",
        "",
        "DELETE FROM `cities`;",
        "",
    ]

    batch_size = 50
    for i in range(0, len(ordered), batch_size):
        batch = ordered[i : i + batch_size]
        lines.append(
            "INSERT INTO `cities` (`id`, `title`, `code`, `parent_id`, `has_children`, `level`) VALUES"
        )
        values = []
        for row in batch:
            parent_id = row.get("parent_id")
            level = row.get("level")
            values.append(
                f"  ({row['id']}, {sql_str(normalize_title(row['title']))}, "
                f"{sql_str(row.get('code'))}, "
                f"{parent_id if parent_id is not None else 'NULL'}, "
                f"{sql_bool(row.get('has_children'))}, "
                f"{level if level is not None else 'NULL'})"
            )
        lines.append(",\n".join(values) + ";")
        lines.append("")

    max_id = max(row["id"] for row in ordered)
    lines.extend([f"ALTER TABLE `cities` AUTO_INCREMENT = {max_id + 1};", "", "COMMIT;", ""])
    return "\n".join(lines)


def main() -> None:
    source = DEFAULT_SOURCE
    if len(sys.argv) > 1 and sys.argv[1] == "--source" and len(sys.argv) > 2:
        source = sys.argv[2]

    rows = load_rows(source)
    OUTPUT.write_text(generate_sql(rows), encoding="utf-8")
    print(f"Wrote {OUTPUT} ({len(rows)} rows)")


if __name__ == "__main__":
    main()