AdministrationItDepartmens/.venv/Lib/site-packages/sqlalchemy/sql/dml.py

# sql/dml.py
# Copyright (C) 2009-2024 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: https://www.opensource.org/licenses/mit-license.php
"""
Provide :class:`_expression.Insert`, :class:`_expression.Update` and
:class:`_expression.Delete`.

"""
from __future__ import annotations

import collections.abc as collections_abc
import operator
from typing import Any
from typing import cast
from typing import Dict
from typing import Iterable
from typing import List
from typing import MutableMapping
from typing import NoReturn
from typing import Optional
from typing import overload
from typing import Sequence
from typing import Tuple
from typing import Type
from typing import TYPE_CHECKING
from typing import TypeVar
from typing import Union

from . import coercions
from . import roles
from . import util as sql_util
from ._typing import _TP
from ._typing import _unexpected_kw
from ._typing import is_column_element
from ._typing import is_named_from_clause
from .base import _entity_namespace_key
from .base import _exclusive_against
from .base import _from_objects
from .base import _generative
from .base import _select_iterables
from .base import ColumnCollection
from .base import CompileState
from .base import DialectKWArgs
from .base import Executable
from .base import Generative
from .base import HasCompileState
from .elements import BooleanClauseList
from .elements import ClauseElement
from .elements import ColumnClause
from .elements import ColumnElement
from .elements import Null
from .selectable import Alias
from .selectable import ExecutableReturnsRows
from .selectable import FromClause
from .selectable import HasCTE
from .selectable import HasPrefixes
from .selectable import Join
from .selectable import SelectLabelStyle
from .selectable import TableClause
from .selectable import TypedReturnsRows
from .sqltypes import NullType
from .visitors import InternalTraversal
from .. import exc
from .. import util
from ..util.typing import Self
from ..util.typing import TypeGuard

if TYPE_CHECKING:
    from ._typing import _ColumnExpressionArgument
    from ._typing import _ColumnsClauseArgument
    from ._typing import _DMLColumnArgument
    from ._typing import _DMLColumnKeyMapping
    from ._typing import _DMLTableArgument
    from ._typing import _T0  # noqa
    from ._typing import _T1  # noqa
    from ._typing import _T2  # noqa
    from ._typing import _T3  # noqa
    from ._typing import _T4  # noqa
    from ._typing import _T5  # noqa
    from ._typing import _T6  # noqa
    from ._typing import _T7  # noqa
    from ._typing import _TypedColumnClauseArgument as _TCCA  # noqa
    from .base import ReadOnlyColumnCollection
    from .compiler import SQLCompiler
    from .elements import KeyedColumnElement
    from .selectable import _ColumnsClauseElement
    from .selectable import _SelectIterable
    from .selectable import Select
    from .selectable import Selectable

    def isupdate(dml: DMLState) -> TypeGuard[UpdateDMLState]: ...

    def isdelete(dml: DMLState) -> TypeGuard[DeleteDMLState]: ...

    def isinsert(dml: DMLState) -> TypeGuard[InsertDMLState]: ...

else:
    isupdate = operator.attrgetter("isupdate")
    isdelete = operator.attrgetter("isdelete")
    isinsert = operator.attrgetter("isinsert")


_T = TypeVar("_T", bound=Any)

_DMLColumnElement = Union[str, ColumnClause[Any]]
_DMLTableElement = Union[TableClause, Alias, Join]


class DMLState(CompileState):
    _no_parameters = True
    _dict_parameters: Optional[MutableMapping[_DMLColumnElement, Any]] = None
    _multi_parameters: Optional[
        List[MutableMapping[_DMLColumnElement, Any]]
    ] = None
    _ordered_values: Optional[List[Tuple[_DMLColumnElement, Any]]] = None
    _parameter_ordering: Optional[List[_DMLColumnElement]] = None
    _primary_table: FromClause
    _supports_implicit_returning = True

    isupdate = False
    isdelete = False
    isinsert = False

    statement: UpdateBase

    def __init__(
        self, statement: UpdateBase, compiler: SQLCompiler, **kw: Any
    ):
        raise NotImplementedError()

    @classmethod
    def get_entity_description(cls, statement: UpdateBase) -> Dict[str, Any]:
        return {
            "name": (
                statement.table.name
                if is_named_from_clause(statement.table)
                else None
            ),
            "table": statement.table,
        }

    @classmethod
    def get_returning_column_descriptions(
        cls, statement: UpdateBase
    ) -> List[Dict[str, Any]]:
        return [
            {
                "name": c.key,
                "type": c.type,
                "expr": c,
            }
            for c in statement._all_selected_columns
        ]

    @property
    def dml_table(self) -> _DMLTableElement:
        return self.statement.table

    if TYPE_CHECKING:

        @classmethod
        def get_plugin_class(cls, statement: Executable) -> Type[DMLState]: ...

    @classmethod
    def _get_multi_crud_kv_pairs(
        cls,
        statement: UpdateBase,
        multi_kv_iterator: Iterable[Dict[_DMLColumnArgument, Any]],
    ) -> List[Dict[_DMLColumnElement, Any]]:
        return [
            {
                coercions.expect(roles.DMLColumnRole, k): v
                for k, v in mapping.items()
            }
            for mapping in multi_kv_iterator
        ]

    @classmethod
    def _get_crud_kv_pairs(
        cls,
        statement: UpdateBase,
        kv_iterator: Iterable[Tuple[_DMLColumnArgument, Any]],
        needs_to_be_cacheable: bool,
    ) -> List[Tuple[_DMLColumnElement, Any]]:
        return [
            (
                coercions.expect(roles.DMLColumnRole, k),
                (
                    v
                    if not needs_to_be_cacheable
                    else coercions.expect(
                        roles.ExpressionElementRole,
                        v,
                        type_=NullType(),
                        is_crud=True,
                    )
                ),
            )
            for k, v in kv_iterator
        ]

    def _make_extra_froms(
        self, statement: DMLWhereBase
    ) -> Tuple[FromClause, List[FromClause]]:
        froms: List[FromClause] = []

        all_tables = list(sql_util.tables_from_leftmost(statement.table))
        primary_table = all_tables[0]
        seen = {primary_table}

        consider = statement._where_criteria
        if self._dict_parameters:
            consider += tuple(self._dict_parameters.values())

        for crit in consider:
            for item in _from_objects(crit):
                if not seen.intersection(item._cloned_set):
                    froms.append(item)
                seen.update(item._cloned_set)

        froms.extend(all_tables[1:])
        return primary_table, froms

    def _process_values(self, statement: ValuesBase) -> None:
        if self._no_parameters:
            self._dict_parameters = statement._values
            self._no_parameters = False

    def _process_select_values(self, statement: ValuesBase) -> None:
        assert statement._select_names is not None
        parameters: MutableMapping[_DMLColumnElement, Any] = {
            name: Null() for name in statement._select_names
        }

        if self._no_parameters:
            self._no_parameters = False
            self._dict_parameters = parameters
        else:
            # this condition normally not reachable as the Insert
            # does not allow this construction to occur
            assert False, "This statement already has parameters"

    def _no_multi_values_supported(self, statement: ValuesBase) -> NoReturn:
        raise exc.InvalidRequestError(
            "%s construct does not support "
            "multiple parameter sets." % statement.__visit_name__.upper()
        )

    def _cant_mix_formats_error(self) -> NoReturn:
        raise exc.InvalidRequestError(
            "Can't mix single and multiple VALUES "
            "formats in one INSERT statement; one style appends to a "
            "list while the other replaces values, so the intent is "
            "ambiguous."
        )


@CompileState.plugin_for("default", "insert")
class InsertDMLState(DMLState):
    isinsert = True

    include_table_with_column_exprs = False

    _has_multi_parameters = False

    def __init__(
        self,
        statement: Insert,
        compiler: SQLCompiler,
        disable_implicit_returning: bool = False,
        **kw: Any,
    ):
        self.statement = statement
        self._primary_table = statement.table

        if disable_implicit_returning:
            self._supports_implicit_returning = False

        self.isinsert = True
        if statement._select_names:
            self._process_select_values(statement)
        if statement._values is not None:
            self._process_values(statement)
        if statement._multi_values:
            self._process_multi_values(statement)

    @util.memoized_property
    def _insert_col_keys(self) -> List[str]:
        # this is also done in crud.py -> _key_getters_for_crud_column
        return [
            coercions.expect(roles.DMLColumnRole, col, as_key=True)
            for col in self._dict_parameters or ()
        ]

    def _process_values(self, statement: ValuesBase) -> None:
        if self._no_parameters:
            self._has_multi_parameters = False
            self._dict_parameters = statement._values
            self._no_parameters = False
        elif self._has_multi_parameters:
            self._cant_mix_formats_error()

    def _process_multi_values(self, statement: ValuesBase) -> None:
        for parameters in statement._multi_values:
            multi_parameters: List[MutableMapping[_DMLColumnElement, Any]] = [
                (
                    {
                        c.key: value
                        for c, value in zip(statement.table.c, parameter_set)
                    }
                    if isinstance(parameter_set, collections_abc.Sequence)
                    else parameter_set
                )
                for parameter_set in parameters
            ]

            if self._no_parameters:
                self._no_parameters = False
                self._has_multi_parameters = True
                self._multi_parameters = multi_parameters
                self._dict_parameters = self._multi_parameters[0]
            elif not self._has_multi_parameters:
                self._cant_mix_formats_error()
            else:
                assert self._multi_parameters
                self._multi_parameters.extend(multi_parameters)


@CompileState.plugin_for("default", "update")
class UpdateDMLState(DMLState):
    isupdate = True

    include_table_with_column_exprs = False

    def __init__(self, statement: Update, compiler: SQLCompiler, **kw: Any):
        self.statement = statement

        self.isupdate = True
        if statement._ordered_values is not None:
            self._process_ordered_values(statement)
        elif statement._values is not None:
            self._process_values(statement)
        elif statement._multi_values:
            self._no_multi_values_supported(statement)
        t, ef = self._make_extra_froms(statement)
        self._primary_table = t
        self._extra_froms = ef

        self.is_multitable = mt = ef
        self.include_table_with_column_exprs = bool(
            mt and compiler.render_table_with_column_in_update_from
        )

    def _process_ordered_values(self, statement: ValuesBase) -> None:
        parameters = statement._ordered_values

        if self._no_parameters:
            self._no_parameters = False
            assert parameters is not None
            self._dict_parameters = dict(parameters)
            self._ordered_values = parameters
            self._parameter_ordering = [key for key, value in parameters]
        else:
            raise exc.InvalidRequestError(
                "Can only invoke ordered_values() once, and not mixed "
                "with any other values() call"
            )


@CompileState.plugin_for("default", "delete")
class DeleteDMLState(DMLState):
    isdelete = True

    def __init__(self, statement: Delete, compiler: SQLCompiler, **kw: Any):
        self.statement = statement

        self.isdelete = True
        t, ef = self._make_extra_froms(statement)
        self._primary_table = t
        self._extra_froms = ef
        self.is_multitable = ef


class UpdateBase(
    roles.DMLRole,
    HasCTE,
    HasCompileState,
    DialectKWArgs,
    HasPrefixes,
    Generative,
    ExecutableReturnsRows,
    ClauseElement,
):
    """Form the base for ``INSERT``, ``UPDATE``, and ``DELETE`` statements."""

    __visit_name__ = "update_base"

    _hints: util.immutabledict[Tuple[_DMLTableElement, str], str] = (
        util.EMPTY_DICT
    )
    named_with_column = False

    _label_style: SelectLabelStyle = (
        SelectLabelStyle.LABEL_STYLE_DISAMBIGUATE_ONLY
    )
    table: _DMLTableElement

    _return_defaults = False
    _return_defaults_columns: Optional[Tuple[_ColumnsClauseElement, ...]] = (
        None
    )
    _supplemental_returning: Optional[Tuple[_ColumnsClauseElement, ...]] = None
    _returning: Tuple[_ColumnsClauseElement, ...] = ()

    is_dml = True

    def _generate_fromclause_column_proxies(
        self, fromclause: FromClause
    ) -> None:
        fromclause._columns._populate_separate_keys(
            col._make_proxy(fromclause)
            for col in self._all_selected_columns
            if is_column_element(col)
        )

    def params(self, *arg: Any, **kw: Any) -> NoReturn:
        """Set the parameters for the statement.

        This method raises ``NotImplementedError`` on the base class,
        and is overridden by :class:`.ValuesBase` to provide the
        SET/VALUES clause of UPDATE and INSERT.

        """
        raise NotImplementedError(
            "params() is not supported for INSERT/UPDATE/DELETE statements."
            " To set the values for an INSERT or UPDATE statement, use"
            " stmt.values(**parameters)."
        )

    @_generative
    def with_dialect_options(self, **opt: Any) -> Self:
        """Add dialect options to this INSERT/UPDATE/DELETE object.

        e.g.::

            upd = table.update().dialect_options(mysql_limit=10)

        .. versionadded: 1.4 - this method supersedes the dialect options
           associated with the constructor.


        """
        self._validate_dialect_kwargs(opt)
        return self

    @_generative
    def return_defaults(
        self,
        *cols: _DMLColumnArgument,
        supplemental_cols: Optional[Iterable[_DMLColumnArgument]] = None,
        sort_by_parameter_order: bool = False,
    ) -> Self:
        """Make use of a :term:`RETURNING` clause for the purpose
        of fetching server-side expressions and defaults, for supporting
        backends only.

        .. deepalchemy::

            The :meth:`.UpdateBase.return_defaults` method is used by the ORM
            for its internal work in fetching newly generated primary key
            and server default values, in particular to provide the underyling
            implementation of the :paramref:`_orm.Mapper.eager_defaults`
            ORM feature as well as to allow RETURNING support with bulk
            ORM inserts.  Its behavior is fairly idiosyncratic
            and is not really intended for general use.  End users should
            stick with using :meth:`.UpdateBase.returning` in order to
            add RETURNING clauses to their INSERT, UPDATE and DELETE
            statements.

        Normally, a single row INSERT statement will automatically populate the
        :attr:`.CursorResult.inserted_primary_key` attribute when executed,
        which stores the primary key of the row that was just inserted in the
        form of a :class:`.Row` object with column names as named tuple keys
        (and the :attr:`.Row._mapping` view fully populated as well). The
        dialect in use chooses the strategy to use in order to populate this
        data; if it was generated using server-side defaults and / or SQL
        expressions, dialect-specific approaches such as ``cursor.lastrowid``
        or ``RETURNING`` are typically used to acquire the new primary key
        value.

        However, when the statement is modified by calling
        :meth:`.UpdateBase.return_defaults` before executing the statement,
        additional behaviors take place **only** for backends that support
        RETURNING and for :class:`.Table` objects that maintain the
        :paramref:`.Table.implicit_returning` parameter at its default value of
        ``True``. In these cases, when the :class:`.CursorResult` is returned
        from the statement's execution, not only will
        :attr:`.CursorResult.inserted_primary_key` be populated as always, the
        :attr:`.CursorResult.returned_defaults` attribute will also be
        populated with a :class:`.Row` named-tuple representing the full range
        of server generated
        values from that single row, including values for any columns that
        specify :paramref:`_schema.Column.server_default` or which make use of
        :paramref:`_schema.Column.default` using a SQL expression.

        When invoking INSERT statements with multiple rows using
        :ref:`insertmanyvalues <engine_insertmanyvalues>`, the
        :meth:`.UpdateBase.return_defaults` modifier will have the effect of
        the :attr:`_engine.CursorResult.inserted_primary_key_rows` and
        :attr:`_engine.CursorResult.returned_defaults_rows` attributes being
        fully populated with lists of :class:`.Row` objects representing newly
        inserted primary key values as well as newly inserted server generated
        values for each row inserted. The