from __future__ import annotations from dataclasses import dataclass, field from decimal import Decimal, InvalidOperation from pathlib import Path from openpyxl import load_workbook from sqlalchemy import delete, inspect from app.extensions import db from app.models import CommunicationPart, EngineFamily, EngineVariant, TankSpec, VehicleCost @dataclass class WorkbookImportSummary: workbook: str replaced_existing: bool imported_rows: dict[str, int] = field(default_factory=dict) warnings: list[str] = field(default_factory=list) def to_dict(self) -> dict[str, object]: return { "workbook": self.workbook, "replaced_existing": self.replaced_existing, "imported_rows": self.imported_rows, "warnings": self.warnings, } def import_workbook_data( workbook_path: str | Path, replace_existing: bool = False ) -> WorkbookImportSummary: path = Path(workbook_path) if not path.exists(): raise FileNotFoundError(f"Workbook not found: {path}") _assert_schema_ready() summary = WorkbookImportSummary(workbook=path.name, replaced_existing=replace_existing) workbook = load_workbook(path, data_only=True, read_only=True) try: if replace_existing: _clear_import_tables() elif _database_has_data(): raise ValueError("数据库已有数据。若要重导,请使用 --replace。") family_count, variant_count = _import_engines(workbook["Engine Database"], summary) communication_count = _import_communication_parts(workbook["Communication"]) tank_count = _import_tank_specs(workbook["Tank Chart"]) vehicle_cost_count = _import_vehicle_costs(workbook["KSP Vehicle Cost"]) db.session.commit() except Exception: db.session.rollback() raise finally: workbook.close() summary.imported_rows = { "engine_families": family_count, "engine_variants": variant_count, "communication_parts": communication_count, "tank_specs": tank_count, "vehicle_costs": vehicle_cost_count, } return summary def _assert_schema_ready() -> None: existing_tables = set(inspect(db.engine).get_table_names()) required_tables = { "engine_families", "engine_variants", "communication_parts", "tank_specs", "vehicle_costs", } missing_tables = sorted(required_tables - existing_tables) if missing_tables: missing = ", ".join(missing_tables) raise ValueError(f"数据库缺少表:{missing}。请先执行数据库迁移。") def _database_has_data() -> bool: for model in (EngineFamily, EngineVariant, CommunicationPart, TankSpec, VehicleCost): if db.session.query(model).first() is not None: return True return False def _clear_import_tables() -> None: for model in (EngineVariant, EngineFamily, CommunicationPart, TankSpec, VehicleCost): db.session.execute(delete(model)) db.session.flush() def _import_engines(worksheet: object, summary: WorkbookImportSummary) -> tuple[int, int]: families_by_name: dict[str, EngineFamily] = {} variant_count = 0 for row_number, row in enumerate(worksheet.iter_rows(min_row=2, values_only=True), start=2): if not _row_has_values(row): continue engine_name = _text(row[0]) if engine_name is None: summary.warnings.append(f"Engine Database 第 {row_number} 行缺少 Engine,已跳过。") continue family = families_by_name.get(engine_name) if family is None: family = EngineFamily( engine_name=engine_name, part_name=_text(row[18]), cycle=_text(row[2]), size_m=_decimal(row[12]), entry_cost=_integer(row[14]), ) families_by_name[engine_name] = family db.session.add(family) else: family.part_name = family.part_name or _text(row[18]) family.cycle = family.cycle or _text(row[2]) family.size_m = family.size_m or _decimal(row[12]) family.entry_cost = family.entry_cost or _integer(row[14]) fuel_type = _text(row[1]) if fuel_type is None: summary.warnings.append(f"Engine Database 第 {row_number} 行缺少 Fuel Type,已跳过。") continue min_thrust = _decimal(row[6]) max_thrust = _decimal(row[7]) mass = _decimal(row[8]) variant = EngineVariant( family=family, fuel_type=fuel_type, work_env=_text(row[3]), sl_isp=_decimal(row[4]), vac_isp=_decimal(row[5]), min_thrust_kn=min_thrust, max_thrust_kn=max_thrust, mass_t=mass, twr=_decimal(row[9]) or _calculate_twr(max_thrust, mass), throttle_ratio=_decimal(row[10]) or _calculate_throttle_ratio(min_thrust, max_thrust), tvc_deg=_decimal(row[11]), price_kd=_integer(row[13]), ignitions=_integer(row[15]), has_unlimited_ignitions=_integer(row[15]) is None, mod_source=_text(row[16]), engine_note=_text(row[17]), config_name=_text(row[19]), tech_required=_text(row[20]), config_note=_text(row[21]), source_sheet_row=row_number, ) db.session.add(variant) variant_count += 1 return len(families_by_name), variant_count def _import_communication_parts(worksheet: object) -> int: count = 0 for row in worksheet.iter_rows(min_row=2, values_only=True): if not _row_has_values(row): continue part_name = _text(row[0]) if part_name is None: continue db.session.add( CommunicationPart( part_name=part_name, display_name=_text(row[1]), mass_t=_decimal(row[2]), is_active=_boolean(row[3]), is_deployable=_boolean(row[4]), antenna_type=_text(row[5]), deployed_diameter_m=_decimal(row[6]), range_raw=_decimal(row[7]), range_km=_decimal(row[8]), range_au=_decimal(row[9]), range_light_year=_decimal(row[10]), angle_deg=_decimal(row[11]), speed=_decimal(row[12]), idle_power_watt=_decimal(row[13]), idle_power_text=_text(row[14]), transmitting_power_watt=_decimal(row[15]), transmitting_power_text=_text(row[16]), source=_text(row[17]), entry_cost=_integer(row[18]), cost=_integer(row[19]), description=_text(row[20]), rescale_factor=_decimal(row[21]), tweakscale=_text(row[22]), is_feeder=_boolean(row[23]), tech_required=_text(row[24]), note=_text(row[25]), ) ) count += 1 return count def _import_tank_specs(worksheet: object) -> int: count = 0 for row in worksheet.iter_rows(min_row=2, values_only=True): if not _row_has_values(row): continue tank_name = _text(row[0]) if tank_name is None: continue db.session.add( TankSpec( tank_name=tank_name, fuel_type=_text(row[1]), dry_mass_t=_decimal(row[2]), fuel_mass_t=_decimal(row[3]), wet_mass_t=_decimal(row[4]), tank_volume_l=_decimal(row[5]), mass_ratio=_decimal(row[6]), kiloliters_per_ton=_decimal(row[7]), vehicle_name=_text(row[8]), source=_text(row[9]), note=_text(row[10]), ) ) count += 1 return count def _import_vehicle_costs(worksheet: object) -> int: count = 0 for row in worksheet.iter_rows(min_row=2, values_only=True): if not _row_has_values(row): continue vehicle_name = _text(row[0]) if vehicle_name is None: continue db.session.add( VehicleCost( vehicle_name=vehicle_name, launch_price=_integer(row[1]), source=_text(row[2]), ) ) count += 1 return count def _row_has_values(row: tuple[object, ...]) -> bool: return any(value not in (None, "") for value in row) def _text(value: object) -> str | None: if value in (None, ""): return None cleaned = str(value).strip() return cleaned or None def _decimal(value: object) -> Decimal | None: if value in (None, ""): return None if isinstance(value, Decimal): return value try: return Decimal(str(value)) except (InvalidOperation, ValueError, TypeError): return None def _integer(value: object) -> int | None: decimal_value = _decimal(value) if decimal_value is None: return None return int(decimal_value.to_integral_value()) def _boolean(value: object) -> bool | None: if value in (None, ""): return None if isinstance(value, bool): return value if isinstance(value, (int, float)): return bool(value) text = str(value).strip().lower() if text in {"true", "1", "yes", "y"}: return True if text in {"false", "0", "no", "n"}: return False return None def _calculate_twr(max_thrust: Decimal | None, mass: Decimal | None) -> Decimal | None: if max_thrust is None or mass in (None, Decimal("0")): return None return max_thrust / Decimal("9.80665") / mass def _calculate_throttle_ratio( min_thrust: Decimal | None, max_thrust: Decimal | None ) -> Decimal | None: if min_thrust is None or max_thrust in (None, Decimal("0")): return None return min_thrust / max_thrust