"""Postgres connector — run guarded SQL against an external PostgreSQL database.

The read/write tools, the SQL safety guard, and row caps live in
``sql_base.BaseSqlConnector``; this module only adds the psycopg driver seams and
a libpq connection string. ``postgres.query`` additionally runs in a READ ONLY
transaction so a write cannot slip through even if the text guard is fooled.
"""

from __future__ import annotations

from legate.connectors.sql_base import BaseSqlConnector, _sql_config_schema


class PostgresConnector(BaseSqlConnector):
    type = "postgres"
    title = "PostgreSQL"
    _default_port = 5432
    config_schema = _sql_config_schema(
        extra={"sslmode": {"type": "string", "description": "e.g. require, verify-full, disable."}}
    )

    def _conninfo(self, config: dict, secrets: dict) -> str:
        parts = {
            "host": config.get("host", ""),
            "port": str(config.get("port", self._default_port)),
            "dbname": config.get("database", ""),
            "user": config.get("user", ""),
            "password": secrets.get("password", ""),
        }
        if config.get("sslmode"):
            parts["sslmode"] = config["sslmode"]
        return " ".join(f"{k}={v}" for k, v in parts.items() if v != "")

    def _run_query(
        self, config: dict, secrets: dict, sql: str, params: list | None, max_rows: int
    ) -> dict:  # pragma: no cover - exercised against a live database
        import psycopg

        with psycopg.connect(self._conninfo(config, secrets)) as conn:
            conn.read_only = True
            with conn.cursor() as cur:
                cur.execute(sql, params or None)
                columns = [d.name for d in cur.description] if cur.description else []
                rows = [list(r) for r in cur.fetchmany(max_rows)]
        return {"columns": columns, "rows": rows}

    def _run_exec(
        self, config: dict, secrets: dict, sql: str, params: list | None
    ) -> int:  # pragma: no cover - exercised against a live database
        import psycopg

        with psycopg.connect(self._conninfo(config, secrets)) as conn:
            with conn.cursor() as cur:
                cur.execute(sql, params or None)
                affected = cur.rowcount
            conn.commit()
        return affected
