删除旧文件;更新数据库查询方法
This commit is contained in:
104
postgresql/database.py
Normal file
104
postgresql/database.py
Normal file
@@ -0,0 +1,104 @@
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import AsyncGenerator, Dict, Optional
|
||||
import psycopg_pool
|
||||
from psycopg.rows import dict_row
|
||||
import postgresql_info
|
||||
|
||||
# Configure logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Database:
|
||||
def __init__(self, db_name=None):
|
||||
self.pool = None
|
||||
self.db_name = db_name
|
||||
|
||||
def init_pool(self, db_name=None):
|
||||
"""Initialize the connection pool."""
|
||||
# Use provided db_name, or the one from constructor, or default from config
|
||||
target_db_name = db_name or self.db_name
|
||||
conn_string = postgresql_info.get_pgconn_string(db_name=target_db_name)
|
||||
try:
|
||||
self.pool = psycopg_pool.AsyncConnectionPool(
|
||||
conninfo=conn_string,
|
||||
min_size=1,
|
||||
max_size=20,
|
||||
open=False, # Don't open immediately, wait for startup
|
||||
kwargs={"row_factory": dict_row}, # Return rows as dictionaries
|
||||
)
|
||||
logger.info(f"PostgreSQL connection pool initialized for database: {target_db_name or 'default'}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize postgresql connection pool: {e}")
|
||||
raise
|
||||
|
||||
async def open(self):
|
||||
if self.pool:
|
||||
await self.pool.open()
|
||||
|
||||
async def close(self):
|
||||
"""Close the connection pool."""
|
||||
if self.pool:
|
||||
await self.pool.close()
|
||||
logger.info("PostgreSQL connection pool closed.")
|
||||
|
||||
@asynccontextmanager
|
||||
async def get_connection(self) -> AsyncGenerator:
|
||||
"""Get a connection from the pool."""
|
||||
if not self.pool:
|
||||
raise Exception("Database pool is not initialized.")
|
||||
|
||||
async with self.pool.connection() as conn:
|
||||
yield conn
|
||||
|
||||
|
||||
# 默认数据库实例
|
||||
db = Database()
|
||||
|
||||
# 缓存不同数据库的实例 - 避免重复创建连接池
|
||||
_database_instances: Dict[str, Database] = {}
|
||||
|
||||
def create_database_instance(db_name):
|
||||
"""Create a new Database instance for a specific database."""
|
||||
return Database(db_name=db_name)
|
||||
|
||||
async def get_database_instance(db_name: Optional[str] = None) -> Database:
|
||||
"""Get or create a database instance for the specified database name."""
|
||||
if not db_name:
|
||||
return db # 返回默认数据库实例
|
||||
|
||||
if db_name not in _database_instances:
|
||||
# 创建新的数据库实例
|
||||
instance = create_database_instance(db_name)
|
||||
instance.init_pool()
|
||||
await instance.open()
|
||||
_database_instances[db_name] = instance
|
||||
logger.info(f"Created new database instance for: {db_name}")
|
||||
|
||||
return _database_instances[db_name]
|
||||
|
||||
async def get_db_connection():
|
||||
"""Dependency for FastAPI to get a database connection."""
|
||||
async with db.get_connection() as conn:
|
||||
yield conn
|
||||
|
||||
async def get_database_connection(db_name: Optional[str] = None):
|
||||
"""
|
||||
FastAPI dependency to get database connection with optional database name.
|
||||
使用方法: conn: AsyncConnection = Depends(lambda: get_database_connection("your_db_name"))
|
||||
或在路由函数中: conn: AsyncConnection = Depends(get_database_connection)
|
||||
"""
|
||||
instance = await get_database_instance(db_name)
|
||||
async with instance.get_connection() as conn:
|
||||
yield conn
|
||||
|
||||
async def cleanup_database_instances():
|
||||
"""Clean up all database instances (call this on application shutdown)."""
|
||||
for db_name, instance in _database_instances.items():
|
||||
await instance.close()
|
||||
logger.info(f"Closed database instance for: {db_name}")
|
||||
_database_instances.clear()
|
||||
|
||||
# 关闭默认数据库
|
||||
await db.close()
|
||||
logger.info("All database instances cleaned up.")
|
||||
40
postgresql/router.py
Normal file
40
postgresql/router.py
Normal file
@@ -0,0 +1,40 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from typing import List, Optional
|
||||
from datetime import datetime
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
from .database import get_database_instance
|
||||
from .scada_info import query_pg_scada_info
|
||||
|
||||
router = APIRouter(prefix="/postgresql", tags=["postgresql"])
|
||||
|
||||
|
||||
# 创建支持数据库选择的连接依赖函数
|
||||
async def get_database_connection(
|
||||
db_name: Optional[str] = Query(
|
||||
None, description="指定要连接的数据库名称,为空时使用默认数据库"
|
||||
)
|
||||
):
|
||||
"""获取数据库连接,支持通过查询参数指定数据库名称"""
|
||||
instance = await get_database_instance(db_name)
|
||||
async with instance.get_connection() as conn:
|
||||
yield conn
|
||||
|
||||
|
||||
@router.get("/scada-info")
|
||||
async def get_scada_info_with_connection(
|
||||
conn: AsyncConnection = Depends(get_database_connection),
|
||||
):
|
||||
"""
|
||||
使用连接池查询SCADA信息
|
||||
"""
|
||||
try:
|
||||
# 使用连接查询SCADA信息
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute("SELECT * FROM scada_info")
|
||||
scada_data = await cur.fetchall()
|
||||
return {"success": True, "data": scada_data, "count": len(scada_data)}
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"查询SCADA信息时发生错误: {str(e)}"
|
||||
)
|
||||
35
postgresql/scada_info.py
Normal file
35
postgresql/scada_info.py
Normal file
@@ -0,0 +1,35 @@
|
||||
from typing import List, Optional, Any
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
|
||||
class ScadaRepository:
|
||||
|
||||
@staticmethod
|
||||
async def get_scadas_info(conn: AsyncConnection) -> List[dict]:
|
||||
"""
|
||||
查询pg数据库中,scada_info 的所有记录
|
||||
:param conn: 异步数据库连接
|
||||
:return: 包含所有记录的列表,每条记录为一个字典
|
||||
"""
|
||||
async with conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
"""
|
||||
SELECT id, type, transmission_mode, transmission_frequency, reliability
|
||||
FROM public.scada_info
|
||||
"""
|
||||
)
|
||||
records = await cur.fetchall()
|
||||
|
||||
# 将查询结果转换为字典列表
|
||||
records_list = []
|
||||
for record in records:
|
||||
record_dict = {
|
||||
"id": record[0],
|
||||
"type": record[1],
|
||||
"transmission_mode": record[2],
|
||||
"transmission_frequency": record[3],
|
||||
"reliability": record[4],
|
||||
}
|
||||
records_list.append(record_dict)
|
||||
|
||||
return records_list
|
||||
Reference in New Issue
Block a user