boss_dp/server/models.py

# -*- coding: utf-8 -*-
"""
数据模型：SQLAlchemy ORM 表模型 + Pydantic 请求/响应模型。
"""
from __future__ import annotations

import time
import uuid
from datetime import datetime
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field
from sqlalchemy import (
    Boolean, Column, DateTime, Integer, JSON, String, Text,
    UniqueConstraint, func,
)
from sqlalchemy.orm import DeclarativeBase

from common.protocol import TaskStatus, TaskType


# ══════════════════════════════════════════════════════════════
#  SQLAlchemy ORM 模型
# ══════════════════════════════════════════════════════════════

class Base(DeclarativeBase):
    """SQLAlchemy 声明式基类。"""
    pass


class BossAccount(Base):
    """BOSS 账号登录状态表。"""
    __tablename__ = "boss_account"

    id = Column(Integer, primary_key=True, autoincrement=True)
    worker_id = Column(String(64), nullable=False, comment="Worker 标识")
    browser_id = Column(String(128), nullable=False, default="", comment="比特浏览器窗口 ID")
    browser_name = Column(String(128), default="", comment="比特浏览器窗口名称（环境名）")
    boss_username = Column(String(128), default="", comment="BOSS 直聘用户名")
    is_logged_in = Column(Boolean, default=False, comment="是否已登录")
    checked_at = Column(DateTime, nullable=True, comment="最近一次检测时间")
    created_at = Column(DateTime, default=func.now(), comment="创建时间")
    updated_at = Column(DateTime, default=func.now(), onupdate=func.now(), comment="更新时间")

    __table_args__ = (
        UniqueConstraint("worker_id", "browser_id", name="uk_worker_browser"),
        {"mysql_charset": "utf8mb4", "comment": "BOSS 账号登录状态"},
    )

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "worker_id": self.worker_id,
            "browser_id": self.browser_id,
            "browser_name": self.browser_name,
            "boss_username": self.boss_username,
            "is_logged_in": self.is_logged_in,
            "checked_at": self.checked_at.isoformat() if self.checked_at else None,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }


class TaskLog(Base):
    """任务执行记录表。"""
    __tablename__ = "task_log"

    id = Column(Integer, primary_key=True, autoincrement=True)
    task_id = Column(String(32), nullable=False, unique=True, comment="任务 ID")
    task_type = Column(String(64), nullable=False, comment="任务类型")
    worker_id = Column(String(64), default="", comment="执行的 Worker")
    status = Column(String(32), default="", comment="最终状态")
    params = Column(JSON, nullable=True, comment="任务参数")
    result = Column(JSON, nullable=True, comment="任务结果")
    error = Column(Text, nullable=True, comment="错误信息")
    created_at = Column(DateTime, default=func.now(), comment="创建时间")

    __table_args__ = (
        {"mysql_charset": "utf8mb4", "comment": "任务执行记录"},
    )

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "task_id": self.task_id,
            "task_type": self.task_type,
            "worker_id": self.worker_id,
            "status": self.status,
            "params": self.params,
            "result": self.result,
            "error": self.error,
            "created_at": self.created_at.isoformat() if self.created_at else None,
        }


# ══════════════════════════════════════════════════════════════
#  Pydantic 请求 / 响应模型（API 用）
# ══════════════════════════════════════════════════════════════

# ─── Worker ───

class BrowserProfile(BaseModel):
    """比特浏览器窗口信息（Worker 上报）。"""
    id: str
    name: str = ""
    remark: str = ""


class WorkerInfo(BaseModel):
    """一台 Worker 的运行时信息（内存中保存）。"""
    worker_id: str
    worker_name: str = ""
    browsers: List[BrowserProfile] = []
    online: bool = True
    last_heartbeat: float = Field(default_factory=time.time)
    connected_at: float = Field(default_factory=time.time)
    current_task_id: Optional[str] = None


class WorkerOut(BaseModel):
    """返回给前端的 Worker 信息。"""
    worker_id: str
    worker_name: str
    browsers: List[BrowserProfile]
    online: bool
    current_task_id: Optional[str] = None


# ─── Task ───

class TaskCreate(BaseModel):
    """前端提交任务的请求体。"""
    task_type: TaskType
    worker_id: Optional[str] = None
    account_name: Optional[str] = None
    params: Dict[str, Any] = {}


class TaskInfo(BaseModel):
    """任务完整信息（内存 / 返回前端）。"""
    task_id: str = Field(default_factory=lambda: uuid.uuid4().hex[:12])
    task_type: TaskType
    status: TaskStatus = TaskStatus.PENDING
    worker_id: Optional[str] = None
    account_name: Optional[str] = None
    params: Dict[str, Any] = {}
    progress: Optional[str] = None
    result: Any = None
    error: Optional[str] = None
    created_at: float = Field(default_factory=time.time)
    updated_at: float = Field(default_factory=time.time)


class TaskOut(BaseModel):
    """返回给前端的任务信息。"""
    task_id: str
    task_type: TaskType
    status: TaskStatus
    worker_id: Optional[str] = None
    account_name: Optional[str] = None
    params: Dict[str, Any] = {}
    progress: Optional[str] = None
    result: Any = None
    error: Optional[str] = None
    created_at: float
    updated_at: float


# ─── 简化接口：前端添加环境名 ───

class CheckLoginRequest(BaseModel):
    """前端提交检测登录请求。worker_id 可不传（走绑定关系）。"""
    browser_name: str
    worker_id: Optional[str] = None


class AccountBindRequest(BaseModel):
    """前端添加账号时提交绑定：账号环境名 + 归属电脑。"""
    browser_name: str
    worker_id: str