ddl_references.py 7.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254
  1. """
  2. Helpers to manipulate deferred DDL statements that might need to be adjusted or
  3. discarded within when executing a migration.
  4. """
  5. from copy import deepcopy
  6. class Reference:
  7. """Base class that defines the reference interface."""
  8. def references_table(self, table):
  9. """
  10. Return whether or not this instance references the specified table.
  11. """
  12. return False
  13. def references_column(self, table, column):
  14. """
  15. Return whether or not this instance references the specified column.
  16. """
  17. return False
  18. def rename_table_references(self, old_table, new_table):
  19. """
  20. Rename all references to the old_name to the new_table.
  21. """
  22. pass
  23. def rename_column_references(self, table, old_column, new_column):
  24. """
  25. Rename all references to the old_column to the new_column.
  26. """
  27. pass
  28. def __repr__(self):
  29. return "<%s %r>" % (self.__class__.__name__, str(self))
  30. def __str__(self):
  31. raise NotImplementedError(
  32. "Subclasses must define how they should be converted to string."
  33. )
  34. class Table(Reference):
  35. """Hold a reference to a table."""
  36. def __init__(self, table, quote_name):
  37. self.table = table
  38. self.quote_name = quote_name
  39. def references_table(self, table):
  40. return self.table == table
  41. def rename_table_references(self, old_table, new_table):
  42. if self.table == old_table:
  43. self.table = new_table
  44. def __str__(self):
  45. return self.quote_name(self.table)
  46. class TableColumns(Table):
  47. """Base class for references to multiple columns of a table."""
  48. def __init__(self, table, columns):
  49. self.table = table
  50. self.columns = columns
  51. def references_column(self, table, column):
  52. return self.table == table and column in self.columns
  53. def rename_column_references(self, table, old_column, new_column):
  54. if self.table == table:
  55. for index, column in enumerate(self.columns):
  56. if column == old_column:
  57. self.columns[index] = new_column
  58. class Columns(TableColumns):
  59. """Hold a reference to one or many columns."""
  60. def __init__(self, table, columns, quote_name, col_suffixes=()):
  61. self.quote_name = quote_name
  62. self.col_suffixes = col_suffixes
  63. super().__init__(table, columns)
  64. def __str__(self):
  65. def col_str(column, idx):
  66. col = self.quote_name(column)
  67. try:
  68. suffix = self.col_suffixes[idx]
  69. if suffix:
  70. col = "{} {}".format(col, suffix)
  71. except IndexError:
  72. pass
  73. return col
  74. return ", ".join(
  75. col_str(column, idx) for idx, column in enumerate(self.columns)
  76. )
  77. class IndexName(TableColumns):
  78. """Hold a reference to an index name."""
  79. def __init__(self, table, columns, suffix, create_index_name):
  80. self.suffix = suffix
  81. self.create_index_name = create_index_name
  82. super().__init__(table, columns)
  83. def __str__(self):
  84. return self.create_index_name(self.table, self.columns, self.suffix)
  85. class IndexColumns(Columns):
  86. def __init__(self, table, columns, quote_name, col_suffixes=(), opclasses=()):
  87. self.opclasses = opclasses
  88. super().__init__(table, columns, quote_name, col_suffixes)
  89. def __str__(self):
  90. def col_str(column, idx):
  91. # Index.__init__() guarantees that self.opclasses is the same
  92. # length as self.columns.
  93. col = "{} {}".format(self.quote_name(column), self.opclasses[idx])
  94. try:
  95. suffix = self.col_suffixes[idx]
  96. if suffix:
  97. col = "{} {}".format(col, suffix)
  98. except IndexError:
  99. pass
  100. return col
  101. return ", ".join(
  102. col_str(column, idx) for idx, column in enumerate(self.columns)
  103. )
  104. class ForeignKeyName(TableColumns):
  105. """Hold a reference to a foreign key name."""
  106. def __init__(
  107. self,
  108. from_table,
  109. from_columns,
  110. to_table,
  111. to_columns,
  112. suffix_template,
  113. create_fk_name,
  114. ):
  115. self.to_reference = TableColumns(to_table, to_columns)
  116. self.suffix_template = suffix_template
  117. self.create_fk_name = create_fk_name
  118. super().__init__(
  119. from_table,
  120. from_columns,
  121. )
  122. def references_table(self, table):
  123. return super().references_table(table) or self.to_reference.references_table(
  124. table
  125. )
  126. def references_column(self, table, column):
  127. return super().references_column(
  128. table, column
  129. ) or self.to_reference.references_column(table, column)
  130. def rename_table_references(self, old_table, new_table):
  131. super().rename_table_references(old_table, new_table)
  132. self.to_reference.rename_table_references(old_table, new_table)
  133. def rename_column_references(self, table, old_column, new_column):
  134. super().rename_column_references(table, old_column, new_column)
  135. self.to_reference.rename_column_references(table, old_column, new_column)
  136. def __str__(self):
  137. suffix = self.suffix_template % {
  138. "to_table": self.to_reference.table,
  139. "to_column": self.to_reference.columns[0],
  140. }
  141. return self.create_fk_name(self.table, self.columns, suffix)
  142. class Statement(Reference):
  143. """
  144. Statement template and formatting parameters container.
  145. Allows keeping a reference to a statement without interpolating identifiers
  146. that might have to be adjusted if they're referencing a table or column
  147. that is removed
  148. """
  149. def __init__(self, template, **parts):
  150. self.template = template
  151. self.parts = parts
  152. def references_table(self, table):
  153. return any(
  154. hasattr(part, "references_table") and part.references_table(table)
  155. for part in self.parts.values()
  156. )
  157. def references_column(self, table, column):
  158. return any(
  159. hasattr(part, "references_column") and part.references_column(table, column)
  160. for part in self.parts.values()
  161. )
  162. def rename_table_references(self, old_table, new_table):
  163. for part in self.parts.values():
  164. if hasattr(part, "rename_table_references"):
  165. part.rename_table_references(old_table, new_table)
  166. def rename_column_references(self, table, old_column, new_column):
  167. for part in self.parts.values():
  168. if hasattr(part, "rename_column_references"):
  169. part.rename_column_references(table, old_column, new_column)
  170. def __str__(self):
  171. return self.template % self.parts
  172. class Expressions(TableColumns):
  173. def __init__(self, table, expressions, compiler, quote_value):
  174. self.compiler = compiler
  175. self.expressions = expressions
  176. self.quote_value = quote_value
  177. columns = [
  178. col.target.column
  179. for col in self.compiler.query._gen_cols([self.expressions])
  180. ]
  181. super().__init__(table, columns)
  182. def rename_table_references(self, old_table, new_table):
  183. if self.table != old_table:
  184. return
  185. self.expressions = self.expressions.relabeled_clone({old_table: new_table})
  186. super().rename_table_references(old_table, new_table)
  187. def rename_column_references(self, table, old_column, new_column):
  188. if self.table != table:
  189. return
  190. expressions = deepcopy(self.expressions)
  191. self.columns = []
  192. for col in self.compiler.query._gen_cols([expressions]):
  193. if col.target.column == old_column:
  194. col.target.column = new_column
  195. self.columns.append(col.target.column)
  196. self.expressions = expressions
  197. def __str__(self):
  198. sql, params = self.compiler.compile(self.expressions)
  199. params = map(self.quote_value, params)
  200. return sql % tuple(params)