Spaces:
Running
Running
File size: 4,402 Bytes
89157f5 55ae875 89157f5 55ae875 89157f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | from __future__ import annotations
import asyncio
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any
from app.core.logger import get_logger
_logger = get_logger(__name__)
@dataclass(frozen=True)
class ConnectionConfig:
db_type: str
host: str
port: int
database: str
username: str
password: str
selected_schema: str = "public"
ssl_enabled: bool = False
ssl_ca_cert: str | None = None
ssl_cert: str | None = None
ssl_key: str | None = None
query_timeout_seconds: float = 30.0
connection_timeout_seconds: float = 10.0
max_rows: int = 10000
@property
def pool_key(self) -> str:
return f"{self.db_type}:{self.username}@{self.host}:{self.port}/{self.database}"
@property
def safe_repr(self) -> str:
return (
f"ConnectionConfig(db_type={self.db_type}, host={self.host}, "
f"port={self.port}, database={self.database}, "
f"selected_schema={self.selected_schema}, username={self.username}, ssl={self.ssl_enabled})"
)
@dataclass
class StatementResult:
success: bool
rows: int = 0
data: list[dict[str, Any]] = field(default_factory=list)
error: str | None = None
error_code: str | None = None
class BaseExecutor(ABC):
def __init__(self, config: ConnectionConfig) -> None:
self._config = config
self._pool: Any = None
self._closed = False
self._lock = asyncio.Lock()
self._max_connection_retries = 3
@abstractmethod
async def _create_pool(self) -> Any:
...
@abstractmethod
async def _execute_queries(
self, pool: Any, queries: list[Any], use_transaction: bool
) -> list[StatementResult]:
...
async def execute(
self, queries: list[Any], use_transaction: bool = True
) -> list[StatementResult]:
if self._closed:
raise RuntimeError("Executor has been closed")
pool = await self._get_or_create_pool()
return await self._execute_queries(pool, queries, use_transaction)
async def _get_or_create_pool(self) -> Any:
async with self._lock:
if self._pool is not None:
return self._pool
last_exc: Exception | None = None
for attempt in range(self._max_connection_retries):
try:
self._pool = await asyncio.wait_for(
self._create_pool(),
timeout=self._config.connection_timeout_seconds,
)
_logger.info(
"Created pool for %s (attempt %d)",
self._config.safe_repr, attempt + 1,
)
return self._pool
except asyncio.TimeoutError:
last_exc = TimeoutError(
f"Connection timed out after {self._config.connection_timeout_seconds}s"
)
_logger.warning(
"Pool creation timeout for %s (attempt %d)",
self._config.safe_repr, attempt + 1,
)
except Exception as exc:
last_exc = exc
_logger.warning(
"Pool creation failed for %s (attempt %d): %s",
self._config.safe_repr, attempt + 1, exc,
)
if attempt < self._max_connection_retries - 1:
wait = 0.1 * (2**attempt)
await asyncio.sleep(wait)
msg = (
f"Failed to create connection pool after {self._max_connection_retries} attempts"
)
if last_exc is not None:
msg += f": {last_exc}"
raise RuntimeError(msg) from last_exc
async def close(self) -> None:
async with self._lock:
if self._closed:
return
self._closed = True
if self._pool is not None:
await self._close_pool(self._pool)
self._pool = None
_logger.info("Closed pool for %s", self._config.safe_repr)
@abstractmethod
async def _close_pool(self, pool: Any) -> None:
...
@property
def is_closed(self) -> bool:
return self._closed
|