feat: 增加远程自动调度处理

This commit is contained in:
2026-05-25 17:48:32 +08:00
parent 89fe45c813
commit 8e891964d5
14 changed files with 1676 additions and 20 deletions
+42 -6
View File
@@ -1,7 +1,7 @@
from datetime import datetime
from datetime import date, datetime
from pathlib import Path
from threading import Thread
from typing import Any
from typing import Any, Callable, Iterable
from fastapi import APIRouter, Body, HTTPException
@@ -27,6 +27,29 @@ async def test_remote_connection(config: dict[str, Any] | None = Body(None)):
@router.post("/api/remote/start")
async def start_remote_processing():
return start_remote_processing_job(source="manual")
@router.get("/api/remote/scheduler/status")
def get_scheduler_status():
if state.auto_scheduler is None:
return {"enabled": False, "running": False, "message": "自动调度器未启动"}
return state.auto_scheduler.get_status()
@router.post("/api/remote/scheduler/trigger")
def trigger_scheduler_check():
if state.auto_scheduler is None:
raise HTTPException(status_code=503, detail="自动调度器未启动")
return state.auto_scheduler.check_and_run(manual=True)
def start_remote_processing_job(
*,
source: str = "manual",
on_finish: Callable[[str, str], None] | None = None,
target_dates: Iterable[date] | None = None,
) -> dict[str, Any]:
if state.global_task_lock["locked"]:
raise HTTPException(status_code=409, detail="已有任务在运行,请等待当前任务完成")
@@ -70,14 +93,14 @@ async def start_remote_processing():
thread = Thread(
target=_run_remote_processing,
args=(task_id, work_dir, app_config, remote_config, logger),
args=(task_id, work_dir, app_config, remote_config, logger, source, on_finish, target_dates),
daemon=True,
)
thread.start()
return {
"success": True,
"message": "远程下载处理任务已启动",
"message": "自动调度远程下载处理任务已启动" if source == "scheduler" else "远程下载处理任务已启动",
"task_id": task_id,
"stage": "downloading",
}
@@ -99,14 +122,18 @@ def _run_remote_processing(
app_config: AppConfig,
remote_config: RemoteDataConfig,
logger: ProcessLogger,
source: str = "manual",
on_finish: Callable[[str, str], None] | None = None,
target_dates: Iterable[date] | None = None,
) -> None:
final_status = "failed"
try:
logger.info(
f"开始远程下载,协议: {remote_config.protocol.upper()},"
f"服务器: {remote_config.host}:{remote_config.port},目录: {remote_config.remote_dir}"
)
downloader = RemoteDataDownloader(remote_config, logger.info)
download_result = downloader.download_to(work_dir)
download_result = downloader.download_to(work_dir, target_dates=target_dates)
logger.success(
f"远程下载完成,共 {download_result.file_count} 个文件,"
f"{_format_bytes(download_result.total_bytes)}"
@@ -127,8 +154,12 @@ def _run_remote_processing(
try:
deleted_count = downloader.delete_source_files(download_result.remote_files)
logger.success(f"远程源文件清理完成,共删除 {deleted_count} 个文件,目录已保留")
final_status = status
except Exception as exc:
logger.warning(f"远程源文件清理失败,数据处理结果已保留: {exc}")
final_status = "source_cleanup_failed"
else:
final_status = status
state.history_manager.update(
task_id,
@@ -145,7 +176,10 @@ def _run_remote_processing(
}
except Exception as exc:
error_detail = exc.to_detail() if isinstance(exc, LicenseError) else None
logger.error(f"远程自动化任务失败: {exc}")
if source == "scheduler":
logger.error(f"自动调度任务失败: {exc}")
else:
logger.error(f"远程自动化任务失败: {exc}")
state.history_manager.update(task_id, status="failed", error=str(exc))
state.processing_tasks[task_id] = {
"logs": state.history_manager.get_logs(task_id),
@@ -160,6 +194,8 @@ def _run_remote_processing(
except Exception as exc:
print(f"自动清理处理历史失败: {exc}")
state.reset_task_lock()
if on_finish:
on_finish(task_id, final_status)
def _format_bytes(size: int) -> str:
+60 -2
View File
@@ -34,6 +34,59 @@ class MySQLConfig:
dbname: str = "CapacityReport"
@dataclass
class AutoSchedulerConfig:
enabled: bool = False
check_interval_hours: int = 1
expected_directories: List[str] = field(default_factory=list)
week_offset: int = 0
def normalized(self) -> "AutoSchedulerConfig":
try:
check_interval_hours = int(self.check_interval_hours)
except (TypeError, ValueError):
check_interval_hours = 1
try:
week_offset = int(self.week_offset)
except (TypeError, ValueError):
week_offset = 0
directories = []
seen = set()
for directory in self.expected_directories or []:
normalized = str(directory).replace("\\", "/").strip().strip("/")
if normalized and normalized not in seen:
directories.append(normalized)
seen.add(normalized)
return AutoSchedulerConfig(
enabled=bool(self.enabled),
check_interval_hours=max(check_interval_hours, 1),
expected_directories=directories,
week_offset=week_offset,
)
def to_dict(self) -> Dict[str, Any]:
normalized = self.normalized()
return {
"enabled": normalized.enabled,
"check_interval_hours": normalized.check_interval_hours,
"expected_directories": normalized.expected_directories,
"week_offset": normalized.week_offset,
}
@classmethod
def from_dict(cls, data: Dict[str, Any] | None) -> "AutoSchedulerConfig":
data = data or {}
directories = data.get("expected_directories", [])
return cls(
enabled=bool(data.get("enabled", False)),
check_interval_hours=data.get("check_interval_hours", 1),
expected_directories=directories if isinstance(directories, list) else [],
week_offset=data.get("week_offset", 0),
).normalized()
@dataclass
class RemoteDataConfig:
enabled: bool = False
@@ -46,6 +99,7 @@ class RemoteDataConfig:
passive: bool = True
timeout: int = 30
auto_delete_source: bool = False
auto_scheduler: AutoSchedulerConfig = field(default_factory=AutoSchedulerConfig)
def normalized(self) -> "RemoteDataConfig":
protocol = self.protocol.lower().strip()
@@ -53,8 +107,9 @@ class RemoteDataConfig:
protocol = "sftp"
port = self.port or (22 if protocol == "sftp" else 21)
scheduler = self.auto_scheduler.normalized()
return RemoteDataConfig(
enabled=bool(self.enabled),
enabled=bool(self.enabled) or scheduler.enabled,
protocol=protocol,
host=self.host.strip(),
port=port,
@@ -63,7 +118,8 @@ class RemoteDataConfig:
remote_dir=(self.remote_dir or "/").strip() or "/",
passive=bool(self.passive),
timeout=max(int(self.timeout or 30), 1),
auto_delete_source=bool(self.auto_delete_source),
auto_delete_source=bool(self.auto_delete_source) or scheduler.enabled,
auto_scheduler=scheduler,
)
def to_dict(self, include_password: bool = False) -> Dict[str, Any]:
@@ -77,6 +133,7 @@ class RemoteDataConfig:
"passive": self.passive,
"timeout": self.timeout,
"auto_delete_source": self.auto_delete_source,
"auto_scheduler": self.auto_scheduler.normalized().to_dict(),
}
if include_password:
data["passwd"] = self.passwd
@@ -106,6 +163,7 @@ class RemoteDataConfig:
passive=bool(data.get("passive", True)),
timeout=timeout,
auto_delete_source=bool(data.get("auto_delete_source", False)),
auto_scheduler=AutoSchedulerConfig.from_dict(data.get("auto_scheduler")),
).normalized()
+28
View File
@@ -1,4 +1,5 @@
import argparse
from contextlib import asynccontextmanager
from pathlib import Path
import uvicorn
@@ -25,6 +26,7 @@ from app.api.routers import (
from app.auth import extract_access_token, resolve_access_context, resolve_login_context
from app.config import BASE_DIR
from app.services.api_tokens import touch_token_usage
from app.services.auto_scheduler import AutoScheduler
APP_VERSION = "3.0.0"
@@ -139,6 +141,14 @@ OPENAPI_OPERATION_DOCS = {
"summary": "远程下载并处理",
"description": "从已配置的 FTP/SFTP 目录递归下载源数据,然后自动执行完整处理流程。",
},
("get", "/api/remote/scheduler/status"): {
"summary": "查询远程自动调度状态",
"description": "返回自动调度启用状态、目标周、就绪标识、下次检查时间和各远程目录的日期覆盖情况。",
},
("post", "/api/remote/scheduler/trigger"): {
"summary": "手动触发自动调度检查",
"description": "立即执行一次远程目录就绪检查;如果已存在就绪标识,会直接触发远程下载并处理。",
},
("post", "/api/history"): {
"summary": "查询处理历史",
"description": "按最近时间返回处理历史记录。",
@@ -249,6 +259,12 @@ OPENAPI_OPERATION_DOCS = {
"passive": True,
"timeout": 30,
"auto_delete_source": False,
"auto_scheduler": {
"enabled": False,
"check_interval_hours": 1,
"expected_directories": ["4G/FDD", "4G/900", "5G/2.6", "5G/700"],
"week_offset": 0,
},
},
},
("post", "/api/config/history-retention"): {
@@ -312,6 +328,17 @@ OPENAPI_OPERATION_DOCS = {
}
@asynccontextmanager
async def app_lifespan(app: FastAPI):
state.auto_scheduler = AutoScheduler()
state.auto_scheduler.start()
try:
yield
finally:
if state.auto_scheduler is not None:
state.auto_scheduler.stop()
def create_app() -> FastAPI:
app = FastAPI(
title="CapacityReport",
@@ -320,6 +347,7 @@ def create_app() -> FastAPI:
docs_url=None,
redoc_url=None,
openapi_url=None,
lifespan=app_lifespan,
)
app.openapi = lambda: custom_openapi(app) # type: ignore[method-assign]
+41 -3
View File
@@ -18,6 +18,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
from app.config import AppConfig, SQL_SCRIPT
from app.database import DatabaseManager
from app.utils.file_dates import DirectoryDateSelection, select_recent_items_by_directory
class ProcessLogger:
@@ -246,7 +247,7 @@ class DataProcessor:
def _unzip_files(self):
"""解压所有 ZIP 文件(支持中文文件名)"""
self.logger.info("正在解压 ZIP 文件...")
zip_files = list(self.work_dir.rglob("*.zip"))
zip_files = self._filter_recent_files(list(self.work_dir.rglob("*.zip")), "ZIP")
zip_count = 0
for zip_file in zip_files:
@@ -338,6 +339,43 @@ class DataProcessor:
for ext in extensions:
for file in directory.rglob(f"*{ext}"):
yield file
def _filter_recent_files(self, files: list[Path], label: str, root: Path | None = None) -> list[Path]:
if not files:
return files
base = (root or self.work_dir).resolve()
def parent_key(file_path: Path) -> str:
try:
parent = file_path.parent.resolve().relative_to(base)
except ValueError:
parent = file_path.parent
parent_text = str(parent).replace("\\", "/")
return "" if parent_text == "." else parent_text
selected, summaries = select_recent_items_by_directory(
files,
parent_key=parent_key,
name_key=lambda file_path: file_path.name,
)
self._log_recent_file_selection(label, summaries)
return sorted(selected)
def _log_recent_file_selection(self, label: str, summaries: list[DirectoryDateSelection]) -> None:
skipped_total = sum(summary.skipped_count for summary in summaries)
if not skipped_total:
return
for summary in summaries:
if not summary.skipped_count or not summary.start_date or not summary.max_date:
continue
self.logger.info(
f"{label}目录 {summary.directory or '.'}: 仅处理 "
f"{summary.start_date.isoformat()} 至 {summary.max_date.isoformat()} "
f"的 {summary.selected_count}/{summary.total_count} 个文件,"
f"跳过 {summary.skipped_count} 个旧文件"
)
def _process_single_excel(self, excel_file: Path, sheet_filter: set) -> int:
"""处理单个 Excel 文件(用于并行)"""
@@ -368,7 +406,7 @@ class DataProcessor:
def _process_excel_files_parallel(self):
"""并行处理 Excel 文件"""
self.logger.info("正在并行处理 Excel 文件...")
excel_files = list(self._scan_files(self.work_dir, ['.xlsx', '.xls']))
excel_files = self._filter_recent_files(list(self._scan_files(self.work_dir, ['.xlsx', '.xls'])), "Excel")
self.logger.info(f"找到 {len(excel_files)} 个 Excel 文件")
if not excel_files:
@@ -794,7 +832,7 @@ class DataProcessor:
self.db.drop_table(table_name)
# 处理该目录下的所有 CSV
csv_files = list(self._scan_files(subdir, ['.csv']))
csv_files = self._filter_recent_files(list(self._scan_files(subdir, ['.csv'])), "CSV", root=subdir)
self.logger.info(f"找到 {len(csv_files)} 个 CSV 文件")
total_rows = 0
+408
View File
@@ -0,0 +1,408 @@
from __future__ import annotations
import json
import threading
from collections import defaultdict
from dataclasses import dataclass
from datetime import date, datetime, timedelta
from typing import Any
from fastapi import HTTPException
from app import state
from app.config import CACHE_DIR, AutoSchedulerConfig, RemoteDataConfig
from app.services.remote_download import RemoteDataDownloader, RemoteFileInfo
from app.utils.file_dates import extract_file_date, required_week_days
READY_DIR = CACHE_DIR / "auto_scheduler"
READY_FLAG = READY_DIR / "ready.flag"
DISABLED_CHECK_SECONDS = 60
STARTUP_CHECK_SECONDS = 5
FAILURE_RESULTS = {"scan_failed", "trigger_failed", "source_cleanup_failed", "failed", "invalid_flag"}
@dataclass(frozen=True)
class DirectoryReadyStatus:
directory: str
ready: bool
found_days: list[date]
missing_days: list[date]
file_count: int
error: str | None = None
def to_dict(self) -> dict[str, Any]:
return {
"ready": self.ready,
"found_days": [item.isoformat() for item in self.found_days],
"missing_days": [item.isoformat() for item in self.missing_days],
"found_count": len(self.found_days),
"required_count": len(self.found_days) + len(self.missing_days),
"file_count": self.file_count,
"error": self.error,
}
class AutoScheduler:
def __init__(self) -> None:
self._stop_event = threading.Event()
self._check_lock = threading.Lock()
self._status_lock = threading.Lock()
self._thread: threading.Thread | None = None
self._running = False
self._last_check_at: datetime | None = None
self._next_check_at: datetime | None = None
self._last_result = "not_started"
self._last_message = "自动调度器尚未检查"
self._directory_status: dict[str, dict[str, Any]] = {}
self._task_running = False
self._task_id: str | None = None
self._failure_count = 0
def start(self) -> None:
if self._thread and self._thread.is_alive():
return
self._running = True
self._stop_event.clear()
self._set_next_check(datetime.now() + timedelta(seconds=STARTUP_CHECK_SECONDS))
self._thread = threading.Thread(target=self._run_loop, name="auto-scheduler", daemon=True)
self._thread.start()
def stop(self) -> None:
self._running = False
self._stop_event.set()
if self._thread and self._thread.is_alive():
self._thread.join(timeout=5)
def get_status(self) -> dict[str, Any]:
config = state.current_config().remote_data.normalized()
scheduler = config.auto_scheduler.normalized()
target_days = required_week_days(scheduler.week_offset)
ready_flag = self._read_ready_flag()
with self._status_lock:
return {
"enabled": scheduler.enabled,
"running": self._running,
"check_interval_hours": scheduler.check_interval_hours,
"expected_directories": scheduler.expected_directories,
"week_offset": scheduler.week_offset,
"auto_delete_source": config.auto_delete_source,
"next_check_at": self._format_dt(self._next_check_at),
"last_check_at": self._format_dt(self._last_check_at),
"last_result": self._last_result,
"last_message": self._last_message,
"failure_count": self._failure_count,
"task_running": self._task_running,
"task_id": self._task_id,
"ready_flag": ready_flag,
"target_week": {
"start": target_days[0].isoformat(),
"end": target_days[-1].isoformat(),
"days": [item.isoformat() for item in target_days],
},
"directory_status": self._directory_status,
}
def check_and_run(self, manual: bool = False) -> dict[str, Any]:
if not self._check_lock.acquire(blocking=False):
return self._finish_check("busy", "自动调度器正在检查中", manual)
try:
return self._check_and_run_locked(manual)
finally:
self._check_lock.release()
def _run_loop(self) -> None:
while not self._stop_event.wait(self._seconds_until_next_check()):
self.check_and_run(manual=False)
def _check_and_run_locked(self, manual: bool) -> dict[str, Any]:
now = datetime.now()
remote_config = state.current_config().remote_data.normalized()
scheduler = remote_config.auto_scheduler.normalized()
self._last_check_at = now
if not scheduler.enabled:
self._set_next_check(now + timedelta(seconds=DISABLED_CHECK_SECONDS))
return self._finish_check("disabled", "自动调度未启用", manual)
self._set_next_check(now + timedelta(hours=scheduler.check_interval_hours))
if not remote_config.enabled:
return self._finish_check("remote_disabled", "远程数据源未启用", manual)
if state.global_task_lock["locked"]:
return self._finish_check("task_running", "已有任务在运行,本轮自动调度跳过", manual)
ready_flag = self._read_ready_flag()
if ready_flag["exists"]:
if ready_flag.get("invalid"):
self._clear_ready_flag()
return self._finish_check("invalid_flag", "就绪标识格式错误,已清除,本轮不触发处理", manual)
target_dates = self._target_dates_from_flag(ready_flag, scheduler)
return self._trigger_processing(target_dates, ready_flag, manual)
target_days = required_week_days(scheduler.week_offset)
directory_status = self._check_remote_ready(remote_config, scheduler, target_days)
ready = bool(directory_status) and all(item.ready for item in directory_status.values())
self._set_directory_status(directory_status)
error_count = sum(1 for item in directory_status.values() if item.error)
if error_count:
return self._finish_check(
"scan_failed",
f"远程目录扫描失败,{error_count}/{len(directory_status)} 个目录无法访问或扫描失败",
manual,
)
if ready:
self._mark_ready(target_days, directory_status)
return self._finish_check(
"marked_ready",
"远程数据已满足目标周 7 天,已写入就绪标识,下次检查将自动处理",
manual,
)
ready_count = sum(1 for item in directory_status.values() if item.ready)
return self._finish_check(
"waiting",
f"远程数据未就绪,{ready_count}/{len(directory_status)} 个目录满足目标周 7 天",
manual,
)
def _check_remote_ready(
self,
remote_config: RemoteDataConfig,
scheduler: AutoSchedulerConfig,
target_days: list[date],
) -> dict[str, DirectoryReadyStatus]:
downloader = RemoteDataDownloader(remote_config)
expected_directories = scheduler.expected_directories
if expected_directories:
return {
directory: self._directory_ready_status(
directory,
self._safe_list_remote_zip_files(downloader, directory),
target_days,
)
for directory in expected_directories
}
files = self._safe_list_remote_zip_files(downloader, None)
if files is None:
return {
".": DirectoryReadyStatus(
directory=".",
ready=False,
found_days=[],
missing_days=target_days,
file_count=0,
error="远程目录扫描失败",
)
}
grouped: dict[str, list[RemoteFileInfo]] = defaultdict(list)
for remote_file in files:
grouped[remote_file.parent or "."].append(remote_file)
if not grouped:
return {
".": DirectoryReadyStatus(
directory=".",
ready=False,
found_days=[],
missing_days=target_days,
file_count=0,
error="远程目录未找到 ZIP 文件",
)
}
return {
directory: self._directory_ready_status(directory, directory_files, target_days)
for directory, directory_files in sorted(grouped.items(), key=lambda item: item[0])
}
def _safe_list_remote_zip_files(
self,
downloader: RemoteDataDownloader,
directory: str | None,
) -> list[RemoteFileInfo] | None:
try:
return downloader.list_remote_zip_files(directory)
except Exception:
return None
def _directory_ready_status(
self,
directory: str,
files: list[RemoteFileInfo] | None,
target_days: list[date],
) -> DirectoryReadyStatus:
if files is None:
return DirectoryReadyStatus(
directory=directory,
ready=False,
found_days=[],
missing_days=target_days,
file_count=0,
error="远程目录不存在或无法访问",
)
required = set(target_days)
found = {
file_date
for remote_file in files
if (file_date := extract_file_date(remote_file.name)) in required
}
missing = [item for item in target_days if item not in found]
return DirectoryReadyStatus(
directory=directory,
ready=not missing,
found_days=sorted(found),
missing_days=missing,
file_count=len(files),
)
def _trigger_processing(
self,
target_dates: list[date],
ready_flag: dict[str, Any],
manual: bool,
) -> dict[str, Any]:
from app.api.routers.remote import start_remote_processing_job
try:
result = start_remote_processing_job(
source="scheduler",
on_finish=self._on_processing_finish,
target_dates=target_dates,
)
except HTTPException as exc:
return self._finish_check("trigger_failed", str(exc.detail), manual)
except Exception as exc:
return self._finish_check("trigger_failed", f"自动调度触发失败: {exc}", manual)
task_id = result.get("task_id")
with self._status_lock:
self._task_running = True
self._task_id = str(task_id) if task_id else None
self._last_result = "triggered"
self._last_message = "已根据就绪标识触发远程下载并处理"
return {
"success": True,
"result": "triggered",
"message": "已根据就绪标识触发远程下载并处理",
"task_id": task_id,
"ready_flag": ready_flag,
"status": self.get_status(),
}
def _mark_ready(
self,
target_days: list[date],
directory_status: dict[str, DirectoryReadyStatus],
) -> None:
READY_DIR.mkdir(parents=True, exist_ok=True)
payload = {
"ready_at": datetime.now().isoformat(timespec="seconds"),
"week_start": target_days[0].isoformat(),
"week_end": target_days[-1].isoformat(),
"target_dates": [item.isoformat() for item in target_days],
"directories": {
name: item.to_dict()
for name, item in directory_status.items()
},
}
READY_FLAG.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
def _clear_ready_flag(self) -> None:
try:
READY_FLAG.unlink(missing_ok=True)
except OSError:
pass
def _read_ready_flag(self) -> dict[str, Any]:
if not READY_FLAG.exists():
return {"exists": False}
try:
data = json.loads(READY_FLAG.read_text(encoding="utf-8"))
if isinstance(data, dict):
return {"exists": True, **data}
except Exception as exc:
return {"exists": True, "invalid": True, "error": str(exc)}
return {"exists": True, "invalid": True, "error": "ready.flag 格式错误"}
def _target_dates_from_flag(
self,
ready_flag: dict[str, Any],
scheduler: AutoSchedulerConfig,
) -> list[date]:
raw_dates = ready_flag.get("target_dates")
if isinstance(raw_dates, list):
dates: list[date] = []
for raw_date in raw_dates:
try:
dates.append(date.fromisoformat(str(raw_date)))
except ValueError:
continue
if dates:
return sorted(dates)
return required_week_days(scheduler.week_offset)
def _on_processing_finish(self, task_id: str, status: str) -> None:
if status == "completed":
self._clear_ready_flag()
message = "自动调度任务处理成功,已清除就绪标识"
result = "completed"
else:
message = "自动调度任务未完成,就绪标识已保留,下次检查会重试"
result = status or "failed"
with self._status_lock:
self._task_running = False
self._task_id = task_id
self._last_result = result
self._last_message = message
self._failure_count = 0 if result == "completed" else self._failure_count + 1
def _finish_check(self, result: str, message: str, manual: bool) -> dict[str, Any]:
with self._status_lock:
self._last_result = result
self._last_message = message
if result in FAILURE_RESULTS:
self._failure_count += 1
elif result not in {"busy", "task_running"}:
self._failure_count = 0
if result not in {"triggered", "task_running"}:
self._task_running = False
return {
"success": result not in FAILURE_RESULTS,
"manual": manual,
"result": result,
"message": message,
"status": self.get_status(),
}
def _set_directory_status(self, directory_status: dict[str, DirectoryReadyStatus]) -> None:
with self._status_lock:
self._directory_status = {
name: item.to_dict()
for name, item in directory_status.items()
}
def _set_next_check(self, value: datetime) -> None:
with self._status_lock:
self._next_check_at = value
def _seconds_until_next_check(self) -> float:
with self._status_lock:
next_check_at = self._next_check_at
if not next_check_at:
return STARTUP_CHECK_SECONDS
return max((next_check_at - datetime.now()).total_seconds(), 0)
@staticmethod
def _format_dt(value: datetime | None) -> str | None:
return value.isoformat(timespec="seconds") if value else None
+181 -2
View File
@@ -1,12 +1,15 @@
from __future__ import annotations
import stat
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import date
from ftplib import FTP
from pathlib import Path
from typing import Callable, Iterable
from app.config import RemoteDataConfig
from app.utils.file_dates import extract_file_date, select_recent_items_by_directory
LogFn = Callable[[str], None]
@@ -19,6 +22,15 @@ class RemoteDownloadResult:
remote_files: list[str] = field(default_factory=list)
@dataclass(frozen=True)
class RemoteFileInfo:
path: str
relative_path: str
parent: str
name: str
size: int = 0
class RemoteDownloadError(RuntimeError):
pass
@@ -42,13 +54,71 @@ class RemoteDataDownloader:
finally:
ssh.close()
def download_to(self, destination: Path) -> RemoteDownloadResult:
def download_to(self, destination: Path, target_dates: Iterable[date] | None = None) -> RemoteDownloadResult:
self._validate_config()
destination.mkdir(parents=True, exist_ok=True)
zip_files = self.list_remote_zip_files()
date_filter = set(target_dates or [])
if date_filter:
selected_files = self._select_files_by_dates(zip_files, date_filter)
return self._download_selected_files(destination, selected_files)
if zip_files:
selected_files, summaries = select_recent_items_by_directory(
zip_files,
parent_key=lambda item: item.parent,
name_key=lambda item: item.name,
)
for summary in summaries:
if summary.max_date and summary.start_date and summary.skipped_count:
self._log(
f"远程目录 {summary.directory or '.'}: 仅下载 "
f"{summary.start_date.isoformat()} 至 {summary.max_date.isoformat()} "
f"的 {summary.selected_count}/{summary.total_count} 个 ZIP 文件,"
f"跳过 {summary.skipped_count} 个旧文件"
)
return self._download_selected_files(destination, selected_files)
if self.config.protocol == "ftp":
return self._download_ftp(destination)
return self._download_sftp(destination)
def _select_files_by_dates(self, zip_files: list[RemoteFileInfo], target_dates: set[date]) -> list[RemoteFileInfo]:
selected_files: list[RemoteFileInfo] = []
grouped: dict[str, list[RemoteFileInfo]] = defaultdict(list)
for remote_file in zip_files:
grouped[remote_file.parent].append(remote_file)
for parent, files in sorted(grouped.items(), key=lambda item: item[0]):
selected = [
remote_file
for remote_file in files
if extract_file_date(remote_file.name) in target_dates
]
selected_files.extend(selected)
skipped_count = len(files) - len(selected)
if skipped_count:
first_day = min(target_dates).isoformat()
last_day = max(target_dates).isoformat()
self._log(
f"远程目录 {parent or '.'}: 仅下载调度目标日期 "
f"{first_day} 至 {last_day} 的 {len(selected)}/{len(files)} 个 ZIP 文件,"
f"跳过 {skipped_count} 个非目标日期文件"
)
return selected_files
def list_remote_zip_files(self, directory: str | None = None) -> list[RemoteFileInfo]:
self._validate_config()
remote_dir = (
self._join_remote_path(self.config.remote_dir, directory.strip("/"))
if directory
else self.config.remote_dir
)
if self.config.protocol == "ftp":
return self._list_ftp_zip_files(remote_dir)
return self._list_sftp_zip_files(remote_dir)
def delete_source_files(self, remote_files: Iterable[str] | None = None) -> int:
self._validate_config()
source_files = list(remote_files or [])
@@ -164,6 +234,59 @@ class RemoteDataDownloader:
result.total_bytes += expected_size or local_path.stat().st_size
result.remote_files.append(remote_path)
def _download_selected_files(self, destination: Path, remote_files: list[RemoteFileInfo]) -> RemoteDownloadResult:
result = RemoteDownloadResult()
if not remote_files:
return result
if self.config.protocol == "ftp":
with self._ftp_client() as ftp:
self._log(f"已连接 FTP: {self.config.host}:{self.config.port}")
for remote_file in remote_files:
self._download_ftp_file(
ftp,
remote_file.path,
destination / remote_file.relative_path,
result,
remote_file.size,
)
return result
ssh = self._sftp_ssh_client()
try:
with ssh.open_sftp() as sftp:
self._log(f"已连接 SFTP: {self.config.host}:{self.config.port}")
for remote_file in remote_files:
self._download_sftp_file(
sftp,
remote_file.path,
destination / remote_file.relative_path,
result,
remote_file.size,
)
finally:
ssh.close()
return result
def _list_ftp_zip_files(self, remote_dir: str) -> list[RemoteFileInfo]:
with self._ftp_client() as ftp:
return self._collect_ftp_zip_files(ftp, remote_dir)
def _collect_ftp_zip_files(self, ftp: FTP, remote_dir: str) -> list[RemoteFileInfo]:
files: list[RemoteFileInfo] = []
for name, entry_type, size in self._list_ftp_entries(ftp, remote_dir):
if name in {".", ".."}:
continue
remote_path = self._join_remote_path(remote_dir, name)
if entry_type == "dir" or (entry_type == "unknown" and self._ftp_is_dir(ftp, remote_path)):
files.extend(self._collect_ftp_zip_files(ftp, remote_path))
continue
if name.lower().endswith(".zip"):
files.append(self._remote_file_info(remote_path, size))
return files
def _delete_ftp_source_files(self) -> int:
with self._ftp_client() as ftp:
self._log(f"开始清理 FTP 源文件: {self.config.remote_dir}")
@@ -246,13 +369,50 @@ class RemoteDataDownloader:
)
return
self._download_sftp_file(sftp, remote_path, local_path, result, int(getattr(attrs, "st_size", 0) or 0))
def _download_sftp_file(
self,
sftp,
remote_path: str,
local_path: Path,
result: RemoteDownloadResult,
expected_size: int = 0,
) -> None:
local_path.parent.mkdir(parents=True, exist_ok=True)
self._log(f"下载: {remote_path}")
sftp.get(remote_path, str(local_path))
result.file_count += 1
result.total_bytes += int(getattr(attrs, "st_size", 0) or local_path.stat().st_size)
result.total_bytes += expected_size or local_path.stat().st_size
result.remote_files.append(remote_path)
def _list_sftp_zip_files(self, remote_dir: str) -> list[RemoteFileInfo]:
ssh = self._sftp_ssh_client()
try:
with ssh.open_sftp() as sftp:
return self._collect_sftp_zip_files(sftp, remote_dir)
finally:
ssh.close()
def _collect_sftp_zip_files(self, sftp, remote_path: str) -> list[RemoteFileInfo]:
attrs = sftp.stat(remote_path)
if not stat.S_ISDIR(attrs.st_mode):
name = Path(remote_path.replace("\\", "/")).name
if name.lower().endswith(".zip"):
return [self._remote_file_info(remote_path, int(getattr(attrs, "st_size", 0) or 0))]
return []
files: list[RemoteFileInfo] = []
for item in sftp.listdir_attr(remote_path):
if item.filename in {".", ".."}:
continue
child_path = self._join_remote_path(remote_path, item.filename)
if stat.S_ISDIR(item.st_mode):
files.extend(self._collect_sftp_zip_files(sftp, child_path))
elif item.filename.lower().endswith(".zip"):
files.append(self._remote_file_info(child_path, int(getattr(item, "st_size", 0) or 0)))
return files
def _delete_sftp_source_files(self) -> int:
ssh = self._sftp_ssh_client()
try:
@@ -301,3 +461,22 @@ class RemoteDataDownloader:
if parent == "/":
return f"/{child}"
return f"{parent}/{child}"
def _remote_file_info(self, remote_path: str, size: int = 0) -> RemoteFileInfo:
normalized_root = self.config.remote_dir.replace("\\", "/").rstrip("/")
normalized_path = remote_path.replace("\\", "/")
if normalized_root and normalized_root != "/" and normalized_path.startswith(f"{normalized_root}/"):
relative_path = normalized_path[len(normalized_root) + 1 :]
else:
relative_path = normalized_path.lstrip("/")
relative = Path(relative_path)
parent = str(relative.parent).replace("\\", "/")
if parent == ".":
parent = ""
return RemoteFileInfo(
path=remote_path,
relative_path=relative_path,
parent=parent,
name=relative.name,
size=size,
)
+1
View File
@@ -8,6 +8,7 @@ config = AppConfig.load()
history_manager = HistoryManager()
processing_tasks: dict[str, dict[str, Any]] = {}
upload_sessions: dict[str, dict[str, Any]] = {}
auto_scheduler: Any | None = None
global_task_lock: dict[str, Any] = {
"locked": False,
+108
View File
@@ -0,0 +1,108 @@
"""Filename date parsing and per-directory date filtering."""
from __future__ import annotations
import re
from dataclasses import dataclass
from datetime import date, datetime, timedelta
from pathlib import Path
from typing import Callable, Iterable, TypeVar
FILENAME_DATE_RE = re.compile(r"(?<!\d)(20\d{10}(?:\d{2})?)(?!\d)")
T = TypeVar("T")
@dataclass(frozen=True)
class DirectoryDateSelection:
directory: str
total_count: int
selected_count: int
skipped_count: int
undated_count: int
max_date: date | None = None
start_date: date | None = None
def extract_file_date(filename: str | Path) -> date | None:
"""Extract the first timestamp date from a filename."""
match = FILENAME_DATE_RE.search(Path(str(filename)).name)
if not match:
return None
return parse_file_timestamp(match.group(1))
def parse_file_timestamp(value: str) -> date | None:
fmt = "%Y%m%d%H%M%S" if len(value) == 14 else "%Y%m%d%H%M"
try:
return datetime.strptime(value, fmt).date()
except ValueError:
return None
def required_week_days(week_offset: int = 0, today: date | None = None) -> list[date]:
"""Return Monday-Sunday dates for the target natural week.
week_offset=0 means last week, -1 means the week before last.
"""
base_day = today or date.today()
this_monday = base_day - timedelta(days=base_day.weekday())
target_monday = this_monday - timedelta(days=7 * (1 - week_offset))
return [target_monday + timedelta(days=offset) for offset in range(7)]
def select_recent_items_by_directory(
items: Iterable[T],
*,
parent_key: Callable[[T], str],
name_key: Callable[[T], str],
days: int = 7,
) -> tuple[list[T], list[DirectoryDateSelection]]:
"""Select the latest N natural days of dated files in each directory.
If a directory has no dated files at all, all files in that directory are
kept to preserve compatibility with legacy/manual data.
"""
safe_days = max(int(days), 1)
groups: dict[str, list[T]] = {}
for item in items:
groups.setdefault(parent_key(item), []).append(item)
selected: list[T] = []
summaries: list[DirectoryDateSelection] = []
for directory, group_items in sorted(groups.items(), key=lambda pair: pair[0]):
dated_items = [(item, extract_file_date(name_key(item))) for item in group_items]
valid_dates = [file_date for _, file_date in dated_items if file_date]
if not valid_dates:
selected.extend(group_items)
summaries.append(
DirectoryDateSelection(
directory=directory,
total_count=len(group_items),
selected_count=len(group_items),
skipped_count=0,
undated_count=len(group_items),
)
)
continue
max_date = max(valid_dates)
start_date = max_date - timedelta(days=safe_days - 1)
group_selected = [
item
for item, file_date in dated_items
if file_date is not None and start_date <= file_date <= max_date
]
selected.extend(group_selected)
summaries.append(
DirectoryDateSelection(
directory=directory,
total_count=len(group_items),
selected_count=len(group_selected),
skipped_count=len(group_items) - len(group_selected),
undated_count=sum(1 for _, file_date in dated_items if file_date is None),
max_date=max_date,
start_date=start_date,
)
)
return selected, summaries