206 lines
7.2 KiB
Python
206 lines
7.2 KiB
Python
# -*- 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,
|
||
}
|
||
|
||
|
||
class AuthToken(Base):
|
||
"""登录 token 表：每个用户名仅保留当前有效 token。"""
|
||
__tablename__ = "auth_token"
|
||
|
||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||
username = Column(String(64), nullable=False, unique=True, comment="用户名")
|
||
token = Column(String(64), nullable=False, comment="当前有效 token")
|
||
created_at = Column(DateTime, default=func.now(), comment="创建时间")
|
||
|
||
__table_args__ = (
|
||
{"mysql_charset": "utf8mb4", "comment": "登录 token"},
|
||
)
|
||
|
||
|
||
# ══════════════════════════════════════════════════════════════
|
||
# 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
|
||
|
||
|
||
class LoginRequest(BaseModel):
|
||
"""登录请求：用户名 + 密码。"""
|
||
username: str
|
||
password: str
|
||
|
||
|
||
class LoginResponse(BaseModel):
|
||
"""登录成功响应：返回 token。"""
|
||
token: str
|