629 lines
19 KiB
Python
629 lines
19 KiB
Python
# mypy: allow-untyped-defs, allow-incomplete-defs, allow-untyped-calls
|
|
# mypy: no-warn-return-any, allow-any-generics
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterable
|
|
from collections.abc import Iterator
|
|
import contextlib
|
|
from itertools import chain
|
|
import re
|
|
from typing import Any
|
|
from typing import Callable
|
|
from typing import NamedTuple
|
|
from typing import Protocol
|
|
from typing import TYPE_CHECKING
|
|
from typing import TypeGuard
|
|
from typing import TypeVar
|
|
from typing import Union
|
|
|
|
from sqlalchemy import __version__
|
|
from sqlalchemy import exc as sa_exc
|
|
from sqlalchemy import schema
|
|
from sqlalchemy import sql
|
|
from sqlalchemy import types as sqltypes
|
|
from sqlalchemy.schema import CheckConstraint
|
|
from sqlalchemy.schema import Column
|
|
from sqlalchemy.schema import ForeignKeyConstraint
|
|
from sqlalchemy.schema import Table
|
|
from sqlalchemy.sql import visitors
|
|
from sqlalchemy.sql.base import _NoneName
|
|
from sqlalchemy.sql.base import DialectKWArgs
|
|
from sqlalchemy.sql.elements import BindParameter
|
|
from sqlalchemy.sql.elements import ColumnClause
|
|
from sqlalchemy.sql.elements import conv
|
|
from sqlalchemy.sql.elements import TextClause
|
|
from sqlalchemy.sql.elements import UnaryExpression
|
|
from sqlalchemy.sql.naming import _constraint_name_for_table
|
|
from sqlalchemy.sql.naming import _NONE_NAME as _NONE_NAME # type: ignore[attr-defined] # noqa: E501
|
|
from sqlalchemy.sql.visitors import traverse
|
|
|
|
from .messaging import warn
|
|
|
|
if TYPE_CHECKING:
|
|
from sqlalchemy import ClauseElement
|
|
from sqlalchemy import Identity
|
|
from sqlalchemy import Index
|
|
from sqlalchemy.engine import Connection
|
|
from sqlalchemy.engine import Dialect
|
|
from sqlalchemy.engine import Transaction
|
|
from sqlalchemy.sql.base import ColumnCollection
|
|
from sqlalchemy.sql.compiler import SQLCompiler
|
|
from sqlalchemy.sql.elements import ColumnElement
|
|
from sqlalchemy.sql.schema import Constraint
|
|
from sqlalchemy.sql.schema import ForeignKey
|
|
from sqlalchemy.sql.schema import SchemaItem
|
|
|
|
_CE = TypeVar("_CE", bound=Union["ColumnElement[Any]", "SchemaItem"])
|
|
|
|
|
|
class _CompilerProtocol(Protocol):
|
|
def __call__(self, element: Any, compiler: Any, **kw: Any) -> str: ...
|
|
|
|
|
|
def _safe_int(value: str) -> int | str:
|
|
try:
|
|
return int(value)
|
|
except:
|
|
return value
|
|
|
|
|
|
_vers = tuple(
|
|
[_safe_int(x) for x in re.findall(r"(\d+|[abc]\d)", __version__)]
|
|
)
|
|
sqla_2_0_25 = _vers >= (2, 25)
|
|
sqla_2_1 = _vers >= (2, 1)
|
|
sqlalchemy_version = __version__
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
def compiles(
|
|
element: type[ClauseElement], *dialects: str
|
|
) -> Callable[[_CompilerProtocol], _CompilerProtocol]: ...
|
|
|
|
else:
|
|
from sqlalchemy.ext.compiler import compiles # noqa: I100,I202
|
|
|
|
|
|
identity_has_dialect_kwargs = issubclass(schema.Identity, DialectKWArgs)
|
|
|
|
|
|
def _get_identity_options_dict(
|
|
identity: Identity | schema.Sequence | None,
|
|
dialect_kwargs: bool = False,
|
|
) -> dict[str, Any]:
|
|
if identity is None:
|
|
return {}
|
|
elif identity_has_dialect_kwargs:
|
|
assert hasattr(identity, "_as_dict")
|
|
as_dict = identity._as_dict()
|
|
if dialect_kwargs:
|
|
assert isinstance(identity, DialectKWArgs)
|
|
as_dict.update(identity.dialect_kwargs)
|
|
else:
|
|
as_dict = {}
|
|
if isinstance(identity, schema.Identity):
|
|
# always=None means something different than always=False
|
|
as_dict["always"] = identity.always
|
|
if identity.on_null is not None:
|
|
as_dict["on_null"] = identity.on_null
|
|
# attributes common to Identity and Sequence
|
|
attrs = (
|
|
"start",
|
|
"increment",
|
|
"minvalue",
|
|
"maxvalue",
|
|
"nominvalue",
|
|
"nomaxvalue",
|
|
"cycle",
|
|
"cache",
|
|
"order",
|
|
)
|
|
as_dict.update(
|
|
{
|
|
key: getattr(identity, key, None)
|
|
for key in attrs
|
|
if getattr(identity, key, None) is not None
|
|
}
|
|
)
|
|
return as_dict
|
|
|
|
|
|
_ConstraintName = Union[None, str, _NoneName]
|
|
_ConstraintNameDefined = Union[str, _NoneName]
|
|
|
|
|
|
def constraint_name_defined(
|
|
name: _ConstraintName,
|
|
) -> TypeGuard[_ConstraintNameDefined]:
|
|
return name is _NONE_NAME or isinstance(name, (str, _NoneName))
|
|
|
|
|
|
def constraint_name_string(name: _ConstraintName) -> TypeGuard[str]:
|
|
return isinstance(name, str)
|
|
|
|
|
|
def constraint_name_or_none(name: _ConstraintName) -> str | None:
|
|
return name if constraint_name_string(name) else None
|
|
|
|
|
|
AUTOINCREMENT_DEFAULT = "auto"
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _ensure_scope_for_ddl(
|
|
connection: Connection | None,
|
|
) -> Iterator[None]:
|
|
try:
|
|
in_transaction = connection.in_transaction # type: ignore[union-attr]
|
|
except AttributeError:
|
|
# catch for MockConnection, None
|
|
in_transaction = None
|
|
pass
|
|
|
|
# yield outside the catch
|
|
if in_transaction is None:
|
|
yield
|
|
else:
|
|
if not in_transaction():
|
|
assert connection is not None
|
|
with connection.begin():
|
|
yield
|
|
else:
|
|
yield
|
|
|
|
|
|
def _safe_begin_connection_transaction(
|
|
connection: Connection,
|
|
) -> Transaction:
|
|
transaction = connection.get_transaction()
|
|
if transaction:
|
|
return transaction
|
|
else:
|
|
return connection.begin()
|
|
|
|
|
|
def _safe_commit_connection_transaction(
|
|
connection: Connection,
|
|
) -> None:
|
|
transaction = connection.get_transaction()
|
|
if transaction:
|
|
transaction.commit()
|
|
|
|
|
|
def _safe_rollback_connection_transaction(
|
|
connection: Connection,
|
|
) -> None:
|
|
transaction = connection.get_transaction()
|
|
if transaction:
|
|
transaction.rollback()
|
|
|
|
|
|
def _get_connection_in_transaction(connection: Connection | None) -> bool:
|
|
try:
|
|
in_transaction = connection.in_transaction # type: ignore
|
|
except AttributeError:
|
|
# catch for MockConnection
|
|
return False
|
|
else:
|
|
return in_transaction()
|
|
|
|
|
|
def _idx_table_bound_expressions(idx: Index) -> Iterable[ColumnElement[Any]]:
|
|
return idx.expressions # type: ignore
|
|
|
|
|
|
def _copy(schema_item: _CE, **kw) -> _CE:
|
|
if hasattr(schema_item, "_copy"):
|
|
return schema_item._copy(**kw)
|
|
else:
|
|
return schema_item.copy(**kw) # type: ignore[union-attr]
|
|
|
|
|
|
def _connectable_has_table(
|
|
connectable: Connection, tablename: str, schemaname: str | None
|
|
) -> bool:
|
|
return connectable.dialect.has_table(connectable, tablename, schemaname)
|
|
|
|
|
|
def _exec_on_inspector(inspector, statement, **params):
|
|
with inspector._operation_context() as conn:
|
|
return conn.execute(statement, params)
|
|
|
|
|
|
def _nullability_might_be_unset(metadata_column):
|
|
from sqlalchemy.sql import schema
|
|
|
|
return metadata_column._user_defined_nullable is schema.NULL_UNSPECIFIED
|
|
|
|
|
|
def _server_default_is_computed(*server_default) -> bool:
|
|
return any(isinstance(sd, schema.Computed) for sd in server_default)
|
|
|
|
|
|
def _server_default_is_identity(*server_default) -> bool:
|
|
return any(isinstance(sd, schema.Identity) for sd in server_default)
|
|
|
|
|
|
def _table_for_constraint(constraint: Constraint) -> Table:
|
|
if isinstance(constraint, ForeignKeyConstraint):
|
|
table = constraint.parent
|
|
assert table is not None
|
|
return table # type: ignore[return-value]
|
|
elif isinstance(constraint, CheckConstraint):
|
|
parent = constraint.parent
|
|
if not isinstance(parent, Table):
|
|
table = parent.table
|
|
assert table is not None
|
|
return table
|
|
else:
|
|
return parent
|
|
else:
|
|
return constraint.table
|
|
|
|
|
|
def _columns_for_constraint(constraint):
|
|
if isinstance(constraint, ForeignKeyConstraint):
|
|
return [fk.parent for fk in constraint.elements]
|
|
elif isinstance(constraint, CheckConstraint):
|
|
return _find_columns(constraint.sqltext)
|
|
else:
|
|
return list(constraint.columns)
|
|
|
|
|
|
def _resolve_for_variant(type_, dialect):
|
|
if _type_has_variants(type_):
|
|
base_type, mapping = _get_variant_mapping(type_)
|
|
return mapping.get(dialect.name, base_type)
|
|
else:
|
|
return type_
|
|
|
|
|
|
if hasattr(sqltypes.TypeEngine, "_variant_mapping"): # 2.0
|
|
|
|
def _type_has_variants(type_):
|
|
return bool(type_._variant_mapping)
|
|
|
|
def _get_variant_mapping(type_):
|
|
return type_, type_._variant_mapping
|
|
|
|
else:
|
|
|
|
def _type_has_variants(type_):
|
|
return type(type_) is sqltypes.Variant
|
|
|
|
def _get_variant_mapping(type_):
|
|
return type_.impl, type_.mapping
|
|
|
|
|
|
def _get_table_key(name: str, schema: str | None) -> str:
|
|
if schema is None:
|
|
return name
|
|
else:
|
|
return schema + "." + name
|
|
|
|
|
|
def _fk_spec(constraint: ForeignKeyConstraint) -> Any:
|
|
if TYPE_CHECKING:
|
|
assert constraint.columns is not None
|
|
assert constraint.elements is not None
|
|
assert isinstance(constraint.parent, Table)
|
|
|
|
source_columns = [
|
|
constraint.columns[key].name for key in constraint.column_keys
|
|
]
|
|
|
|
source_table = constraint.parent.name
|
|
source_schema = constraint.parent.schema
|
|
target_schema = constraint.elements[0].column.table.schema
|
|
target_table = constraint.elements[0].column.table.name
|
|
target_columns = [element.column.name for element in constraint.elements]
|
|
ondelete = constraint.ondelete
|
|
onupdate = constraint.onupdate
|
|
deferrable = constraint.deferrable
|
|
initially = constraint.initially
|
|
return (
|
|
source_schema,
|
|
source_table,
|
|
source_columns,
|
|
target_schema,
|
|
target_table,
|
|
target_columns,
|
|
onupdate,
|
|
ondelete,
|
|
deferrable,
|
|
initially,
|
|
)
|
|
|
|
|
|
class _ForeignKeyTarget(NamedTuple):
|
|
"""Stand-in for SQLAlchemy 2.1's ``ForeignKeyTarget`` named tuple."""
|
|
|
|
schema: str | None
|
|
table_name: str
|
|
column_name: str | None
|
|
|
|
|
|
def _fk_target_tokens(fk: ForeignKey) -> _ForeignKeyTarget:
|
|
"""Return the ``(schema, table_name, column_name)`` names that identify
|
|
the target of a :class:`.ForeignKey`.
|
|
|
|
These names are the representation SQLAlchemy computes against; unlike
|
|
the dotted string form of ``ForeignKey._get_colspec()`` they remain
|
|
unambiguous when a name itself contains a dot.
|
|
|
|
"""
|
|
fk_any: Any = fk
|
|
if sqla_2_1:
|
|
return _ForeignKeyTarget(*fk_any.target_tokens)
|
|
else:
|
|
return _ForeignKeyTarget(*fk_any._column_tokens)
|
|
|
|
|
|
def _fk_target_table_key(fk: ForeignKey) -> str | None:
|
|
"""Return the :attr:`.MetaData.tables` key under which the target of a
|
|
:class:`.ForeignKey` is, or would be, registered.
|
|
|
|
Returns ``None`` when the target was given as a :class:`.Column` that is
|
|
not associated with a :class:`.Table`, in which case no key can be named.
|
|
|
|
"""
|
|
fk_any: Any = fk
|
|
if sqla_2_1:
|
|
return fk_any.target_table_key
|
|
else:
|
|
return fk_any._table_key()
|
|
|
|
|
|
def _fk_target_is_named(fk: ForeignKey) -> bool:
|
|
"""Return True if a :class:`.ForeignKey` identifies its target by name
|
|
rather than by a :class:`.Column` object.
|
|
|
|
"""
|
|
fk_any: Any = fk
|
|
if sqla_2_1:
|
|
return fk_any.target_column is None
|
|
else:
|
|
return isinstance(fk_any._colspec, str)
|
|
|
|
|
|
def _fk_is_self_referential(constraint: ForeignKeyConstraint) -> bool:
|
|
assert constraint.parent is not None
|
|
return (
|
|
_fk_target_table_key(constraint.elements[0]) == constraint.parent.key
|
|
)
|
|
|
|
|
|
def _is_type_bound(constraint: Constraint) -> bool:
|
|
# this deals with SQLAlchemy #3260, don't copy CHECK constraints
|
|
# that will be generated by the type.
|
|
# new feature added for #3260
|
|
return constraint._type_bound
|
|
|
|
|
|
def _conv_constraint_names(table: Table) -> None:
|
|
"""Mark the names of constraints and indexes on ``table`` as final.
|
|
|
|
A name that came back from reflection is the name that's actually in
|
|
the database, so it should not be elaborated further by a naming
|
|
convention; :class:`sqlalchemy.schema.conv` is the marker that
|
|
indicates this.
|
|
|
|
"""
|
|
for const in table.constraints:
|
|
if constraint_name_string(const.name) and not isinstance(
|
|
const.name, conv
|
|
):
|
|
const.name = conv(const.name)
|
|
|
|
for idx in table.indexes:
|
|
if constraint_name_string(idx.name) and not isinstance(idx.name, conv):
|
|
idx.name = conv(idx.name)
|
|
|
|
|
|
def _name_type_bound_constraints(table: Table, source_table: Table) -> None:
|
|
"""Apply the naming convention on ``source_table.metadata`` to the
|
|
type-bound constraints on ``table``.
|
|
|
|
``table`` here is a copy of ``source_table`` under a different name, so
|
|
the convention is resolved against ``source_table`` in order for the
|
|
``%(table_name)s`` token to produce the real table name.
|
|
|
|
"""
|
|
if not source_table.metadata.naming_convention:
|
|
return
|
|
|
|
for const in table.constraints:
|
|
if not _is_type_bound(const):
|
|
continue
|
|
try:
|
|
newname = _constraint_name_for_table(const, source_table)
|
|
except sa_exc.InvalidRequestError:
|
|
# the convention includes %(constraint_name)s but the type that
|
|
# generates this constraint has no name of its own, so there's
|
|
# nothing to interpolate. the constraint is emitted without a
|
|
# name, which is what it was before the naming convention was
|
|
# taken into account here at all.
|
|
colnames = ", ".join(
|
|
c.name for c in _columns_for_constraint(const)
|
|
)
|
|
warn(
|
|
f"The CHECK constraint generated by the type of column "
|
|
f"'{source_table.name}.{colnames}' has no name, while the "
|
|
f"naming convention in use requires one; the constraint is "
|
|
f"emitted without a name. Boolean and Enum datatypes should "
|
|
f"be given a name, e.g. Boolean(name='...'), in order for "
|
|
f"their CHECK constraints to participate fully in batch "
|
|
f"mode."
|
|
)
|
|
continue
|
|
if newname:
|
|
const.name = newname
|
|
|
|
|
|
def _find_columns(clause):
|
|
"""locate Column objects within the given expression."""
|
|
|
|
cols: set[ColumnElement[Any]] = set()
|
|
traverse(clause, {}, {"column": cols.add})
|
|
return cols
|
|
|
|
|
|
def _remove_column_from_collection(
|
|
collection: ColumnCollection, column: Column[Any] | ColumnClause[Any]
|
|
) -> None:
|
|
"""remove a column from a ColumnCollection."""
|
|
|
|
# workaround for older SQLAlchemy, remove the
|
|
# same object that's present
|
|
assert column.key is not None
|
|
to_remove = collection[column.key]
|
|
|
|
# SQLAlchemy 2.0 will use more ReadOnlyColumnCollection
|
|
# (renamed from ImmutableColumnCollection)
|
|
if hasattr(collection, "_immutable") or hasattr(collection, "_readonly"):
|
|
collection._parent.remove(to_remove)
|
|
else:
|
|
collection.remove(to_remove)
|
|
|
|
|
|
def _textual_index_column(
|
|
table: Table, text_: str | TextClause | ColumnElement[Any]
|
|
) -> ColumnElement[Any] | Column[Any]:
|
|
"""a workaround for the Index construct's severe lack of flexibility"""
|
|
if isinstance(text_, str):
|
|
c = Column(text_, sqltypes.NULLTYPE)
|
|
table.append_column(c)
|
|
return c
|
|
elif isinstance(text_, TextClause):
|
|
return _textual_index_element(table, text_)
|
|
elif isinstance(text_, _textual_index_element):
|
|
return _textual_index_column(table, text_.text)
|
|
elif isinstance(text_, sql.ColumnElement):
|
|
return _copy_expression(text_, table)
|
|
else:
|
|
raise ValueError("String or text() construct expected")
|
|
|
|
|
|
def _copy_expression(expression: _CE, target_table: Table) -> _CE:
|
|
def replace(col):
|
|
if (
|
|
isinstance(col, Column)
|
|
and col.table is not None
|
|
and col.table is not target_table
|
|
):
|
|
if col.name in target_table.c:
|
|
return target_table.c[col.name]
|
|
else:
|
|
c = _copy(col)
|
|
target_table.append_column(c)
|
|
return c
|
|
else:
|
|
return None
|
|
|
|
return visitors.replacement_traverse( # type: ignore[call-overload]
|
|
expression, {}, replace
|
|
)
|
|
|
|
|
|
class _textual_index_element(sql.ColumnElement):
|
|
"""Wrap around a sqlalchemy text() construct in such a way that
|
|
we appear like a column-oriented SQL expression to an Index
|
|
construct.
|
|
|
|
The issue here is that currently the Postgresql dialect, the biggest
|
|
recipient of functional indexes, keys all the index expressions to
|
|
the corresponding column expressions when rendering CREATE INDEX,
|
|
so the Index we create here needs to have a .columns collection that
|
|
is the same length as the .expressions collection. Ultimately
|
|
SQLAlchemy should support text() expressions in indexes.
|
|
|
|
See SQLAlchemy issue 3174.
|
|
|
|
"""
|
|
|
|
__visit_name__ = "_textual_idx_element"
|
|
|
|
def __init__(self, table: Table, text: TextClause) -> None:
|
|
self.table = table
|
|
self.text = text
|
|
self.key = text.text
|
|
self.fake_column = schema.Column(self.text.text, sqltypes.NULLTYPE)
|
|
table.append_column(self.fake_column)
|
|
|
|
def get_children(self, **kw):
|
|
return [self.fake_column]
|
|
|
|
|
|
@compiles(_textual_index_element)
|
|
def _render_textual_index_column(
|
|
element: _textual_index_element, compiler: SQLCompiler, **kw
|
|
) -> str:
|
|
return compiler.process(element.text, **kw)
|
|
|
|
|
|
class _literal_bindparam(BindParameter):
|
|
pass
|
|
|
|
|
|
@compiles(_literal_bindparam)
|
|
def _render_literal_bindparam(
|
|
element: _literal_bindparam, compiler: SQLCompiler, **kw
|
|
) -> str:
|
|
return compiler.render_literal_bindparam(element, **kw)
|
|
|
|
|
|
def _get_constraint_final_name(
|
|
constraint: Index | Constraint, dialect: Dialect | None
|
|
) -> str | None:
|
|
if constraint.name is None:
|
|
return None
|
|
assert dialect is not None
|
|
# make use of SQLAlchemy API to return what would be the final
|
|
# compiled form of the name for this dialect.
|
|
return dialect.identifier_preparer.format_constraint(
|
|
constraint, _alembic_quote=False
|
|
)
|
|
|
|
|
|
def _constraint_is_named(
|
|
constraint: Constraint | Index, dialect: Dialect | None
|
|
) -> bool:
|
|
if constraint.name is None:
|
|
return False
|
|
assert dialect is not None
|
|
name = dialect.identifier_preparer.format_constraint(
|
|
constraint, _alembic_quote=False
|
|
)
|
|
return name is not None
|
|
|
|
|
|
def is_expression_index(index: Index) -> bool:
|
|
for expr in index.expressions:
|
|
if is_expression(expr):
|
|
return True
|
|
return False
|
|
|
|
|
|
def is_expression(expr: Any) -> bool:
|
|
while isinstance(expr, UnaryExpression):
|
|
expr = expr.element
|
|
if not isinstance(expr, ColumnClause) or expr.is_literal:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _inherit_schema_deprecated() -> bool:
|
|
# at some point in 2.1 inherit_schema was replaced with a property
|
|
# so that's preset at the class level, while before it wasn't.
|
|
return sqla_2_1 and hasattr(sqltypes.Enum, "inherit_schema")
|
|
|
|
|
|
def all_table_check_constraints(table: Table) -> set[CheckConstraint]:
|
|
"""Returns all check constraint including type bound ones and
|
|
those on columns."""
|
|
candidates = chain(
|
|
table.constraints, *(c.constraints for c in table.columns)
|
|
)
|
|
return {ck for ck in candidates if isinstance(ck, CheckConstraint)}
|