192 lines
7.0 KiB
Python
192 lines
7.0 KiB
Python
#!/usr/bin/env python3
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import re
|
||
import shutil
|
||
from dataclasses import dataclass
|
||
from datetime import UTC, datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import clickhouse_connect
|
||
import yaml
|
||
from openpyxl import load_workbook
|
||
|
||
|
||
@dataclass
|
||
class Config:
|
||
clickhouse: dict[str, Any]
|
||
archive_dir: str | None
|
||
delete_after_load: bool
|
||
min_file_age_seconds: int
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
p = argparse.ArgumentParser(description="Load company registry xlsx into ClickHouse")
|
||
p.add_argument("--config", required=True)
|
||
p.add_argument("--landing", required=True)
|
||
return p.parse_args()
|
||
|
||
|
||
def load_config(path: str) -> Config:
|
||
raw = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
|
||
return Config(
|
||
clickhouse=raw["clickhouse"],
|
||
archive_dir=raw.get("archive_dir"),
|
||
delete_after_load=bool(raw.get("delete_after_load", False)),
|
||
min_file_age_seconds=int(raw.get("min_file_age_seconds", 180)),
|
||
)
|
||
|
||
|
||
def ch_client(conf: Config):
|
||
return clickhouse_connect.get_client(
|
||
host=conf.clickhouse["host"],
|
||
port=conf.clickhouse.get("port", 8123),
|
||
username=conf.clickhouse.get("username", "default"),
|
||
password=conf.clickhouse.get("password", ""),
|
||
database=conf.clickhouse.get("database", "analytics_1c"),
|
||
)
|
||
|
||
|
||
def normalize_company_key(value: str) -> str:
|
||
text = (value or "").upper().replace("Ё", "Е")
|
||
text = re.sub(r"(^|\s)20\d{2}($|\s)", " ", text)
|
||
text = re.sub(r"[^0-9A-ZА-Я]+", " ", text)
|
||
text = re.sub(r"\s+", " ", text)
|
||
return text.strip()
|
||
|
||
|
||
def as_text(value: Any) -> str:
|
||
if value is None:
|
||
return ""
|
||
return str(value).strip()
|
||
|
||
|
||
def archive_or_delete(conf: Config, path: Path) -> None:
|
||
if conf.archive_dir:
|
||
archive_root = Path(conf.archive_dir) / "registry"
|
||
archive_root.mkdir(parents=True, exist_ok=True)
|
||
shutil.move(str(path), archive_root / path.name)
|
||
return
|
||
if conf.delete_after_load:
|
||
path.unlink(missing_ok=True)
|
||
|
||
|
||
def parse_registry(path: Path) -> list[dict[str, Any]]:
|
||
wb = load_workbook(path, read_only=True, data_only=True)
|
||
now = datetime.now(UTC).replace(tzinfo=None, microsecond=0)
|
||
|
||
tax_map: dict[str, tuple[str, str]] = {}
|
||
if "Лист2" in wb.sheetnames:
|
||
ws = wb["Лист2"]
|
||
for row in ws.iter_rows(min_row=3, values_only=True):
|
||
company_name = as_text(row[1] if len(row) > 1 else "")
|
||
if not company_name:
|
||
continue
|
||
tax_map[normalize_company_key(company_name)] = (
|
||
as_text(row[2] if len(row) > 2 else ""),
|
||
as_text(row[3] if len(row) > 3 else ""),
|
||
)
|
||
|
||
rows: list[dict[str, Any]] = []
|
||
if "ОСНОВНОЙ" not in wb.sheetnames:
|
||
return rows
|
||
|
||
ws = wb["ОСНОВНОЙ"]
|
||
top_headers = [as_text(v) for v in next(ws.iter_rows(min_row=1, max_row=1, values_only=True))]
|
||
manager_headers = [as_text(v) for v in next(ws.iter_rows(min_row=2, max_row=2, values_only=True))]
|
||
|
||
col_specs: list[dict[str, Any]] = []
|
||
current_manager: dict[str, Any] | None = None
|
||
for idx in range(1, len(manager_headers)):
|
||
manager = manager_headers[idx]
|
||
top = top_headers[idx] if idx < len(top_headers) else ""
|
||
if manager:
|
||
current_manager = {
|
||
"col": idx,
|
||
"assignee_name": manager,
|
||
"meta_col": None,
|
||
"meta_label": "",
|
||
"registry_status": "active",
|
||
}
|
||
col_specs.append(current_manager)
|
||
continue
|
||
if top and "исключ" in top.lower():
|
||
current_manager = {
|
||
"col": idx,
|
||
"assignee_name": "",
|
||
"meta_col": None,
|
||
"meta_label": "",
|
||
"registry_status": "excluded",
|
||
}
|
||
col_specs.append(current_manager)
|
||
continue
|
||
if current_manager is not None:
|
||
current_manager["meta_col"] = idx
|
||
current_manager["meta_label"] = top or "meta"
|
||
current_manager = None
|
||
|
||
for row in ws.iter_rows(min_row=3, values_only=True):
|
||
for spec in col_specs:
|
||
company_name = as_text(row[spec["col"]] if spec["col"] < len(row) else "")
|
||
if not company_name:
|
||
continue
|
||
meta_value = ""
|
||
if spec.get("meta_col") is not None:
|
||
meta_value = as_text(row[spec["meta_col"]] if spec["meta_col"] < len(row) else "")
|
||
company_key = normalize_company_key(company_name)
|
||
inn, kpp = tax_map.get(company_key, ("", ""))
|
||
rows.append(
|
||
{
|
||
"ts": now,
|
||
"source_file": path.name,
|
||
"source_sheet": "ОСНОВНОЙ",
|
||
"company_name": company_name,
|
||
"company_key": company_key,
|
||
"assignee_name": spec["assignee_name"],
|
||
"registry_status": spec["registry_status"],
|
||
"share_text": meta_value if "ключ" not in spec.get("meta_label", "").lower() else "",
|
||
"key_contour": 1 if meta_value.upper() == "ЕСТЬ" and "ключ" in spec.get("meta_label", "").lower() else 0,
|
||
"inn": inn,
|
||
"kpp": kpp,
|
||
}
|
||
)
|
||
|
||
return rows
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
conf = load_config(args.config)
|
||
client = ch_client(conf)
|
||
landing = Path(args.landing)
|
||
if not landing.exists():
|
||
return 0
|
||
|
||
for path in sorted(p for p in landing.iterdir() if p.is_file() and p.suffix.lower() == ".xlsx"):
|
||
age_seconds = max(0, int((datetime.now(UTC) - datetime.fromtimestamp(path.stat().st_mtime, UTC)).total_seconds()))
|
||
if age_seconds < conf.min_file_age_seconds:
|
||
print(f"skip registry: {path.name} age={age_seconds}s < min_file_age_seconds={conf.min_file_age_seconds}")
|
||
continue
|
||
rows = parse_registry(path)
|
||
if rows:
|
||
client.insert(
|
||
"analytics_1c.raw_1c_company_registry",
|
||
[[row["source_file"], row["source_sheet"], json.dumps({k: (v.isoformat() if isinstance(v, datetime) else v) for k, v in row.items()}, ensure_ascii=False)] for row in rows],
|
||
column_names=["source_file", "source_sheet", "payload"],
|
||
)
|
||
client.insert(
|
||
"analytics_1c.company_registry",
|
||
[[row["ts"], row["source_file"], row["source_sheet"], row["company_name"], row["company_key"], row["assignee_name"], row["registry_status"], row["share_text"], row["key_contour"], row["inn"], row["kpp"]] for row in rows],
|
||
column_names=["ts", "source_file", "source_sheet", "company_name", "company_key", "assignee_name", "registry_status", "share_text", "key_contour", "inn", "kpp"],
|
||
)
|
||
print(f"loaded registry: {path.name} rows={len(rows)}")
|
||
archive_or_delete(conf, path)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|