Fix Alembic migration to handle existing tables
This commit is contained in:
Automated Action
parent
cae8527a4b
commit
33cac31475
@ -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')
|
||||
# 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")
|
||||
Loading…
x
Reference in New Issue
Block a user