"""Cluster connection helpers — primary (writes) and secondary (reads).

The trust ceremony has three legs (see docs/architecture.md):

1. Resolve the leader URL from a SIGNED DNS TXT record.

2. Verify the cluster server's TDX-attested cert (currently NoopVerifier;

production switches to DcapVerifier when the Python wheel ships).

3. Present a TDX-attested CLIENT cert from the dstack guest agent so the

cluster sidecar can DCAP-verify our quote.

`psycopg-ra-tls`'s `connect_via_manifest` does (1) and (3) and accepts the

verifier for (2).

Lifecycle:

- Each connection (primary, secondary) lives for the FastAPI process

lifetime. One mutual-RA-TLS handshake at startup, NOT one per query.

- Every connection brings its own localhost forwarder — psycopg-ra-tls

starts a daemon thread internally. Closing the connection closes the

forwarder. Don't fight that contract.

- On `OperationalError` (leader rotation, network glitch, sidecar

restart) the caller closes the dead connection and calls

`open_primary` / `open_secondary` again to re-handshake. That's the

only place we pay the RA-TLS cost in steady state.

Two consumer shapes:

`with connect(settings, role=...) as conn:` — short-lived, one query

or migration. Forwarder dies with the block.

`conn = open_primary(settings)` / `open_secondary(settings)` —

long-lived, the FastAPI app stores it on `app.state` and re-uses

it for every request / poll. Caller is responsible for `.close()`.

"""

from __future__ import annotations

import logging

from collections.abc import Iterator

from contextlib import contextmanager

import psycopg

from psycopg.rows import dict_row

from psycopg_ratls import connect_via_manifest

from ra_tls_verify import NoopVerifier

from .config import Settings

log = logging.getLogger(__name__)

def _build_dsn(settings: Settings, role: str) -> str:

"""Construct the DSN handed to psycopg-ra-tls.

The host/port are placeholders — psycopg-ra-tls rewrites them to point

at the localhost forwarder it spins up internally. The driver itself

sees `127.0.0.1:<ephemeral>` with `sslmode=disable`. The `secret` is

the cluster-wide allowlist token; the cluster sidecar substitutes the

KMS-derived password on the wire.

"""

return (

f"postgresql://{role}:{settings.cluster_secret}"

f"@placeholder:5433/{settings.database}"

)

def _open(settings: Settings, role: str) -> psycopg.Connection:

"""Open a fresh psycopg connection through a fresh psycopg-ra-tls forwarder.

The forwarder lives as long as the returned connection. Caller owns

the lifetime — we DO NOT wrap in a context manager.

Verifier choice: see CLAUDE.md "Hard rules" #3 for the trade-off.

Production should switch to DcapVerifier once ra-tls-verify ships

the Python wheel of dcap-qvl. Until then, the manifest signature is

the trust anchor (the dns-controller's TEE signs that this leader

URL is a real cluster member).

"""

log.info(

"opening connection: role=%s cluster=%s db=%s",

role,

settings.cluster_uuid,

settings.database,

)

conn = connect_via_manifest(

cluster_domain=settings.cluster_domain,

dsn_template=_build_dsn(settings, role),

manifest_signer=settings.manifest_signer_bytes,

verifier=NoopVerifier(),

allow_simulator=settings.allow_simulator,

)

# dict_row everywhere — matches hivemind-core, reads better in

# downstream code, and makes mocking easier.

conn.row_factory = dict_row # type: ignore[assignment]

return conn

def open_primary(settings: Settings) -> psycopg.Connection:

"""Open a long-lived connection to the cluster's primary (read-write).

Used by the FastAPI app's startup hook; stored on `app.state.primary_conn`

and reused for every `POST /events` request. The caller owns shutdown

— call `.close()` on app shutdown to release the forwarder cleanly.

"""

return _open(settings, role=settings.primary_role)

def open_secondary(settings: Settings) -> psycopg.Connection:

"""Open a long-lived read-only connection to the cluster.

Used by the poller task. The current cluster's manifest TXT only

advertises the leader URL, so this connection physically lands on

the primary too — but it authenticates as a read-only role

(`teesql_read` by default), so any accidental write attempt is

rejected by the cluster's grants.

When the cluster's manifest gains explicit replica routing, this

helper will switch to picking a replica URL. Callers shouldn't have

to change.

"""

return _open(settings, role=settings.secondary_role)

@contextmanager

def connect(settings: Settings, role: str | None = None) -> Iterator[psycopg.Connection]:

"""Short-lived helper: yield a connection, close it (and its forwarder)

on exit.

Use this for one-shot operations like `migrate`, smoke tests, or any

code path where you do exactly one set of queries and then exit.

For long-running services, use `open_primary` / `open_secondary` and

keep the connection alive — see the module docstring.

"""

conn = _open(settings, role=role or settings.primary_role)

try:

yield conn

finally:

try:

conn.close()

except Exception as e: # noqa: BLE001

log.warning("error closing connection: %s", e)