introspection.py

from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

from django.db.backends.postgresql.introspection import (  # type: ignore[import]
    DatabaseIntrospection,
)

from psqlextra.types import PostgresPartitioningMethod

from . import base_impl

PARTITIONING_STRATEGY_TO_METHOD = {
    "r": PostgresPartitioningMethod.RANGE,
    "l": PostgresPartitioningMethod.LIST,
    "h": PostgresPartitioningMethod.HASH,
}


@dataclass
class PostgresIntrospectedPartitionTable:
    """Data container for information about a partition."""

    name: str
    full_name: str
    comment: Optional[str]


@dataclass
class PostgresIntrospectedPartitonedTable:
    """Data container for information about a partitioned table."""

    name: str
    method: PostgresPartitioningMethod
    key: List[str]
    partitions: List[PostgresIntrospectedPartitionTable]

    def partition_by_name(
        self, name: str
    ) -> Optional[PostgresIntrospectedPartitionTable]:
        """Finds the partition with the specified name."""

        return next(
            (
                partition
                for partition in self.partitions
                if partition.name == name
            ),
            None,
        )


if TYPE_CHECKING:

    class Introspection(DatabaseIntrospection):
        pass

else:
    Introspection = base_impl.introspection()


class PostgresIntrospection(Introspection):
    """Adds introspection features specific to PostgreSQL."""

    # TODO: This class is a mess, both here and in the
    # the base.
    #
    # Some methods return untyped dicts, some named tuples,
    # some flat lists of strings. It's horribly inconsistent.
    #
    # Most methods are poorly named. For example; `get_table_description`
    # does not return a complete table description. It merely returns
    # the columns.
    #
    # We do our best in this class to stay consistent with
    # the base in Django by respecting its naming scheme
    # and commonly used return types. Creating an API that
    # matches the look&feel from the Django base class
    # is more important than fixing those issues.

    def get_partitioned_tables(
        self, cursor
    ) -> List[PostgresIntrospectedPartitonedTable]:
        """Gets a list of partitioned tables."""

        cursor.execute(
            """
            SELECT
                pg_class.relname,
                pg_partitioned_table.partstrat
            FROM
                pg_partitioned_table
            JOIN
                pg_class
            ON
                pg_class.oid = pg_partitioned_table.partrelid
        """
        )

        return [
            PostgresIntrospectedPartitonedTable(
                name=row[0],
                method=PARTITIONING_STRATEGY_TO_METHOD[row[1]],
                key=self.get_partition_key(cursor, row[0]),
                partitions=self.get_partitions(cursor, row[0]),
            )
            for row in cursor.fetchall()
        ]

    def get_partitioned_table(self, cursor, table_name: str):
        """Gets a single partitioned table."""

        return next(
            (
                table
                for table in self.get_partitioned_tables(cursor)
                if table.name == table_name
            ),
            None,
        )

    def get_partitions(
        self, cursor, table_name
    ) -> List[PostgresIntrospectedPartitionTable]:
        """Gets a list of partitions belonging to the specified partitioned
        table."""

        sql = """
            SELECT
                child.relname,
                pg_description.description
            FROM pg_inherits
            JOIN
                pg_class parent
            ON
                pg_inherits.inhparent = parent.oid
            JOIN
                pg_class child
            ON
                pg_inherits.inhrelid = child.oid
            JOIN
                pg_namespace nmsp_parent
            ON
                nmsp_parent.oid = parent.relnamespace
            JOIN
                pg_namespace nmsp_child
            ON
                nmsp_child.oid = child.relnamespace
            LEFT JOIN
                pg_description
            ON
                pg_description.objoid = child.oid
            WHERE
                parent.relname = %s
        """

        cursor.execute(sql, (table_name,))

        return [
            PostgresIntrospectedPartitionTable(
                name=row[0].replace(f"{table_name}_", ""),
                full_name=row[0],
                comment=row[1] or None,
            )
            for row in cursor.fetchall()
        ]

    def get_partition_key(self, cursor, table_name: str) -> List[str]:
        """Gets the partition key for the specified partitioned table.

        Returns:
            A list of column names that are part of the
            partition key.
        """

        sql = """
            SELECT
                col.column_name
            FROM
                (SELECT partrelid,
                        partnatts,
                        CASE partstrat
                            WHEN 'l' THEN 'list'
                            WHEN 'r' THEN 'range'
                            WHEN 'h' THEN 'hash'
                        END AS partition_strategy,
                        Unnest(partattrs) column_index
                 FROM pg_partitioned_table) pt
            JOIN
                pg_class par
            ON par.oid = pt.partrelid
            JOIN
                information_schema.COLUMNS col
            ON
                col.table_schema = par.relnamespace :: regnamespace :: text
                AND col.table_name = par.relname
                AND ordinal_position = pt.column_index
            WHERE
                table_name = %s
        """

        cursor.execute(sql, (table_name,))
        return [row[0] for row in cursor.fetchall()]

    def get_columns(self, cursor, table_name: str):
        return self.get_table_description(cursor, table_name)

    def get_schema_list(self, cursor) -> List[str]:
        """A flat list of available schemas."""

        cursor.execute(
            """
            SELECT
                schema_name
            FROM
                information_schema.schemata
            """,
            tuple(),
        )

        return [name for name, in cursor.fetchall()]

    def get_constraints(self, cursor, table_name: str):
        """Retrieve any constraints or keys (unique, pk, fk, check, index)
        across one or more columns.

        Also retrieve the definition of expression-based indexes.
        """

        constraints = super().get_constraints(cursor, table_name)

        # standard Django implementation does not return the definition
        # for indexes, only for constraints, let's patch that up
        cursor.execute(
            "SELECT indexname, indexdef FROM pg_indexes WHERE schemaname = current_schema() AND tablename = %s",
            (table_name,),
        )
        for index_name, definition in cursor.fetchall():
            # PostgreSQL 13 or older won't give a definition if the
            # index is actually a primary key.
            constraint = constraints.get(index_name)
            if not constraint:
                continue

            if constraint.get("definition") is None:
                constraint["definition"] = definition

        return constraints

    def get_table_locks(self, cursor) -> List[Tuple[str, str, str]]:
        cursor.execute(
            """
            SELECT
                n.nspname,
                t.relname,
                l.mode
            FROM pg_locks l
            INNER JOIN pg_class t ON t.oid = l.relation
            INNER JOIN pg_namespace n ON n.oid = t.relnamespace
            WHERE t.relnamespace >= 2200
            ORDER BY n.nspname, t.relname, l.mode
        """
        )

        return cursor.fetchall()

    def get_storage_settings(self, cursor, table_name: str) -> Dict[str, str]:
        sql = """
            SELECT
                unnest(c.reloptions || array(select 'toast.' || x from pg_catalog.unnest(tc.reloptions) x))
            FROM
                pg_catalog.pg_class c
            LEFT JOIN
                pg_catalog.pg_class tc ON (c.reltoastrelid = tc.oid)
            LEFT JOIN
                pg_catalog.pg_am am ON (c.relam = am.oid)
            WHERE
                c.relname::text = %s
                AND pg_catalog.pg_table_is_visible(c.oid)
        """

        cursor.execute(sql, (table_name,))

        storage_settings = {}
        for row in cursor.fetchall():
            # It's hard to believe, but storage settings are really
            # represented as `key=value` strings in Postgres.
            # See: https://www.postgresql.org/docs/current/catalog-pg-class.html
            name, value = row[0].split("=")
            storage_settings[name] = value

        return storage_settings

    def get_relations(self, cursor, table_name: str):
        """Gets a dictionary {field_name: (field_name_other_table,
        other_table)} representing all relations in the specified table.

        This is overriden because the query in Django does not handle
        relations between tables in different schemas properly.
        """

        cursor.execute(
            """
            SELECT a1.attname, c2.relname, a2.attname
            FROM pg_constraint con
            LEFT JOIN pg_class c1 ON con.conrelid = c1.oid
            LEFT JOIN pg_class c2 ON con.confrelid = c2.oid
            LEFT JOIN pg_attribute a1 ON c1.oid = a1.attrelid AND a1.attnum = con.conkey[1]
            LEFT JOIN pg_attribute a2 ON c2.oid = a2.attrelid AND a2.attnum = con.confkey[1]
            WHERE
                con.conrelid = %s::regclass AND
                con.contype = 'f' AND
                pg_catalog.pg_table_is_visible(c1.oid)
        """,
            [table_name],
        )
        return {row[0]: (row[2], row[1]) for row in cursor.fetchall()}