ADD file via upload

Db.py

@@ -0,0 +1,138 @@
+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()