添加数据库连接池

This commit is contained in:
wzy-warehouse
2026-06-06 13:34:18 +08:00
parent 39b46b58fd
commit a8737f1ec8
3 changed files with 25 additions and 19 deletions
+1 -3
View File
@@ -218,7 +218,7 @@ class EarthquakeDBN:
lat = point.get('lat')
source_type = point.get('source_type')
logger.info(f"地震预测点 ID={point_id}, source_type={source_type}")
logger.debug(f"地震预测点 ID={point_id}, source_type={source_type}")
# 计算震中距(如果未直接提供)
if epicenter_distance is None:
@@ -226,7 +226,6 @@ class EarthquakeDBN:
epicenter_distance = self._haversine_distance(
lon, lat, epicenter_lon, epicenter_lat
)
logger.info(f"计算震中距: {epicenter_distance:.1f} km")
else:
logger.warning("未提供震中距或震中坐标,使用默认值 100km")
epicenter_distance = 100.0
@@ -234,7 +233,6 @@ class EarthquakeDBN:
# 估算地震烈度(如果未直接提供)
if seismic_intensity is None:
seismic_intensity = self.estimate_seismic_intensity(magnitude, epicenter_distance, depth)
logger.info(f"估算地震烈度: {seismic_intensity:.1f}")
# 获取静态因子数据
raw_factors = point.get('static_factors', {})
+1 -1
View File
@@ -234,7 +234,7 @@ class RainfallDBN:
lat = point.get('lat')
source_type = point.get('source_type')
logger.info(f"预测点 ID={point_id}, source_type={source_type}")
logger.debug(f"预测点 ID={point_id}, source_type={source_type}")
# 获取降雨数据
if rainfall is not None and duration is not None:
+23 -15
View File
@@ -1,8 +1,9 @@
"""
PostgreSQL 数据库工具类
提供增删改查方法
提供增删改查方法,内置连接池
"""
import psycopg2
from psycopg2 import pool
from psycopg2.extras import RealDictCursor
from typing import List, Dict, Any, Optional, Tuple
from contextlib import contextmanager
@@ -10,10 +11,10 @@ from config import settings
class PostgresSQLHelper:
"""PostgreSQL 数据库帮助类"""
"""PostgreSQL 数据库帮助类(连接池版)"""
def __init__(self):
"""初始化数据库连接配置"""
"""初始化数据库连接"""
self.db_config = {
'host': settings.DB_HOST,
'port': settings.DB_PORT,
@@ -21,31 +22,38 @@ class PostgresSQLHelper:
'password': settings.DB_PASSWORD,
'database': settings.DB_NAME,
}
self._pool = None
def _ensure_pool(self):
"""延迟初始化连接池"""
if self._pool is None:
self._pool = pool.ThreadedConnectionPool(
minconn=2,
maxconn=20,
**self.db_config
)
@contextmanager
def get_connection(self):
"""
获取数据库连接的上下文管理器
自动管理连接的开启和关闭
从连接池获取连接(复用TCP连接,省去握手开销)
"""
conn = None
self._ensure_pool()
conn = self._pool.getconn()
try:
conn = psycopg2.connect(**self.db_config)
yield conn
conn.commit()
except Exception as e:
if conn:
conn.rollback()
conn.rollback()
raise e
finally:
if conn:
conn.close()
self._pool.putconn(conn)
@contextmanager
def get_cursor(self, dict_cursor=False):
"""
获取数据库游标的上下文管理器
Args:
dict_cursor: 是否使用字典游标(返回字典格式结果)
"""