82 lines
2.2 KiB
Python
82 lines
2.2 KiB
Python
"""Initialize empty Tenant codes
|
|
|
|
Revision ID: 867deef0888b
|
|
Revises: cab899dbb213
|
|
Create Date: 2025-04-03 09:01:20.446536
|
|
|
|
"""
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
import uuid
|
|
from sqlalchemy.sql import table, column, select, or_
|
|
from sqlalchemy.orm import Session
|
|
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision = '867deef0888b'
|
|
down_revision = 'cab899dbb213'
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def upgrade():
|
|
# Create a reference to the tenant table
|
|
tenant_table = table('tenant',
|
|
column('id', sa.Integer),
|
|
column('code', sa.String(50)),
|
|
schema='public' # Assuming the table is in the 'public' schema
|
|
)
|
|
|
|
# Get a connection
|
|
connection = op.get_bind()
|
|
session = Session(bind=connection)
|
|
|
|
try:
|
|
# Find all tenants with empty or null code
|
|
# Note the updated select syntax for SQLAlchemy 2.0
|
|
query = select(tenant_table.c.id).where(
|
|
or_(
|
|
tenant_table.c.code == None,
|
|
tenant_table.c.code == ''
|
|
)
|
|
)
|
|
|
|
results = connection.execute(query)
|
|
|
|
# Update each tenant with a UUID-based code
|
|
for row in results:
|
|
tenant_id = row[0]
|
|
code = f"TENANT-{str(uuid.uuid4())}"
|
|
|
|
# Update the tenant record
|
|
update_stmt = tenant_table.update().where(
|
|
tenant_table.c.id == tenant_id
|
|
).values(
|
|
code=code
|
|
)
|
|
|
|
connection.execute(update_stmt)
|
|
|
|
# Commit changes
|
|
session.commit()
|
|
|
|
# Log how many records were updated
|
|
count_query = select(sa.func.count()).select_from(tenant_table).where(
|
|
tenant_table.c.code.like('TENANT-%')
|
|
)
|
|
|
|
updated_count = connection.execute(count_query).scalar()
|
|
|
|
print(f"Updated {updated_count} tenant records with UUID-based codes")
|
|
|
|
except Exception as e:
|
|
session.rollback()
|
|
print(f"Error updating tenant codes: {str(e)}")
|
|
raise e
|
|
finally:
|
|
session.close()
|
|
|
|
|
|
def downgrade():
|
|
# No downgrade needed for this data migration
|
|
pass |