修复mongodb ssh tunnel 超时问题

sql/engines/mongo.py

@@ -831,20 +831,24 @@ def get_all_databases(self):
             result.rows = conn.list_database_names()
         except OperationFailure:
             result.rows = [self.db_name]
+        conn.close()
         return result
 
     def get_all_tables(self, db_name, **kwargs):
         result = ResultSet()
         conn = self.get_connection()
         db = conn[db_name]
         result.rows = db.list_collection_names()
+        conn.close()
         return result
 
     def get_all_columns_by_tb(self, db_name, tb_name, **kwargs):
         """获取所有字段, 返回一个ResultSet"""
         # https://github.com/getredash/redash/blob/master/redash/query_runner/mongodb.py
         result = ResultSet()
-        db = self.get_connection()[db_name]
+        conn = self.get_connection()
+        db = conn[db_name]
+        conn.close()
         collection_name = tb_name
         documents_sample = []
         if "viewOn" in db[collection_name].options():
@@ -1072,6 +1076,7 @@ def query(self, db_name=None, sql="", limit_num=0, close_conn=True, **kwargs):
             conn = self.get_connection()
             db = conn[db_name]
             collection = db[collection_name]
+            conn.close()
 
             # 执行语句
             logger.debug(find_cmd)
@@ -1247,6 +1252,7 @@ def current_op(self, command_type):
         except Exception as e:
             logger.warning(f"mongodb获取连接信息错误，错误信息{traceback.format_exc()}")
             result_set.error = str(e)
+        conn.close()
 
         return result_set
 
@@ -1267,7 +1273,7 @@ def get_kill_command(self, opids):
                 kill_command = kill_command + "db.killOp({});".format(opid)
             else:
                 kill_command = kill_command + 'db.killOp("{}");'.format(opid)
-
+        conn.close()
         return kill_command
 
     def kill_op(self, opids):
@@ -1288,6 +1294,7 @@ def kill_op(self, opids):
                     f"{self.name}语句执行killOp报错，语句：db.runCommand({sql}) ，错误信息{traceback.format_exc()}"
                 )
                 result.error = str(e)
+        conn.close()
         return result
 
     def get_all_databases_summary(self):
@@ -1314,6 +1321,7 @@ def get_all_databases_summary(self):
                 }
                 rows.append(row)
             query_result.rows = rows
+            conn.close()
         return query_result
 
     def get_instance_users_summary(self):
@@ -1339,6 +1347,7 @@ def get_instance_users_summary(self):
                         }
                     )
             query_result.rows = rows
+            conn.close()
         return query_result
 
     def create_instance_user(self, **kwargs):
@@ -1360,6 +1369,7 @@ def create_instance_user(self, **kwargs):
                     "remark": remark,
                 }
             ]
+            conn.close()
         except Exception as e:
             exec_result.error = str(e)
         return exec_result
@@ -1373,6 +1383,7 @@ def drop_instance_user(self, db_name_user: str, **kwarg):
         try:
             conn = self.get_connection()
             conn[db_name].command("dropUser", user)
+            conn.close()
         except Exception as e:
             exec_result.error = str(e)
         return exec_result
@@ -1386,6 +1397,7 @@ def reset_instance_user_pwd(self, db_name_user: str, reset_pwd: str, **kwargs):
         try:
             conn = self.get_connection()
             conn[db_name].command("updateUser", user, pwd=reset_pwd)
+            conn.close()
         except Exception as e:
             exec_result.error = str(e)
         return exec_result

sql/instance.py

@@ -348,6 +348,10 @@ def instance_resource(request):
         result["status"] = 1
         result["msg"] = str(msg)
     else:
+        if instance.tunnel:
+            query_engine.close()
+            query_engine.ssh.server.close()
+            del query_engine.ssh
         if resource.error:
             result["status"] = 1
             result["msg"] = resource.error

sql/query.py

@@ -60,11 +60,19 @@ def query(request):
             # 引擎内部判断为 bad_query
             result["status"] = 1
             result["msg"] = query_check_info.get("msg")
+            if instance.tunnel:
+                query_engine.close()
+                query_engine.ssh.server.close()
+                del query_engine.ssh
             return HttpResponse(json.dumps(result), content_type="application/json")
         if query_check_info.get("has_star") and config.get("disable_star") is True:
             # 引擎内部判断为有 * 且禁止 * 选项打开
             result["status"] = 1
             result["msg"] = query_check_info.get("msg")
+            if instance.tunnel:
+                query_engine.close()
+                query_engine.ssh.server.close()
+                del query_engine.ssh
             return HttpResponse(json.dumps(result), content_type="application/json")
         sql_content = query_check_info["filtered_sql"]
 
@@ -78,6 +86,10 @@ def query(request):
         else:
             result["status"] = priv_check_info["status"]
             result["msg"] = priv_check_info["msg"]
+            if instance.tunnel:
+                query_engine.close()
+                query_engine.ssh.server.close()
+                del query_engine.ssh
             return HttpResponse(json.dumps(result), content_type="application/json")
         # explain的limit_num设置为0
         limit_num = 0 if re.match(r"^explain", sql_content.lower()) else limit_num
@@ -188,9 +200,17 @@ def query(request):
         )
         result["status"] = 1
         result["msg"] = f"查询异常报错，错误信息：{e}"
+        if instance.tunnel:
+            query_engine.close()
+            query_engine.ssh.server.close()
+            del query_engine.ssh
         return HttpResponse(json.dumps(result), content_type="application/json")
     # 返回查询结果
     try:
+        if instance.tunnel:
+            query_engine.close()
+            query_engine.ssh.server.close()
+            del query_engine.ssh
         return HttpResponse(
             json.dumps(
                 result,
@@ -202,6 +222,10 @@ def query(request):
         )
     # 虽然能正常返回，但是依然会乱码
     except UnicodeDecodeError:
+        if instance.tunnel:
+            query_engine.close()
+            query_engine.ssh.server.close()
+            del query_engine.ssh
         return HttpResponse(
             json.dumps(result, default=str, bigint_as_string=True, encoding="latin1"),
             content_type="application/json",