pz-cli 5.0.0__py3-none-any.whl

pz_foundation/db/postgresql.py

@@ -0,0 +1,409 @@
+from __future__ import annotations
+
+import contextlib
+from dataclasses import dataclass
+import importlib
+import re
+from typing import Any
+from typing import Optional
+from urllib.parse import quote, urlencode
+
+
+_SEARCH_PATH_OPTION_RE = re.compile(r"search_path=([A-Za-z0-9_]+)")
+
+
+@dataclass(frozen=True, slots=True)
+class ClusterNode:
+    host: str
+    port: int
+    role: str = "unknown"
+
+
+class PostgreSQLClusterManager:
+    def __init__(
+        self,
+        *,
+        nodes: list[str],
+        db_user: str,
+        db_password: str,
+        db_name: str,
+        ssl_mode: str = "",
+        ssl_root_cert: str = "",
+        ssl_cert: str = "",
+        ssl_key: str = "",
+        ssl_password: str = "",
+    ) -> None:
+        self.nodes = list(nodes)
+        self.db_user = db_user
+        self.db_password = db_password
+        self.db_name = db_name
+        self.ssl_mode = ssl_mode
+        self.ssl_root_cert = ssl_root_cert
+        self.ssl_cert = ssl_cert
+        self.ssl_key = ssl_key
+        self.ssl_password = ssl_password
+
+    async def find_master(self) -> ClusterNode | None:
+        for node in self.nodes:
+            candidate = self._probe_node(node)
+            if candidate is not None and candidate.role == "master":
+                return candidate
+        if not self.nodes:
+            return None
+        host, port = _split_node(self.nodes[0])
+        return ClusterNode(host=host, port=port, role="master")
+
+    async def find_read_replica(self) -> ClusterNode | None:
+        for node in self.nodes:
+            candidate = self._probe_node(node)
+            if candidate is not None and candidate.role == "replica":
+                return candidate
+        if len(self.nodes) < 2:
+            return None
+        host, port = _split_node(self.nodes[1])
+        return ClusterNode(host=host, port=port, role="replica")
+
+    def connection_urls(self) -> dict[str, str]:
+        out: dict[str, str] = {}
+        for idx, node in enumerate(self.nodes):
+            host, port = _split_node(node)
+            label = "default" if idx == 0 else f"read_only_{idx}"
+            out[label] = _build_postgresql_url(
+                host=host,
+                port=port,
+                user=self.db_user,
+                password=self.db_password,
+                database=self.db_name,
+                ssl_mode=self.ssl_mode,
+                ssl_root_cert=self.ssl_root_cert,
+                ssl_cert=self.ssl_cert,
+                ssl_key=self.ssl_key,
+                ssl_password=self.ssl_password,
+            )
+        return out
+
+    def _probe_node(self, node: str) -> ClusterNode | None:
+        try:
+            psycopg = importlib.import_module("psycopg")
+        except Exception:
+            return None
+        host, port = _split_node(node)
+        conn = None
+        try:
+            conn = psycopg.connect(
+                f"postgresql://{self.db_user}:{self.db_password}@{host}:{port}/{self.db_name}",
+                connect_timeout=2,
+                autocommit=True,
+                **_build_postgresql_ssl_credentials(self),
+            )
+            cursor = conn.cursor()
+            try:
+                cursor.execute("SELECT pg_is_in_recovery()")
+                row = cursor.fetchone()
+            finally:
+                with contextlib.suppress(Exception):
+                    cursor.close()
+            in_recovery = bool(row[0]) if row else False
+            return ClusterNode(host=host, port=port, role="replica" if in_recovery else "master")
+        except Exception:
+            return None
+        finally:
+            if conn is not None:
+                with contextlib.suppress(Exception):
+                    conn.close()
+
+
+def build_postgresql_connections(settings) -> dict[str, Any]:
+    hosts, ports = _split_hosts_ports(settings)
+    runtime_credentials = _build_postgresql_runtime_credentials(settings)
+    writer_target_session_attrs = _writer_target_session_attrs(settings, hosts)
+    writer_conninfo = _build_writer_conninfo(settings, hosts, ports)
+    default_credentials = {
+        "host": hosts[0],
+        "port": int(ports[0] or 5432),
+        "user": getattr(settings, "db_user", ""),
+        "password": getattr(settings, "db_password", ""),
+        "database": getattr(settings, "db_name", ""),
+        **runtime_credentials,
+    }
+    if writer_target_session_attrs:
+        default_credentials["target_session_attrs"] = writer_target_session_attrs
+    if writer_conninfo:
+        default_credentials["conninfo"] = writer_conninfo
+    base = {
+        "engine": "pz_foundation.db.tortoise_psycopg",
+        "credentials": default_credentials,
+    }
+    out: dict[str, Any] = {"default": base}
+    if len(hosts) >= 2 and not writer_conninfo:
+        read_only_credentials = dict(default_credentials)
+        read_only_credentials.pop("conninfo", None)
+        read_only_credentials.pop("target_session_attrs", None)
+        read_only_credentials["host"] = hosts[1]
+        read_only_credentials["port"] = int(ports[1] or ports[0] or 5432)
+        out["read_only"] = {
+            "engine": "pz_foundation.db.tortoise_psycopg",
+            "credentials": read_only_credentials,
+        }
+    return out
+
+
+def _build_writer_conninfo(settings: Any, hosts: list[str] | tuple[str, ...], ports: list[str] | tuple[str, ...]) -> str:
+    target_session_attrs = _writer_target_session_attrs(settings, hosts)
+    if len(hosts) <= 1 and not target_session_attrs and not tuple(getattr(settings, "db_urls", ()) or ()):
+        return ""
+    query = _conninfo_query(settings, target_session_attrs=target_session_attrs)
+    port_values = _expand_ports(ports, count=len(hosts), fallback=5432)
+    return _build_postgresql_url(
+        host=",".join(_format_node_host(host) for host in hosts),
+        port=0,
+        user=getattr(settings, "db_user", ""),
+        password=getattr(settings, "db_password", ""),
+        database=getattr(settings, "db_name", ""),
+        query=query,
+        host_port_pairs=tuple(zip(hosts, port_values)),
+    )
+
+
+def _writer_target_session_attrs(settings: Any, hosts: list[str] | tuple[str, ...]) -> str:
+    target_session_attrs = str(getattr(settings, "target_session_attrs", "") or "").strip()
+    if not target_session_attrs and len(hosts) > 1:
+        target_session_attrs = "read-write"
+    return target_session_attrs
+
+
+def _conninfo_query(settings: Any, *, target_session_attrs: str) -> dict[str, str]:
+    query: dict[str, str] = {}
+    for key, value in _build_postgresql_ssl_credentials(settings).items():
+        query[key] = str(value)
+    if target_session_attrs:
+        query["target_session_attrs"] = target_session_attrs
+    for key, value in dict(getattr(settings, "options", {}) or {}).items():
+        if key == "options" or key in query:
+            continue
+        if isinstance(value, (str, int, float, bool)) and str(value).strip():
+            query[key] = str(value)
+    return query
+
+
+async def find_cluster_master(
+    nodes: list[str],
+    db_user: str,
+    db_password: str,
+    db_name: str,
+    ssl_mode: str = "",
+    ssl_root_cert: str = "",
+    ssl_cert: str = "",
+    ssl_key: str = "",
+    ssl_password: str = "",
+) -> Optional[tuple[str, int]]:
+    manager = PostgreSQLClusterManager(
+        nodes=nodes,
+        db_user=db_user,
+        db_password=db_password,
+        db_name=db_name,
+        ssl_mode=ssl_mode,
+        ssl_root_cert=ssl_root_cert,
+        ssl_cert=ssl_cert,
+        ssl_key=ssl_key,
+        ssl_password=ssl_password,
+    )
+    master = await manager.find_master()
+    if master is None:
+        return None
+    return master.host, master.port
+
+
+def _split_node(node: str) -> tuple[str, int]:
+    text = str(node or "").strip()
+    if not text:
+        return "", 5432
+    if text.startswith("["):
+        closing = text.find("]")
+        if closing > 0:
+            host = text[1:closing].strip()
+            remainder = text[closing + 1 :].strip()
+            if remainder.startswith(":"):
+                try:
+                    return host, int(remainder[1:] or 5432)
+                except Exception:
+                    return host, 5432
+            return host, 5432
+    if text.count(":") == 1:
+        host, raw_port = text.rsplit(":", 1)
+        try:
+            return host.strip(), int(raw_port or 5432)
+        except Exception:
+            return host.strip(), 5432
+    return text, 5432
+
+
+def _format_node_host(host: str) -> str:
+    value = str(host or "").strip()
+    if ":" in value and not (value.startswith("[") and value.endswith("]")):
+        return f"[{value}]"
+    return value
+
+
+def _pool_credentials(settings) -> dict[str, Any]:
+    minsize = max(1, int(getattr(settings, "db_min_size", 2) or 1))
+    maxsize = max(minsize, int(getattr(settings, "db_max_size", 10) or minsize))
+    creds: dict[str, Any] = {"minsize": minsize, "maxsize": maxsize}
+    try:
+        acquire_timeout = float(getattr(settings, "db_pool_acquire_timeout", 0) or 0)
+    except Exception:
+        acquire_timeout = 0
+    try:
+        connect_timeout = float(getattr(settings, "db_connection_timeout", 0) or 0)
+    except Exception:
+        connect_timeout = 0
+    if acquire_timeout > 0:
+        creds["timeout"] = acquire_timeout
+    if connect_timeout > 0:
+        creds["connect_timeout"] = int(max(1, connect_timeout))
+    target_session_attrs = str(getattr(settings, "target_session_attrs", "") or "").strip()
+    if target_session_attrs:
+        creds["target_session_attrs"] = target_session_attrs
+    for source_name, credential_name in (
+        ("db_pool_max_lifetime", "max_lifetime"),
+        ("db_pool_max_idle", "max_idle"),
+        ("db_reconnect_timeout", "reconnect_timeout"),
+    ):
+        try:
+            value = float(getattr(settings, source_name, 0) or 0)
+        except Exception:
+            value = 0
+        if value > 0:
+            creds[credential_name] = value
+    return creds
+
+
+def _build_postgresql_ssl_credentials(settings) -> dict[str, Any]:
+    credentials: dict[str, Any] = {}
+    ssl_mode = str(getattr(settings, "ssl_mode", "") or "").strip()
+    ssl_root_cert = str(getattr(settings, "ssl_root_cert", "") or "").strip()
+    ssl_cert = str(getattr(settings, "ssl_cert", "") or "").strip()
+    ssl_key = str(getattr(settings, "ssl_key", "") or "").strip()
+    ssl_password = str(getattr(settings, "ssl_password", "") or "").strip()
+    if ssl_mode:
+        credentials["sslmode"] = ssl_mode
+    if ssl_root_cert:
+        credentials["sslrootcert"] = ssl_root_cert
+    if ssl_cert:
+        credentials["sslcert"] = ssl_cert
+    if ssl_key:
+        credentials["sslkey"] = ssl_key
+    if ssl_password:
+        credentials["sslpassword"] = ssl_password
+    return credentials
+
+
+def _build_postgresql_runtime_credentials(settings) -> dict[str, Any]:
+    credentials: dict[str, Any] = {
+        **_pool_credentials(settings),
+        **_build_postgresql_ssl_credentials(settings),
+    }
+    option_tokens: list[str] = []
+    statement_timeout_ms = _positive_int(getattr(settings, "db_statement_timeout_ms", 0), 0)
+    idle_transaction_timeout_ms = _positive_int(
+        getattr(settings, "db_idle_in_transaction_session_timeout_ms", 0),
+        0,
+    )
+    if statement_timeout_ms > 0:
+        option_tokens.extend(("-c", f"statement_timeout={statement_timeout_ms}"))
+    if idle_transaction_timeout_ms > 0:
+        option_tokens.extend(("-c", f"idle_in_transaction_session_timeout={idle_transaction_timeout_ms}"))
+    option_payload = dict(getattr(settings, "options", {}) or {})
+    schema = _search_path_schema(str(option_payload.get("options", "") or ""))
+    if schema:
+        option_tokens.extend(("-c", f"search_path={schema}"))
+    if option_tokens:
+        credentials["server_options"] = " ".join(option_tokens)
+    return credentials
+
+
+def _positive_int(value: Any, default: int) -> int:
+    try:
+        parsed = int(value)
+    except Exception:
+        return default
+    return parsed if parsed > 0 else default
+
+
+def _search_path_schema(raw_options: str) -> str:
+    match = _SEARCH_PATH_OPTION_RE.search(str(raw_options or "").strip())
+    if match is None:
+        return ""
+    return str(match.group(1) or "").strip()
+
+
+def _split_hosts_ports(settings) -> tuple[list[str], list[str]]:
+    hosts = [h.strip() for h in str(getattr(settings, "db_host", "") or "").split(",") if h.strip()]
+    ports = [p.strip() for p in str(getattr(settings, "db_port", "") or "").split(",") if p.strip()]
+    if not hosts:
+        raise ValueError("foundation postgresql host is not configured")
+    if not ports:
+        ports = [str(getattr(settings, "db_port", 5432) or 5432)]
+    if len(ports) < len(hosts):
+        ports = [ports[0]] * len(hosts)
+    return hosts, ports
+
+
+def _build_postgresql_url(
+    *,
+    host: str,
+    port: int,
+    user: str,
+    password: str,
+    database: str,
+    query: dict[str, str] | None = None,
+    host_port_pairs: tuple[tuple[str, int], ...] = (),
+    ssl_mode: str = "",
+    ssl_root_cert: str = "",
+    ssl_cert: str = "",
+    ssl_key: str = "",
+    ssl_password: str = "",
+) -> str:
+    auth = ""
+    if user or password:
+        auth = f"{quote(user or '', safe='')}:{quote(password or '', safe='')}@"
+    if host_port_pairs:
+        hostspec = ",".join(f"{_format_node_host(item_host)}:{int(item_port or 5432)}" for item_host, item_port in host_port_pairs)
+    else:
+        hostspec = f"{_format_node_host(host)}:{int(port or 5432)}"
+    base = f"postgresql://{auth}{hostspec}/{quote(database or '', safe='')}"
+    query = dict(query or {})
+    if ssl_mode:
+        query["sslmode"] = ssl_mode
+    if ssl_root_cert:
+        query["sslrootcert"] = ssl_root_cert
+    if ssl_cert:
+        query["sslcert"] = ssl_cert
+    if ssl_key:
+        query["sslkey"] = ssl_key
+    if ssl_password:
+        query["sslpassword"] = ssl_password
+    if not query:
+        return base
+    return f"{base}?{urlencode(query)}"
+
+
+def _expand_ports(ports: list[str] | tuple[str, ...], *, count: int, fallback: int) -> tuple[int, ...]:
+    parsed: list[int] = []
+    for item in ports:
+        try:
+            value = int(item)
+        except Exception:
+            continue
+        if value > 0:
+            parsed.append(value)
+    if len(parsed) >= count:
+        return tuple(parsed[:count])
+    if len(parsed) == 1:
+        return tuple(parsed[0] for _ in range(count))
+    if parsed:
+        return tuple([*parsed, *(parsed[0] for _ in range(count - len(parsed)))])
+    return tuple(int(fallback) for _ in range(count))
+
+
+__all__ = ["ClusterNode", "PostgreSQLClusterManager", "build_postgresql_connections", "find_cluster_master"]

pz_foundation/db/production_pool_fallback.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable
+
+import pz_settings_center.public as settings_center
+
+from ..observability import get_logger
+
+
+logger = get_logger(__name__)
+
+Reconnect = Callable[[str], Awaitable[None]]
+LifecycleReader = Callable[[], Any]
+
+
+@dataclass(frozen=True, slots=True)
+class ProductionPoolFallbackResult:
+    attempted: bool
+    recovered: bool
+    reason: str
+    lifecycle: Any | None = None
+    error: BaseException | None = None
+
+
+def production_pool_fallback_enabled() -> bool:
+    try:
+        pz_env = settings_center.current_pz_env(default="develop")
+    except Exception:
+        return False
+    return str(pz_env or "").strip().lower() == "product"
+
+
+def should_apply_production_pool_replenishment_fallback(
+    *,
+    diagnosis: dict[str, Any],
+    independent_probe: dict[str, Any],
+) -> tuple[bool, str]:
+    if not production_pool_fallback_enabled():
+        return False, "not_product_env"
+    if diagnosis.get("category") != "pool_replenishment_stalled":
+        return False, "category_not_pool_replenishment_stalled"
+    if independent_probe.get("attempted") and independent_probe.get("reachable") is False:
+        return False, "database_not_reachable"
+    if independent_probe.get("attempted") and independent_probe.get("status") not in (None, "ok"):
+        return False, "independent_probe_not_ok"
+    return True, "enabled"
+
+
+async def maybe_apply_production_pool_replenishment_fallback(
+    *,
+    alias: str,
+    error: BaseException,
+    lifecycle: Any,
+    pool_status: dict[str, Any],
+    independent_probe: dict[str, Any],
+    diagnosis: dict[str, Any],
+    reconnect: Reconnect,
+    lifecycle_reader: LifecycleReader,
+) -> ProductionPoolFallbackResult:
+    should_apply, reason = should_apply_production_pool_replenishment_fallback(
+        diagnosis=diagnosis,
+        independent_probe=independent_probe,
+    )
+    if not should_apply:
+        return ProductionPoolFallbackResult(attempted=False, recovered=False, reason=reason)
+
+    logger.error(
+        "生产数据库连接池强兜底已触发: alias=%s error=%s lifecycle=%s pool_status=%s "
+        "independent_probe=%s diagnosis=%s action=alias_reconnect",
+        alias,
+        _safe_message(error),
+        lifecycle,
+        pool_status,
+        independent_probe,
+        diagnosis,
+    )
+
+    try:
+        await reconnect(alias)
+        rebuilt_lifecycle = lifecycle_reader()
+    except Exception as exc:  # noqa: BLE001
+        logger.error(
+            "生产数据库连接池强兜底重建失败: alias=%s diagnosis=%s",
+            alias,
+            diagnosis,
+            exc_info=True,
+        )
+        return ProductionPoolFallbackResult(
+            attempted=True,
+            recovered=False,
+            reason="reconnect_failed",
+            error=exc,
+        )
+
+    if not _alias_ready(rebuilt_lifecycle, alias):
+        logger.error(
+            "生产数据库连接池强兜底重建后生命周期仍不可用: alias=%s lifecycle=%s diagnosis=%s",
+            alias,
+            rebuilt_lifecycle,
+            diagnosis,
+        )
+        return ProductionPoolFallbackResult(
+            attempted=True,
+            recovered=False,
+            reason="lifecycle_not_ready",
+            lifecycle=rebuilt_lifecycle,
+        )
+
+    return ProductionPoolFallbackResult(
+        attempted=True,
+        recovered=True,
+        reason="recovered",
+        lifecycle=rebuilt_lifecycle,
+    )
+
+
+def _alias_ready(lifecycle: Any, alias: str) -> bool:
+    alias_status = getattr(lifecycle, "alias_status", None)
+    if callable(alias_status):
+        try:
+            return alias_status(alias) == "ok"
+        except Exception:
+            return False
+    if alias == "default":
+        return bool(getattr(lifecycle, "default_ready", False))
+    if alias == "read_only":
+        return bool(getattr(lifecycle, "read_only_ready", False))
+    return False
+
+
+def _safe_message(error: BaseException) -> str:
+    text = str(error or "").replace("\n", " ").strip()
+    if len(text) > 300:
+        return f"{text[:300]}..."
+    return text
+
+
+__all__ = [
+    "ProductionPoolFallbackResult",
+    "maybe_apply_production_pool_replenishment_fallback",
+    "production_pool_fallback_enabled",
+    "should_apply_production_pool_replenishment_fallback",
+]

pz_foundation/db/public.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from . import (
+    describe_foundation_database_binding,
+    connection_lifecycle_status,
+    DBConnectionLifecycleStatus,
+    DBConnectionReadinessProbeResult,
+    FoundationDatabase,
+    FoundationDatabaseSettings,
+    MySQLFoundationDatabase,
+    normalize_foundation_database_settings,
+    PostgreSQLFoundationDatabase,
+    probe_connection_readiness,
+    probe_foundation_database_binding,
+    SQLiteFoundationDatabase,
+    get_foundation_database,
+    get_foundation_database_if_enabled,
+    load_foundation_database_settings,
+    get_connection_status,
+    tortoise_manager,
+)
+
+STABLE_EXPORTS: tuple[str, ...] = (
+    "describe_foundation_database_binding",
+    "connection_lifecycle_status",
+    "DBConnectionLifecycleStatus",
+    "DBConnectionReadinessProbeResult",
+    "FoundationDatabase",
+    "FoundationDatabaseSettings",
+    "MySQLFoundationDatabase",
+    "normalize_foundation_database_settings",
+    "PostgreSQLFoundationDatabase",
+    "probe_connection_readiness",
+    "probe_foundation_database_binding",
+    "SQLiteFoundationDatabase",
+    "get_connection_status",
+    "get_foundation_database",
+    "get_foundation_database_if_enabled",
+    "load_foundation_database_settings",
+    "tortoise_manager",
+)
+
+COMPAT_EXPORTS: tuple[str, ...] = ()
+
+__all__ = [*STABLE_EXPORTS, *COMPAT_EXPORTS]