|
|
from typing import List, Tuple, Optional
|
|
|
import sqlite3
|
|
|
|
|
|
|
|
|
class Db_Operation:
|
|
|
def __init__(self, db_name: str = "rainbow_chain.db", table_name: str = "rainbow_table"):
|
|
|
"""
|
|
|
初始化数据库连接和表结构，并进行必要的优化配置。
|
|
|
:param db_name: 数据库文件名
|
|
|
:param table_name: 表名
|
|
|
"""
|
|
|
self.db_name = db_name
|
|
|
self.table_name = table_name
|
|
|
self.conn = sqlite3.connect(self.db_name, timeout=30) # 设置较高的超时时间以应对高并发
|
|
|
self.cursor = self.conn.cursor()
|
|
|
if not self._initialize_table():
|
|
|
self.close()
|
|
|
return
|
|
|
if not self._configure_database():
|
|
|
self.close()
|
|
|
return
|
|
|
if not self._create_indexes():
|
|
|
self.close()
|
|
|
return
|
|
|
|
|
|
def _initialize_table(self) -> bool:
|
|
|
"""
|
|
|
创建表结构，如果表不存在。
|
|
|
|
|
|
:return: 如果成功返回 True，否则返回 False
|
|
|
"""
|
|
|
try:
|
|
|
create_table_query = f"""
|
|
|
CREATE TABLE IF NOT EXISTS {self.table_name} (
|
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
|
start_str TEXT NOT NULL,
|
|
|
end_str TEXT NOT NULL
|
|
|
)
|
|
|
"""
|
|
|
self.cursor.execute(create_table_query)
|
|
|
self.conn.commit()
|
|
|
return True
|
|
|
except sqlite3.Error:
|
|
|
return False
|
|
|
|
|
|
def _create_indexes(self) -> bool:
|
|
|
"""
|
|
|
为 end_str 列创建索引，以加速查询。
|
|
|
|
|
|
:return: 如果成功返回 True，否则返回 False
|
|
|
"""
|
|
|
try:
|
|
|
create_index_query = f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_end_str ON {self.table_name} " \
|
|
|
f"(end_str) "
|
|
|
self.cursor.execute(create_index_query)
|
|
|
self.conn.commit()
|
|
|
return True
|
|
|
except sqlite3.Error:
|
|
|
return False
|
|
|
|
|
|
def _configure_database(self) -> bool:
|
|
|
"""
|
|
|
配置SQLite参数以优化性能。
|
|
|
|
|
|
:return: 如果成功返回 True，否则返回 False
|
|
|
"""
|
|
|
try:
|
|
|
# 启用WAL模式，提高并发性能
|
|
|
self.cursor.execute("PRAGMA journal_mode=WAL;")
|
|
|
# 增加缓存大小，以减少磁盘I/O
|
|
|
self.cursor.execute("PRAGMA cache_size = 10000;") # 根据需要调整
|
|
|
# 关闭同步以提高写入性能（注意：存在数据丢失风险）
|
|
|
self.cursor.execute("PRAGMA synchronous = NORMAL;")
|
|
|
# 启用内存中的临时存储，以提升性能
|
|
|
self.cursor.execute("PRAGMA temp_store = MEMORY;")
|
|
|
# 提高SQLITE的锁定超时，从默认的5秒延长到30秒
|
|
|
self.cursor.execute("PRAGMA busy_timeout = 30000;")
|
|
|
self.conn.commit()
|
|
|
return True
|
|
|
except sqlite3.Error:
|
|
|
return False
|
|
|
|
|
|
def write_to_database(self, data: List[Tuple[str, str]]) -> bool:
|
|
|
"""
|
|
|
批量写入数据到数据库。使用事务来提升性能。
|
|
|
|
|
|
:param data: 包含 (start_str, end_str) 元组的列表
|
|
|
:return: 如果成功返回 True，否则返回 False
|
|
|
"""
|
|
|
if not data:
|
|
|
return True # 如果没有数据需要写入，直接返回 True
|
|
|
|
|
|
try:
|
|
|
self.cursor.execute("BEGIN TRANSACTION;")
|
|
|
insert_query = f"INSERT INTO {self.table_name} (start_str, end_str) VALUES (?, ?)"
|
|
|
self.cursor.executemany(insert_query, data)
|
|
|
self.conn.commit()
|
|
|
return True
|
|
|
except sqlite3.Error:
|
|
|
self.conn.rollback()
|
|
|
return False
|
|
|
|
|
|
def check_string_in_end_str(self, target_string: str) -> Optional[Tuple[int, str, str]]:
|
|
|
"""
|
|
|
查询 end_str 列中是否存在目标字符串。
|
|
|
|
|
|
:param target_string: 要查询的字符串
|
|
|
:return: 如果有匹配记录，返回一条数据的元组 (id, start_str, end_str)；否则返回 None
|
|
|
"""
|
|
|
try:
|
|
|
query = f"SELECT id, start_str, end_str FROM {self.table_name} WHERE end_str = ? LIMIT 1"
|
|
|
self.cursor.execute(query, (target_string,))
|
|
|
result = self.cursor.fetchone()
|
|
|
return result
|
|
|
except sqlite3.Error:
|
|
|
return None
|
|
|
|
|
|
def close(self):
|
|
|
"""
|
|
|
关闭数据库连接。
|
|
|
"""
|
|
|
try:
|
|
|
self.cursor.close()
|
|
|
self.conn.close()
|
|
|
except sqlite3.Error:
|
|
|
pass
|
|
|
|
|
|
def __enter__(self):
|
|
|
"""
|
|
|
支持 with 语句，通过返回自身来管理资源。
|
|
|
"""
|
|
|
return self
|
|
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
|
"""
|
|
|
支持 with 语句，确保关闭连接。
|
|
|
"""
|
|
|
self.close()
|
