218 lines
6.8 KiB
Python
218 lines
6.8 KiB
Python
# Oracle 连接池。thick client 进程内只初始化一次。
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from contextlib import contextmanager
|
||
from collections.abc import Iterator
|
||
|
||
from .errors import RepoUnavailable
|
||
|
||
_client_ready = False
|
||
# 班次池只有一条连接。空闲达到这个秒数才探活，坏了就丢掉再借。
|
||
PING_INTERVAL = 60
|
||
|
||
|
||
class OraclePool:
|
||
"""oracledb 连接池的薄封装，借还连接、对外异常固定文案。"""
|
||
|
||
def __init__(self, raw_pool):
|
||
"""包住 oracledb.ConnectionPool。调用方不要直接碰 raw_pool。"""
|
||
self._raw = raw_pool
|
||
|
||
@contextmanager
|
||
def acquire(self) -> Iterator:
|
||
"""借一条活连接。正常退出归还；未提交失败则丢掉，不放回池里。
|
||
|
||
探活失败或执行前已死时丢掉，再借一次。班次池只有一条，靠探活和丢弃换连接。
|
||
借还失败时抛 RepoUnavailable，消息不含 DSN。
|
||
"""
|
||
conn = self._borrow_live()
|
||
try:
|
||
yield conn
|
||
except Exception:
|
||
self._drop(conn)
|
||
raise
|
||
else:
|
||
self._release(conn)
|
||
|
||
def run(self, fn) -> None:
|
||
"""执行一笔。未提交失败则丢掉连接，再借一次，只重试这一笔。
|
||
|
||
fn 收到连接。这一笔里已经调用过 commit 时不再重试，已提交的语句不重做。
|
||
班次池只有一条，靠探活和丢弃换连接。
|
||
"""
|
||
conn = self._borrow_live()
|
||
committed = {"done": False}
|
||
try:
|
||
_invoke(conn, fn, committed)
|
||
except Exception:
|
||
self._drop(conn)
|
||
if committed["done"]:
|
||
raise
|
||
else:
|
||
self._release(conn)
|
||
return
|
||
retry = self._borrow_live()
|
||
try:
|
||
_invoke(retry, fn, {"done": False})
|
||
except Exception:
|
||
self._drop(retry)
|
||
raise
|
||
else:
|
||
self._release(retry)
|
||
|
||
def _borrow_live(self):
|
||
"""借一条连接。探活失败或执行前已死则丢掉，再借一次。"""
|
||
try:
|
||
conn = self._raw.acquire()
|
||
except Exception as exc:
|
||
raise RepoUnavailable("数据库不可用") from exc
|
||
if _alive(conn):
|
||
return conn
|
||
self._drop(conn)
|
||
try:
|
||
conn = self._raw.acquire()
|
||
except Exception as exc:
|
||
raise RepoUnavailable("数据库不可用") from exc
|
||
if _alive(conn):
|
||
return conn
|
||
self._drop(conn)
|
||
raise RepoUnavailable("数据库不可用")
|
||
|
||
def _drop(self, conn) -> None:
|
||
"""丢掉坏连接。池上没有 drop 时退回 close，避免占住唯一的槽。"""
|
||
drop = getattr(self._raw, "drop", None)
|
||
if drop is not None:
|
||
try:
|
||
drop(conn)
|
||
return
|
||
except Exception:
|
||
pass
|
||
self._release(conn)
|
||
|
||
def _release(self, conn) -> None:
|
||
"""把连接归还池里。归还失败不再抛出。"""
|
||
close = getattr(conn, "close", None)
|
||
if close is None:
|
||
return
|
||
try:
|
||
close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# 返回连接池下限和上限。
|
||
#
|
||
# oms 默认 1–4，班次（job）默认 1–1。RPA_ORACLE_POOL_MIN / MAX 可覆盖。
|
||
# max 小于 min 时抬到与 min 相同。
|
||
def pool_limits(role: str = "oms") -> tuple[int, int]:
|
||
if role == "job":
|
||
default_min, default_max = 1, 1
|
||
else:
|
||
default_min, default_max = 1, 4
|
||
try:
|
||
min_size = int(os.getenv("RPA_ORACLE_POOL_MIN", str(default_min)))
|
||
max_size = int(os.getenv("RPA_ORACLE_POOL_MAX", str(default_max)))
|
||
except ValueError as exc:
|
||
raise RepoUnavailable("数据库不可用") from exc
|
||
if max_size < min_size:
|
||
max_size = min_size
|
||
return min_size, max_size
|
||
|
||
|
||
# 建班次连接池。默认大小 1。缺 HOST / 用户 / 密码抛「数据库不可用」。
|
||
def job_pool() -> OraclePool:
|
||
return create_pool(role="job")
|
||
|
||
|
||
def create_pool(*, role: str = "oms") -> OraclePool:
|
||
"""按环境变量建一个连接池。
|
||
|
||
缺 HOST / USER / PASSWORD、thick 初始化失败或建池失败时，
|
||
抛 RepoUnavailable「数据库不可用」，消息不含 DSN 和库路径。
|
||
ping_interval 为 60 秒。班次池只有一条连接，靠探活和丢弃换连接，大小仍是 min=1, max=1。
|
||
"""
|
||
user, password, dsn = _credentials()
|
||
_ensure_client()
|
||
min_size, max_size = pool_limits(role)
|
||
try:
|
||
import oracledb
|
||
|
||
raw = oracledb.create_pool(
|
||
user=user,
|
||
password=password,
|
||
dsn=dsn,
|
||
min=min_size,
|
||
max=max_size,
|
||
increment=1,
|
||
ping_interval=PING_INTERVAL,
|
||
)
|
||
except RepoUnavailable:
|
||
raise
|
||
except Exception as exc:
|
||
raise RepoUnavailable("数据库不可用") from exc
|
||
return OraclePool(raw)
|
||
|
||
|
||
# 读用户、密码、DSN。缺任一项抛固定文案，不把 host 拼进消息。
|
||
def _credentials() -> tuple[str, str, str]:
|
||
user = os.getenv("RPA_ORACLE_USER", "")
|
||
password = os.getenv("RPA_ORACLE_PASSWORD", "")
|
||
host = os.getenv("RPA_ORACLE_HOST", "")
|
||
if not user or not password or not host:
|
||
raise RepoUnavailable("数据库不可用")
|
||
port = os.getenv("RPA_ORACLE_PORT", "1521")
|
||
service = os.getenv("RPA_ORACLE_SERVICE", "rpa")
|
||
return user, password, f"{host}:{port}/{service}"
|
||
|
||
|
||
def _invoke(conn, fn, committed: dict) -> None:
|
||
"""调用 fn。包住 commit，记下这一笔是否已经提交。"""
|
||
original = getattr(conn, "commit", None)
|
||
|
||
def commit():
|
||
"""先记下已提交，再执行原来的 commit。失败也不再重做这一笔。"""
|
||
committed["done"] = True
|
||
if original is not None:
|
||
original()
|
||
|
||
if callable(original):
|
||
conn.commit = commit
|
||
try:
|
||
fn(conn)
|
||
finally:
|
||
if callable(original):
|
||
conn.commit = original
|
||
|
||
|
||
def _alive(conn) -> bool:
|
||
"""执行前探活。没有 ping 时看 dead 标记，两者都没有视为可用。"""
|
||
if getattr(conn, "dead", False):
|
||
return False
|
||
ping = getattr(conn, "ping", None)
|
||
if not callable(ping):
|
||
return True
|
||
try:
|
||
ping()
|
||
except Exception:
|
||
return False
|
||
return True
|
||
|
||
|
||
# 进程内只调用一次 init_oracle_client。失败时文案不含 Instant Client 路径。
|
||
def _ensure_client() -> None:
|
||
global _client_ready
|
||
if _client_ready:
|
||
return
|
||
import oracledb
|
||
|
||
if os.getenv("RPA_ORACLE_THICK", "1") == "1":
|
||
lib_dir = os.getenv("RPA_ORACLE_CLIENT_LIB", "")
|
||
kwargs = {"lib_dir": lib_dir} if lib_dir else {}
|
||
try:
|
||
oracledb.init_oracle_client(**kwargs)
|
||
except Exception as exc:
|
||
raise RepoUnavailable("数据库不可用") from exc
|
||
_client_ready = True
|