#!/usr/bin/env python3 """ import_to_eia.py — 统一入库脚本 ============================= 所有爬虫的输出最终都通过此脚本写入 eia.db。 支持三种输入： 1. JSONL 文件（crawl_*.py 输出格式） 2. pandas DataFrame（批量） 3. 直接从 list_crawl.db 迁移 用法: python3 import_to_eia.py --from-listdb # 从 list_crawl.db 迁移 python3 import_to_eia.py --file data.jsonl # 导入 JSONL 文件 python3 import_to_eia.py --stats # 查看 eia.db 统计 """ import json, os, re, sqlite3, sys, hashlib, time from datetime import datetime BASE_DIR = os.path.dirname(os.path.abspath(__file__)) EIA_DB = os.path.join(BASE_DIR, "eia.db") LIST_DB = os.path.join(BASE_DIR, "list_crawl.db") def init_db(): """创建 eia.db 表结构""" conn = sqlite3.connect(EIA_DB) conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") conn.execute("PRAGMA cache_size=-64000") # 主表 conn.execute(""" CREATE TABLE IF NOT EXISTS projects ( id INTEGER PRIMARY KEY AUTOINCREMENT, site_name TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '', url TEXT UNIQUE, content TEXT DEFAULT '', pub_date TEXT DEFAULT '', summary TEXT DEFAULT '', source_url TEXT DEFAULT '', tags TEXT DEFAULT '', status TEXT DEFAULT 'active', created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) # FTS5 全文索引 conn.execute(""" CREATE VIRTUAL TABLE IF NOT EXISTS projects_fts USING fts5( title, content, summary, site_name, content='projects', content_rowid='id', tokenize='unicode61' ) """) # 站点元信息 conn.execute(""" CREATE TABLE IF NOT EXISTS sites ( site_key TEXT PRIMARY KEY, name TEXT, url TEXT, method TEXT DEFAULT '', total_items INTEGER DEFAULT 0, last_sync TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) # 自动同步触发器 triggers = [ ("projects_ai", "AFTER INSERT", """ INSERT INTO projects_fts(rowid, title, content, summary, site_name) VALUES (new.id, new.title, new.content, new.summary, new.site_name); """), ("projects_ad", "AFTER DELETE", """ INSERT INTO projects_fts(projects_fts, rowid, title, content, summary, site_name) VALUES ('delete', old.id, old.title, old.content, old.summary, old.site_name); """), ("projects_au", "AFTER UPDATE", """ INSERT INTO projects_fts(projects_fts, rowid, title, content, summary, site_name) VALUES ('delete', old.id, old.title, old.content, old.summary, old.site_name); INSERT INTO projects_fts(rowid, title, content, summary, site_name) VALUES (new.id, new.title, new.content, new.summary, new.site_name); """), ] for name, timing, body in triggers: conn.execute(f"DROP TRIGGER IF EXISTS {name}") conn.execute(f"CREATE TRIGGER {name} {timing} ON projects BEGIN {body} END;") conn.commit() conn.close() return EIA_DB def extract_summary(content, max_len=300): """从正文HTML提取纯文本摘要""" if not content: return "" text = re.sub(r'<[^>]+>', ' ', content) text = re.sub(r'\s+', ' ', text).strip() return text[:max_len] def import_item(conn, item): """导入单条记录""" try: conn.execute(""" INSERT OR IGNORE INTO projects (site_name, title, url, content, pub_date, summary, source_url, tags) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( item.get("site_name", "")[:200], item.get("title", "")[:500], item.get("url", ""), item.get("content", ""), item.get("pub_date", "")[:10], item.get("summary", "")[:500], item.get("source_url", ""), item.get("tags", "")[:100], )) return conn.total_changes > 0 except: return False def import_jsonl(filepath): """从 JSONL 文件导入""" conn = sqlite3.connect(EIA_DB) total, new, skip = 0, 0, 0 site_stats = {} with open(filepath, encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue try: item = json.loads(line) except: skip += 1 continue # 自动生成摘要 if not item.get("summary") and item.get("content"): item["summary"] = extract_summary(item["content"]) if import_item(conn, item): new += 1 site_name = item.get("site_name", "未知") site_stats[site_name] = site_stats.get(site_name, 0) + 1 else: skip += 1 total += 1 conn.commit() # 更新站点统计 for name, cnt in site_stats.items(): site_key = hashlib.md5(name.encode()).hexdigest()[:12] conn.execute( "INSERT OR REPLACE INTO sites (site_key, name, total_items, last_sync) VALUES (?, ?, COALESCE((SELECT total_items FROM sites WHERE site_key=?),0)+?, CURRENT_TIMESTAMP)", (site_key, name[:200], site_key, cnt) ) conn.commit() conn.close() return total, new, skip def migrate_from_listdb(): """从 list_crawl.db 迁移数据到 eia.db""" if not os.path.exists(LIST_DB): print(" ❌ list_crawl.db 不存在") return conn_src = sqlite3.connect(LIST_DB) conn_dst = sqlite3.connect(EIA_DB) # 读取所有 completed 和有 list_only 但已有正文的 rows = conn_src.execute(""" SELECT site_name, title, url, content, pub_date FROM items WHERE status IN ('completed', 'date_fixed') ORDER BY id """).fetchall() print(f" 📋 待迁移: {len(rows)} 条") new = 0 skip = 0 for site_name, title, url, content, pub_date in rows: summary = extract_summary(content) if content else "" try: conn_dst.execute(""" INSERT OR IGNORE INTO projects (site_name, title, url, content, pub_date, summary) VALUES (?, ?, ?, ?, ?, ?) """, (site_name, title, url, content or "", pub_date or "", summary)) if conn_dst.total_changes > 0: new += 1 else: skip += 1 except: skip += 1 conn_dst.commit() conn_src.close() conn_dst.close() print(f" ✅ 迁移完成: 新增{new}, 跳过重复{skip}") return new def print_stats(): """打印 eia.db 统计""" if not os.path.exists(EIA_DB): print(" eia.db 尚未创建") return conn = sqlite3.connect(EIA_DB) conn.execute("PRAGMA query_only=1") total = conn.execute("SELECT COUNT(*) FROM projects").fetchone()[0] with_content = conn.execute("SELECT COUNT(*) FROM projects WHERE content != ''").fetchone()[0] with_date = conn.execute("SELECT COUNT(*) FROM projects WHERE pub_date != ''").fetchone()[0] sites_count = conn.execute("SELECT COUNT(*) FROM sites").fetchone()[0] print(f"\n{'='*50}") print(f" 📊 eia.db 统计") print(f"{'='*50}") print(f" 项目总数: {total}") print(f" 已有正文: {with_content}") print(f" 已有日期: {with_date}") print(f" 来源站点数: {sites_count}") if sites_count > 0: print(f"\n 各站数据量:") rows = conn.execute("SELECT name, total_items FROM sites ORDER BY total_items DESC LIMIT 15").fetchall() for r in rows: print(f" {r[0][:50]:50s} {r[1]:4d}条") # 最近导入 rows = conn.execute("SELECT site_name, created_at FROM projects ORDER BY id DESC LIMIT 5").fetchall() if rows: print(f"\n 最近5条:") for r in rows: print(f" {r[0][:40]:40s} {r[1]}") conn.close() def main(): import argparse parser = argparse.ArgumentParser(description="统一入库脚本 → eia.db") parser.add_argument("--from-listdb", action="store_true", help="从 list_crawl.db 迁移") parser.add_argument("--file", help="导入 JSONL 文件") parser.add_argument("--stats", action="store_true", help="查看统计") parser.add_argument("--init", action="store_true", help="初始化数据库（建表）") args = parser.parse_args() if args.init: init_db() print(f" ✅ eia.db 已初始化: {EIA_DB}") print_stats() return if args.stats: print_stats() return if args.from_listdb: init_db() migrate_from_listdb() print_stats() return if args.file: init_db() total, new, skip = import_jsonl(args.file) print(f" ✅ 导入完成: {total}条处理, {new}新增, {skip}跳过") print_stats() return parser.print_help() if __name__ == "__main__": main()