db.py
1.16 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
"""
数据库连接管理模块
统一管理 MySQL 数据库连接
"""
import mysql.connector
from mysql.connector import MySQLConnection
from ..config import get_settings
def get_connection() -> MySQLConnection:
"""
获取数据库连接
Returns:
MySQLConnection: MySQL 数据库连接
"""
settings = get_settings()
return mysql.connector.connect(
host=settings.mysql_host,
port=settings.mysql_port,
user=settings.mysql_user,
password=settings.mysql_password,
database=settings.mysql_database,
)
class DatabaseConnection:
"""
数据库连接上下文管理器
使用示例:
with DatabaseConnection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM table")
rows = cursor.fetchall()
"""
def __init__(self):
self._conn: MySQLConnection = None
def __enter__(self) -> MySQLConnection:
self._conn = get_connection()
return self._conn
def __exit__(self, exc_type, exc_val, exc_tb):
if self._conn and self._conn.is_connected():
self._conn.close()