优化redis在线执行功能

app/dao/__init__.py

@@ -1,21 +1,6 @@
 from sqlalchemy import select
 
 from app.models import Base, engine, async_session, DatabaseHelper
-from app.models.user import User
-from app.models.project import Project
-from app.models.project_role import ProjectRole
-from app.models.test_case import TestCase
-from app.models.testcase_asserts import TestCaseAsserts
-from app.models.environment import Environment
-from app.models.gconfig import GConfig
-from app.models.constructor import Constructor
-from app.models.report import PityReport
-from app.models.result import PityTestResult
-from app.models.database import PityDatabase
-from app.models.testcase_directory import PityTestcaseDirectory
-from app.models.testcase_data import PityTestcaseData
-from app.models.test_plan import PityTestPlan
-from app.models.redis_config import PityRedis
 
 Base.metadata.create_all(engine)
 
@@ -129,6 +114,9 @@ async def delete_record_by_id(cls, user, id):
                     if original is None:
                         raise Exception("记录不存在")
                     DatabaseHelper.delete_model(original, user)
+                    await session.flush()
+                    session.expunge(original)
+                    return original
         except Exception as e:
             cls.log.error(f"删除{cls.model}记录失败, error: {e}")
             raise Exception(f"删除记录失败")

app/middleware/RedisManager.py

@@ -1,7 +1,8 @@
 """
-redis客户端，基于aredis(支持集群，aioredis不支持集群)
+redis客户端Manager
 """
-from aredis import StrictRedisCluster, ClusterConnectionPool, ConnectionPool, StrictRedis
+from redis import ConnectionPool, StrictRedis
+from rediscluster import RedisCluster, ClusterConnectionPool
 
 from app.excpetions.RedisException import RedisException
 
@@ -12,6 +13,19 @@ class PityRedisManager(object):
     _cluster_pool = dict()
     _pool = dict()
 
+    @staticmethod
+    def delete_client(redis_id: int, cluster: bool):
+        """
+        根据redis_id和是否是集群删除客户端
+        :param redis_id:
+        :param cluster:
+        :return:
+        """
+        if cluster:
+            PityRedisManager._cluster_pool.pop(redis_id)
+        else:
+            PityRedisManager._pool.pop(redis_id)
+
     @staticmethod
     def get_cluster_client(redis_id: int, addr: str):
         """
@@ -28,7 +42,7 @@ def get_cluster_client(redis_id: int, addr: str):
         return client
 
     @staticmethod
-    def get_single_node_client(redis_id: int, addr: str, password: str, db: str):
+    def get_single_node_client(redis_id: int, addr: str, password: str, db: int):
         """
         获取redis单实例客户端
         :param redis_id:
@@ -37,14 +51,14 @@ def get_single_node_client(redis_id: int, addr: str, password: str, db: str):
         :param db:
         :return:
         """
-        node = PityRedisManager._cluster_pool.get(redis_id)
+        node = PityRedisManager._pool.get(redis_id)
         if node is not None:
             return node
         host, port = addr.split(":")
         pool = ConnectionPool(host=host, port=port, db=db, max_connections=100, password=password,
                               decode_responses=True)
         client = StrictRedis(connection_pool=pool)
-        PityRedisManager._pool[redis_id] = PityRedisManager.get_cluster(addr)
+        PityRedisManager._pool[redis_id] = client
         return client
 
     @staticmethod
@@ -78,7 +92,7 @@ def get_cluster(addr: str):
             nodes = addr.split(',')
             startup_nodes = [{"host": n.split(":")[0], "port": n.split(":")[1]} for n in nodes]
             pool = ClusterConnectionPool(startup_nodes=startup_nodes, max_connections=100, decode_responses=True)
-            client = StrictRedisCluster(connection_pool=pool, decode_responses=True)
+            client = RedisCluster(connection_pool=pool, decode_responses=True)
             return client
         except Exception as e:
             raise RedisException(f"获取Redis连接失败, {e}")

app/routers/config/redis_config.py

@@ -1,4 +1,5 @@
 from fastapi import Depends
+from starlette.background import BackgroundTasks
 
 from app.dao.config.RedisConfigDao import PityRedisConfigDao
 from app.handler.fatcory import PityResponse
@@ -42,19 +43,27 @@ async def insert_redis_config(form: RedisConfigForm,
 
 @router.post("/redis/update")
 async def update_redis_config(form: RedisConfigForm,
+                              background_tasks: BackgroundTasks,
                               user_info=Depends(Permission(Config.ADMIN))):
     try:
         result = await PityRedisConfigDao.update_record_by_id(user_info['id'], form)
+        if result.cluster:
+            background_tasks.add_task(PityRedisManager.refresh_redis_cluster, *(result.id, result.addr))
+        else:
+            background_tasks.add_task(PityRedisManager.refresh_redis_client,
+                                      *(result.id, result.addr, result.password, result.db))
         return PityResponse.success(data=PityResponse.model_to_dict(result))
     except Exception as err:
         return PityResponse.failed(err)
 
 
 @router.get("/redis/delete")
-async def delete_redis_config(id: int,
+async def delete_redis_config(id: int, background_tasks: BackgroundTasks,
                               user_info=Depends(Permission(Config.ADMIN))):
     try:
-        await PityRedisConfigDao.delete_record_by_id(user_info['id'], id)
+        ans = await PityRedisConfigDao.delete_record_by_id(user_info['id'], id)
+        # 更新缓存
+        background_tasks.add_task(PityRedisManager.delete_client, *(id, ans.cluster))
         return PityResponse.success()
     except Exception as err:
         return PityResponse.failed(err)
@@ -64,10 +73,12 @@ async def delete_redis_config(id: int,
 async def test_redis_command(form: OnlineRedisForm):
     try:
         redis_config = await PityRedisConfigDao.query_record(id=form.id)
-        client = PityRedisManager.get_single_node_client(redis_config.id, redis_config.addr,
-                                                         redis_config.password, redis_config.db)
-        res = await client.execute_command(form.command)
-        # res = await client.set("name", "李相赫")
+        if not redis_config.cluster:
+            client = PityRedisManager.get_single_node_client(redis_config.id, redis_config.addr,
+                                                             redis_config.password, redis_config.db)
+        else:
+            client = PityRedisManager.get_cluster_client(redis_config.id, redis_config.addr)
+        res = client.execute_command(form.command)
         return PityResponse.success(res)
     except Exception as err:
         return PityResponse.failed(err)

requirements.txt

@@ -15,5 +15,5 @@ Jinja2
 aiofiles
 psycopg2-binary
 APScheduler
-aredis
-hiredis
+redis
+redis-py-cluster