yu-zhi-ran/platform/backend/app/database.py

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import os
from pathlib import Path
from datetime import datetime, timezone, timedelta

PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent

USE_POSTGRES = os.getenv('USE_POSTGRES', 'true').lower() == 'true'

if USE_POSTGRES:
    POSTGRES_CONFIG = {
        'host': os.getenv('PG_HOST', '127.0.0.1'),
        'port': os.getenv('PG_PORT', '5432'),
        'database': os.getenv('PG_DATABASE', 'yzr_nr'),
        'user': os.getenv('PG_USER', 'yzr_nr'),
        'password': os.getenv('PG_PASSWORD', 'aTX3WKKnPfRnM5PC')
    }
    SQLALCHEMY_DATABASE_URL = (
        f"postgresql://{POSTGRES_CONFIG['user']}:{POSTGRES_CONFIG['password']}"
        f"@{POSTGRES_CONFIG['host']}:{POSTGRES_CONFIG['port']}/{POSTGRES_CONFIG['database']}"
    )
    engine = create_engine(
        SQLALCHEMY_DATABASE_URL,
        pool_pre_ping=True,
        pool_size=10,
        max_overflow=20
    )
else:
    DATA_DIR = os.getenv('DATA_DIR', str(PROJECT_ROOT / 'data'))
    os.makedirs(DATA_DIR, exist_ok=True)
    DB_PATH = os.path.join(DATA_DIR, 'yzr.db')
    SQLALCHEMY_DATABASE_URL = f"sqlite:///{DB_PATH}"
    engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()

def init_db():
    Base.metadata.create_all(bind=engine)
    from sqlalchemy import text
    try:
        with engine.connect() as conn:
            # 迁移：为已有表添加列
            conn.execute(text("ALTER TABLE users ADD COLUMN IF NOT EXISTS last_login TIMESTAMP"))
            conn.execute(text("ALTER TABLE llm_configs ADD COLUMN IF NOT EXISTS provider VARCHAR DEFAULT 'opencode-go'"))
            conn.execute(text("ALTER TABLE llm_configs ADD COLUMN IF NOT EXISTS base_url VARCHAR"))
            conn.execute(text("ALTER TABLE llm_configs ADD COLUMN IF NOT EXISTS api_key VARCHAR"))
            conn.execute(text("ALTER TABLE llm_configs ADD COLUMN IF NOT EXISTS is_default BOOLEAN DEFAULT FALSE"))
            conn.execute(text("ALTER TABLE llm_configs ADD COLUMN IF NOT EXISTS rate_limit INTEGER DEFAULT 0"))
            conn.execute(text("ALTER TABLE llm_configs ADD COLUMN IF NOT EXISTS rate_limit_window_minutes INTEGER DEFAULT 300"))
            try:
                conn.execute(text("ALTER TABLE articles ADD COLUMN IF NOT EXISTS images JSON DEFAULT '{}'::json"))
            except Exception:
                conn.execute(text("ALTER TABLE articles ADD COLUMN IF NOT EXISTS images TEXT DEFAULT '{}'"))
            for table, col, typ in [
                ("users", "org_id", "VARCHAR DEFAULT 'default'"),
                ("articles", "title", "VARCHAR"),
                ("articles", "content", "TEXT"),
                ("articles", "updated_at", "TIMESTAMP"),
                ("media_assets", "org_id", "VARCHAR DEFAULT 'default'"),
                ("topics", "org_id", "VARCHAR DEFAULT 'default'"),
                ("topics", "reviewed_at", "TIMESTAMP"),
                ("platform_configs", "requires_image", "BOOLEAN DEFAULT FALSE"),
                ("platform_configs", "image_count_min", "INTEGER DEFAULT 0"),
                ("platform_configs", "image_count_max", "INTEGER DEFAULT 0"),
                ("platform_configs", "image_width", "INTEGER DEFAULT 0"),
                ("platform_configs", "image_height", "INTEGER DEFAULT 0"),
                ("platform_configs", "min_words", "INTEGER DEFAULT 0"),
                ("platform_configs", "max_words", "INTEGER DEFAULT 0"),
                ("platform_configs", "website_url", "VARCHAR"),
                ("task_logs", "module_id", "VARCHAR"),
                ("task_logs", "error_trace", "TEXT"),
                ("task_logs", "triggered_by", "VARCHAR DEFAULT 'scheduler'"),
                ("task_logs", "result_data", "JSON DEFAULT '{}'::json"),
                ("task_logs", "next_run_time", "TIMESTAMP"),
                ("task_configs", "module_id", "VARCHAR UNIQUE"),
                ("task_configs", "enabled", "BOOLEAN DEFAULT TRUE"),
                ("task_configs", "params", "JSON DEFAULT '{}'::json"),
                ("task_configs", "schedule", "VARCHAR"),
                ("task_configs", "last_modified_by", "VARCHAR"),
                ("prompt_configs", "key", "VARCHAR UNIQUE"),
                ("prompt_configs", "module_id", "VARCHAR"),
                ("prompt_configs", "category", "VARCHAR DEFAULT 'prompt'"),
                ("prompt_configs", "version", "VARCHAR DEFAULT 'v1'"),
                ("prompt_configs", "content", "TEXT"),
                ("prompt_configs", "variables", "JSON DEFAULT '[]'::json"),
                ("prompt_configs", "description", "VARCHAR"),
                ("prompt_configs", "enabled", "BOOLEAN DEFAULT TRUE"),
                ("prompt_configs", "temperature", "FLOAT"),
                ("prompt_configs", "max_tokens", "INTEGER"),
                ("prompt_configs", "created_by", "VARCHAR"),
                ("search_providers", "console_url", "VARCHAR"),
                ("keyword_domain_map", "id", "INTEGER PRIMARY KEY"),
                ("keyword_domain_map", "pattern", "VARCHAR"),
                ("keyword_domain_map", "domain", "VARCHAR"),
                ("keyword_domain_map", "sort_order", "INTEGER DEFAULT 0"),
                ("keyword_domain_map", "is_active", "BOOLEAN DEFAULT TRUE"),
                ("sensitive_words", "id", "INTEGER PRIMARY KEY"),
                ("sensitive_words", "word", "VARCHAR"),
                ("sensitive_words", "category", "VARCHAR DEFAULT 'general'"),
                ("sensitive_words", "is_active", "BOOLEAN DEFAULT TRUE"),
                ("sensitive_words", "added_by", "VARCHAR"),
                ("content_clean_rules", "id", "INTEGER PRIMARY KEY"),
                ("content_clean_rules", "rule_type", "VARCHAR"),
                ("content_clean_rules", "pattern", "TEXT"),
                ("content_clean_rules", "description", "VARCHAR"),
                ("content_clean_rules", "is_active", "BOOLEAN DEFAULT TRUE"),
                ("content_clean_rules", "sort_order", "INTEGER DEFAULT 0"),
                ("collector_categories", "pain_template", "TEXT"),
                ("trend_field_mappings", "id", "INTEGER PRIMARY KEY"),
                ("trend_field_mappings", "trend_keyword", "VARCHAR"),
                ("trend_field_mappings", "field_name", "VARCHAR"),
                ("trend_field_mappings", "sort_order", "INTEGER DEFAULT 0"),
                ("trend_field_mappings", "is_active", "BOOLEAN DEFAULT TRUE"),
                ("roles", "id", "INTEGER PRIMARY KEY"),
                ("roles", "name", "VARCHAR UNIQUE"),
                ("roles", "description", "VARCHAR DEFAULT ''"),
                ("roles", "is_system", "BOOLEAN DEFAULT FALSE"),
                ("menus", "id", "INTEGER PRIMARY KEY"),
                ("menus", "parent_id", "INTEGER REFERENCES menus(id)"),
                ("menus", "name", "VARCHAR"),
                ("menus", "path", "VARCHAR"),
                ("menus", "icon", "VARCHAR DEFAULT ''"),
                ("menus", "sort_order", "INTEGER DEFAULT 0"),
                ("menus", "roles", "JSON DEFAULT '[]'::json"),
                ("menus", "is_active", "BOOLEAN DEFAULT TRUE"),
            ]:
                try:
                    conn.execute(text(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS {col} {typ}"))
                except Exception:
                    try:
                        conn.execute(text(f"ALTER TABLE {table} ADD COLUMN {col} {typ}"))
                    except Exception:
                        pass
            # Create roles, menus, and search_providers tables if they don't exist
            for tbl_sql in [
                "CREATE TABLE IF NOT EXISTS roles (id SERIAL PRIMARY KEY, name VARCHAR UNIQUE NOT NULL, description VARCHAR DEFAULT '', is_system BOOLEAN DEFAULT FALSE, created_at TIMESTAMP WITH TIME ZONE DEFAULT now())",
                "CREATE TABLE IF NOT EXISTS menus (id SERIAL PRIMARY KEY, parent_id INTEGER REFERENCES menus(id), name VARCHAR NOT NULL, path VARCHAR NOT NULL, icon VARCHAR DEFAULT '', sort_order INTEGER DEFAULT 0, roles JSON DEFAULT '[]'::json, is_active BOOLEAN DEFAULT TRUE, created_at TIMESTAMP WITH TIME ZONE DEFAULT now())",
                "CREATE TABLE IF NOT EXISTS search_providers (id SERIAL PRIMARY KEY, name VARCHAR NOT NULL, provider_type VARCHAR NOT NULL, api_key VARCHAR, api_url VARCHAR, priority INTEGER DEFAULT 1, enabled BOOLEAN DEFAULT TRUE, daily_limit INTEGER DEFAULT 1500, usage_today INTEGER DEFAULT 0, extra_config JSON DEFAULT '{}'::json, last_used_at TIMESTAMP, created_at TIMESTAMP WITH TIME ZONE DEFAULT now(), updated_at TIMESTAMP WITH TIME ZONE DEFAULT now())",
            ]:
                try:
                    conn.execute(text(tbl_sql))
                except Exception:
                    pass
            conn.commit()
    except Exception:
        pass  # SQLite 不支持 IF NOT EXISTS，但 create_all 对 SQLite 够用，这里仅为 PostgreSQL 迁移

    # 清理卡死的 running 任务（超过 30 分钟无更新的标记为失败）
    try:
        with SessionLocal() as sess:
            from sqlalchemy import text as sa_text
            cutoff = datetime.now(timezone.utc) - timedelta(minutes=30)
            sess.execute(
                sa_text("UPDATE task_logs SET status='failed', message='系统启动时清理卡死任务', finished_at=now() "
                        "WHERE status='running' AND started_at < :cutoff AND module_id != 'scheduled_task_monitor'"),
                {"cutoff": cutoff}
            )
            sess.commit()
    except Exception:
        pass

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()