From f0e88dadc8502ea05b1a00f3155bee7d1cf32249 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 14 May 2026 13:08:53 +0900 Subject: [PATCH] refac --- .../38d63c18f30f_add_oauth_session_table.py | 81 ++++--- .../b10670c03dd5_update_user_table.py | 223 +++++++++++------- 2 files changed, 181 insertions(+), 123 deletions(-) diff --git a/backend/open_webui/migrations/versions/38d63c18f30f_add_oauth_session_table.py b/backend/open_webui/migrations/versions/38d63c18f30f_add_oauth_session_table.py index d5442026e3..8fd7656633 100644 --- a/backend/open_webui/migrations/versions/38d63c18f30f_add_oauth_session_table.py +++ b/backend/open_webui/migrations/versions/38d63c18f30f_add_oauth_session_table.py @@ -19,50 +19,63 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: - # Ensure 'id' column in 'user' table is unique and primary key (ForeignKey constraint) inspector = sa.inspect(op.get_bind()) - columns = inspector.get_columns('user') + existing_tables = set(inspector.get_table_names()) + # ── Fix user.id PK only when it is genuinely wrong ──────────────── + # On databases migrated from the old Peewee layer the PK may have + # drifted. Fresh installs (init migration 7e5b5dc7342b) already + # have a correct PK on 'id', so we must NOT attempt to re-create it + # — PostgreSQL will reject a duplicate PRIMARY KEY constraint. pk_columns = inspector.get_pk_constraint('user')['constrained_columns'] - id_column = next((col for col in columns if col['name'] == 'id'), None) - if id_column and not id_column.get('unique', False): + if pk_columns != ['id']: + columns = inspector.get_columns('user') + id_column = next((col for col in columns if col['name'] == 'id'), None) unique_constraints = inspector.get_unique_constraints('user') unique_columns = {tuple(u['column_names']) for u in unique_constraints} - with op.batch_alter_table('user') as batch_op: - # If primary key is wrong, drop it - if pk_columns and pk_columns != ['id']: - batch_op.drop_constraint(inspector.get_pk_constraint('user')['name'], type_='primary') + if id_column: + with op.batch_alter_table('user') as batch_op: + if pk_columns: + batch_op.drop_constraint( + inspector.get_pk_constraint('user')['name'], + type_='primary', + ) + if ('id',) not in unique_columns: + batch_op.create_unique_constraint('uq_user_id', ['id']) + batch_op.create_primary_key('pk_user_id', ['id']) - # Add unique constraint if missing - if ('id',) not in unique_columns: - batch_op.create_unique_constraint('uq_user_id', ['id']) + # ── Create oauth_session table (idempotent) ─────────────────────── + if 'oauth_session' not in existing_tables: + op.create_table( + 'oauth_session', + sa.Column('id', sa.Text(), primary_key=True, nullable=False, unique=True), + sa.Column( + 'user_id', + sa.Text(), + sa.ForeignKey('user.id', ondelete='CASCADE'), + nullable=False, + ), + sa.Column('provider', sa.Text(), nullable=False), + sa.Column('token', sa.Text(), nullable=False), + sa.Column('expires_at', sa.BigInteger(), nullable=False), + sa.Column('created_at', sa.BigInteger(), nullable=False), + sa.Column('updated_at', sa.BigInteger(), nullable=False), + ) - # Re-create correct primary key - batch_op.create_primary_key('pk_user_id', ['id']) + # Create indexes (idempotent — no-ops when table was just created + # with the columns above, and safe to call if indexes already exist). + existing_indexes = { + idx['name'] for idx in inspector.get_indexes('oauth_session') + } if 'oauth_session' in existing_tables else set() - # Create oauth_session table - op.create_table( - 'oauth_session', - sa.Column('id', sa.Text(), primary_key=True, nullable=False, unique=True), - sa.Column( - 'user_id', - sa.Text(), - sa.ForeignKey('user.id', ondelete='CASCADE'), - nullable=False, - ), - sa.Column('provider', sa.Text(), nullable=False), - sa.Column('token', sa.Text(), nullable=False), - sa.Column('expires_at', sa.BigInteger(), nullable=False), - sa.Column('created_at', sa.BigInteger(), nullable=False), - sa.Column('updated_at', sa.BigInteger(), nullable=False), - ) - - # Create indexes for better performance - op.create_index('idx_oauth_session_user_id', 'oauth_session', ['user_id']) - op.create_index('idx_oauth_session_expires_at', 'oauth_session', ['expires_at']) - op.create_index('idx_oauth_session_user_provider', 'oauth_session', ['user_id', 'provider']) + if 'idx_oauth_session_user_id' not in existing_indexes: + op.create_index('idx_oauth_session_user_id', 'oauth_session', ['user_id']) + if 'idx_oauth_session_expires_at' not in existing_indexes: + op.create_index('idx_oauth_session_expires_at', 'oauth_session', ['expires_at']) + if 'idx_oauth_session_user_provider' not in existing_indexes: + op.create_index('idx_oauth_session_user_provider', 'oauth_session', ['user_id', 'provider']) def downgrade() -> None: diff --git a/backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py b/backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py index 556b3a8eac..a609426ce8 100644 --- a/backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py +++ b/backend/open_webui/migrations/versions/b10670c03dd5_update_user_table.py @@ -20,20 +20,49 @@ down_revision: Union[str, None] = '2f1211949ecc' branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None +# ── Ad-hoc table references for Core DML ───────────────────────────────── +# These are lightweight table() / column() references used only inside this +# migration for SELECT / UPDATE / INSERT — they do NOT create or alter +# anything on disk. + +_user = sa.table( + 'user', + sa.column('id', sa.Text), + sa.column('oauth_sub', sa.Text), + sa.column('oauth', sa.JSON), + sa.column('api_key', sa.Text), + sa.column('info', sa.Text), + sa.column('settings', sa.Text), +) + +_api_key = sa.table( + 'api_key', + sa.column('id', sa.Text), + sa.column('user_id', sa.Text), + sa.column('key', sa.Text), + sa.column('created_at', sa.BigInteger), + sa.column('updated_at', sa.BigInteger), +) + 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. + SQLite requires manual removal of any user-created indexes referencing + a column before ALTER TABLE ... DROP COLUMN can succeed. + + NOTE: PRAGMAs have no Core equivalent — raw text is unavoidable here. """ 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 + # Skip system-managed autoindexes (PK / UNIQUE constraints) — they + # cannot be dropped directly and will disappear when the column is + # removed via batch_alter_table. + if index_name.startswith('sqlite_autoindex_'): + continue idx_info = conn.execute(sa.text(f"PRAGMA index_info('{index_name}')")).fetchall() - - indexed_cols = [row[2] for row in idx_info] # col names + indexed_cols = [row[2] for row in idx_info] if column_name in indexed_cols: conn.execute(sa.text(f'DROP INDEX IF EXISTS {index_name}')) @@ -42,33 +71,31 @@ def _convert_column_to_json(table: str, column: str): conn = op.get_bind() dialect = conn.dialect.name + t = sa.table(table, sa.column('id', sa.Text), sa.column(column, sa.Text)) + t_json = sa.column(f'{column}_json', sa.JSON) + # 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() + rows = conn.execute(sa.select(t.c.id, t.c[column])).fetchall() - for row in rows: - uid, raw = row + for uid, raw in rows: if raw is None: parsed = None else: try: parsed = json.loads(raw) except Exception: - parsed = None # fallback safe behavior + parsed = None conn.execute( - sa.text(f'UPDATE "{table}" SET {column}_json = :val WHERE id = :id'), - {'val': json.dumps(parsed) if parsed else None, 'id': uid}, + sa.update(sa.table(table, sa.column('id'), t_json)) + .where(sa.column('id') == uid) + .values({f'{column}_json': json.dumps(parsed) if parsed else None}) ) - # 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: @@ -85,15 +112,19 @@ def _convert_column_to_text(table: str, column: str): conn = op.get_bind() dialect = conn.dialect.name + t = sa.table(table, sa.column('id', sa.Text), sa.column(column)) + t_text = sa.column(f'{column}_text', sa.Text) + 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() + rows = conn.execute(sa.select(t.c.id, t.c[column])).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}, + sa.update(sa.table(table, sa.column('id'), t_text)) + .where(sa.column('id') == uid) + .values({f'{column}_text': json.dumps(raw) if raw else None}) ) op.drop_column(table, column) @@ -109,88 +140,102 @@ def _convert_column_to_text(table: str, column: str): 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)) + conn = op.get_bind() + inspector = sa.inspect(conn) + existing_tables = set(inspector.get_table_names()) + user_columns = {c['name'] for c in inspector.get_columns('user')} - 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)) + # ── Add new columns (idempotent) ────────────────────────────────── + for col_name, col_type in [ + ('profile_banner_image_url', sa.Text()), + ('timezone', sa.String()), + ('presence_state', sa.String()), + ('status_emoji', sa.String()), + ('status_message', sa.Text()), + ('status_expires_at', sa.BigInteger()), + ('oauth', sa.JSON()), + ]: + if col_name not in user_columns: + op.add_column('user', sa.Column(col_name, col_type, 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), - ) + # ── Create api_key table (idempotent) ───────────────────────────── + if 'api_key' not in existing_tables: + 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() + # ── Migrate oauth_sub → oauth JSON (only if old column still exists) + if 'oauth_sub' in user_columns: + rows = conn.execute( + sa.select(_user.c.id, _user.c.oauth_sub) + .where(_user.c.oauth_sub.is_not(None)) + ).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 + for uid, oauth_sub in rows: + if oauth_sub: + provider, sub = (oauth_sub.split('@', 1) if '@' in oauth_sub + else ('oidc', oauth_sub)) + conn.execute( + sa.update(_user) + .where(_user.c.id == uid) + .values(oauth=json.dumps({provider: {'sub': 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}, - ) + # ── Migrate api_key column → api_key table (only if old column still exists) + if 'api_key' in user_columns: + rows = conn.execute( + sa.select(_user.c.id, _user.c.api_key) + .where(_user.c.api_key.is_not(None)) + ).fetchall() + now = int(time.time()) - 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, key_val in rows: + if key_val: + conn.execute( + sa.insert(_api_key).values( + id=f'key_{uid}', + user_id=uid, + key=key_val, + created_at=now, + updated_at=now, + ) + ) - 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, - }, - ) + # ── Drop legacy columns (idempotent) ────────────────────────────── + cols_to_drop = {'api_key', 'oauth_sub'} & user_columns + if cols_to_drop: + if conn.dialect.name == 'sqlite': + for col in cols_to_drop: + _drop_sqlite_indexes_for_column('user', col, conn) - 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') + with op.batch_alter_table('user') as batch_op: + for col in cols_to_drop: + batch_op.drop_column(col) 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() + rows = conn.execute( + sa.select(_user.c.id, _user.c.oauth) + .where(_user.c.oauth.is_not(None)) + ).fetchall() - for uid, oauth in users: + for uid, oauth in rows: try: data = json.loads(oauth) provider = list(data.keys())[0] @@ -200,24 +245,24 @@ def downgrade() -> None: oauth_sub = None conn.execute( - sa.text('UPDATE "user" SET oauth_sub = :oauth_sub WHERE id = :id'), - {'oauth_sub': oauth_sub, 'id': uid}, + sa.update(_user) + .where(_user.c.id == uid) + .values(oauth_sub=oauth_sub) ) op.drop_column('user', 'oauth') - # --- 2. Restore api_key field --- + # --- 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() + keys = conn.execute(sa.select(_api_key.c.user_id, _api_key.c.key)).fetchall() for uid, key in keys: conn.execute( - sa.text('UPDATE "user" SET api_key = :key WHERE id = :id'), - {'key': key, 'id': uid}, + sa.update(_user) + .where(_user.c.id == uid) + .values(api_key=key) ) - # Drop new table op.drop_table('api_key') with op.batch_alter_table('user') as batch_op: