193 lines
5.9 KiB
Python
193 lines
5.9 KiB
Python
"""
|
||
数据库连接
|
||
|
||
使用 psycopg2 直连 PostgreSQL，不引入 ORM。
|
||
连接参数从环境变量读取（经 config 模块加载）。
|
||
|
||
提供两类连接：
|
||
- get_connection()：zqyy_app 读写连接（用户/队列/调度等业务数据）
|
||
- get_etl_readonly_connection(site_id)：etl_feiqiu 只读连接（数据库查看器），
|
||
自动设置 RLS site_id 隔离
|
||
|
||
当 DEV_TRACE_ENABLED=true 且存在活跃 TraceContext 时，
|
||
get_connection() 返回 TracedConnection 包装，自动记录 DB_CONN / DB_QUERY / DB_CONN_RELEASE span。
|
||
"""
|
||
|
||
import time
|
||
|
||
import psycopg2
|
||
from psycopg2.extensions import connection as PgConnection
|
||
|
||
from app.config import (
|
||
APP_DB_NAME,
|
||
DB_HOST,
|
||
DB_PASSWORD,
|
||
DB_PORT,
|
||
DB_USER,
|
||
ETL_DB_HOST,
|
||
ETL_DB_NAME,
|
||
ETL_DB_PASSWORD,
|
||
ETL_DB_PORT,
|
||
ETL_DB_USER,
|
||
)
|
||
|
||
# TCP keepalive 参数：防止长期运行的后台服务连接被网络设备/PostgreSQL 回收
|
||
_KEEPALIVE_KWARGS = {
|
||
"keepalives": 1,
|
||
"keepalives_idle": 60, # 空闲 60 秒后开始探测
|
||
"keepalives_interval": 10, # 每 10 秒探测一次
|
||
"keepalives_count": 3, # 连续 3 次失败判定断开
|
||
}
|
||
|
||
# 连接重试参数：应对 PostgreSQL 瞬时不可用（Tailscale 网络抖动等）
|
||
_CONNECT_MAX_RETRIES = 3
|
||
_CONNECT_RETRY_DELAY = 1.0 # 秒
|
||
|
||
|
||
def _connect_with_retry(connect_fn, max_retries=_CONNECT_MAX_RETRIES):
|
||
"""带重试的数据库连接，应对服务端瞬时拒绝连接。"""
|
||
last_exc = None
|
||
for attempt in range(max_retries):
|
||
try:
|
||
return connect_fn()
|
||
except psycopg2.OperationalError as e:
|
||
last_exc = e
|
||
if attempt < max_retries - 1:
|
||
time.sleep(_CONNECT_RETRY_DELAY * (attempt + 1))
|
||
raise last_exc
|
||
|
||
|
||
def get_connection() -> PgConnection:
|
||
"""
|
||
获取 zqyy_app 数据库连接。
|
||
|
||
调用方负责关闭连接（推荐配合 contextmanager 或 try/finally 使用）。
|
||
当 trace 启用且有活跃 TraceContext 时，返回 TracedConnection 包装，
|
||
自动记录 DB_CONN span（连接获取耗时）并拦截后续 SQL 执行。
|
||
"""
|
||
# CHANGE 2026-03-22 | task 8.2 | 集成 trace db_wrapper，仅 trace 启用时包装
|
||
from app.trace.config import get_trace_config
|
||
from app.trace.context import SpanType, TraceSpan, get_current_trace
|
||
|
||
config = get_trace_config()
|
||
should_trace = config.enabled and get_current_trace() is not None
|
||
|
||
start = time.perf_counter() if should_trace else 0.0
|
||
|
||
conn = _connect_with_retry(lambda: psycopg2.connect(
|
||
host=DB_HOST,
|
||
port=DB_PORT,
|
||
user=DB_USER,
|
||
password=DB_PASSWORD,
|
||
dbname=APP_DB_NAME,
|
||
**_KEEPALIVE_KWARGS,
|
||
))
|
||
|
||
if should_trace:
|
||
from datetime import datetime
|
||
from app.trace.db_wrapper import traced_connection
|
||
|
||
elapsed_ms = (time.perf_counter() - start) * 1000
|
||
ctx = get_current_trace()
|
||
# ctx 不为 None（上面已检查）
|
||
ctx.add_span(TraceSpan(
|
||
span_type=SpanType.DB_CONN,
|
||
module="app.database",
|
||
function="get_connection",
|
||
description_zh=f"获取数据库连接，耗时 {elapsed_ms:.1f}ms",
|
||
description_en=f"Acquired database connection in {elapsed_ms:.1f}ms",
|
||
params={},
|
||
result_summary=f"{elapsed_ms:.1f}ms",
|
||
duration_ms=elapsed_ms,
|
||
timestamp=datetime.now().isoformat(),
|
||
))
|
||
return traced_connection(conn)
|
||
|
||
return conn
|
||
|
||
|
||
def get_etl_global_readonly_connection() -> PgConnection:
|
||
"""
|
||
获取 ETL 数据库的全局只读连接（不设 RLS）。
|
||
|
||
用于系统管理后台等不需要门店隔离的场景（如 ETL 状态监控）。
|
||
"""
|
||
conn = _connect_with_retry(lambda: psycopg2.connect(
|
||
host=ETL_DB_HOST,
|
||
port=ETL_DB_PORT,
|
||
user=ETL_DB_USER,
|
||
password=ETL_DB_PASSWORD,
|
||
dbname=ETL_DB_NAME,
|
||
**_KEEPALIVE_KWARGS,
|
||
))
|
||
try:
|
||
conn.autocommit = False
|
||
with conn.cursor() as cur:
|
||
cur.execute("SET default_transaction_read_only = on")
|
||
conn.commit()
|
||
except Exception:
|
||
conn.close()
|
||
raise
|
||
return conn
|
||
|
||
|
||
def get_etl_readonly_connection(site_id: int | str) -> PgConnection:
|
||
"""
|
||
获取 ETL 数据库（etl_feiqiu）的只读连接。
|
||
|
||
连接建立后自动执行：
|
||
1. SET default_transaction_read_only = on — 禁止写操作
|
||
2. SET LOCAL app.current_site_id = '{site_id}' — 启用 RLS 门店隔离
|
||
|
||
调用方负责关闭连接。典型用法::
|
||
|
||
conn = get_etl_readonly_connection(site_id)
|
||
try:
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT ...")
|
||
finally:
|
||
conn.close()
|
||
"""
|
||
conn = _connect_with_retry(lambda: psycopg2.connect(
|
||
host=ETL_DB_HOST,
|
||
port=ETL_DB_PORT,
|
||
user=ETL_DB_USER,
|
||
password=ETL_DB_PASSWORD,
|
||
dbname=ETL_DB_NAME,
|
||
**_KEEPALIVE_KWARGS,
|
||
))
|
||
try:
|
||
conn.autocommit = False
|
||
with conn.cursor() as cur:
|
||
# 会话级只读：防止任何写操作
|
||
cur.execute("SET default_transaction_read_only = on")
|
||
# 事务级 RLS 隔离：设置当前门店 ID
|
||
cur.execute(
|
||
"SET LOCAL app.current_site_id = %s", (str(site_id),)
|
||
)
|
||
conn.commit()
|
||
except Exception:
|
||
conn.close()
|
||
raise
|
||
return conn
|
||
|
||
|
||
def get_etl_write_connection() -> PgConnection:
|
||
"""
|
||
获取 ETL 数据库（etl_feiqiu）的可写连接。
|
||
|
||
仅用于后端需要回写 ETL 汇总表的场景（如 task_generator 回写关系指数统计）。
|
||
不设置 RLS 隔离，调用方需在 SQL 中显式指定 site_id。
|
||
调用方负责关闭连接。
|
||
"""
|
||
conn = _connect_with_retry(lambda: psycopg2.connect(
|
||
host=ETL_DB_HOST,
|
||
port=ETL_DB_PORT,
|
||
user=ETL_DB_USER,
|
||
password=ETL_DB_PASSWORD,
|
||
dbname=ETL_DB_NAME,
|
||
**_KEEPALIVE_KWARGS,
|
||
))
|
||
conn.autocommit = False
|
||
return conn
|