删除旧文件；更新数据库查询方法

timescaledb/database.py

@@ -1,6 +1,6 @@
 import logging
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Dict, Optional
 import psycopg_pool
 from psycopg.rows import dict_row
 import timescaledb.timescaledb_info as timescaledb_info
@@ -9,12 +9,15 @@ import timescaledb.timescaledb_info as timescaledb_info
 logger = logging.getLogger(__name__)
 
 class Database:
-    def __init__(self):
+    def __init__(self, db_name=None):
         self.pool = None
+        self.db_name = db_name
 
-    def init_pool(self):
+    def init_pool(self, db_name=None):
         """Initialize the connection pool."""
-        conn_string = timescaledb_info.get_pgconn_string()
+        # Use provided db_name, or the one from constructor, or default from config
+        target_db_name = db_name or self.db_name
+        conn_string = timescaledb_info.get_pgconn_string(db_name=target_db_name)
         try:
             self.pool = psycopg_pool.AsyncConnectionPool(
                 conninfo=conn_string,
@@ -23,7 +26,7 @@ class Database:
                 open=False, # Don't open immediately, wait for startup
                 kwargs={"row_factory": dict_row} # Return rows as dictionaries
             )
-            logger.info("TimescaleDB connection pool initialized.")
+            logger.info(f"TimescaleDB connection pool initialized for database: {target_db_name or 'default'}")
         except Exception as e:
             logger.error(f"Failed to initialize TimescaleDB connection pool: {e}")
             raise
@@ -47,9 +50,53 @@ class Database:
         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.")