476 lines
19 KiB
Python
476 lines
19 KiB
Python
"""
|
|
数据库模型和操作
|
|
"""
|
|
import sqlite3
|
|
import json
|
|
from datetime import datetime
|
|
from contextlib import contextmanager
|
|
from config import Config
|
|
|
|
class Database:
|
|
def __init__(self, db_path=None):
|
|
self.db_path = db_path or Config.DATABASE
|
|
self.init_db()
|
|
|
|
@contextmanager
|
|
def get_connection(self):
|
|
"""获取数据库连接"""
|
|
conn = sqlite3.connect(self.db_path)
|
|
conn.row_factory = sqlite3.Row
|
|
try:
|
|
yield conn
|
|
finally:
|
|
conn.close()
|
|
|
|
def init_db(self):
|
|
"""初始化数据库"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# 内容库表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS articles (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
product_names TEXT NOT NULL,
|
|
search_title TEXT,
|
|
category TEXT,
|
|
keywords TEXT,
|
|
summary TEXT,
|
|
content TEXT,
|
|
source TEXT,
|
|
url TEXT,
|
|
fetch_date DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# 为旧表添加字段（如果不存在）
|
|
try:
|
|
cursor.execute('ALTER TABLE articles ADD COLUMN search_title TEXT')
|
|
except:
|
|
pass
|
|
|
|
# 待处理产品列表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS pending_products (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
product_name TEXT NOT NULL UNIQUE,
|
|
category TEXT,
|
|
subcategory TEXT,
|
|
priority INTEGER DEFAULT 0,
|
|
source TEXT,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# 处理中产品列表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS processing_products (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
product_name TEXT NOT NULL UNIQUE,
|
|
category TEXT,
|
|
subcategory TEXT,
|
|
status TEXT DEFAULT 'processing',
|
|
started_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
error_message TEXT
|
|
)
|
|
''')
|
|
|
|
# 处理历史记录
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS process_history (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
product_name TEXT NOT NULL,
|
|
category TEXT,
|
|
subcategory TEXT,
|
|
status TEXT,
|
|
review_id TEXT,
|
|
submitted_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
details TEXT
|
|
)
|
|
''')
|
|
|
|
# 任务配置表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS task_configs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT NOT NULL UNIQUE,
|
|
config TEXT,
|
|
enabled INTEGER DEFAULT 1,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# 系统配置表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS system_config (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# 失败的URL记录表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS failed_urls (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
url TEXT NOT NULL,
|
|
title TEXT,
|
|
error_message TEXT,
|
|
retry_count INTEGER DEFAULT 0,
|
|
status TEXT DEFAULT 'failed',
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
last_retry_at DATETIME,
|
|
source TEXT DEFAULT 'search'
|
|
)
|
|
''')
|
|
|
|
# 搜索缓存表
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS search_cache (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
keyword TEXT NOT NULL,
|
|
engine TEXT DEFAULT 'bing_cn',
|
|
results TEXT NOT NULL,
|
|
result_count INTEGER,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
expires_at DATETIME,
|
|
UNIQUE(keyword, engine)
|
|
)
|
|
''')
|
|
|
|
conn.commit()
|
|
|
|
# ========== 内容库操作 ==========
|
|
def add_article(self, product_names, category, keywords, summary, content, source, url=None, search_title=None):
|
|
"""添加文章到内容库"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT INTO articles (product_names, search_title, category, keywords, summary, content, source, url)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
''', (json.dumps(product_names, ensure_ascii=False), search_title,
|
|
category, json.dumps(keywords, ensure_ascii=False), summary, content, source, url))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
|
|
def search_articles(self, keyword, category=None):
|
|
"""搜索文章"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
if category:
|
|
cursor.execute('''
|
|
SELECT * FROM articles
|
|
WHERE (product_names LIKE ? OR keywords LIKE ? OR summary LIKE ? OR content LIKE ? OR url LIKE ?)
|
|
AND category = ?
|
|
ORDER BY fetch_date DESC
|
|
''', (f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', category))
|
|
else:
|
|
cursor.execute('''
|
|
SELECT * FROM articles
|
|
WHERE product_names LIKE ? OR keywords LIKE ? OR summary LIKE ? OR content LIKE ? OR url LIKE ?
|
|
ORDER BY fetch_date DESC
|
|
''', (f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', f'%{keyword}%'))
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_article_by_id(self, article_id):
|
|
"""获取文章详情"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT * FROM articles WHERE id = ?', (article_id,))
|
|
row = cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_all_articles(self, limit=100, offset=0):
|
|
"""获取所有文章"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT * FROM articles ORDER BY fetch_date DESC LIMIT ? OFFSET ?', (limit, offset))
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def delete_article(self, article_id):
|
|
"""删除文章"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM articles WHERE id = ?', (article_id,))
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def get_articles_count(self):
|
|
"""获取文章总数"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT COUNT(*) as count FROM articles')
|
|
row = cursor.fetchone()
|
|
return row['count'] if row else 0
|
|
|
|
# ========== 待处理产品操作 ==========
|
|
def add_pending_product(self, product_name, category=None, subcategory=None, priority=0, source='manual'):
|
|
"""添加待处理产品"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
try:
|
|
cursor.execute('''
|
|
INSERT INTO pending_products (product_name, category, subcategory, priority, source)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
''', (product_name, category, subcategory, priority, source))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
except sqlite3.IntegrityError:
|
|
# 产品已存在，更新优先级
|
|
cursor.execute('''
|
|
UPDATE pending_products
|
|
SET priority = MAX(priority, ?), updated_at = CURRENT_TIMESTAMP
|
|
WHERE product_name = ?
|
|
''', (priority, product_name))
|
|
conn.commit()
|
|
return None
|
|
|
|
def get_pending_products(self, limit=10, order_by='priority'):
|
|
"""获取待处理产品列表"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
if order_by == 'priority':
|
|
cursor.execute('SELECT * FROM pending_products ORDER BY priority DESC, created_at ASC LIMIT ?', (limit,))
|
|
else:
|
|
cursor.execute('SELECT * FROM pending_products ORDER BY created_at ASC LIMIT ?', (limit,))
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_pending_count(self):
|
|
"""获取待处理产品数量"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT COUNT(*) FROM pending_products')
|
|
return cursor.fetchone()[0]
|
|
|
|
def remove_pending_product(self, product_name):
|
|
"""从待处理列表移除产品"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM pending_products WHERE product_name = ?', (product_name,))
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
# ========== 处理中产品操作 ==========
|
|
def start_processing(self, product_name, category, subcategory):
|
|
"""开始处理产品"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
try:
|
|
cursor.execute('''
|
|
INSERT INTO processing_products (product_name, category, subcategory, status)
|
|
VALUES (?, ?, ?, 'processing')
|
|
''', (product_name, category, subcategory))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
except sqlite3.IntegrityError:
|
|
return None
|
|
|
|
def finish_processing(self, product_name, status='completed', error_message=None):
|
|
"""完成处理"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM processing_products WHERE product_name = ?', (product_name,))
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def get_processing_products(self):
|
|
"""获取处理中的产品"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT * FROM processing_products')
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
# ========== 处理历史操作 ==========
|
|
def add_process_history(self, product_name, category, subcategory, status, review_id=None, details=None):
|
|
"""添加处理历史"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT INTO process_history (product_name, category, subcategory, status, review_id, details)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
''', (product_name, category, subcategory, status, review_id,
|
|
json.dumps(details, ensure_ascii=False) if details else None))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
|
|
def get_process_history(self, limit=100):
|
|
"""获取处理历史"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT * FROM process_history ORDER BY submitted_at DESC LIMIT ?', (limit,))
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_history_by_product(self, product_name):
|
|
"""获取指定产品的处理历史"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT * FROM process_history WHERE product_name = ? ORDER BY submitted_at DESC', (product_name,))
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
# ========== 任务配置操作 ==========
|
|
def save_task_config(self, name, config):
|
|
"""保存任务配置"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT OR REPLACE INTO task_configs (name, config, updated_at)
|
|
VALUES (?, ?, CURRENT_TIMESTAMP)
|
|
''', (name, json.dumps(config, ensure_ascii=False)))
|
|
conn.commit()
|
|
|
|
def get_task_config(self, name):
|
|
"""获取任务配置"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT * FROM task_configs WHERE name = ?', (name,))
|
|
row = cursor.fetchone()
|
|
if row:
|
|
result = dict(row)
|
|
result['config'] = json.loads(result['config'])
|
|
return result
|
|
return None
|
|
|
|
# ========== 系统配置操作 ==========
|
|
def get_system_config(self, key, default=None):
|
|
"""获取系统配置"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT value FROM system_config WHERE key = ?', (key,))
|
|
row = cursor.fetchone()
|
|
return row['value'] if row else default
|
|
|
|
def set_system_config(self, key, value):
|
|
"""设置系统配置"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT OR REPLACE INTO system_config (key, value, updated_at)
|
|
VALUES (?, ?, CURRENT_TIMESTAMP)
|
|
''', (key, value))
|
|
conn.commit()
|
|
|
|
# ========== 失败URL操作 ==========
|
|
def add_failed_url(self, url, title=None, error_message=None, source='search'):
|
|
"""添加失败的URL"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
# 先检查是否已存在
|
|
cursor.execute('SELECT id, retry_count FROM failed_urls WHERE url = ?', (url,))
|
|
existing = cursor.fetchone()
|
|
|
|
if existing:
|
|
# 更新重试次数和错误信息
|
|
cursor.execute('''
|
|
UPDATE failed_urls
|
|
SET error_message = ?, last_retry_at = CURRENT_TIMESTAMP, retry_count = retry_count + 1
|
|
WHERE url = ?
|
|
''', (error_message, url))
|
|
else:
|
|
# 新增失败记录
|
|
cursor.execute('''
|
|
INSERT INTO failed_urls (url, title, error_message, source)
|
|
VALUES (?, ?, ?, ?)
|
|
''', (url, title, error_message, source))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
|
|
def get_failed_urls(self, limit=100, status='failed'):
|
|
"""获取失败的URL列表"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
SELECT * FROM failed_urls
|
|
WHERE status = ?
|
|
ORDER BY created_at DESC
|
|
LIMIT ?
|
|
''', (status, limit))
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_failed_url_count(self):
|
|
"""获取失败URL数量"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('SELECT COUNT(*) FROM failed_urls WHERE status = "failed"')
|
|
return cursor.fetchone()[0]
|
|
|
|
def mark_url_success(self, url):
|
|
"""标记URL为成功（已处理）"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
UPDATE failed_urls
|
|
SET status = 'success', last_retry_at = CURRENT_TIMESTAMP
|
|
WHERE url = ?
|
|
''', (url,))
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def delete_failed_url(self, url_id):
|
|
"""删除失败URL记录"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM failed_urls WHERE id = ?', (url_id,))
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def clear_failed_urls(self):
|
|
"""清空所有失败URL记录"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM failed_urls WHERE status = "failed"')
|
|
conn.commit()
|
|
return cursor.rowcount
|
|
|
|
# ========== 搜索缓存操作 ==========
|
|
def save_search_cache(self, keyword, engine, results, expire_days=7):
|
|
"""保存搜索结果缓存"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT OR REPLACE INTO search_cache (keyword, engine, results, result_count, created_at, expires_at)
|
|
VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP, datetime('now', '+' || ? || ' days'))
|
|
''', (keyword, engine, json.dumps(results, ensure_ascii=False), len(results), expire_days))
|
|
conn.commit()
|
|
return cursor.lastrowid
|
|
|
|
def get_search_cache(self, keyword, engine='bing_cn'):
|
|
"""获取搜索结果缓存"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
SELECT results, result_count, created_at, expires_at
|
|
FROM search_cache
|
|
WHERE keyword = ? AND engine = ? AND expires_at > datetime('now')
|
|
''', (keyword, engine))
|
|
row = cursor.fetchone()
|
|
if row:
|
|
return {
|
|
'results': json.loads(row['results']),
|
|
'count': row['result_count'],
|
|
'cached_at': row['created_at'],
|
|
'expires_at': row['expires_at']
|
|
}
|
|
return None
|
|
|
|
def clear_expired_cache(self):
|
|
"""清理过期缓存"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM search_cache WHERE expires_at <= datetime("now")')
|
|
conn.commit()
|
|
return cursor.rowcount
|
|
|
|
def clear_all_cache(self):
|
|
"""清空所有缓存"""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('DELETE FROM search_cache')
|
|
conn.commit()
|
|
return cursor.rowcount
|
|
|
|
# 全局数据库实例
|
|
db = Database() |