53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
"""Shared controlled PostgreSQL connection helper for ARR processes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, Optional
|
|
|
|
|
|
def read_controlled_database(path: Path) -> Dict[str, object]:
|
|
values: Dict[str, str] = {}
|
|
for raw in path.read_text(encoding="utf-8").splitlines():
|
|
line = raw.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
key, value = line.split("=", 1)
|
|
values[key.strip()] = value.strip().strip('"').strip("'")
|
|
required = (
|
|
"ARR_DB_HOST",
|
|
"ARR_DB_PORT",
|
|
"ARR_DB_USER",
|
|
"ARR_DB_PASSWORD",
|
|
"ARR_DB_NAME",
|
|
)
|
|
if any(not values.get(key) for key in required):
|
|
raise ValueError("controlled database configuration is incomplete")
|
|
return {
|
|
"host": values["ARR_DB_HOST"],
|
|
"port": int(values["ARR_DB_PORT"]),
|
|
"user": values["ARR_DB_USER"],
|
|
"password": values["ARR_DB_PASSWORD"],
|
|
"dbname": values["ARR_DB_NAME"],
|
|
}
|
|
|
|
|
|
def controlled_connect(
|
|
config_path: Path,
|
|
driver_path: Optional[Path],
|
|
) -> Callable[[str], Any]:
|
|
if driver_path is not None:
|
|
resolved = driver_path.expanduser().resolve()
|
|
if not resolved.is_dir():
|
|
raise ValueError("database driver path is unavailable")
|
|
sys.path.insert(0, str(resolved))
|
|
import psycopg # type: ignore[import-not-found]
|
|
|
|
parameters = read_controlled_database(config_path.expanduser().resolve())
|
|
|
|
def connect(_dsn: str) -> Any:
|
|
return psycopg.connect(**parameters, autocommit=False)
|
|
|
|
return connect
|