From c8cb5ce58c65b6a440690079b88847a1a1dacf22 Mon Sep 17 00:00:00 2001 From: wzy-warehouse <18135009705@163.com> Date: Fri, 8 May 2026 16:33:57 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=95=B0=E6=8D=AE=E5=BA=93?= =?UTF-8?q?=E7=AE=A1=E7=90=86=E7=B1=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/utils/__init__.py | 3 + app/utils/db_helper.py | 150 +++++++++++++++++++++++++++++++++++++++++ requirements.txt | 3 +- 3 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 app/utils/db_helper.py diff --git a/app/utils/__init__.py b/app/utils/__init__.py index 63b5914..1b92493 100644 --- a/app/utils/__init__.py +++ b/app/utils/__init__.py @@ -1,3 +1,6 @@ """ Utility functions package """ +from app.utils.db_helper import db_helper, PostgresSQLHelper + +__all__ = ['db_helper', 'PostgresSQLHelper'] diff --git a/app/utils/db_helper.py b/app/utils/db_helper.py new file mode 100644 index 0000000..9a9647c --- /dev/null +++ b/app/utils/db_helper.py @@ -0,0 +1,150 @@ +""" +PostgreSQL 数据库工具类 +提供增删改查方法 +""" +import psycopg2 +from psycopg2.extras import RealDictCursor +from typing import List, Dict, Any, Optional, Tuple +from contextlib import contextmanager +from config import settings + + +class PostgresSQLHelper: + """PostgreSQL 数据库帮助类""" + + def __init__(self): + """初始化数据库连接配置""" + self.db_config = { + 'host': settings.DB_HOST, + 'port': settings.DB_PORT, + 'user': settings.DB_USER, + 'password': settings.DB_PASSWORD, + 'database': settings.DB_NAME, + } + + @contextmanager + def get_connection(self): + """ + 获取数据库连接的上下文管理器 + 自动管理连接的开启和关闭 + """ + conn = None + try: + conn = psycopg2.connect(**self.db_config) + yield conn + conn.commit() + except Exception as e: + if conn: + conn.rollback() + raise e + finally: + if conn: + conn.close() + + @contextmanager + def get_cursor(self, dict_cursor=False): + """ + 获取数据库游标的上下文管理器 + + Args: + dict_cursor: 是否使用字典游标(返回字典格式结果) + """ + with self.get_connection() as conn: + cursor = None + try: + if dict_cursor: + cursor = conn.cursor(cursor_factory=RealDictCursor) + else: + cursor = conn.cursor() + yield cursor + finally: + if cursor: + cursor.close() + + def execute_query(self, sql: str, params: tuple = None) -> List[Dict[str, Any]]: + """ + 执行查询语句,返回字典列表 + + Args: + sql: SQL 查询语句 + params: 参数元组 + + Returns: + 查询结果列表,每个元素为字典 + """ + with self.get_cursor(dict_cursor=True) as cursor: + cursor.execute(sql, params) + results = cursor.fetchall() + # 将 RealDictRow 转换为普通字典 + return [dict(row) for row in results] + + def execute_query_one(self, sql: str, params: tuple = None) -> Optional[Dict[str, Any]]: + """ + 执行查询语句,返回单条记录 + + Args: + sql: SQL 查询语句 + params: 参数元组 + + Returns: + 单条记录的字典,如果没有结果则返回 None + """ + with self.get_cursor(dict_cursor=True) as cursor: + cursor.execute(sql, params) + result = cursor.fetchone() + return dict(result) if result else None + + def execute_update(self, sql: str, params: tuple = None) -> int: + """ + 执行更新/插入/删除语句 + + Args: + sql: SQL 语句 + params: 参数元组 + + Returns: + 影响的行数 + """ + with self.get_cursor() as cursor: + cursor.execute(sql, params) + return cursor.rowcount + + def execute_insert(self, sql: str, params: tuple = None, returning: str = None) -> Any: + """ + 执行插入语句并返回新生成的 ID 或指定字段 + + Args: + sql: INSERT SQL 语句 + params: 参数元组 + returning: RETURNING 子句指定的字段名,默认为 'id' + + Returns: + 新生成的 ID 或指定字段的值 + """ + if returning and not sql.upper().endswith(f'RETURNING {returning}'.upper()): + sql = f"{sql} RETURNING {returning}" + + with self.get_cursor() as cursor: + cursor.execute(sql, params) + result = cursor.fetchone() + return result[0] if result else None + + def execute_many(self, sql: str, params_list: List[tuple]) -> int: + """ + 批量执行 SQL 语句 + + Args: + sql: SQL 语句 + params_list: 参数列表,每个元素为参数元组 + + Returns: + 影响的总行数 + """ + with self.get_cursor() as cursor: + cursor.executemany(sql, params_list) + return cursor.rowcount + + +# 创建全局实例 +db_helper = PostgresSQLHelper() + diff --git a/requirements.txt b/requirements.txt index ad70ac5..3c40ee6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ -dynaconf == 3.2.13 \ No newline at end of file +dynaconf == 3.2.13 +psycopg2-binary == 2.9.12 \ No newline at end of file