"""add instruction tables Revision ID: f9c3a857d9e6 Revises: b61d5f5a1ded 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. revision = 'f9c3a857d9e6' down_revision = 'b61d5f5a1ded' 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(): # 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") # 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") # 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')] 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(): # 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 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 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 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")