mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-09 01:50:44 +08:00
82 lines
2.3 KiB
Python
82 lines
2.3 KiB
Python
"""psycopg connection surface compatible with sqlite3 usage in SqliteStore."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import Any, Iterable, Sequence
|
|
|
|
import psycopg
|
|
from psycopg.rows import dict_row
|
|
|
|
from svc.persistence.pg_compat import adapt_sql_for_postgres
|
|
|
|
|
|
def normalize_psycopg_conninfo(url: str) -> str:
|
|
"""Strip SQLAlchemy driver suffix so :func:`psycopg.connect` accepts the URI."""
|
|
u = str(url or "").strip()
|
|
for prefix in (
|
|
"postgresql+psycopg://",
|
|
"postgresql+psycopg2://",
|
|
"postgres+psycopg://",
|
|
"postgres+psycopg2://",
|
|
):
|
|
if u.startswith(prefix):
|
|
rest = u.split("://", 1)[1]
|
|
return "postgresql://" + rest
|
|
return u
|
|
|
|
|
|
class PgCursorShim:
|
|
def __init__(self, raw: Any) -> None:
|
|
self._raw = raw
|
|
|
|
def fetchone(self) -> Any:
|
|
return self._raw.fetchone()
|
|
|
|
def fetchall(self) -> list[Any]:
|
|
return self._raw.fetchall()
|
|
|
|
@property
|
|
def lastrowid(self) -> int:
|
|
return 0
|
|
|
|
@property
|
|
def rowcount(self) -> int:
|
|
return int(self._raw.rowcount or 0)
|
|
|
|
def __iter__(self) -> Iterable[Any]:
|
|
return iter(self._raw)
|
|
|
|
|
|
class PgConnShim:
|
|
def __init__(self, raw: psycopg.Connection) -> None:
|
|
self._raw = raw
|
|
|
|
def execute(self, sql: str, params: Sequence[Any] | None = None) -> PgCursorShim:
|
|
adapted = adapt_sql_for_postgres(sql)
|
|
cur = self._raw.execute(adapted, params or ())
|
|
return PgCursorShim(cur)
|
|
|
|
def executemany(self, sql: str, seq_of_params: Sequence[Sequence[Any]]) -> None:
|
|
adapted = adapt_sql_for_postgres(sql)
|
|
self._raw.executemany(adapted, seq_of_params)
|
|
|
|
|
|
def _pg_connect_timeout_s() -> int:
|
|
raw = str(os.getenv("AIA_ASSISTANT_PG_CONNECT_TIMEOUT") or "10").strip()
|
|
try:
|
|
return max(1, min(int(raw), 120))
|
|
except ValueError:
|
|
return 10
|
|
|
|
|
|
def connect_postgres(url: str) -> psycopg.Connection:
|
|
"""Open a new PostgreSQL connection (short ``connect_timeout`` to avoid hanging on bad hosts)."""
|
|
return psycopg.connect(
|
|
normalize_psycopg_conninfo(url),
|
|
row_factory=dict_row,
|
|
connect_timeout=_pg_connect_timeout_s(),
|
|
)
|
|
|
|
|
|
__all__ = ["PgConnShim", "PgCursorShim", "connect_postgres", "normalize_psycopg_conninfo"]
|