118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
"""Add users table and update todos
|
|
|
|
Revision ID: 2_add_users_table_and_update_todos
|
|
Revises: 1_initial_create_todos_table
|
|
Create Date: 2023-07-20 10:00:00.000000
|
|
|
|
"""
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision = '2_add_users_table_and_update_todos'
|
|
down_revision = '1_initial_create_todos_table'
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def upgrade():
|
|
# Create users table if it doesn't exist
|
|
from sqlalchemy.exc import OperationalError
|
|
try:
|
|
# Use inspector to check if the table exists
|
|
from sqlalchemy import inspect
|
|
conn = op.get_bind()
|
|
inspector = inspect(conn)
|
|
if 'users' not in inspector.get_table_names():
|
|
op.create_table(
|
|
'users',
|
|
sa.Column('id', sa.Integer(), nullable=False),
|
|
sa.Column('email', sa.String(), nullable=False),
|
|
sa.Column('username', sa.String(), nullable=False),
|
|
sa.Column('hashed_password', sa.String(), nullable=False),
|
|
sa.Column('is_active', sa.Boolean(), nullable=True, default=True),
|
|
sa.Column('is_superuser', sa.Boolean(), nullable=True, default=False),
|
|
sa.Column(
|
|
'created_at',
|
|
sa.DateTime(timezone=True),
|
|
server_default=sa.text('(CURRENT_TIMESTAMP)'),
|
|
nullable=True,
|
|
),
|
|
sa.Column('updated_at', sa.DateTime(timezone=True), nullable=True),
|
|
sa.PrimaryKeyConstraint('id'),
|
|
)
|
|
op.create_index(op.f('ix_users_email'), 'users', ['email'], unique=True)
|
|
op.create_index(op.f('ix_users_id'), 'users', ['id'], unique=False)
|
|
op.create_index(op.f('ix_users_username'), 'users', ['username'], unique=True)
|
|
except OperationalError:
|
|
# Table already exists, skip creation
|
|
pass
|
|
|
|
# Add owner_id column to todos table if it doesn't exist
|
|
from sqlalchemy.exc import OperationalError
|
|
|
|
# Check if owner_id column exists in todos table
|
|
conn = op.get_bind()
|
|
inspector = inspect(conn)
|
|
columns = [col['name'] for col in inspector.get_columns('todos')]
|
|
|
|
if 'owner_id' not in columns:
|
|
try:
|
|
op.add_column('todos', sa.Column('owner_id', sa.Integer(), nullable=True))
|
|
|
|
# Create foreign key from todos to users
|
|
op.create_foreign_key(
|
|
'fk_todos_owner_id_users',
|
|
'todos', 'users',
|
|
['owner_id'], ['id'],
|
|
)
|
|
|
|
# Update all existing todos to have owner_id = 1 (first user)
|
|
op.execute("UPDATE todos SET owner_id = 1")
|
|
|
|
# Make owner_id non-nullable after updating existing todos
|
|
op.alter_column('todos', 'owner_id', nullable=False)
|
|
except OperationalError as e:
|
|
# Log the error but continue
|
|
print(f"Operation failed with error: {str(e)}")
|
|
# If this is because the column already exists, we can continue
|
|
pass
|
|
|
|
|
|
def downgrade():
|
|
from sqlalchemy import inspect
|
|
from sqlalchemy.exc import OperationalError
|
|
|
|
conn = op.get_bind()
|
|
inspector = inspect(conn)
|
|
|
|
# Check if the constraint exists before trying to drop it
|
|
try:
|
|
constraints = inspector.get_foreign_keys('todos')
|
|
constraint_name = 'fk_todos_owner_id_users'
|
|
fk_exists = any(constraint['name'] == constraint_name for constraint in constraints)
|
|
if fk_exists:
|
|
op.drop_constraint('fk_todos_owner_id_users', 'todos', type_='foreignkey')
|
|
except OperationalError:
|
|
pass
|
|
|
|
# Check if owner_id column exists before trying to drop it
|
|
columns = [col['name'] for col in inspector.get_columns('todos')]
|
|
if 'owner_id' in columns:
|
|
op.drop_column('todos', 'owner_id')
|
|
|
|
# Check if users table exists before trying to drop it
|
|
if 'users' in inspector.get_table_names():
|
|
# Check for each index before dropping
|
|
indices = inspector.get_indexes('users')
|
|
index_names = [index['name'] for index in indices]
|
|
|
|
if 'ix_users_username' in index_names:
|
|
op.drop_index(op.f('ix_users_username'), table_name='users')
|
|
if 'ix_users_id' in index_names:
|
|
op.drop_index(op.f('ix_users_id'), table_name='users')
|
|
if 'ix_users_email' in index_names:
|
|
op.drop_index(op.f('ix_users_email'), table_name='users')
|
|
|
|
op.drop_table('users')
|