96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
|
|
import data
|
|
|
|
|
|
class PostgresServiceError(Exception):
|
|
pass
|
|
|
|
|
|
def _load_driver():
|
|
try:
|
|
import psycopg
|
|
return "psycopg", psycopg
|
|
except Exception:
|
|
pass
|
|
try:
|
|
import psycopg2
|
|
return "psycopg2", psycopg2
|
|
except Exception as e:
|
|
raise PostgresServiceError("Nie je nainstalovany PostgreSQL driver psycopg ani psycopg2.") from e
|
|
|
|
|
|
def is_configured(conn: data.PostgresConnection) -> bool:
|
|
return bool(
|
|
conn.host.strip()
|
|
and conn.database.strip()
|
|
and conn.user.strip()
|
|
and int(conn.port or 0) > 0
|
|
)
|
|
|
|
|
|
def connection_kwargs(conn: data.PostgresConnection) -> dict:
|
|
kwargs = {
|
|
"host": conn.host,
|
|
"port": int(conn.port or 5432),
|
|
"dbname": conn.database,
|
|
"user": conn.user,
|
|
"password": conn.password,
|
|
"connect_timeout": int(conn.connect_timeout or 5),
|
|
}
|
|
if conn.sslmode:
|
|
kwargs["sslmode"] = conn.sslmode
|
|
return kwargs
|
|
|
|
|
|
def test_connection(conn: data.PostgresConnection) -> None:
|
|
if not conn.enabled:
|
|
raise PostgresServiceError("PostgreSQL pripojenie nie je povolene pre instalaciu.")
|
|
if not is_configured(conn):
|
|
raise PostgresServiceError("PostgreSQL pripojenie nie je vyplnene.")
|
|
driver_name, driver = _load_driver()
|
|
try:
|
|
pg = driver.connect(**connection_kwargs(conn))
|
|
try:
|
|
cur = pg.cursor()
|
|
try:
|
|
cur.execute("SELECT 1")
|
|
cur.fetchone()
|
|
finally:
|
|
cur.close()
|
|
finally:
|
|
pg.close()
|
|
except Exception as e:
|
|
raise PostgresServiceError(f"PostgreSQL spojenie zlyhalo ({driver_name}): {e}") from e
|
|
|
|
|
|
def open_connection(conn: data.PostgresConnection):
|
|
if not conn.enabled:
|
|
raise PostgresServiceError("PostgreSQL pripojenie nie je povolene pre instalaciu.")
|
|
if not is_configured(conn):
|
|
raise PostgresServiceError("PostgreSQL pripojenie nie je vyplnene.")
|
|
driver_name, driver = _load_driver()
|
|
try:
|
|
return driver.connect(**connection_kwargs(conn))
|
|
except Exception as e:
|
|
raise PostgresServiceError(f"PostgreSQL spojenie zlyhalo ({driver_name}): {e}") from e
|
|
|
|
|
|
@contextmanager
|
|
def connect(conn: data.PostgresConnection):
|
|
if not conn.enabled:
|
|
raise PostgresServiceError("PostgreSQL pripojenie nie je povolene pre instalaciu.")
|
|
if not is_configured(conn):
|
|
raise PostgresServiceError("PostgreSQL pripojenie nie je vyplnene.")
|
|
driver_name, driver = _load_driver()
|
|
try:
|
|
pg = driver.connect(**connection_kwargs(conn))
|
|
except Exception as e:
|
|
raise PostgresServiceError(f"PostgreSQL spojenie zlyhalo ({driver_name}): {e}") from e
|
|
try:
|
|
yield pg
|
|
finally:
|
|
pg.close()
|