butubb
d937ff5fe6
上一个提交加了 4 个调度字段，却没在 _ADDED_COLUMNS 里登记，后果是生产环境： - 任务列表接口报 Unknown column，UI 打不开任务列表 - 调度器每 20 秒 tick 一次炸一次，定时任务完全不会触发 最阴险的是启动完全正常——能连库、能起来，只是随后每条查询都失败。 - _ensure_columns 不再遍历手写清单，改为遍历 MonitorBase.metadata.sorted_tables， 从根上消掉「加了字段忘了登记」这类漏 - 新增 _column_ddl：用 CreateColumn 渲染类型与可空性，并给 NOT NULL 列补 DEFAULT。 模型的 default= 是 ORM 侧行为、不会进 DDL，而给已有数据的表加 NOT NULL 列必须有值， 否则能否成功取决于服务端 sql_mode - 主键列跳过：MySQL 不允许 AUTO_INCREMENT 与 DEFAULT 共存 新增 tests/test_monitor_column_migration.py 守住：每个 NOT NULL 列都必须能生成带 DEFAULT 的合法 ALTER。
343 lines
12 KiB
Python
343 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) 2025 [email protected]
|
|
#
|
|
# This file is part of MediaCrawler project.
|
|
# Repository: https://github.com/NanmiCoder/MediaCrawler/blob/main/api/monitor/db.py
|
|
# GitHub: https://github.com/NanmiCoder
|
|
# Licensed under NON-COMMERCIAL LEARNING LICENSE 1.1
|
|
#
|
|
# 声明：本代码仅供学习和研究目的使用。使用者应遵守以下原则：
|
|
# 1. 不得用于任何商业用途。
|
|
# 2. 使用时应遵守目标平台的使用条款和robots.txt规则。
|
|
# 3. 不得进行大规模爬取或对平台造成运营干扰。
|
|
# 4. 应合理控制请求频率，避免给目标平台带来不必要的负担。
|
|
# 5. 不得用于任何非法或不当的用途。
|
|
#
|
|
# 详细许可条款请参阅项目根目录下的LICENSE文件。
|
|
# 使用本代码即表示您同意遵守上述原则和LICENSE中的所有条款。
|
|
|
|
"""Database engine for the monitoring layer.
|
|
|
|
**MySQL** by default (see ``config/db_config.py`` and ``.env``), with SQLite kept
|
|
as an option so the test suite can run without a reachable server.
|
|
|
|
Three things here exist because of specific MySQL 5.7 behaviour:
|
|
|
|
* **utf8mb4 is forced per table.** This instance's server *and* the target schema
|
|
default to ``latin1``; relying on either would mangle or reject Chinese text.
|
|
The charset is set on every table rather than on the database, so it holds no
|
|
matter what the schema default is.
|
|
* **Connections are recycled.** The monitor runs for weeks, and MySQL drops idle
|
|
connections after ``wait_timeout`` (8h by default). Without ``pool_recycle`` and
|
|
``pool_pre_ping`` the first query after a quiet night fails with "server has
|
|
gone away".
|
|
* **The connected schema is asserted at startup.** A misconfigured database name
|
|
is caught immediately instead of silently writing to the wrong schema.
|
|
|
|
Only the configured schema is ever touched: no ``CREATE DATABASE``, no ``USE``,
|
|
no cross-schema query.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import AsyncIterator, Optional
|
|
|
|
from sqlalchemy import Column, event, text
|
|
from sqlalchemy.dialects import mysql
|
|
from sqlalchemy.schema import CreateColumn
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
|
|
from .models import MonitorBase
|
|
|
|
PROJECT_ROOT = Path(__file__).parent.parent.parent
|
|
DATA_DIR = PROJECT_ROOT / "data"
|
|
DEFAULT_SQLITE_PATH = DATA_DIR / "monitor.db"
|
|
|
|
# Load .env here as well as in api/main.py: this module is imported directly by
|
|
# scripts and tests, and a configuration that only applies when the server is the
|
|
# entry point is a trap. load_dotenv does not override real environment variables.
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv(PROJECT_ROOT / ".env")
|
|
|
|
# Kept identical to config/db_config.py's defaults so one .env drives both the
|
|
# monitor database and the crawler's own DB output.
|
|
MYSQL_HOST = lambda: os.getenv("MYSQL_DB_HOST", "localhost") # noqa: E731
|
|
MYSQL_PORT = lambda: int(os.getenv("MYSQL_DB_PORT", "3306")) # noqa: E731
|
|
MYSQL_USER = lambda: os.getenv("MYSQL_DB_USER", "root") # noqa: E731
|
|
MYSQL_PWD = lambda: os.getenv("MYSQL_DB_PWD", "") # noqa: E731
|
|
MYSQL_DB_NAME = lambda: os.getenv("MYSQL_DB_NAME", "mediacrawler") # noqa: E731
|
|
|
|
_engine: Optional[AsyncEngine] = None
|
|
_session_factory: Optional[async_sessionmaker[AsyncSession]] = None
|
|
# None means "resolve from the environment" (MySQL). Tests set a SQLite URL.
|
|
_db_url: Optional[str] = None
|
|
_expected_schema: Optional[str] = None
|
|
|
|
|
|
def resolve_db_url() -> str:
|
|
"""Build the connection URL. MySQL unless overridden."""
|
|
if _db_url is not None:
|
|
return _db_url
|
|
|
|
from urllib.parse import quote_plus
|
|
|
|
user = quote_plus(MYSQL_USER())
|
|
password = quote_plus(MYSQL_PWD())
|
|
host = MYSQL_HOST()
|
|
port = MYSQL_PORT()
|
|
name = MYSQL_DB_NAME()
|
|
return f"mysql+aiomysql://{user}:{password}@{host}:{port}/{name}?charset=utf8mb4"
|
|
|
|
|
|
def is_mysql() -> bool:
|
|
return resolve_db_url().startswith("mysql")
|
|
|
|
|
|
def set_sqlite_path(path: Path) -> None:
|
|
"""Point the layer at SQLite. Used by the test suite only."""
|
|
global _db_url, _engine, _session_factory, _expected_schema
|
|
_db_url = f"sqlite+aiosqlite:///{Path(path)}"
|
|
_engine = None
|
|
_session_factory = None
|
|
_expected_schema = None
|
|
|
|
|
|
def set_db_url(url: str, expected_schema: Optional[str] = None) -> None:
|
|
"""Point the layer at an explicit URL. ``expected_schema`` enables the guard."""
|
|
global _db_url, _engine, _session_factory, _expected_schema
|
|
_db_url = url
|
|
_engine = None
|
|
_session_factory = None
|
|
_expected_schema = expected_schema
|
|
|
|
|
|
def expected_schema() -> Optional[str]:
|
|
"""The schema the connection must be using, if the guard applies."""
|
|
if _expected_schema is not None:
|
|
return _expected_schema
|
|
return MYSQL_DB_NAME() if is_mysql() else None
|
|
|
|
|
|
def get_engine() -> AsyncEngine:
|
|
global _engine
|
|
if _engine is None:
|
|
url = resolve_db_url()
|
|
kwargs: dict = {"future": True}
|
|
|
|
if url.startswith("mysql"):
|
|
# Recycle well inside MySQL's default 8h wait_timeout, and verify a
|
|
# pooled connection before handing it out.
|
|
kwargs.update(pool_recycle=3600, pool_pre_ping=True, pool_size=5, max_overflow=5)
|
|
kwargs["connect_args"] = {"charset": "utf8mb4"}
|
|
else:
|
|
Path(url.split("///", 1)[-1]).parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
_engine = create_async_engine(url, **kwargs)
|
|
|
|
if url.startswith("sqlite"):
|
|
|
|
@event.listens_for(_engine.sync_engine, "connect")
|
|
def _set_sqlite_pragmas(dbapi_connection, _connection_record): # pragma: no cover
|
|
cursor = dbapi_connection.cursor()
|
|
cursor.execute("PRAGMA journal_mode=WAL")
|
|
cursor.execute("PRAGMA foreign_keys=ON")
|
|
cursor.close()
|
|
|
|
return _engine
|
|
|
|
|
|
def get_session_factory() -> async_sessionmaker[AsyncSession]:
|
|
global _session_factory
|
|
if _session_factory is None:
|
|
_session_factory = async_sessionmaker(
|
|
bind=get_engine(),
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
return _session_factory
|
|
|
|
|
|
@asynccontextmanager
|
|
async def get_session() -> AsyncIterator[AsyncSession]:
|
|
"""Transactional session. Commits on success, rolls back on error."""
|
|
factory = get_session_factory()
|
|
async with factory() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
|
|
|
|
async def _assert_correct_schema(conn) -> None:
|
|
"""Refuse to run against anything but the configured schema.
|
|
|
|
A guard, not the guarantee: the real protection is a MySQL account scoped to
|
|
this one schema (see UPSTREAM.md). This catches the ordinary mistake of a
|
|
wrong database name in configuration, before a single row is written.
|
|
"""
|
|
if not is_mysql():
|
|
return
|
|
|
|
expected = expected_schema()
|
|
if not expected:
|
|
return
|
|
|
|
current = (await conn.execute(text("SELECT DATABASE()"))).scalar()
|
|
if current is None:
|
|
raise RuntimeError(
|
|
f"数据库连接未选定 schema，期望 {expected!r}。请检查 MYSQL_DB_NAME。"
|
|
)
|
|
# lower_case_table_names=1 makes names case-insensitive server-side.
|
|
if current.lower() != expected.lower():
|
|
raise RuntimeError(
|
|
f"连接的库是 {current!r}，但配置要求 {expected!r}。"
|
|
f"为避免误写其它库，已拒绝启动。"
|
|
)
|
|
print(f"[monitor.db] 已连接 MySQL schema: {current}", flush=True)
|
|
|
|
|
|
async def init_db() -> None:
|
|
"""Create missing tables, then run the small in-place migrations."""
|
|
engine = get_engine()
|
|
async with engine.begin() as conn:
|
|
await _assert_correct_schema(conn)
|
|
await conn.run_sync(MonitorBase.metadata.create_all)
|
|
await _ensure_columns(conn)
|
|
await _migrate_setting_keys(conn)
|
|
|
|
|
|
# ``create_all`` creates missing *tables* but never adds *columns* to a table that
|
|
# already exists, so those need an explicit ALTER TABLE.
|
|
#
|
|
# Which columns those are is derived from the ORM metadata, not kept by hand. The
|
|
# hand-kept version was a trap: forgetting to register a new column there still let
|
|
# the app start -- it connects fine, then fails on every query and every scheduler
|
|
# tick. Which is exactly what happened when the scheduling columns were added.
|
|
|
|
|
|
def _implicit_default(column: Column) -> Optional[str]:
|
|
"""A literal to seed existing rows with when a NOT NULL column is added."""
|
|
default = column.default
|
|
if default is not None and getattr(default, "is_scalar", False):
|
|
value = default.arg
|
|
if isinstance(value, bool):
|
|
return "1" if value else "0"
|
|
if isinstance(value, (int, float)):
|
|
return str(value)
|
|
return "'" + str(value).replace("'", "''") + "'"
|
|
|
|
# No scalar default on the model. Fall back to the type's zero value, so that
|
|
# adding the column cannot depend on the server's sql_mode.
|
|
try:
|
|
python_type = column.type.python_type
|
|
except NotImplementedError:
|
|
return None
|
|
if python_type in (bool, int, float):
|
|
return "0"
|
|
if python_type is str:
|
|
return "''"
|
|
return None
|
|
|
|
|
|
def _column_ddl(column: Column) -> str:
|
|
"""One column as MySQL DDL for ``ALTER TABLE ... ADD COLUMN``.
|
|
|
|
``CreateColumn`` renders the name, type and nullability. The default is added
|
|
separately because a model's ``default=`` is applied by the ORM and never
|
|
reaches the DDL -- and a NOT NULL column added to a populated table needs a
|
|
value for the rows already sitting there.
|
|
"""
|
|
ddl = str(CreateColumn(column).compile(dialect=mysql.dialect()))
|
|
if not column.nullable and column.server_default is None:
|
|
seed = _implicit_default(column)
|
|
if seed is not None:
|
|
ddl += f" DEFAULT {seed}"
|
|
return ddl
|
|
|
|
|
|
async def _existing_columns(conn, table: str) -> set[str]:
|
|
if is_mysql():
|
|
rows = await conn.execute(
|
|
text(
|
|
"SELECT COLUMN_NAME FROM information_schema.COLUMNS "
|
|
"WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = :t"
|
|
),
|
|
{"t": table},
|
|
)
|
|
return {row[0] for row in rows}
|
|
|
|
rows = await conn.execute(text(f"PRAGMA table_info({table})"))
|
|
return {row[1] for row in rows}
|
|
|
|
|
|
async def _ensure_columns(conn) -> None:
|
|
"""Add every model column the live table is missing."""
|
|
for table in MonitorBase.metadata.sorted_tables:
|
|
existing = await _existing_columns(conn, table.name)
|
|
if not existing:
|
|
# Table did not exist before this run; create_all built it complete.
|
|
continue
|
|
for column in table.columns:
|
|
# Primary keys are always present, and MySQL rejects AUTO_INCREMENT
|
|
# alongside the DEFAULT this helper appends -- so skip them rather
|
|
# than emit DDL that could never run.
|
|
if column.name in existing or column.primary_key:
|
|
continue
|
|
print(f"[monitor.db] 补齐缺失字段 {table.name}.{column.name}", flush=True)
|
|
await conn.execute(
|
|
text(f"ALTER TABLE {table.name} ADD COLUMN {_column_ddl(column)}")
|
|
)
|
|
|
|
|
|
async def _migrate_setting_keys(conn) -> None:
|
|
"""Move pre-namespacing setting keys to their scoped names.
|
|
|
|
Idempotent: the legacy row is only renamed when the new key is absent, so an
|
|
operator's later value is never overwritten.
|
|
"""
|
|
from .models import LEGACY_SETTING_KEY_RENAMES
|
|
|
|
for legacy, scoped in LEGACY_SETTING_KEY_RENAMES.items():
|
|
exists = (
|
|
await conn.execute(
|
|
text("SELECT 1 FROM monitor_setting WHERE `key` = :k"), {"k": legacy}
|
|
)
|
|
).first()
|
|
if not exists:
|
|
continue
|
|
|
|
already = (
|
|
await conn.execute(
|
|
text("SELECT 1 FROM monitor_setting WHERE `key` = :k"), {"k": scoped}
|
|
)
|
|
).first()
|
|
if already:
|
|
# Both present: the scoped one is authoritative; drop the stale row.
|
|
await conn.execute(
|
|
text("DELETE FROM monitor_setting WHERE `key` = :k"), {"k": legacy}
|
|
)
|
|
continue
|
|
|
|
await conn.execute(
|
|
text("UPDATE monitor_setting SET `key` = :new WHERE `key` = :old"),
|
|
{"new": scoped, "old": legacy},
|
|
)
|
|
|
|
|
|
async def dispose_engine() -> None:
|
|
global _engine, _session_factory
|
|
if _engine is not None:
|
|
await _engine.dispose()
|
|
_engine = None
|
|
_session_factory = None
|