Python企业级 PostgreSQL 数据库创建函数实现
# 企业级 PostgreSQL 数据库创建函数实现
以下是一个企业级的 Python 函数实现，用于创建 PostgreSQL 数据库，包含错误处理、日志记录和配置管理等最佳实践：
```python
import logging
import psycopg2
from psycopg2 import sql
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
from typing import Optional, Dict, Any
import os
from dataclasses import dataclass
from datetime import datetime
# 配置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@dataclass
class DatabaseConfig:
"""数据库配置数据类"""
host: str = "localhost"
port: int = 5432
user: str = "postgres"
password: str = ""
maintenance_db: str = "postgres" # 用于连接以创建新数据库的维护数据库
timeout: int = 30 # 连接超时时间(秒)
class DatabaseCreator:
"""企业级PostgreSQL数据库创建器"""
def __init__(self, config: DatabaseConfig):
self.config = config
self.connection = None
def __enter__(self):
"""上下文管理器入口"""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""上下文管理器退出 - 确保连接关闭"""
self.close_connection()
def create_connection(self, dbname: Optional[str] = None) -> bool:
"""
创建到PostgreSQL服务器的连接
Args:
dbname: 要连接的数据库名称，None表示使用维护数据库
Returns:
bool: 连接是否成功
"""
try:
# 构建连接参数
connect_dbname = dbname or self.config.maintenance_db
conn_params = {
"host": self.config.host,
"port": self.config.port,
"user": self.config.user,
"password": self.config.password,
"dbname": connect_dbname,
"connect_timeout": self.config.timeout
}
# 建立连接
self.connection = psycopg2.connect(**conn_params)
self.connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
logger.info(f"成功连接到数据库服务器 {self.config.host}:{self.config.port}")
return True
except psycopg2.Error as e:
logger.error(f"数据库连接失败: {e}")
return False
def close_connection(self) -> None:
"""安全关闭数据库连接"""
if self.connection and not self.connection.closed:
self.connection.close()
logger.info("数据库连接已关闭")
def database_exists(self, dbname: str) -> bool:
"""
检查数据库是否已存在
Args:
dbname: 要检查的数据库名称
Returns:
bool: 数据库是否存在
"""
if not self.connection or self.connection.closed:
if not self.create_connection():
return False
try:
with self.connection.cursor() as cursor:
# 查询pg_database系统目录以检查数据库是否存在
query = sql.SQL("SELECT 1 FROM pg_database WHERE datname = %s")
cursor.execute(query, (dbname,))
return cursor.fetchone() is not None
except psycopg2.Error as e:
logger.error(f"检查数据库存在性时出错: {e}")
return False
def create_database(
self,
dbname: str,
owner: str = "postgres",
encoding: str = "UTF8",
lc_collate: str = "Chinese (Simplified)_China.936",
lc_ctype: str = "Chinese (Simplified)_China.936",
locale_provider: str = "libc",
tablespace: str = "pg_default",
connection_limit: int = -1,
is_template: bool = False,
comment: Optional[str] = None,
if_not_exists: bool = True,
drop_if_exists: bool = False
) -> bool:
"""
创建PostgreSQL数据库
Args:
dbname: 数据库名称
owner: 数据库所有者
encoding: 字符编码
lc_collate: 排序规则
lc_ctype: 字符分类
locale_provider: 区域提供者
tablespace: 表空间
connection_limit: 连接限制(-1表示无限制)
is_template: 是否作为模板数据库
comment: 数据库注释
if_not_exists: 如果数据库已存在是否跳过创建
drop_if_exists: 如果数据库已存在是否删除后重新创建
Returns:
bool: 数据库创建是否成功
"""
# 首先连接到维护数据库
if not self.create_connection():
return False
# 检查数据库是否已存在
db_exists = self.database_exists(dbname)
if db_exists:
if if_not_exists and not drop_if_exists:
logger.info(f"数据库 '{dbname}' 已存在，跳过创建")
return True
elif drop_if_exists:
if not self.drop_database(dbname):
logger.error(f"无法删除已存在的数据库 '{dbname}'")
return False
try:
with self.connection.cursor() as cursor:
# 构建CREATE DATABASE语句
base_query = sql.SQL("CREATE DATABASE {}")
if if_not_exists and not drop_if_exists:
base_query = sql.SQL("CREATE DATABASE IF NOT EXISTS {}")
# 添加选项
options = [
sql.SQL("OWNER = {}").format(sql.Identifier(owner)),
sql.SQL("ENCODING = {}").format(sql.Literal(encoding)),
sql.SQL("LC_COLLATE = {}").format(sql.Literal(lc_collate)),
sql.SQL("LC_CTYPE = {}").format(sql.Literal(lc_ctype)),
sql.SQL("LOCALE_PROVIDER = {}").format(sql.Literal(locale_provider)),
sql.SQL("TABLESPACE = {}").format(sql.Identifier(tablespace)),
sql.SQL("CONNECTION LIMIT = {}").format(sql.Literal(connection_limit)),
sql.SQL("IS_TEMPLATE = {}").format(sql.Literal(is_template))
]
# 组合完整的SQL语句
query = sql.SQL("{} WITH {}").format(
base_query.format(sql.Identifier(dbname)),
sql.SQL(", ").join(options)
)
# 执行创建语句
cursor.execute(query)
logger.info(f"数据库 '{dbname}' 创建成功")
# 添加注释（如果提供）
if comment:
comment_query = sql.SQL("COMMENT ON DATABASE {} IS {}").format(
sql.Identifier(dbname),
sql.Literal(comment)
)
cursor.execute(comment_query)
logger.info(f"已为数据库 '{dbname}' 添加注释")
return True
except psycopg2.Error as e:
logger.error(f"创建数据库时出错: {e}")
return False
def drop_database(self, dbname: str) -> bool:
"""
删除数据库
Args:
dbname: 要删除的数据库名称
Returns:
bool: 删除是否成功
"""
if not self.connection or self.connection.closed:
if not self.create_connection():
return False
try:
with self.connection.cursor() as cursor:
# 确保没有其他会话连接到要删除的数据库
terminate_query = sql.SQL(
"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = %s"
)
cursor.execute(terminate_query, (dbname,))
# 执行删除操作
drop_query = sql.SQL("DROP DATABASE IF EXISTS {}").format(sql.Identifier(dbname))
cursor.execute(drop_query)
logger.info(f"数据库 '{dbname}' 已删除")
return True
except psycopg2.Error as e:
logger.error(f"删除数据库时出错: {e}")
return False
def create_safe_sentry_database() -> bool:
"""
创建safe_sentry数据库的专用函数
Returns:
bool: 数据库创建是否成功
"""
# 从环境变量获取配置或使用默认值
config = DatabaseConfig(
host=os.getenv("DB_HOST", "localhost"),
port=int(os.getenv("DB_PORT", "5432")),
user=os.getenv("DB_USER", "postgres"),
password=os.getenv("DB_PASSWORD", ""),
maintenance_db=os.getenv("DB_MAINTENANCE_DB", "postgres")
)
# 创建注释，包含当前日期
current_date = datetime.now().strftime("%Y年%m月%d日 %A")
comment = f"{current_date} 创建。"
# 使用上下文管理器确保资源正确释放
with DatabaseCreator(config) as creator:
return creator.create_database(
dbname="safe_sentry",
owner="postgres",
encoding="UTF8",
lc_collate="Chinese (Simplified)_China.936",
lc_ctype="Chinese (Simplified)_China.936",
locale_provider="libc",
tablespace="pg_default",
connection_limit=-1,
is_template=False,
comment=comment,
if_not_exists=True,
drop_if_exists=False # 设置为True将删除并重新创建已存在的数据库
)
# 使用示例
if __name__ == "__main__":
# 创建safe_sentry数据库
success = create_safe_sentry_database()
if success:
print("数据库创建成功")
else:
print("数据库创建失败")
exit(1)
```
## 企业级特性说明
1. **配置管理**:
- 使用数据类存储配置信息
- 支持环境变量配置，提高安全性
2. **错误处理**:
- 全面的异常捕获和处理
- 详细的日志记录
3. **资源管理**:
- 使用上下文管理器确保连接正确关闭
- 防止资源泄漏
4. **灵活性**:
- 参数化所有数据库选项
- 支持条件创建(IF NOT EXISTS)
- 支持删除并重新创建现有数据库
5. **安全性**:
- 使用参数化查询防止SQL注入
- 使用psycopg2的sql模块安全构建查询
6. **可维护性**:
- 清晰的代码结构
- 详细的文档字符串
- 模块化设计
## 使用说明
1. 安装依赖: `pip install psycopg2-binary`
2. 设置环境变量(可选):
- `DB_HOST`: 数据库主机
- `DB_PORT`: 数据库端口
- `DB_USER`: 数据库用户
- `DB_PASSWORD`: 数据库密码
3. 直接运行脚本或导入使用
这个实现遵循了企业级应用的最佳实践，适合在生产环境中使用。
浙公网安备 33010602011771号