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()