sqlite3优化

server/jumpserver.py

@@ -13,7 +13,7 @@ from urllib.parse import urlencode
 import hashlib
 import requests
 
-from flask import request, Flask, redirect, session, render_template
+from flask import request, Flask, redirect, session, render_template, g
 from flask_cors import CORS
 
 app = Flask(__name__,
@@ -27,6 +27,9 @@ gHost = ""
 gPort = 0
 gDebug = 0
 
+# 数据库文件
+gSqlite3File = "/usr/local/jumpserver/jumpserver.db"
+
 # 默认ssh管理账号
 gDefaultSSHAdmin = "ec2-user"
 
@@ -193,6 +196,17 @@ def checkCookie(request):
     # return True
 
 
+@app.before_request
+def before_request():
+    g.db = connect_db()
+ 
+
+@app.after_request
+def after_request(response):
+    g.db.close()
+    return response
+
+
 # sso回调接口
 @app.route('/sso', methods=['GET', 'POST'])
 def sso():
@@ -316,20 +330,6 @@ def host_deluser():
     return do_host_deluser(request)
 
 
-@app.route('/jump/host/sudouser', methods=['GET', 'POST'])
-def host_sudouser():
-    if not checkCookie(request):
-        return redirect_sso()
-    return do_host_sudouser(request)
-
-
-@app.route('/jump/host/unsudouser', methods=['GET', 'POST'])
-def host_unsudouser():
-    if not checkCookie(request):
-        return redirect_sso()
-    return do_host_unsudouser(request)
-
-
 @app.route('/jump/host/modifyuser', methods=['GET', 'POST'])
 def host_modifyuser():
     if not checkCookie(request):
@@ -370,12 +370,7 @@ def do_login(request):
 
 # get主机列表
 def do_hostlist(request):
-    conn, cur = get_db()
-    cur.execute("select id,name,ip,port,desc,date from hosts where isdelete=0;")
-    hosts = cur.fetchall()
-    cur.close()
-    conn.close()
-
+    hosts = g.db.execute("select id,name,ip,port,desc,date from hosts where isdelete=0;").fetchall()
     resp = []
     for host in hosts:
         res = {}
@@ -391,12 +386,7 @@ def do_hostlist(request):
 
 # 获取用户列表
 def do_userlist(request):
-    conn, cur = get_db()
-    cur.execute("select id,name,sudo,desc,date from users where isdelete=0;")
-    users = cur.fetchall()
-    cur.close()
-    conn.close()
-
+    users = g.db.execute("select id,name,sudo,desc,date from users where isdelete=0;").fetchall()
     resp = []
     for user in users:
         res = {}
@@ -417,9 +407,7 @@ def do_add_user(request):
         name = request.form["name"]
         desc = request.form["desc"]
 
-    conn, cur = get_db()
-    cur.execute("select count(1) from users where name='%s'" % name)
-    ret = cur.fetchone()
+    ret = g.db.execute("select count(1) from users where name='%s'" % name).fetchone()
     if (len(ret) > 0 and ret[0]) >= 1:
         return "user %s exists" % name
 
@@ -430,11 +418,9 @@ def do_add_user(request):
         return "error %s" % output
 
     # 新增用户 sql
-    cur.execute("insert into users(name,desc) values('%s',\"%s\")" %
+    g.db.execute("insert into users(name,desc) values('%s',\"%s\")" %
                 (name, desc))
-    conn.commit()
-    cur.close()
-    conn.close()
+    g.db.commit()
     return "ok"
 
 
@@ -453,12 +439,9 @@ def do_del_user(request):
         print("output=%s" % output)
         return "error %s" % output
 
-    conn, cur = get_db()
-    cur.execute("delete from users where name='%s'" % username)
-    cur.execute("delete from hostuser where username='%s'" % username)
-    conn.commit()
-    cur.close()
-    conn.close()
+    g.db.execute("delete from users where name='%s'" % username)
+    g.db.execute("delete from hostuser where username='%s'" % username)
+    g.db.commit()
     return "del user %s ok" % username
 
 
@@ -479,12 +462,8 @@ def do_modify_user(request):
     if sudo != 0 and sudo != 1:
         return "invalid request sudo:%d param" % sudo
 
-    # 需求数据库
-    conn, cur = get_db()
-
     # 检查
-    cur.execute("select sudo,desc from users where name='%s'" % username)
-    users = cur.fetchall()
+    users = g.db.execute("select sudo,desc from users where name='%s'" % username).fetchall()
     if len(users) == 0:
         print("user(%s) not exitst" % username)
         return "user(%s) not exitst" % username
@@ -515,13 +494,10 @@ def do_modify_user(request):
     if change:
         print("update users set sudo=%d,desc=\"%s\" where name='%s'" %
               (sudo, desc, username))
-        cur.execute("update users set sudo=%d,desc=\"%s\" where name='%s'" % (
+        g.db.execute("update users set sudo=%d,desc=\"%s\" where name='%s'" % (
             sudo, desc, username))
-        conn.commit()
+        g.db.commit()
 
-    # 清理数据库
-    cur.close()
-    conn.close()
     print("modify user:%s successful [output: %s]" % (username, output))
     return "modify user:%s successful [output: %s]" % (username, output)
 
@@ -539,22 +515,17 @@ def do_add_host(request):
         port = int(request.form["port"])
         desc = request.form["desc"]
 
-    conn, cur = get_db()
-    cur.execute(
-        "select count(1) from hosts where name='%s' or ip='%s'" % (name, ip))
-    ret = cur.fetchone()
+    ret = g.db.execute(
+        "select count(1) from hosts where name='%s' or ip='%s'" % (name, ip)).fetchone()
     if (len(ret) > 0 and ret[0]) >= 1:
         return "alias name(%s) or ip(%s) is exists" % (name, ip)
 
     # 新增用户 sql
-    cur.execute("insert into hosts(name,ip,port,desc) values('%s','%s',%d,'%s')" % (
+    g.db.execute("insert into hosts(name,ip,port,desc) values('%s','%s',%d,'%s')" % (
         name, ip, port, desc))
-    cur.execute("insert into hostuser(hostname,username,sudo) values('%s','%s','%d')" % (
+    g.db.execute("insert into hostuser(hostname,username,sudo) values('%s','%s','%d')" % (
         name, gDefaultSSHAdmin, 1))
-    conn.commit()
-
-    cur.close()
-    conn.close()
+    g.db.commit()
     return "add host %s:%s ok" % (name, ip)
 
 
@@ -569,14 +540,10 @@ def do_del_host(request):
     else:
         return "invalid request for del host"
 
-    conn, cur = get_db()
-    cur.execute("delete from hosts where name='%s' and ip='%s'" %
+    g.db.execute("delete from hosts where name='%s' and ip='%s'" %
                 (hostname, ip))
-    cur.execute("delete from hostuser where hostname='%s'" % hostname)
-    conn.commit()
-
-    cur.close()
-    conn.close()
+    g.db.execute("delete from hostuser where hostname='%s'" % hostname)
+    g.db.commit()
     return "delete host %s:%s ok" % (hostname, ip)
 
 
@@ -591,19 +558,15 @@ def do_host_adduser(request):
     else:
         return "invalid request for add user to host"
 
-    # 需求数据库
-    conn, cur = get_db()
     # 检查
-    cur.execute("select count(1) from hostuser where hostname='%s' and username='%s'" % (
-        hostname, username))
-    ret = cur.fetchone()
+    ret = g.db.execute("select count(1) from hostuser where hostname='%s' and username='%s'" % (
+        hostname, username)).fetchone()
     if (len(ret) > 0 and ret[0]) >= 1:
         print("user(%s) exitst on host(%s)" % (username, hostname))
         return "user(%s) exitst on host(%s)" % (username, hostname)
 
     # 检查
-    cur.execute("select ip,port from hosts where name='%s'" % (hostname))
-    hostips = cur.fetchone()
+    hostips = g.db.execute("select ip,port from hosts where name='%s'" % (hostname)).fetchone()
     if hostips == None:
         print("host(%s) not exitst on hosts" % hostname)
         return "host(%s) not exitst on hosts" % hostname
@@ -631,14 +594,10 @@ def do_host_adduser(request):
         return "error: remote add user host(%s) user(%s) failed! => output=%s" % (hostname, username, output)
 
     # 记录在数据库中
-    cur.execute("insert into hostuser(hostname,username) values('%s','%s')" % (
-        hostname, username))
-    cur.fetchone()
-    conn.commit()
+    g.db.execute("insert into hostuser(hostname,username) values('%s','%s')" % (
+        hostname, username)).fetchone()
+    g.db.commit()
 
-    # 清理数据库
-    cur.close()
-    conn.close()
     print("host remote =>> add user:%s to host:%s successful [output: %s]" % (
         username, hostname, output))
     return "host remote =>> add user:%s to host:%s successful" % (username, hostname)
@@ -655,20 +614,15 @@ def do_host_deluser(request):
     else:
         return "invalid request for add user to host"
 
-    # 需求数据库
-    conn, cur = get_db()
-
     # 检查
-    cur.execute("select count(1) from hostuser where hostname='%s' and username='%s'" % (
-        hostname, username))
-    ret = cur.fetchone()
+    ret = g.db.execute("select count(1) from hostuser where hostname='%s' and username='%s'" % (
+        hostname, username)).fetchone()
     if (len(ret) > 0 and ret[0]) == 0:
         print("user(%s) not exitst on host(%s)" % (username, hostname))
         return "user(%s) not exitst on host(%s)" % (username, hostname)
 
     # 检查
-    cur.execute("select ip,port from hosts where name='%s'" % (hostname))
-    hostips = cur.fetchone()
+    hostips = g.db.execute("select ip,port from hosts where name='%s'" % (hostname)).fetchone()
     if hostips == None:
         print("host(%s) not exitst on hosts" % hostname)
         return "host(%s) not exitst on hosts" % hostname
@@ -687,13 +641,10 @@ def do_host_deluser(request):
         return "error: remote del user host(%s) user(%s) failed! => output=%s" % (hostname, username, output)
 
     # 记录在数据库中
-    cur.execute("delete from hostuser where hostname='%s' and username='%s'" % (
+    g.db.execute("delete from hostuser where hostname='%s' and username='%s'" % (
         hostname, username))
-    conn.commit()
+    g.db.commit()
 
-    # 清理数据库
-    cur.close()
-    conn.close()
     print("host remote =>> del user:%s from host:%s successful [output: %s]" % (
         username, hostname, output))
     return "host remote =>> del user:%s from host:%s successful" % (username, hostname)
@@ -713,19 +664,14 @@ def do_host_modifyuser(request):
     else:
         return "invalid request for add user to host"
 
-    # 需求数据库
-    conn, cur = get_db()
-
     # 检查
-    cur.execute("select sudo,desc from hostuser where hostname='%s' and username='%s' and isdelete=0" % (
-        hostname, username))
-    hostusers = cur.fetchone()
+    hostusers = g.db.execute("select sudo,desc from hostuser where hostname='%s' and username='%s' and isdelete=0" % (
+        hostname, username)).fetchone()
     if len(hostusers) > 0 and hostusers[0] != None:
         print("user(%s) not exitst on host(%s)" % (username, hostname))
         return "user(%s) not exitst on host(%s)" % (username, hostname)
 
     hostuser = hostusers[0]
-
     user_sudo = hostuser[0]
     user_desc = hostuser[1]
     change = False
@@ -738,8 +684,7 @@ def do_host_modifyuser(request):
         change = True
 
         # 检查
-        cur.execute("select ip,port from hosts where name='%s'" % (hostname))
-        hostips = cur.fetchone()
+        hostips = g.db.execute("select ip,port from hosts where name='%s'" % (hostname)).fetchone()
         if hostips == None:
             print("host(%s) not exitst on hosts" % hostname)
             return "host(%s) not exitst on hosts" % hostname
@@ -760,13 +705,10 @@ def do_host_modifyuser(request):
 
     if change:
         # 记录在数据库中
-        cur.execute("update hostuser set sudo=%d,desc=%s where hostname='%s' and username='%s'" % (
+        g.db.execute("update hostuser set sudo=%d,desc=%s where hostname='%s' and username='%s'" % (
             sudo, desc, hostname, username))
-        conn.commit()
+        g.db.commit()
 
-    # 清理数据库
-    cur.close()
-    conn.close()
     print("host remote =>> %s user:%s from host:%s successful [output: %s]" % (
         opParam, username, hostname, output))
     return "host remote =>> %s user:%s from host:%s successful [output: %s]" % (opParam, username, hostname, output)
@@ -781,12 +723,7 @@ def do_userhostlist(request):
     else:
         return "invalid request for getting user host list"
 
-    conn, cur = get_db()
-    cur.execute("select id,name,ip,port,desc,date from hosts where isdelete=0 and name in (select hostname from hostuser where username='%s')" % username)
-    hosts = cur.fetchall()
-    cur.close()
-    conn.close()
-
+    hosts = g.db.execute("select id,name,ip,port,desc,date from hosts where isdelete=0 and name in (select hostname from hostuser where username='%s')" % username).fetchall()
     resp = []
     for host in hosts:
         res = {}
@@ -809,13 +746,8 @@ def do_hostuserlist(request):
     else:
         return "invalid request for getting host user list"
 
-    conn, cur = get_db()
-    cur.execute(
-        "select id,username,sudo,date from hostuser where isdelete=0 and hostname='%s'" % hostname)
-    users = cur.fetchall()
-    cur.close()
-    conn.close()
-
+    users = g.db.execute(
+        "select id,username,sudo,date from hostuser where isdelete=0 and hostname='%s'" % hostname).fetchall()
     resp = []
     for user in users:
         res = {}
@@ -829,13 +761,8 @@ def do_hostuserlist(request):
 
 # 获取所有的用户主机列表
 def do_hostuserall(request):
-    conn, cur = get_db()
-    cur.execute(
-        "select id,hostname,username,sudo,date from hostuser where isdelete=0")
-    users = cur.fetchall()
-    cur.close()
-    conn.close()
-
+    users = g.db.execute(
+        "select id,hostname,username,sudo,date from hostuser where isdelete=0").fetchall()
     resp = []
     for user in users:
         res = {}
@@ -849,23 +776,22 @@ def do_hostuserall(request):
 
 
 # 连接数据库
-def get_db():
-    conn = sqlite3.connect('/usr/local/jumpserver/jumpserver.db')
-    curr = conn.cursor()
-    return conn, curr
+def connect_db():
+    return sqlite3.connect(gSqlite3File)
 
 
 # 初始化表
 def init_db():
     # 连接数据库
-    conn, _ = get_db()
+    conn = connect_db()
     if conn == None:
         sys.exit(1)
 
-    # 初始化表
+    # 初始化用户表
     conn.execute(gUsersTableSql)
+    # 初始化主机表
     conn.execute(gHostsTableSql)
-    # 创建主机数据库
+    # 创建主机用户表
     conn.execute(gHostUserSql)