modify: 调整项目结构，更新调度器

pyproject.toml

@@ -19,6 +19,10 @@ readme = "README.md"
 requires-python = ">=3.7"
 version = "0.1.0"
 
+[project.optional-dependencies]
+pyside6 = ["PySide6"]
+pyqt6 = ["PyQt6"]
+
 [[tool.uv.index]]
 default = true
 url = "https://pypi.tuna.tsinghua.edu.cn/simple"

src/queue_sqlite/core/__init__.pyi

@@ -5,48 +5,47 @@ from ..constant import MessageStatus
 class core:
     class QueueOperation:
         def __init__(self, db_path: str):
+            """
+            Initialize the queue operation.
+
+            Args:
+                db_path (str): The path to the database file.
+
+            """
             ...
 
         def init_db(self):
+            """
+            Initialize the database.
+            """
             ...
 
         def enqueue(self, message_dict: dict) -> str:
+            """
+            Enqueue a message.
+
+            Args:
+                message_dict (dict): The message to be enqueued.
+
+            Returns:
+                str: The message id.
+
+                Raises:
+                    ValueError: If the message is not valid.
+            """
             ...
 
-        def dequeue(self, size: int = 1) -> List[dict]:
-            ...
-
-        def get_queue_length(self) -> int:
-            ...
-
-        def get_completed_messages(self) -> List[dict]:
-            ...
-
-        def get_result(self, message_id: str):
-            ...
-
-        def update_status(self, message_id: str, status: MessageStatus):
-            ...
-
-        def update_result(self, message_id: str, result: str):
-            ...
-
-        def delete_message(self, message_id: str):
-            ...
-
-        def clean_old_messages(self, days: int):
-            ...
-
-        def clean_expired_messages(self):
-            ...
+        def dequeue(self, size: int = 1) -> List[dict]: ...
+        def get_queue_length(self) -> int: ...
+        def get_completed_messages(self) -> List[dict]: ...
+        def get_result(self, message_id: str): ...
+        def update_status(self, message_id: str, status: MessageStatus): ...
+        def update_result(self, message_id: str, result: str): ...
+        def delete_message(self, message_id: str): ...
+        def clean_old_messages(self, days: int): ...
+        def clean_expired_messages(self): ...
 
     class TaskMounter:
-        def __init__(self, task_class: type[TaskMounter]):
-            ...
-        
-        def get_task_list(self) -> List[str]:
-            ...
-
-        def get_task_function(self, name: str) -> Callable:
-            ...
-            
+        def __init__(self, task_class: type[TaskMounter]): ...
+        def get_task_list(self) -> List[str]: ...
+        def get_task_function(self, name: str) -> Callable: ...

src/queue_sqlite/model/__init__.py

@@ -1,3 +1,4 @@
 from .message_item import MessageItem
 
+
 __all__ = ["MessageItem"]

src/queue_sqlite/queue_operation/listen_operation.py

@@ -9,6 +9,13 @@ class ListenOperation:
         self.listen_fields = ListenMounter.get_Listener_list()
         self.create_table()
 
+    def _get_connection(self):
+        """获取数据库连接"""
+        return sqlite3.connect(
+            self.db_dir,
+            check_same_thread=False,
+        )
+
     def create_table(self):
         if len(self.listen_fields) == 0:
             return
@@ -37,25 +44,22 @@ class ListenOperation:
                 VALUES ('listen_table', NEW.id, 'key', OLD.key, NEW.key);
             END;
         """
-        self.conn = sqlite3.connect(
-            self.db_dir,
-            check_same_thread = False,
-        )
-        self.conn.executescript(sql)
-        self.conn.execute("PRAGMA journal_mode=WAL;")
-        self.conn.execute("PRAGMA synchronous=NORMAL;")
-        self.conn.execute("PRAGMA cache_size=-20000;")
-        self.conn.execute("PRAGMA mmap_size=1073741824;")
-        self.conn.execute("PRAGMA temp_store=MEMORY;")
-        self.conn.execute("PRAGMA busy_timeout=5000;")
-        self.conn.commit()
+        conn = self._get_connection()
+        conn.executescript(sql)
+        conn.execute("PRAGMA journal_mode=WAL;")
+        conn.execute("PRAGMA synchronous=NORMAL;")
+        conn.execute("PRAGMA cache_size=-20000;")
+        conn.execute("PRAGMA mmap_size=1073741824;")
+        conn.execute("PRAGMA temp_store=MEMORY;")
+        conn.execute("PRAGMA busy_timeout=5000;")
+        conn.commit()
 
         # 删除原有数据
         # sql = f"""
         #     DELETE FROM listen_table
         # """
-        # self.conn.execute(sql)
-        # self.conn.commit()
+        # conn.execute(sql)
+        # conn.commit()
 
         for listen_field in self.listen_fields:
             sql = f"""
@@ -64,37 +68,40 @@ class ListenOperation:
                 VALUES 
                     (?, ?)
             """
-            self.conn.execute(sql, (listen_field, "null"))
-            self.conn.commit()
+            conn.execute(sql, (listen_field, "null"))
+            conn.commit()
 
     def listen_data(self) -> Tuple[bool, Union[List[Tuple], str]]:
         sql = f"""
             SELECT * FROM change_log where is_delete = 0 ORDER BY id DESC LIMIT 100
         """
-        result = self.conn.execute(sql).fetchall()
+        conn = self._get_connection()
+        result = conn.execute(sql).fetchall()
         if len(result) == 0:
             return False, "No data found"
         return True, result
 
-        
     def delete_change_log(self, delete_id):
         sql = f"""
             DELETE FROM change_log WHERE id = {delete_id}
         """
-        self.conn.execute(sql)
+        conn = self._get_connection()
+        conn.execute(sql)
 
     def update_listen_data(self, key, value):
         sql = f"""
             UPDATE listen_table SET value = '{value}' WHERE key = '{key}'
         """
-        self.conn.execute(sql)
-        self.conn.commit()
+        conn = self._get_connection()
+        conn.execute(sql)
+        conn.commit()
 
     def get_value(self, key):
         sql = f"""
             SELECT value FROM listen_table WHERE key = '{key}'
         """
-        result = self.conn.execute(sql).fetchone()
+        conn = self._get_connection()
+        result = conn.execute(sql).fetchone()
         if result is None:
             return None
         return result[0]
@@ -103,5 +110,6 @@ class ListenOperation:
         sql = f"""
             SELECT key, value FROM listen_table
         """
-        result = self.conn.execute(sql).fetchall()
+        conn = self._get_connection()
+        result = conn.execute(sql).fetchall()
         return result

src/queue_sqlite/queue_operation/queue_operation.py

@@ -1,5 +1,4 @@
 import os
-import json
 import hashlib
 from typing import Tuple, Union, List
 import threading

src/queue_sqlite/scheduler/__init__.py

@@ -1,49 +1,56 @@
-from .receive_scheduler import ReceiveScheduler
-from .task_scheduler import TaskScheduler
-from .listen_data_scheduler import ListenDataScheduler
+from .standard import StandardQueueScheduler
+from .base import BaseScheduler
 from ..model import MessageItem
 from typing import Callable
-from ..queue_operation import QueueOperation
-from ..queue_operation.listen_operation import ListenOperation
-from .cleanup_scheduler import CleanupScheduler
-import os
 
-class QueueScheduler:
-    def __init__(self, receive_thread_num=1, task_thread_num=1, shard_num=4):
-        self.queue_operation = QueueOperation(shard_num)
-        self.listen_operation = ListenOperation(os.path.join(self.queue_operation.db_dir, "listen.db"))
-        self.listen_operation.create_table()
 
-        self.listen_scheduler = ListenDataScheduler(self.listen_operation)
-        self.receive_scheduler = ReceiveScheduler(self.queue_operation, receive_thread_num)
-        self.task_scheduler = TaskScheduler(self.queue_operation, task_thread_num)
-        self.cleanup_scheduler = CleanupScheduler(self.queue_operation)
+class QueueScheduler(BaseScheduler):
+    def __init__(
+        self,
+        receive_thread_num=1,
+        task_thread_num=1,
+        shard_num=4,
+        scheduler_type="standard",
+    ):
+        self.scheduler = None
+        if scheduler_type == "standard":
+            self.scheduler = StandardQueueScheduler(
+                receive_thread_num, task_thread_num, shard_num
+            )
+        else:
+            pass
+        if self.scheduler:
+            self.queue_operation = self.scheduler.queue_operation
 
     def send_message(self, message: MessageItem, callback: Callable):
-        self.receive_scheduler.send_message(message, callback)
+        if not self.scheduler:
+            raise Exception("Scheduler not initialized")
+        self.scheduler.send_message(message, callback)
+
+    def start(self):
+        if not self.scheduler:
+            raise Exception("Scheduler not initialized")
+        self.scheduler.start()
+
+    def stop(self):
+        if not self.scheduler:
+            raise Exception("Scheduler not initialized")
+        self.scheduler.stop()
 
     def update_listen_data(self, key, value):
-        self.listen_operation.update_listen_data(key, value)
+        if not self.scheduler:
+            raise Exception("Scheduler not initialized")
+        self.scheduler.update_listen_data(key, value)
 
     def get_listen_datas(self):
-        return self.listen_operation.get_values()
+        if not self.scheduler:
+            raise Exception("Scheduler not initialized")
+        return self.scheduler.get_listen_datas()
 
     def get_listen_data(self, key):
-        return self.listen_operation.get_value(key)
+        if not self.scheduler:
+            raise Exception("Scheduler not initialized")
+        return self.scheduler.get_listen_data(key)
 
-    def start_queue_scheduler(self):
-        self.receive_scheduler.start_receive_thread()
-        self.task_scheduler.start_task_thread()
-        self.cleanup_scheduler.start_cleanup()
-        self.listen_scheduler.start_listen_data()
 
-    def stop_queue_scheduler(self):
-        self.receive_scheduler.stop_receive_thread()
-        self.listen_scheduler.stop_listen_data()
-        self.task_scheduler.stop_task_thread()
-        self.cleanup_scheduler.stop_cleanup()
-        # self.queue_operation.close_all_connections()
-
-__all__ = [
-    "QueueScheduler"
-]
+__all__ = ["QueueScheduler"]

src/queue_sqlite/scheduler/base/__init__.py

@@ -0,0 +1,4 @@
+from .base_scheduler import BaseScheduler
+
+
+__all__ = ["BaseScheduler"]

src/queue_sqlite/scheduler/base/base_scheduler.py

@@ -0,0 +1,42 @@
+from abc import ABC, abstractmethod
+from typing import Callable
+from ...model import MessageItem
+
+
+class BaseScheduler(ABC):
+    """调度器抽象类"""
+
+    @abstractmethod
+    def send_message(self, message: MessageItem, callback: Callable):
+        """发送消息到队列
+
+        Args:
+            message (MessageItem): 消息对象
+            callback (Callable): 发送完成后的回调函数
+        """
+        pass
+
+    @abstractmethod
+    def start(self):
+        """启动调度器"""
+        pass
+
+    @abstractmethod
+    def stop(self):
+        """停止调度器"""
+        pass
+
+    @abstractmethod
+    def update_listen_data(self, key, value):
+        """更新监听数据"""
+        pass
+
+    @abstractmethod
+    def get_listen_datas(self) -> list:
+        """获取监听数据"""
+        pass
+
+    @abstractmethod
+    def get_listen_data(self, key):
+        """获取单个监听数据"""
+        pass

src/queue_sqlite/scheduler/standard/__init__.py

@@ -0,0 +1,54 @@
+from .receive_scheduler import ReceiveScheduler
+from .task_scheduler import TaskScheduler
+from .listen_data_scheduler import ListenDataScheduler
+from ...model import MessageItem
+from typing import Callable
+from ...queue_operation import QueueOperation
+from ...queue_operation.listen_operation import ListenOperation
+from .cleanup_scheduler import CleanupScheduler
+import os
+from ..base import BaseScheduler
+
+
+class StandardQueueScheduler(BaseScheduler):
+    def __init__(self, receive_thread_num=1, task_thread_num=1, shard_num=4):
+        self.queue_operation = QueueOperation(shard_num)
+        self.listen_operation = ListenOperation(
+            os.path.join(self.queue_operation.db_dir, "listen.db")
+        )
+        self.listen_operation.create_table()
+
+        self.listen_scheduler = ListenDataScheduler(self.listen_operation)
+        self.receive_scheduler = ReceiveScheduler(
+            self.queue_operation, receive_thread_num
+        )
+        self.task_scheduler = TaskScheduler(self.queue_operation, task_thread_num)
+        self.cleanup_scheduler = CleanupScheduler(self.queue_operation)
+
+    def send_message(self, message: MessageItem, callback: Callable):
+        self.receive_scheduler.send_message(message, callback)
+
+    def update_listen_data(self, key, value):
+        self.listen_operation.update_listen_data(key, value)
+
+    def get_listen_datas(self):
+        return self.listen_operation.get_values()
+
+    def get_listen_data(self, key):
+        return self.listen_operation.get_value(key)
+
+    def start(self):
+        self.receive_scheduler.start_receive_thread()
+        self.task_scheduler.start_task_thread()
+        self.cleanup_scheduler.start_cleanup()
+        self.listen_scheduler.start_listen_data()
+
+    def stop(self):
+        self.receive_scheduler.stop_receive_thread()
+        self.listen_scheduler.stop_listen_data()
+        self.task_scheduler.stop_task_thread()
+        self.cleanup_scheduler.stop_cleanup()
+        # self.queue_operation.close_all_connections()
+
+
+__all__ = ["StandardQueueScheduler"]

src/queue_sqlite/scheduler/standard/cleanup_scheduler.py

@@ -1,7 +1,7 @@
 # 新建清理调度器
 import threading
 import time
-from ..queue_operation import QueueOperation
+from ...queue_operation import QueueOperation
 
 
 class CleanupScheduler:
@@ -37,7 +37,9 @@ class CleanupScheduler:
             return
 
         self.is_running = True
-        self.cleanup_thread = threading.Thread(target=self.cleanup_expired_messages, daemon=True)
+        self.cleanup_thread = threading.Thread(
+            target=self.cleanup_expired_messages, daemon=True
+        )
         self.cleanup_thread.start()
 
     def stop_cleanup(self):

src/queue_sqlite/scheduler/standard/listen_data_scheduler.py

@@ -1,6 +1,6 @@
-from ..queue_operation.listen_operation import ListenOperation
+from ...queue_operation.listen_operation import ListenOperation
 from concurrent.futures import ThreadPoolExecutor
-from ..mounter.listen_mounter import ListenMounter
+from ...mounter.listen_mounter import ListenMounter
 import threading
 
 
@@ -29,7 +29,9 @@ class ListenDataScheduler:
                     key = data[6]
                     new_value = data[7]
                     delete_id = data[0]
-                    self.executor.submit(self._process_listen_data, key, new_value, delete_id)
+                    self.executor.submit(
+                        self._process_listen_data, key, new_value, delete_id
+                    )
 
     def start_listen_data(self):
         if self.is_running:
@@ -46,4 +48,3 @@ class ListenDataScheduler:
             self.listen_thread.join(timeout=2)
 
         self.executor.shutdown(wait=True)
-    

src/queue_sqlite/scheduler/standard/receive_scheduler.py

@@ -1,5 +1,5 @@
-from ..queue_operation import QueueOperation
-from ..model import MessageItem
+from ...queue_operation import QueueOperation
+from ...model import MessageItem
 from typing import Callable
 from concurrent.futures import ThreadPoolExecutor
 import threading
@@ -11,12 +11,14 @@ class ReceiveScheduler:
         self.callbacks = dict()
         self.receive_thread_num = receive_thread_num
         self.is_running = False
-        self.executor = ThreadPoolExecutor(max_workers=receive_thread_num)  # 并行执行回调
+        self.executor = ThreadPoolExecutor(
+            max_workers=receive_thread_num
+        )  # 并行执行回调
         self.lock = threading.Lock()
         self.queue_operation = queue_operation
         self.receive_thread = None  # 单一轮询线程
 
-    def send_message(self, message: MessageItem, callback: Callable = None): # type: ignore
+    def send_message(self, message: MessageItem, callback: Callable = None):  # type: ignore
         if callback is None:
             callback = lambda message: print(message)
         # self.queue_operation.enqueue(message.to_dict())
@@ -77,4 +79,3 @@ class ReceiveScheduler:
 
         # 关闭线程池
         self.executor.shutdown(wait=True)
-        

src/queue_sqlite/scheduler/standard/task_scheduler.py

@@ -1,8 +1,8 @@
-from ..queue_operation import QueueOperation
-from ..model import MessageItem
-from ..constant import MessageStatus
-from ..cycle.task_cycle import TaskCycle
-from ..mounter.task_mounter import TaskMounter
+from ...queue_operation import QueueOperation
+from ...model import MessageItem
+from ...constant import MessageStatus
+from ...cycle.task_cycle import TaskCycle
+from ...mounter.task_mounter import TaskMounter
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 import time

tests/test_core.py

@@ -10,8 +10,13 @@ class TestCore:
     def test_task_mounter(cls):
         task_mounter = core.TaskMounter(TaskMounter)
         task_mounter.get_task_list()
-        print(task_mounter.get_task_function("task")({"num": 1})(lambda x: x+1).meta["num"])
+        print(
+            task_mounter.get_task_function("task")({"num": 1})(lambda x: x + 1).meta[
+                "num"
+            ]
+        )
         from tasks import example
+
         task_mounter.get_task_list()
         print(task_mounter.get_task_function("<lambda>"))
 
@@ -21,7 +26,7 @@ class TestCore:
 
     @classmethod
     def test_queue_operation_enqueue(cls):
-        message_item = MessageItem(content={"num": 1}, destination="test")
+        message_item = MessageItem(content={"num": 1}, destination="test")  # type: ignore
         cls.queue_operation.enqueue(message_item.to_dict_by_core())
 
     @classmethod

tests/test_listen.py

@@ -6,7 +6,7 @@ import time
 class TestListen:
     def test_listen_data(self):
         scheduler = QueueScheduler()
-        scheduler.start_queue_scheduler()
+        scheduler.start()
         scheduler.update_listen_data("key_1", "value_1")
         time.sleep(0.001)
         scheduler.update_listen_data("key_1", "value_2")
@@ -15,6 +15,6 @@ class TestListen:
         time.sleep(0.001)
         scheduler.update_listen_data("key_1", "value_4")
         time.sleep(0.001)
-        scheduler.stop_queue_scheduler()
+        scheduler.stop()
         print(scheduler.get_listen_datas())
         print(scheduler.get_listen_data("key_1"))

tests/test_stress.py

@@ -53,7 +53,9 @@ class TestStress:
         cls.is_monitoring = True
         monitor_thread = threading.Thread(target=cls.monitor_resources, daemon=True)
         monitor_thread.start()
-        queue_scheduler = QueueScheduler(receive_thread_num=4, task_thread_num=8, shard_num=12)
+        queue_scheduler = QueueScheduler(
+            receive_thread_num=4, task_thread_num=8, shard_num=12
+        )
 
         # threads = []
         # for message_item in messages:
@@ -63,7 +65,7 @@ class TestStress:
         # 添加性能计时
         start_time = time.perf_counter()
 
-        queue_scheduler.start_queue_scheduler()
+        queue_scheduler.start()
 
         for message_item in messages:  # 增加任务数量
             message_item: MessageItem
@@ -75,12 +77,15 @@ class TestStress:
 
         # # 动态等待任务完成
         start_wait = time.time()
-        while (queue_scheduler.queue_operation.get_queue_length() > 0 or cls.callback_counter < TASK_NUM):  # 30秒超时
+        while (
+            queue_scheduler.queue_operation.get_queue_length() > 0
+            or cls.callback_counter < TASK_NUM
+        ):  # 30秒超时
             time.sleep(0.5)
             print(f"当前队列长度：{queue_scheduler.queue_operation.get_queue_length()}")
             print(f"当前回调数量：{cls.callback_counter}")
 
-        queue_scheduler.stop_queue_scheduler()
+        queue_scheduler.stop()
 
         print(f"当前队列长度：{queue_scheduler.queue_operation.get_queue_length()}")
         print(f"测试函数：斐波那契数列")
@@ -89,4 +94,3 @@ class TestStress:
 
         cls.is_monitoring = False
         monitor_thread.join(timeout=1)
-