Files
meshview/alembic/versions/add_import_time_us_columns.py
T
2025-11-03 14:11:42 -08:00

63 lines
2.7 KiB
Python

"""add import_time_us columns
Revision ID: add_time_us_cols
Revises: c88468b7ab0b
Create Date: 2025-11-03 14:10:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = 'add_time_us_cols'
down_revision: Union[str, None] = 'c88468b7ab0b'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# Check if columns already exist, add them if they don't
conn = op.get_bind()
inspector = sa.inspect(conn)
# Add import_time_us to packet table
packet_columns = [col['name'] for col in inspector.get_columns('packet')]
if 'import_time_us' not in packet_columns:
with op.batch_alter_table('packet', schema=None) as batch_op:
batch_op.add_column(sa.Column('import_time_us', sa.BigInteger(), nullable=True))
op.create_index('idx_packet_import_time_us', 'packet', [sa.text('import_time_us DESC')], unique=False)
op.create_index('idx_packet_from_node_time_us', 'packet', ['from_node_id', sa.text('import_time_us DESC')], unique=False)
# Add import_time_us to packet_seen table
packet_seen_columns = [col['name'] for col in inspector.get_columns('packet_seen')]
if 'import_time_us' not in packet_seen_columns:
with op.batch_alter_table('packet_seen', schema=None) as batch_op:
batch_op.add_column(sa.Column('import_time_us', sa.BigInteger(), nullable=True))
op.create_index('idx_packet_seen_import_time_us', 'packet_seen', ['import_time_us'], unique=False)
# Add import_time_us to traceroute table
traceroute_columns = [col['name'] for col in inspector.get_columns('traceroute')]
if 'import_time_us' not in traceroute_columns:
with op.batch_alter_table('traceroute', schema=None) as batch_op:
batch_op.add_column(sa.Column('import_time_us', sa.BigInteger(), nullable=True))
op.create_index('idx_traceroute_import_time_us', 'traceroute', ['import_time_us'], unique=False)
def downgrade() -> None:
# Drop indexes and columns
op.drop_index('idx_traceroute_import_time_us', table_name='traceroute')
with op.batch_alter_table('traceroute', schema=None) as batch_op:
batch_op.drop_column('import_time_us')
op.drop_index('idx_packet_seen_import_time_us', table_name='packet_seen')
with op.batch_alter_table('packet_seen', schema=None) as batch_op:
batch_op.drop_column('import_time_us')
op.drop_index('idx_packet_from_node_time_us', table_name='packet')
op.drop_index('idx_packet_import_time_us', table_name='packet')
with op.batch_alter_table('packet', schema=None) as batch_op:
batch_op.drop_column('import_time_us')