perf: improve sqlite context manager

src/server.py

@@ -16,37 +16,34 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan)
 
-# 连接到数据库（如果数据库不存在，则会自动创建）
+
-conn = sqlite3.connect("server.db")
+def init():
-# 创建游标对象，用于执行SQL语句
+    asyncio.create_task(receive_heartbeat_internal())
-cursor = conn.cursor()
+
-# 创建表: id: int; ip: TEXT
+    conn = sqlite3.connect("server.db")
-cursor.execute(
+    cursor = conn.cursor()
+    # init table: id: int; ip: TEXT
+    cursor.execute(
         """CREATE TABLE IF NOT EXISTS nodes (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
                     ip TEXT NOT NULL,
                     last_heartbeat INTEGER
                 )"""
-)
+    )
-
-
-def init():
-    asyncio.create_task(receive_heartbeat_internal())
 
 
 def clean_env():
     clear_database()
-    # 关闭游标和连接
-    cursor.close()
-    conn.close()
 
 
 @app.get("/server/show_nodes")
 async def show_nodes() -> list:
     nodes_list = []
+    with sqlite3.connect("server.db") as db:
         # 查询数据
-    cursor.execute("SELECT * FROM nodes")
+        cursor = db.execute("SELECT * FROM nodes")
         rows = cursor.fetchall()
+
     for row in rows:
         nodes_list.append(row)
     return nodes_list
@@ -71,12 +68,13 @@ async def get_node(ip: str) -> int:
     current_time = int(time.time())
     print("当前时间: ", current_time)
 
+    with sqlite3.connect("server.db") as db:
         # 插入数据
-    cursor.execute(
+        db.execute(
             "INSERT INTO nodes (id, ip, last_heartbeat) VALUES (?, ?, ?)",
             (ip_int, ip, current_time),
         )
-    conn.commit()
+        db.commit()
 
     return ip_int
 
@@ -89,13 +87,15 @@ async def delete_node(ip: str) -> None:
     return:
     None
     """
+    with sqlite3.connect("server.db") as db:
         # 查询要删除的节点
-    cursor.execute("SELECT * FROM nodes WHERE ip=?", (ip,))
+        cursor = db.execute("SELECT * FROM nodes WHERE ip=?", (ip,))
         row = cursor.fetchone()
     if row is not None:
+        with sqlite3.connect("server.db") as db:
             # 执行删除操作
-        cursor.execute("DELETE FROM nodes WHERE ip=?", (ip,))
+            db.execute("DELETE FROM nodes WHERE ip=?", (ip,))
-        conn.commit()
+            db.commit()
         print(f"Node with IP {ip} deleted successfully.")
     else:
         print(f"Node with IP {ip} not found.")
@@ -105,7 +105,8 @@ async def delete_node(ip: str) -> None:
 @app.get("/server/heartbeat")
 async def receive_heartbeat(ip: str):
     print("收到来自", ip, "的心跳包")
-    cursor.execute(
+    with sqlite3.connect("server.db") as db:
+        db.execute(
             "UPDATE nodes SET last_heartbeat = ? WHERE ip = ?", (time.time(), ip)
         )
     return {"status": "received"}
@@ -114,11 +115,12 @@ async def receive_heartbeat(ip: str):
 async def receive_heartbeat_internal():
     while 1:
         timeout = 70
+        with sqlite3.connect("server.db") as db:
             # 删除超时的节点
-        cursor.execute(
+            db.execute(
                 "DELETE FROM nodes WHERE last_heartbeat < ?", (time.time() - timeout,)
             )
-        conn.commit()
+            db.commit()
         await asyncio.sleep(timeout)
 
 
@@ -133,9 +135,11 @@ async def send_nodes_list(count: int) -> list:
     """
     nodes_list = []
 
+    with sqlite3.connect("server.db") as db:
         # 查询数据库中的节点数据
-    cursor.execute("SELECT * FROM nodes LIMIT ?", (count,))
+        cursor = db.execute("SELECT * FROM nodes LIMIT ?", (count,))
         rows = cursor.fetchall()
+        
     for row in rows:
         id, ip, last_heartbeat = row
         nodes_list.append(ip)
@@ -147,8 +151,9 @@ async def send_nodes_list(count: int) -> list:
 
 # @app.get("/server/clear_database")
 def clear_database() -> None:
-    cursor.execute("DELETE FROM nodes")
+    with sqlite3.connect("server.db") as db:
-    conn.commit()
+        db.execute("DELETE FROM nodes")
+        db.commit()
 
 
 if __name__ == "__main__":