schema.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198
  1. # testing/schema.py
  2. # Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
  3. # <see AUTHORS file>
  4. #
  5. # This module is part of SQLAlchemy and is released under
  6. # the MIT License: https://www.opensource.org/licenses/mit-license.php
  7. # mypy: ignore-errors
  8. from __future__ import annotations
  9. import sys
  10. from . import config
  11. from . import exclusions
  12. from .. import event
  13. from .. import schema
  14. from .. import types as sqltypes
  15. from ..orm import mapped_column as _orm_mapped_column
  16. from ..util import OrderedDict
  17. __all__ = ["Table", "Column"]
  18. table_options = {}
  19. def Table(*args, **kw) -> schema.Table:
  20. """A schema.Table wrapper/hook for dialect-specific tweaks."""
  21. # pop out local options; these are not used at the moment
  22. _ = {k: kw.pop(k) for k in list(kw) if k.startswith("test_")}
  23. kw.update(table_options)
  24. return schema.Table(*args, **kw)
  25. def mapped_column(*args, **kw):
  26. """An orm.mapped_column wrapper/hook for dialect-specific tweaks."""
  27. return _schema_column(_orm_mapped_column, args, kw)
  28. def Column(*args, **kw):
  29. """A schema.Column wrapper/hook for dialect-specific tweaks."""
  30. return _schema_column(schema.Column, args, kw)
  31. def _schema_column(factory, args, kw):
  32. test_opts = {k: kw.pop(k) for k in list(kw) if k.startswith("test_")}
  33. if not config.requirements.foreign_key_ddl.enabled_for_config(config):
  34. args = [arg for arg in args if not isinstance(arg, schema.ForeignKey)]
  35. construct = factory(*args, **kw)
  36. if factory is schema.Column:
  37. col = construct
  38. else:
  39. col = construct.column
  40. if test_opts.get("test_needs_autoincrement", False) and kw.get(
  41. "primary_key", False
  42. ):
  43. if col.default is None and col.server_default is None:
  44. col.autoincrement = True
  45. # allow any test suite to pick up on this
  46. col.info["test_needs_autoincrement"] = True
  47. # hardcoded rule for oracle; this should
  48. # be moved out
  49. if exclusions.against(config._current, "oracle"):
  50. def add_seq(c, tbl):
  51. c._init_items(
  52. schema.Sequence(
  53. _truncate_name(
  54. config.db.dialect, tbl.name + "_" + c.name + "_seq"
  55. ),
  56. optional=True,
  57. )
  58. )
  59. event.listen(col, "after_parent_attach", add_seq, propagate=True)
  60. return construct
  61. class eq_type_affinity:
  62. """Helper to compare types inside of datastructures based on affinity.
  63. E.g.::
  64. eq_(
  65. inspect(connection).get_columns("foo"),
  66. [
  67. {
  68. "name": "id",
  69. "type": testing.eq_type_affinity(sqltypes.INTEGER),
  70. "nullable": False,
  71. "default": None,
  72. "autoincrement": False,
  73. },
  74. {
  75. "name": "data",
  76. "type": testing.eq_type_affinity(sqltypes.NullType),
  77. "nullable": True,
  78. "default": None,
  79. "autoincrement": False,
  80. },
  81. ],
  82. )
  83. """
  84. def __init__(self, target):
  85. self.target = sqltypes.to_instance(target)
  86. def __eq__(self, other):
  87. return self.target._type_affinity is other._type_affinity
  88. def __ne__(self, other):
  89. return self.target._type_affinity is not other._type_affinity
  90. class eq_compile_type:
  91. """similar to eq_type_affinity but uses compile"""
  92. def __init__(self, target):
  93. self.target = target
  94. def __eq__(self, other):
  95. return self.target == other.compile()
  96. def __ne__(self, other):
  97. return self.target != other.compile()
  98. class eq_clause_element:
  99. """Helper to compare SQL structures based on compare()"""
  100. def __init__(self, target):
  101. self.target = target
  102. def __eq__(self, other):
  103. return self.target.compare(other)
  104. def __ne__(self, other):
  105. return not self.target.compare(other)
  106. def _truncate_name(dialect, name):
  107. if len(name) > dialect.max_identifier_length:
  108. return (
  109. name[0 : max(dialect.max_identifier_length - 6, 0)]
  110. + "_"
  111. + hex(hash(name) % 64)[2:]
  112. )
  113. else:
  114. return name
  115. def pep435_enum(name):
  116. # Implements PEP 435 in the minimal fashion needed by SQLAlchemy
  117. __members__ = OrderedDict()
  118. def __init__(self, name, value, alias=None):
  119. self.name = name
  120. self.value = value
  121. self.__members__[name] = self
  122. value_to_member[value] = self
  123. setattr(self.__class__, name, self)
  124. if alias:
  125. self.__members__[alias] = self
  126. setattr(self.__class__, alias, self)
  127. value_to_member = {}
  128. @classmethod
  129. def get(cls, value):
  130. return value_to_member[value]
  131. someenum = type(
  132. name,
  133. (object,),
  134. {"__members__": __members__, "__init__": __init__, "get": get},
  135. )
  136. # getframe() trick for pickling I don't understand courtesy
  137. # Python namedtuple()
  138. try:
  139. module = sys._getframe(1).f_globals.get("__name__", "__main__")
  140. except (AttributeError, ValueError):
  141. pass
  142. if module is not None:
  143. someenum.__module__ = module
  144. return someenum