87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
"""第 4-1 课示例：从 TOML 配置读取参数并连接 PostgreSQL。"""
|
||
|
||
from pathlib import Path
|
||
import tomllib
|
||
|
||
import psycopg
|
||
|
||
|
||
# 使用当前 Python 文件的位置定位配置，避免程序依赖 PowerShell 的工作目录。
|
||
CONFIG_PATH = Path(__file__).with_name("config.toml")
|
||
|
||
|
||
def load_database_config(config_path: Path) -> dict[str, str | int]:
|
||
"""读取并检查 TOML 中的 PostgreSQL 连接配置。"""
|
||
if not config_path.exists():
|
||
raise RuntimeError(
|
||
"未找到 config.toml，请复制 config.example.toml 并填写练习数据库配置。"
|
||
)
|
||
|
||
# tomllib 要求以二进制模式读取 TOML 文件，因此这里使用 "rb"。
|
||
with config_path.open("rb") as config_file:
|
||
config_data = tomllib.load(config_file)
|
||
|
||
postgresql_config = config_data.get("postgresql")
|
||
if not isinstance(postgresql_config, dict):
|
||
raise RuntimeError("config.toml 缺少 [postgresql] 配置节。")
|
||
|
||
required_names = ("host", "port", "dbname", "user", "password")
|
||
missing_names = [
|
||
name for name in required_names if postgresql_config.get(name) in (None, "")
|
||
]
|
||
if missing_names:
|
||
missing_text = "、".join(missing_names)
|
||
raise RuntimeError(f"config.toml 缺少数据库配置：{missing_text}。")
|
||
|
||
return postgresql_config
|
||
|
||
|
||
def query_connection_info(
|
||
database_config: dict[str, str | int],
|
||
) -> tuple[str, str, str]:
|
||
"""连接 PostgreSQL，执行只读参数化查询，并返回连接信息。"""
|
||
message = "Psycopg 连接成功"
|
||
|
||
# ** 会把字典中的键值展开为关键字参数，例如 host="数据库主机"。
|
||
# 连接信息来自本地 config.toml，不会写入环境变量。
|
||
with psycopg.connect(**database_config) as connection:
|
||
# 第二层 with 管理游标。游标同时负责执行 SQL 和读取查询结果。
|
||
with connection.cursor() as cursor:
|
||
# SQL 和参数必须分开传递，不能使用 f-string 拼接用户数据。
|
||
cursor.execute(
|
||
"SELECT current_database(), current_user, %s::text",
|
||
(message,),
|
||
)
|
||
result = cursor.fetchone()
|
||
|
||
# 这条 PostgreSQL 查询必然返回一行。普通业务查询仍要考虑 None。
|
||
if result is None:
|
||
raise RuntimeError("数据库没有返回连接验证结果。")
|
||
|
||
database_name, user_name, returned_message = result
|
||
return database_name, user_name, returned_message
|
||
|
||
|
||
def main() -> None:
|
||
"""组织配置读取、数据库查询和结果输出。"""
|
||
try:
|
||
database_config = load_database_config(CONFIG_PATH)
|
||
database_name, user_name, message = query_connection_info(database_config)
|
||
except (OSError, tomllib.TOMLDecodeError, RuntimeError) as error:
|
||
# 这里处理文件读取、TOML 格式以及课程程序主动检查到的问题。
|
||
print(f"配置读取失败：{error}")
|
||
return
|
||
except psycopg.Error as error:
|
||
# Psycopg 的具体异常信息可以保留英文，前面补充中文场景说明。
|
||
print(f"数据库访问失败：{error}")
|
||
return
|
||
|
||
print("连接成功。")
|
||
print(f"当前数据库：{database_name}")
|
||
print(f"当前用户：{user_name}")
|
||
print(f"参数化查询结果：{message}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|