1. 新增数据库自动连接，临时修复数据库断开问题

5 files changed, 13 insertions, 2 deletions

diff --git a/server/apps/agentcomm.py b/server/apps/agentcomm.py
index 5d2a53c..34705e2 100644
--- a/server/apps/agentcomm.py
+++ b/server/apps/agentcomm.py
@@ -116,6 +116,7 @@ def deliver_task(agent_id, policy, policy_param):
         FROM %s
         WHERE AGENT_ID = '%s'
         """ % (MYSQL_TAB_AGENT, agent_id)
+        da.conn.ping(reconnect=True)
         da.cursor.execute(sql)
         data = da.cursor.fetchall()
 
@@ -157,6 +158,7 @@ def task_ret(json_data):
     """ % (MYSQL_TAB_TASK_LOG, json_data["id"], json_data["level"], json_data["info"],
            json_data["taskpolicy"])
     try:
+        da.conn.ping(reconnect=True)
         da.cursor.execute(sql)
         da.conn.commit()
         return "ok"
@@ -215,6 +217,7 @@ def del_agent(json_data):
     sql = """DELETE FROM %s WHERE AGENT_ID = '%s'
     """ % (MYSQL_TAB_AGENT, json_data["id"])
     try:
+        da.conn.ping(reconnect=True)
         da.cursor.execute(sql)
         da.conn.commit()
         if da.cursor.rowcount == 0:
@@ -252,6 +255,7 @@ def insert_agent(param: dict):
             '%(mem)s',
             %(idle)s);""" % param
     try:
+        da.conn.ping(reconnect=True)
         da.cursor.execute(sql)
         da.conn.commit()
     except Exception as e:
diff --git a/server/apps/policy.py b/server/apps/policy.py
index 93f0ca7..eb941b4 100644
--- a/server/apps/policy.py
+++ b/server/apps/policy.py
@@ -49,6 +49,7 @@ def init_task_policy(ptype, target, task):
     FROM %s
     WHERE P_TYPE IN (%s) AND P_PROTO IN (%s)
     """ % (MYSQL_TAB_POLICY, "\'" + "\',\'".join(policy) + "\'", "\'" + "\',\'".join(proto) + "\'")
+    da.conn.ping(reconnect=True)
     da.cursor.execute(sql)
     policy_list = da.cursor.fetchall()
     # 随机选择一个作为初始策略
diff --git a/server/apps/target.py b/server/apps/target.py
index 7a83264..5e51562 100644
--- a/server/apps/target.py
+++ b/server/apps/target.py
@@ -343,6 +343,7 @@ def filter_info():
     # 查询所有的isp
     isp_sql = """SELECT DISTINCT ISP from %s """ % MYSQL_TAB_TARGETDATA
     # 执行查询
+    da.conn.ping(reconnect=True)
     da.cursor.execute(isp_sql)
     isp_data = da.cursor.fetchall()
     isp = [i.popitem()[1] for i in isp_data]
@@ -350,6 +351,7 @@ def filter_info():
     # 查询所有的国家
     cou_sql = """SELECT DISTINCT COU from %s """ % MYSQL_TAB_TARGETDATA
     # 执行查询
+    da.conn.ping(reconnect=True)
     da.cursor.execute(cou_sql)
     cou_data = da.cursor.fetchall()
     cou = [i.popitem()[1] for i in cou_data]
@@ -411,6 +413,7 @@ def map_info(query_data):
                 MYSQL_TAB_TARGETDATA, " AND ".join(["=".join(condition.popitem()) for _ in range(l)]))
 
         # 执行查询
+        da.conn.ping(reconnect=True)
         da.cursor.execute(sql)
         data = da.cursor.fetchall()
         for d in data:
diff --git a/server/apps/task.py b/server/apps/task.py
index f31e700..0f54754 100644
--- a/server/apps/task.py
+++ b/server/apps/task.py
@@ -208,7 +208,7 @@ def make_task(json_data):
     '%(nodes)s'
     );
     """ % task_info
-    # debug(sql)
+    da.conn.ping(reconnect=True)
     da.cursor.execute(sql)
     da.conn.commit()
 
diff --git a/server/apps/util.py b/server/apps/util.py
index c5892d7..7e50f29 100644
--- a/server/apps/util.py
+++ b/server/apps/util.py
@@ -175,10 +175,12 @@ class DataHandler:
         if len(differ) == 0:
             if not count:
                 sql = """SELECT * FROM %s LIMIT %s, %s;""" % (tabname, offset, offset + limit)
+                da.conn.ping(reconnect=True)
                 self.cursor.execute(sql)
                 return self.cursor.fetchall()
             else:
                 sql = """SELECT count(*) FROM %s""" % (tabname)
+                da.conn.ping(reconnect=True)
                 self.cursor.execute(sql)
                 return dict(self.cursor.fetchall()[0]).popitem()[1]
         else:
@@ -196,12 +198,13 @@ class DataHandler:
             if not count:
                 sql = """SELECT * FROM %s WHERE %s LIMIT %s, %s""" % (
                     tabname, " AND ".join(["=".join(condition.popitem()) for _ in range(l)]), offset, offset + limit)
-                # print(sql)
+                da.conn.ping(reconnect=True)
                 self.cursor.execute(sql)
                 return self.cursor.fetchall()
             else:
                 sql = """SELECT count(*) FROM %s WHERE %s LIMIT %s, %s""" % (
                     tabname, " AND ".join(["=".join(condition.popitem()) for _ in range(l)]), offset, limit)
+                da.conn.ping(reconnect=True)
                 self.cursor.execute(sql)
                 return dict(self.cursor.fetchall()[0]).popitem()[1]