From 33cac314752abbd285262ad589b2eba204dc13d2 Mon Sep 17 00:00:00 2001 From: Automated Action Date: Wed, 28 May 2025 22:48:30 +0000 Subject: [PATCH] Fix Alembic migration to handle existing tables --- migrations/versions/add_instruction_tables.py | 192 ++++++++++++------ 1 file changed, 131 insertions(+), 61 deletions(-) diff --git a/migrations/versions/add_instruction_tables.py b/migrations/versions/add_instruction_tables.py index be32d32..fb92bab 100644 --- a/migrations/versions/add_instruction_tables.py +++ b/migrations/versions/add_instruction_tables.py @@ -7,6 +7,7 @@ Create Date: 2023-10-10 12:00:00.000000 """ from alembic import op import sqlalchemy as sa +from sqlalchemy.engine import reflection # revision identifiers, used by Alembic. @@ -16,74 +17,143 @@ branch_labels = None depends_on = None +# Helper function to check if a table exists +def table_exists(table_name): + bind = op.get_bind() + inspector = reflection.Inspector.from_engine(bind) + tables = inspector.get_table_names() + return table_name in tables + + def upgrade(): - # Create instructions table - op.create_table( - 'instructions', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('transaction_id', sa.Integer(), nullable=True), - sa.Column('program_id', sa.String(), nullable=True), - sa.Column('instruction_type', sa.String(), nullable=True), - sa.Column('instruction_index', sa.Integer(), nullable=True), - sa.Column('accounts', sa.Text(), nullable=True), - sa.Column('data', sa.Text(), nullable=True), - sa.Column('parsed_data', sa.JSON(), nullable=True), - sa.ForeignKeyConstraint(['transaction_id'], ['transactions.id'], ), - sa.PrimaryKeyConstraint('id'), - ) - op.create_index(op.f('ix_instructions_id'), 'instructions', ['id'], unique=False) - op.create_index(op.f('ix_instructions_program_id'), 'instructions', ['program_id'], unique=False) - op.create_index(op.f('ix_instructions_instruction_type'), 'instructions', ['instruction_type'], unique=False) + # Check if instructions table exists + if not table_exists('instructions'): + # Create instructions table + op.create_table( + 'instructions', + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('transaction_id', sa.Integer(), nullable=True), + sa.Column('program_id', sa.String(), nullable=True), + sa.Column('instruction_type', sa.String(), nullable=True), + sa.Column('instruction_index', sa.Integer(), nullable=True), + sa.Column('accounts', sa.Text(), nullable=True), + sa.Column('data', sa.Text(), nullable=True), + sa.Column('parsed_data', sa.JSON(), nullable=True), + sa.ForeignKeyConstraint(['transaction_id'], ['transactions.id'], ), + sa.PrimaryKeyConstraint('id'), + ) + op.create_index(op.f('ix_instructions_id'), 'instructions', ['id'], unique=False) + op.create_index(op.f('ix_instructions_program_id'), 'instructions', ['program_id'], unique=False) + op.create_index(op.f('ix_instructions_instruction_type'), 'instructions', ['instruction_type'], unique=False) + else: + print("Table 'instructions' already exists, skipping creation") - # Create swaps table - op.create_table( - 'swaps', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('transaction_id', sa.Integer(), nullable=True), - sa.Column('instruction_id', sa.Integer(), nullable=True), - sa.Column('program_id', sa.String(), nullable=True), - sa.Column('token_in_address', sa.String(), nullable=True), - sa.Column('token_out_address', sa.String(), nullable=True), - sa.Column('amount_in', sa.Float(), nullable=True), - sa.Column('amount_out', sa.Float(), nullable=True), - sa.Column('user_account', sa.String(), nullable=True), - sa.Column('timestamp', sa.DateTime(), nullable=True), - sa.ForeignKeyConstraint(['instruction_id'], ['instructions.id'], ), - sa.ForeignKeyConstraint(['transaction_id'], ['transactions.id'], ), - sa.PrimaryKeyConstraint('id'), - ) - op.create_index(op.f('ix_swaps_id'), 'swaps', ['id'], unique=False) - op.create_index(op.f('ix_swaps_program_id'), 'swaps', ['program_id'], unique=False) - op.create_index(op.f('ix_swaps_token_in_address'), 'swaps', ['token_in_address'], unique=False) - op.create_index(op.f('ix_swaps_token_out_address'), 'swaps', ['token_out_address'], unique=False) - op.create_index(op.f('ix_swaps_user_account'), 'swaps', ['user_account'], unique=False) + # Check if swaps table exists + if not table_exists('swaps'): + # Create swaps table + op.create_table( + 'swaps', + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('transaction_id', sa.Integer(), nullable=True), + sa.Column('instruction_id', sa.Integer(), nullable=True), + sa.Column('program_id', sa.String(), nullable=True), + sa.Column('token_in_address', sa.String(), nullable=True), + sa.Column('token_out_address', sa.String(), nullable=True), + sa.Column('amount_in', sa.Float(), nullable=True), + sa.Column('amount_out', sa.Float(), nullable=True), + sa.Column('user_account', sa.String(), nullable=True), + sa.Column('timestamp', sa.DateTime(), nullable=True), + sa.ForeignKeyConstraint(['instruction_id'], ['instructions.id'], ), + sa.ForeignKeyConstraint(['transaction_id'], ['transactions.id'], ), + sa.PrimaryKeyConstraint('id'), + ) + op.create_index(op.f('ix_swaps_id'), 'swaps', ['id'], unique=False) + op.create_index(op.f('ix_swaps_program_id'), 'swaps', ['program_id'], unique=False) + op.create_index(op.f('ix_swaps_token_in_address'), 'swaps', ['token_in_address'], unique=False) + op.create_index(op.f('ix_swaps_token_out_address'), 'swaps', ['token_out_address'], unique=False) + op.create_index(op.f('ix_swaps_user_account'), 'swaps', ['user_account'], unique=False) + else: + print("Table 'swaps' already exists, skipping creation") - # Add instruction_id to token_transfers - op.add_column('token_transfers', sa.Column('instruction_id', sa.Integer(), nullable=True)) - op.create_foreign_key(None, 'token_transfers', 'instructions', ['instruction_id'], ['id']) + # Check if instruction_id column exists in token_transfers + bind = op.get_bind() + inspector = reflection.Inspector.from_engine(bind) + columns = [col['name'] for col in inspector.get_columns('token_transfers')] - # Add swap_sequence to arbitrage_events - op.add_column('arbitrage_events', sa.Column('swap_sequence', sa.Text(), nullable=True)) + if 'instruction_id' not in columns: + # Add instruction_id to token_transfers + op.add_column('token_transfers', sa.Column('instruction_id', sa.Integer(), nullable=True)) + op.create_foreign_key(None, 'token_transfers', 'instructions', ['instruction_id'], ['id']) + else: + print("Column 'instruction_id' already exists in 'token_transfers', skipping addition") + + # Check if swap_sequence column exists in arbitrage_events + columns = [col['name'] for col in inspector.get_columns('arbitrage_events')] + + if 'swap_sequence' not in columns: + # Add swap_sequence to arbitrage_events + op.add_column('arbitrage_events', sa.Column('swap_sequence', sa.Text(), nullable=True)) + else: + print("Column 'swap_sequence' already exists in 'arbitrage_events', skipping addition") def downgrade(): - # Remove swap_sequence from arbitrage_events - op.drop_column('arbitrage_events', 'swap_sequence') + # Helper function to safely drop a column if it exists + def safe_drop_column(table_name, column_name): + bind = op.get_bind() + inspector = reflection.Inspector.from_engine(bind) + columns = [col['name'] for col in inspector.get_columns(table_name)] + if column_name in columns: + op.drop_column(table_name, column_name) + else: + print(f"Column '{column_name}' does not exist in '{table_name}', skipping drop") + + # Remove swap_sequence from arbitrage_events if it exists + if table_exists('arbitrage_events'): + safe_drop_column('arbitrage_events', 'swap_sequence') - # Remove instruction_id from token_transfers - op.drop_constraint(None, 'token_transfers', type_='foreignkey') - op.drop_column('token_transfers', 'instruction_id') + # Remove instruction_id from token_transfers if it exists + if table_exists('token_transfers'): + inspector = reflection.Inspector.from_engine(op.get_bind()) + + # Check if foreign key exists before trying to drop it + fk_constraints = inspector.get_foreign_keys('token_transfers') + has_instruction_fk = any( + fk['referred_table'] == 'instructions' and 'instruction_id' in fk['constrained_columns'] + for fk in fk_constraints + ) + + if has_instruction_fk: + # Need to iterate through constraints to find the right one + for fk in fk_constraints: + if fk['referred_table'] == 'instructions' and 'instruction_id' in fk['constrained_columns']: + op.drop_constraint(fk['name'], 'token_transfers', type_='foreignkey') + break + + safe_drop_column('token_transfers', 'instruction_id') - # Drop swaps table - op.drop_index(op.f('ix_swaps_user_account'), table_name='swaps') - op.drop_index(op.f('ix_swaps_token_out_address'), table_name='swaps') - op.drop_index(op.f('ix_swaps_token_in_address'), table_name='swaps') - op.drop_index(op.f('ix_swaps_program_id'), table_name='swaps') - op.drop_index(op.f('ix_swaps_id'), table_name='swaps') - op.drop_table('swaps') + # Drop swaps table if it exists + if table_exists('swaps'): + # Drop indexes first + inspector = reflection.Inspector.from_engine(op.get_bind()) + indexes = inspector.get_indexes('swaps') + + for index in indexes: + op.drop_index(index['name'], table_name='swaps') + + op.drop_table('swaps') + else: + print("Table 'swaps' does not exist, skipping drop") - # Drop instructions table - op.drop_index(op.f('ix_instructions_instruction_type'), table_name='instructions') - op.drop_index(op.f('ix_instructions_program_id'), table_name='instructions') - op.drop_index(op.f('ix_instructions_id'), table_name='instructions') - op.drop_table('instructions') \ No newline at end of file + # Drop instructions table if it exists + if table_exists('instructions'): + # Drop indexes first + inspector = reflection.Inspector.from_engine(op.get_bind()) + indexes = inspector.get_indexes('instructions') + + for index in indexes: + op.drop_index(index['name'], table_name='instructions') + + op.drop_table('instructions') + else: + print("Table 'instructions' does not exist, skipping drop") \ No newline at end of file