501 lines
14 KiB
Python
501 lines
14 KiB
Python
"""
|
|
Information about PostgreSQL types
|
|
|
|
These types allow to read information from the system catalog and provide
|
|
information to the adapters if needed.
|
|
"""
|
|
|
|
# Copyright (C) 2020 The Psycopg Team
|
|
from enum import Enum
|
|
from typing import Any, Dict, Iterator, Optional, overload
|
|
from typing import Sequence, Tuple, Type, TypeVar, Union, TYPE_CHECKING
|
|
from typing_extensions import TypeAlias
|
|
|
|
from . import errors as e
|
|
from .abc import AdaptContext, Query
|
|
from .rows import dict_row
|
|
from ._encodings import conn_encoding
|
|
|
|
if TYPE_CHECKING:
|
|
from .connection import BaseConnection, Connection
|
|
from .connection_async import AsyncConnection
|
|
from .sql import Identifier, SQL
|
|
|
|
T = TypeVar("T", bound="TypeInfo")
|
|
RegistryKey: TypeAlias = Union[str, int, Tuple[type, int]]
|
|
|
|
|
|
class TypeInfo:
|
|
"""
|
|
Hold information about a PostgreSQL base type.
|
|
"""
|
|
|
|
__module__ = "psycopg.types"
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
oid: int,
|
|
array_oid: int,
|
|
*,
|
|
regtype: str = "",
|
|
delimiter: str = ",",
|
|
):
|
|
self.name = name
|
|
self.oid = oid
|
|
self.array_oid = array_oid
|
|
self.regtype = regtype or name
|
|
self.delimiter = delimiter
|
|
|
|
def __repr__(self) -> str:
|
|
return (
|
|
f"<{self.__class__.__qualname__}:"
|
|
f" {self.name} (oid: {self.oid}, array oid: {self.array_oid})>"
|
|
)
|
|
|
|
@overload
|
|
@classmethod
|
|
def fetch(
|
|
cls: Type[T], conn: "Connection[Any]", name: Union[str, "Identifier"]
|
|
) -> Optional[T]:
|
|
...
|
|
|
|
@overload
|
|
@classmethod
|
|
async def fetch(
|
|
cls: Type[T], conn: "AsyncConnection[Any]", name: Union[str, "Identifier"]
|
|
) -> Optional[T]:
|
|
...
|
|
|
|
@classmethod
|
|
def fetch(
|
|
cls: Type[T], conn: "BaseConnection[Any]", name: Union[str, "Identifier"]
|
|
) -> Any:
|
|
"""Query a system catalog to read information about a type."""
|
|
from .sql import Composable
|
|
from .connection import Connection
|
|
from .connection_async import AsyncConnection
|
|
|
|
if isinstance(name, Composable):
|
|
name = name.as_string(conn)
|
|
|
|
if isinstance(conn, Connection):
|
|
return cls._fetch(conn, name)
|
|
elif isinstance(conn, AsyncConnection):
|
|
return cls._fetch_async(conn, name)
|
|
else:
|
|
raise TypeError(
|
|
f"expected Connection or AsyncConnection, got {type(conn).__name__}"
|
|
)
|
|
|
|
@classmethod
|
|
def _fetch(cls: Type[T], conn: "Connection[Any]", name: str) -> Optional[T]:
|
|
# This might result in a nested transaction. What we want is to leave
|
|
# the function with the connection in the state we found (either idle
|
|
# or intrans)
|
|
try:
|
|
with conn.transaction():
|
|
if conn_encoding(conn) == "ascii":
|
|
conn.execute("set local client_encoding to utf8")
|
|
with conn.cursor(row_factory=dict_row) as cur:
|
|
cur.execute(cls._get_info_query(conn), {"name": name})
|
|
recs = cur.fetchall()
|
|
except e.UndefinedObject:
|
|
return None
|
|
|
|
return cls._from_records(name, recs)
|
|
|
|
@classmethod
|
|
async def _fetch_async(
|
|
cls: Type[T], conn: "AsyncConnection[Any]", name: str
|
|
) -> Optional[T]:
|
|
try:
|
|
async with conn.transaction():
|
|
if conn_encoding(conn) == "ascii":
|
|
await conn.execute("set local client_encoding to utf8")
|
|
async with conn.cursor(row_factory=dict_row) as cur:
|
|
await cur.execute(cls._get_info_query(conn), {"name": name})
|
|
recs = await cur.fetchall()
|
|
except e.UndefinedObject:
|
|
return None
|
|
|
|
return cls._from_records(name, recs)
|
|
|
|
@classmethod
|
|
def _from_records(
|
|
cls: Type[T], name: str, recs: Sequence[Dict[str, Any]]
|
|
) -> Optional[T]:
|
|
if len(recs) == 1:
|
|
return cls(**recs[0])
|
|
elif not recs:
|
|
return None
|
|
else:
|
|
raise e.ProgrammingError(f"found {len(recs)} different types named {name}")
|
|
|
|
def register(self, context: Optional[AdaptContext] = None) -> None:
|
|
"""
|
|
Register the type information, globally or in the specified `!context`.
|
|
"""
|
|
if context:
|
|
types = context.adapters.types
|
|
else:
|
|
from . import postgres
|
|
|
|
types = postgres.types
|
|
|
|
types.add(self)
|
|
|
|
if self.array_oid:
|
|
from .types.array import register_array
|
|
|
|
register_array(self, context)
|
|
|
|
@classmethod
|
|
def _get_info_query(cls, conn: "BaseConnection[Any]") -> Query:
|
|
from .sql import SQL
|
|
|
|
return SQL(
|
|
"""\
|
|
SELECT
|
|
typname AS name, oid, typarray AS array_oid,
|
|
oid::regtype::text AS regtype, typdelim AS delimiter
|
|
FROM pg_type t
|
|
WHERE t.oid = {regtype}
|
|
ORDER BY t.oid
|
|
"""
|
|
).format(regtype=cls._to_regtype(conn))
|
|
|
|
@classmethod
|
|
def _has_to_regtype_function(cls, conn: "BaseConnection[Any]") -> bool:
|
|
# to_regtype() introduced in PostgreSQL 9.4 and CockroachDB 22.2
|
|
info = conn.info
|
|
if info.vendor == "PostgreSQL":
|
|
return info.server_version >= 90400
|
|
elif info.vendor == "CockroachDB":
|
|
return info.server_version >= 220200
|
|
else:
|
|
return False
|
|
|
|
@classmethod
|
|
def _to_regtype(cls, conn: "BaseConnection[Any]") -> "SQL":
|
|
# `to_regtype()` returns the type oid or NULL, unlike the :: operator,
|
|
# which returns the type or raises an exception, which requires
|
|
# a transaction rollback and leaves traces in the server logs.
|
|
|
|
from .sql import SQL
|
|
|
|
if cls._has_to_regtype_function(conn):
|
|
return SQL("to_regtype(%(name)s)")
|
|
else:
|
|
return SQL("%(name)s::regtype")
|
|
|
|
def _added(self, registry: "TypesRegistry") -> None:
|
|
"""Method called by the `!registry` when the object is added there."""
|
|
pass
|
|
|
|
|
|
class RangeInfo(TypeInfo):
|
|
"""Manage information about a range type."""
|
|
|
|
__module__ = "psycopg.types.range"
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
oid: int,
|
|
array_oid: int,
|
|
*,
|
|
regtype: str = "",
|
|
subtype_oid: int,
|
|
):
|
|
super().__init__(name, oid, array_oid, regtype=regtype)
|
|
self.subtype_oid = subtype_oid
|
|
|
|
@classmethod
|
|
def _get_info_query(cls, conn: "BaseConnection[Any]") -> Query:
|
|
from .sql import SQL
|
|
|
|
return SQL(
|
|
"""\
|
|
SELECT t.typname AS name, t.oid AS oid, t.typarray AS array_oid,
|
|
t.oid::regtype::text AS regtype,
|
|
r.rngsubtype AS subtype_oid
|
|
FROM pg_type t
|
|
JOIN pg_range r ON t.oid = r.rngtypid
|
|
WHERE t.oid = {regtype}
|
|
"""
|
|
).format(regtype=cls._to_regtype(conn))
|
|
|
|
def _added(self, registry: "TypesRegistry") -> None:
|
|
# Map ranges subtypes to info
|
|
registry._registry[RangeInfo, self.subtype_oid] = self
|
|
|
|
|
|
class MultirangeInfo(TypeInfo):
|
|
"""Manage information about a multirange type."""
|
|
|
|
__module__ = "psycopg.types.multirange"
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
oid: int,
|
|
array_oid: int,
|
|
*,
|
|
regtype: str = "",
|
|
range_oid: int,
|
|
subtype_oid: int,
|
|
):
|
|
super().__init__(name, oid, array_oid, regtype=regtype)
|
|
self.range_oid = range_oid
|
|
self.subtype_oid = subtype_oid
|
|
|
|
@classmethod
|
|
def _get_info_query(cls, conn: "BaseConnection[Any]") -> Query:
|
|
from .sql import SQL
|
|
|
|
if conn.info.server_version < 140000:
|
|
raise e.NotSupportedError(
|
|
"multirange types are only available from PostgreSQL 14"
|
|
)
|
|
|
|
return SQL(
|
|
"""\
|
|
SELECT t.typname AS name, t.oid AS oid, t.typarray AS array_oid,
|
|
t.oid::regtype::text AS regtype,
|
|
r.rngtypid AS range_oid, r.rngsubtype AS subtype_oid
|
|
FROM pg_type t
|
|
JOIN pg_range r ON t.oid = r.rngmultitypid
|
|
WHERE t.oid = {regtype}
|
|
"""
|
|
).format(regtype=cls._to_regtype(conn))
|
|
|
|
def _added(self, registry: "TypesRegistry") -> None:
|
|
# Map multiranges ranges and subtypes to info
|
|
registry._registry[MultirangeInfo, self.range_oid] = self
|
|
registry._registry[MultirangeInfo, self.subtype_oid] = self
|
|
|
|
|
|
class CompositeInfo(TypeInfo):
|
|
"""Manage information about a composite type."""
|
|
|
|
__module__ = "psycopg.types.composite"
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
oid: int,
|
|
array_oid: int,
|
|
*,
|
|
regtype: str = "",
|
|
field_names: Sequence[str],
|
|
field_types: Sequence[int],
|
|
):
|
|
super().__init__(name, oid, array_oid, regtype=regtype)
|
|
self.field_names = field_names
|
|
self.field_types = field_types
|
|
# Will be set by register() if the `factory` is a type
|
|
self.python_type: Optional[type] = None
|
|
|
|
@classmethod
|
|
def _get_info_query(cls, conn: "BaseConnection[Any]") -> Query:
|
|
from .sql import SQL
|
|
|
|
return SQL(
|
|
"""\
|
|
SELECT
|
|
t.typname AS name, t.oid AS oid, t.typarray AS array_oid,
|
|
t.oid::regtype::text AS regtype,
|
|
coalesce(a.fnames, '{{}}') AS field_names,
|
|
coalesce(a.ftypes, '{{}}') AS field_types
|
|
FROM pg_type t
|
|
LEFT JOIN (
|
|
SELECT
|
|
attrelid,
|
|
array_agg(attname) AS fnames,
|
|
array_agg(atttypid) AS ftypes
|
|
FROM (
|
|
SELECT a.attrelid, a.attname, a.atttypid
|
|
FROM pg_attribute a
|
|
JOIN pg_type t ON t.typrelid = a.attrelid
|
|
WHERE t.oid = {regtype}
|
|
AND a.attnum > 0
|
|
AND NOT a.attisdropped
|
|
ORDER BY a.attnum
|
|
) x
|
|
GROUP BY attrelid
|
|
) a ON a.attrelid = t.typrelid
|
|
WHERE t.oid = {regtype}
|
|
"""
|
|
).format(regtype=cls._to_regtype(conn))
|
|
|
|
|
|
class EnumInfo(TypeInfo):
|
|
"""Manage information about an enum type."""
|
|
|
|
__module__ = "psycopg.types.enum"
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
oid: int,
|
|
array_oid: int,
|
|
labels: Sequence[str],
|
|
):
|
|
super().__init__(name, oid, array_oid)
|
|
self.labels = labels
|
|
# Will be set by register_enum()
|
|
self.enum: Optional[Type[Enum]] = None
|
|
|
|
@classmethod
|
|
def _get_info_query(cls, conn: "BaseConnection[Any]") -> Query:
|
|
from .sql import SQL
|
|
|
|
return SQL(
|
|
"""\
|
|
SELECT name, oid, array_oid, array_agg(label) AS labels
|
|
FROM (
|
|
SELECT
|
|
t.typname AS name, t.oid AS oid, t.typarray AS array_oid,
|
|
e.enumlabel AS label
|
|
FROM pg_type t
|
|
LEFT JOIN pg_enum e
|
|
ON e.enumtypid = t.oid
|
|
WHERE t.oid = {regtype}
|
|
ORDER BY e.enumsortorder
|
|
) x
|
|
GROUP BY name, oid, array_oid
|
|
"""
|
|
).format(regtype=cls._to_regtype(conn))
|
|
|
|
|
|
class TypesRegistry:
|
|
"""
|
|
Container for the information about types in a database.
|
|
"""
|
|
|
|
__module__ = "psycopg.types"
|
|
|
|
def __init__(self, template: Optional["TypesRegistry"] = None):
|
|
self._registry: Dict[RegistryKey, TypeInfo]
|
|
|
|
# Make a shallow copy: it will become a proper copy if the registry
|
|
# is edited.
|
|
if template:
|
|
self._registry = template._registry
|
|
self._own_state = False
|
|
template._own_state = False
|
|
else:
|
|
self.clear()
|
|
|
|
def clear(self) -> None:
|
|
self._registry = {}
|
|
self._own_state = True
|
|
|
|
def add(self, info: TypeInfo) -> None:
|
|
self._ensure_own_state()
|
|
if info.oid:
|
|
self._registry[info.oid] = info
|
|
if info.array_oid:
|
|
self._registry[info.array_oid] = info
|
|
self._registry[info.name] = info
|
|
|
|
if info.regtype and info.regtype not in self._registry:
|
|
self._registry[info.regtype] = info
|
|
|
|
# Allow info to customise further their relation with the registry
|
|
info._added(self)
|
|
|
|
def __iter__(self) -> Iterator[TypeInfo]:
|
|
seen = set()
|
|
for t in self._registry.values():
|
|
if id(t) not in seen:
|
|
seen.add(id(t))
|
|
yield t
|
|
|
|
@overload
|
|
def __getitem__(self, key: Union[str, int]) -> TypeInfo:
|
|
...
|
|
|
|
@overload
|
|
def __getitem__(self, key: Tuple[Type[T], int]) -> T:
|
|
...
|
|
|
|
def __getitem__(self, key: RegistryKey) -> TypeInfo:
|
|
"""
|
|
Return info about a type, specified by name or oid
|
|
|
|
:param key: the name or oid of the type to look for.
|
|
|
|
Raise KeyError if not found.
|
|
"""
|
|
if isinstance(key, str):
|
|
if key.endswith("[]"):
|
|
key = key[:-2]
|
|
elif not isinstance(key, (int, tuple)):
|
|
raise TypeError(f"the key must be an oid or a name, got {type(key)}")
|
|
try:
|
|
return self._registry[key]
|
|
except KeyError:
|
|
raise KeyError(f"couldn't find the type {key!r} in the types registry")
|
|
|
|
@overload
|
|
def get(self, key: Union[str, int]) -> Optional[TypeInfo]:
|
|
...
|
|
|
|
@overload
|
|
def get(self, key: Tuple[Type[T], int]) -> Optional[T]:
|
|
...
|
|
|
|
def get(self, key: RegistryKey) -> Optional[TypeInfo]:
|
|
"""
|
|
Return info about a type, specified by name or oid
|
|
|
|
:param key: the name or oid of the type to look for.
|
|
|
|
Unlike `__getitem__`, return None if not found.
|
|
"""
|
|
try:
|
|
return self[key]
|
|
except KeyError:
|
|
return None
|
|
|
|
def get_oid(self, name: str) -> int:
|
|
"""
|
|
Return the oid of a PostgreSQL type by name.
|
|
|
|
:param key: the name of the type to look for.
|
|
|
|
Return the array oid if the type ends with "``[]``"
|
|
|
|
Raise KeyError if the name is unknown.
|
|
"""
|
|
t = self[name]
|
|
if name.endswith("[]"):
|
|
return t.array_oid
|
|
else:
|
|
return t.oid
|
|
|
|
def get_by_subtype(self, cls: Type[T], subtype: Union[int, str]) -> Optional[T]:
|
|
"""
|
|
Return info about a `TypeInfo` subclass by its element name or oid.
|
|
|
|
:param cls: the subtype of `!TypeInfo` to look for. Currently
|
|
supported are `~psycopg.types.range.RangeInfo` and
|
|
`~psycopg.types.multirange.MultirangeInfo`.
|
|
:param subtype: The name or OID of the subtype of the element to look for.
|
|
:return: The `!TypeInfo` object of class `!cls` whose subtype is
|
|
`!subtype`. `!None` if the element or its range are not found.
|
|
"""
|
|
try:
|
|
info = self[subtype]
|
|
except KeyError:
|
|
return None
|
|
return self.get((cls, info.oid))
|
|
|
|
def _ensure_own_state(self) -> None:
|
|
# Time to write! so, copy.
|
|
if not self._own_state:
|
|
self._registry = self._registry.copy()
|
|
self._own_state = True
|