126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
"""
|
|
数据库配置和连接工具
|
|
"""
|
|
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
from sqlalchemy import create_engine
|
|
import pandas as pd
|
|
from loguru import logger
|
|
import os
|
|
from typing import Optional
|
|
|
|
|
|
class DatabaseConfig:
|
|
"""数据库配置类"""
|
|
|
|
def __init__(self, env: str = "online"):
|
|
self.host = "host.docker.internal" if env == "online" else "192.168.0.115"
|
|
self.port = 5432
|
|
self.database = "etf_db"
|
|
self.username = "admin"
|
|
self.password = "admin"
|
|
|
|
@property
|
|
def connection_string(self) -> str:
|
|
"""获取连接字符串"""
|
|
return f"postgresql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
|
|
|
|
@property
|
|
def psycopg2_params(self) -> dict:
|
|
"""获取psycopg2连接参数"""
|
|
return {
|
|
"host": self.host,
|
|
"port": self.port,
|
|
"database": self.database,
|
|
"user": self.username,
|
|
"password": self.password,
|
|
}
|
|
|
|
|
|
class DatabaseManager:
|
|
"""数据库管理类"""
|
|
|
|
def __init__(self, config: DatabaseConfig = None):
|
|
self.config = config or DatabaseConfig()
|
|
self.engine = None
|
|
|
|
def get_engine(self):
|
|
"""获取SQLAlchemy引擎"""
|
|
if self.engine is None:
|
|
self.engine = create_engine(
|
|
self.config.connection_string,
|
|
pool_pre_ping=True,
|
|
pool_recycle=300,
|
|
echo=False,
|
|
)
|
|
return self.engine
|
|
|
|
def get_connection(self):
|
|
"""获取psycopg2连接"""
|
|
return psycopg2.connect(**self.config.psycopg2_params)
|
|
|
|
def test_connection(self) -> bool:
|
|
"""测试数据库连接"""
|
|
try:
|
|
with self.get_connection() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute("SELECT 1")
|
|
result = cursor.fetchone()
|
|
logger.info("数据库连接测试成功")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"数据库连接测试失败: {e}")
|
|
return False
|
|
|
|
def create_table_if_not_exists(self, table_name: str, create_sql: str) -> bool:
|
|
"""创建表（如果不存在）"""
|
|
try:
|
|
with self.get_connection() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(create_sql)
|
|
conn.commit()
|
|
logger.info(f"表 {table_name} 创建成功或已存在")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"创建表 {table_name} 失败: {e}")
|
|
return False
|
|
|
|
def insert_dataframe(
|
|
self, df: pd.DataFrame, table_name: str, if_exists: str = "append"
|
|
) -> bool:
|
|
"""将DataFrame插入到数据库表"""
|
|
try:
|
|
engine = self.get_engine()
|
|
df.to_sql(
|
|
table_name,
|
|
engine,
|
|
if_exists=if_exists,
|
|
index=False,
|
|
method="multi",
|
|
chunksize=1000,
|
|
)
|
|
logger.info(f"成功插入 {len(df)} 条记录到表 {table_name}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"插入数据到表 {table_name} 失败: {e}")
|
|
return False
|
|
|
|
def execute_query(self, query: str, params: tuple = None) -> Optional[list]:
|
|
"""执行查询并返回结果"""
|
|
try:
|
|
with self.get_connection() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(query, params)
|
|
result = cursor.fetchall()
|
|
return result
|
|
except Exception as e:
|
|
logger.error(f"执行查询失败: {e}")
|
|
return None
|
|
|
|
def close(self):
|
|
"""关闭连接"""
|
|
if self.engine:
|
|
self.engine.dispose()
|
|
self.engine = None
|