1"""Psycopg adapter registration for Plain.
 2
 3The `AdaptersMap` returned by `get_adapters_template()` is attached to every
 4psycopg connection we open (via `build_connection_params` in `sources.py`).
 5"""
 6
 7from __future__ import annotations
 8
 9from functools import lru_cache
10from typing import Any
11
12from psycopg import adapt, adapters
13from psycopg.abc import PyFormat
14from psycopg.postgres import types as pg_types
15from psycopg.types.range import BaseRangeDumper, Range, RangeDumper
16from psycopg.types.string import TextLoader
17
18TSRANGE_OID = pg_types["tsrange"].oid
19TSTZRANGE_OID = pg_types["tstzrange"].oid
20
21
22class PlainRangeDumper(RangeDumper):
23    """A Range dumper customized for Plain."""
24
25    def upgrade(self, obj: Range[Any], format: PyFormat) -> BaseRangeDumper:
26        dumper = super().upgrade(obj, format)
27        if dumper is not self and dumper.oid == TSRANGE_OID:
28            dumper.oid = TSTZRANGE_OID
29        return dumper
30
31
32@lru_cache
33def get_adapters_template() -> adapt.AdaptersMap:
34    ctx = adapt.AdaptersMap(adapters)
35    # No-op JSON loader to avoid psycopg3 round trips
36    ctx.register_loader("jsonb", TextLoader)
37    # Treat inet/cidr as text
38    ctx.register_loader("inet", TextLoader)
39    ctx.register_loader("cidr", TextLoader)
40    ctx.register_dumper(Range, PlainRangeDumper)
41    return ctx