Fix migration to handle existing tables and columns

- Update SQLite foreign key constraint implementation to use batch operations
- Add helper function to detect existing constraints
- Use batch_alter_table context manager for SQLite compatibility
- Improve error handling in downgrade operations for SQLite compatibility
- Add better error messages for debugging
This commit is contained in:
Automated Action 2025-05-16 02:20:55 +00:00
parent f1b8b8cbee
commit cddf794105

View File

@ -45,6 +45,15 @@ def index_exists(table_name, index_name):
return index_name in indexes return index_name in indexes
def get_constraint_names(table_name):
"""Get all constraint names for a table"""
conn = op.get_bind()
inspector = Inspector.from_engine(conn)
if not table_exists(table_name):
return []
return [fk.get('name', '') for fk in inspector.get_foreign_keys(table_name)]
def upgrade(): def upgrade():
# Create users table if it doesn't exist # Create users table if it doesn't exist
if not table_exists('users'): if not table_exists('users'):
@ -83,23 +92,42 @@ def upgrade():
if not column_exists('todos', 'owner_id'): if not column_exists('todos', 'owner_id'):
op.add_column('todos', sa.Column('owner_id', sa.Integer(), nullable=True)) op.add_column('todos', sa.Column('owner_id', sa.Integer(), nullable=True))
# Create foreign key constraint if it doesn't exist # For SQLite, we need to use batch operations to add a foreign key constraint
# SQLite doesn't support ALTER TABLE ADD CONSTRAINT, so we need to
# use SQLAlchemy's batch operations which recreate the table
existing_constraints = get_constraint_names('todos')
if 'fk_todos_owner_id_users' not in existing_constraints:
try: try:
op.create_foreign_key('fk_todos_owner_id_users', 'todos', 'users', ['owner_id'], ['id']) # Use batch operations to recreate the table with the foreign key
except (OperationalError, sa.exc.InternalError): with op.batch_alter_table('todos') as batch_op:
pass batch_op.create_foreign_key(
'fk_todos_owner_id_users',
'users',
['owner_id'],
['id']
)
except (OperationalError, sa.exc.InternalError) as e:
# Log error but allow migration to continue
print(f"Warning: Could not create foreign key constraint: {str(e)}")
def downgrade(): def downgrade():
# Try to drop foreign key constraint # Try to drop foreign key constraint using batch operations for SQLite
existing_constraints = get_constraint_names('todos')
if 'fk_todos_owner_id_users' in existing_constraints:
try: try:
op.drop_constraint('fk_todos_owner_id_users', 'todos', type_='foreignkey') with op.batch_alter_table('todos') as batch_op:
except (OperationalError, sa.exc.InternalError): batch_op.drop_constraint('fk_todos_owner_id_users', type_='foreignkey')
pass except (OperationalError, sa.exc.InternalError) as e:
print(f"Warning: Could not drop foreign key constraint: {str(e)}")
# Try to drop owner_id column from todos table # Try to drop owner_id column from todos table
if column_exists('todos', 'owner_id'): if column_exists('todos', 'owner_id'):
op.drop_column('todos', 'owner_id') try:
with op.batch_alter_table('todos') as batch_op:
batch_op.drop_column('owner_id')
except (OperationalError, sa.exc.InternalError) as e:
print(f"Warning: Could not drop column: {str(e)}")
# Try to drop indexes on users table # Try to drop indexes on users table
try: try: