"""Update user table Revision ID: b10670c03dd5 Revises: 2f1211949ecc Create Date: 2025-11-28 04:55:31.737538 """ from typing import Sequence, Union from alembic import op import sqlalchemy as sa import open_webui.internal.db import json import time # revision identifiers, used by Alembic. revision: str = 'b10670c03dd5' down_revision: Union[str, None] = '2f1211949ecc' branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None def _drop_sqlite_indexes_for_column(table_name, column_name, conn): """ SQLite requires manual removal of any indexes referencing a column before ALTER TABLE ... DROP COLUMN can succeed. """ indexes = conn.execute(sa.text(f"PRAGMA index_list('{table_name}')")).fetchall() for idx in indexes: index_name = idx[1] # index name # Get indexed columns idx_info = conn.execute(sa.text(f"PRAGMA index_info('{index_name}')")).fetchall() indexed_cols = [row[2] for row in idx_info] # col names if column_name in indexed_cols: conn.execute(sa.text(f'DROP INDEX IF EXISTS {index_name}')) def _convert_column_to_json(table: str, column: str): conn = op.get_bind() dialect = conn.dialect.name # SQLite cannot ALTER COLUMN → must recreate column if dialect == 'sqlite': # 1. Add temporary column op.add_column(table, sa.Column(f'{column}_json', sa.JSON(), nullable=True)) # 2. Load old data rows = conn.execute(sa.text(f'SELECT id, {column} FROM "{table}"')).fetchall() for row in rows: uid, raw = row if raw is None: parsed = None else: try: parsed = json.loads(raw) except Exception: parsed = None # fallback safe behavior conn.execute( sa.text(f'UPDATE "{table}" SET {column}_json = :val WHERE id = :id'), {'val': json.dumps(parsed) if parsed else None, 'id': uid}, ) # 3. Drop old TEXT column op.drop_column(table, column) # 4. Rename new JSON column → original name op.alter_column(table, f'{column}_json', new_column_name=column) else: # PostgreSQL supports direct CAST op.alter_column( table, column, type_=sa.JSON(), postgresql_using=f'{column}::json', ) def _convert_column_to_text(table: str, column: str): conn = op.get_bind() dialect = conn.dialect.name if dialect == 'sqlite': op.add_column(table, sa.Column(f'{column}_text', sa.Text(), nullable=True)) rows = conn.execute(sa.text(f'SELECT id, {column} FROM "{table}"')).fetchall() for uid, raw in rows: conn.execute( sa.text(f'UPDATE "{table}" SET {column}_text = :val WHERE id = :id'), {'val': json.dumps(raw) if raw else None, 'id': uid}, ) op.drop_column(table, column) op.alter_column(table, f'{column}_text', new_column_name=column) else: op.alter_column( table, column, type_=sa.Text(), postgresql_using=f'to_json({column})::text', ) def upgrade() -> None: op.add_column('user', sa.Column('profile_banner_image_url', sa.Text(), nullable=True)) op.add_column('user', sa.Column('timezone', sa.String(), nullable=True)) op.add_column('user', sa.Column('presence_state', sa.String(), nullable=True)) op.add_column('user', sa.Column('status_emoji', sa.String(), nullable=True)) op.add_column('user', sa.Column('status_message', sa.Text(), nullable=True)) op.add_column('user', sa.Column('status_expires_at', sa.BigInteger(), nullable=True)) op.add_column('user', sa.Column('oauth', sa.JSON(), nullable=True)) # Convert info (TEXT/JSONField) → JSON _convert_column_to_json('user', 'info') # Convert settings (TEXT/JSONField) → JSON _convert_column_to_json('user', 'settings') op.create_table( 'api_key', sa.Column('id', sa.Text(), primary_key=True, unique=True), sa.Column('user_id', sa.Text(), sa.ForeignKey('user.id', ondelete='CASCADE')), sa.Column('key', sa.Text(), unique=True, nullable=False), sa.Column('data', sa.JSON(), nullable=True), sa.Column('expires_at', sa.BigInteger(), nullable=True), sa.Column('last_used_at', sa.BigInteger(), nullable=True), sa.Column('created_at', sa.BigInteger(), nullable=False), sa.Column('updated_at', sa.BigInteger(), nullable=False), ) conn = op.get_bind() users = conn.execute(sa.text('SELECT id, oauth_sub FROM "user" WHERE oauth_sub IS NOT NULL')).fetchall() for uid, oauth_sub in users: if oauth_sub: # Example formats supported: # provider@sub # plain sub (stored as {"oidc": {"sub": sub}}) if '@' in oauth_sub: provider, sub = oauth_sub.split('@', 1) else: provider, sub = 'oidc', oauth_sub oauth_json = json.dumps({provider: {'sub': sub}}) conn.execute( sa.text('UPDATE "user" SET oauth = :oauth WHERE id = :id'), {'oauth': oauth_json, 'id': uid}, ) users_with_keys = conn.execute(sa.text('SELECT id, api_key FROM "user" WHERE api_key IS NOT NULL')).fetchall() now = int(time.time()) for uid, api_key in users_with_keys: if api_key: conn.execute( sa.text(""" INSERT INTO api_key (id, user_id, key, created_at, updated_at) VALUES (:id, :user_id, :key, :created_at, :updated_at) """), { 'id': f'key_{uid}', 'user_id': uid, 'key': api_key, 'created_at': now, 'updated_at': now, }, ) if conn.dialect.name == 'sqlite': _drop_sqlite_indexes_for_column('user', 'api_key', conn) _drop_sqlite_indexes_for_column('user', 'oauth_sub', conn) with op.batch_alter_table('user') as batch_op: batch_op.drop_column('api_key') batch_op.drop_column('oauth_sub') def downgrade() -> None: # --- 1. Restore old oauth_sub column --- op.add_column('user', sa.Column('oauth_sub', sa.Text(), nullable=True)) conn = op.get_bind() users = conn.execute(sa.text('SELECT id, oauth FROM "user" WHERE oauth IS NOT NULL')).fetchall() for uid, oauth in users: try: data = json.loads(oauth) provider = list(data.keys())[0] sub = data[provider].get('sub') oauth_sub = f'{provider}@{sub}' except Exception: oauth_sub = None conn.execute( sa.text('UPDATE "user" SET oauth_sub = :oauth_sub WHERE id = :id'), {'oauth_sub': oauth_sub, 'id': uid}, ) op.drop_column('user', 'oauth') # --- 2. Restore api_key field --- op.add_column('user', sa.Column('api_key', sa.String(), nullable=True)) # Restore values from api_key keys = conn.execute(sa.text('SELECT user_id, key FROM api_key')).fetchall() for uid, key in keys: conn.execute( sa.text('UPDATE "user" SET api_key = :key WHERE id = :id'), {'key': key, 'id': uid}, ) # Drop new table op.drop_table('api_key') with op.batch_alter_table('user') as batch_op: batch_op.drop_column('profile_banner_image_url') batch_op.drop_column('timezone') batch_op.drop_column('presence_state') batch_op.drop_column('status_emoji') batch_op.drop_column('status_message') batch_op.drop_column('status_expires_at') # Convert info (JSON) → TEXT _convert_column_to_text('user', 'info') # Convert settings (JSON) → TEXT _convert_column_to_text('user', 'settings')