toimpl.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162
  1. from . import ops
  2. from . import Operations
  3. from sqlalchemy import schema as sa_schema
  4. @Operations.implementation_for(ops.AlterColumnOp)
  5. def alter_column(operations, operation):
  6. compiler = operations.impl.dialect.statement_compiler(
  7. operations.impl.dialect,
  8. None
  9. )
  10. existing_type = operation.existing_type
  11. existing_nullable = operation.existing_nullable
  12. existing_server_default = operation.existing_server_default
  13. type_ = operation.modify_type
  14. column_name = operation.column_name
  15. table_name = operation.table_name
  16. schema = operation.schema
  17. server_default = operation.modify_server_default
  18. new_column_name = operation.modify_name
  19. nullable = operation.modify_nullable
  20. def _count_constraint(constraint):
  21. return not isinstance(
  22. constraint,
  23. sa_schema.PrimaryKeyConstraint) and \
  24. (not constraint._create_rule or
  25. constraint._create_rule(compiler))
  26. if existing_type and type_:
  27. t = operations.schema_obj.table(
  28. table_name,
  29. sa_schema.Column(column_name, existing_type),
  30. schema=schema
  31. )
  32. for constraint in t.constraints:
  33. if _count_constraint(constraint):
  34. operations.impl.drop_constraint(constraint)
  35. operations.impl.alter_column(
  36. table_name, column_name,
  37. nullable=nullable,
  38. server_default=server_default,
  39. name=new_column_name,
  40. type_=type_,
  41. schema=schema,
  42. existing_type=existing_type,
  43. existing_server_default=existing_server_default,
  44. existing_nullable=existing_nullable,
  45. **operation.kw
  46. )
  47. if type_:
  48. t = operations.schema_obj.table(
  49. table_name,
  50. operations.schema_obj.column(column_name, type_),
  51. schema=schema
  52. )
  53. for constraint in t.constraints:
  54. if _count_constraint(constraint):
  55. operations.impl.add_constraint(constraint)
  56. @Operations.implementation_for(ops.DropTableOp)
  57. def drop_table(operations, operation):
  58. operations.impl.drop_table(
  59. operation.to_table(operations.migration_context)
  60. )
  61. @Operations.implementation_for(ops.DropColumnOp)
  62. def drop_column(operations, operation):
  63. column = operation.to_column(operations.migration_context)
  64. operations.impl.drop_column(
  65. operation.table_name,
  66. column,
  67. schema=operation.schema,
  68. **operation.kw
  69. )
  70. @Operations.implementation_for(ops.CreateIndexOp)
  71. def create_index(operations, operation):
  72. idx = operation.to_index(operations.migration_context)
  73. operations.impl.create_index(idx)
  74. @Operations.implementation_for(ops.DropIndexOp)
  75. def drop_index(operations, operation):
  76. operations.impl.drop_index(
  77. operation.to_index(operations.migration_context)
  78. )
  79. @Operations.implementation_for(ops.CreateTableOp)
  80. def create_table(operations, operation):
  81. table = operation.to_table(operations.migration_context)
  82. operations.impl.create_table(table)
  83. return table
  84. @Operations.implementation_for(ops.RenameTableOp)
  85. def rename_table(operations, operation):
  86. operations.impl.rename_table(
  87. operation.table_name,
  88. operation.new_table_name,
  89. schema=operation.schema)
  90. @Operations.implementation_for(ops.AddColumnOp)
  91. def add_column(operations, operation):
  92. table_name = operation.table_name
  93. column = operation.column
  94. schema = operation.schema
  95. t = operations.schema_obj.table(table_name, column, schema=schema)
  96. operations.impl.add_column(
  97. table_name,
  98. column,
  99. schema=schema
  100. )
  101. for constraint in t.constraints:
  102. if not isinstance(constraint, sa_schema.PrimaryKeyConstraint):
  103. operations.impl.add_constraint(constraint)
  104. for index in t.indexes:
  105. operations.impl.create_index(index)
  106. @Operations.implementation_for(ops.AddConstraintOp)
  107. def create_constraint(operations, operation):
  108. operations.impl.add_constraint(
  109. operation.to_constraint(operations.migration_context)
  110. )
  111. @Operations.implementation_for(ops.DropConstraintOp)
  112. def drop_constraint(operations, operation):
  113. operations.impl.drop_constraint(
  114. operations.schema_obj.generic_constraint(
  115. operation.constraint_name,
  116. operation.table_name,
  117. operation.constraint_type,
  118. schema=operation.schema,
  119. )
  120. )
  121. @Operations.implementation_for(ops.BulkInsertOp)
  122. def bulk_insert(operations, operation):
  123. operations.impl.bulk_insert(
  124. operation.table, operation.rows, multiinsert=operation.multiinsert)
  125. @Operations.implementation_for(ops.ExecuteSQLOp)
  126. def execute_sql(operations, operation):
  127. operations.migration_context.impl.execute(
  128. operation.sqltext,
  129. execution_options=operation.execution_options
  130. )