mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-07 08:27:05 +00:00
fix: guard non-dict JSON, incremental flushing, offline txn parity
This commit is contained in:
parent
05ef5c5223
commit
d93fe55ad1
2 changed files with 21 additions and 6 deletions
|
|
@ -57,6 +57,7 @@ def run_migrations_offline() -> None:
|
|||
target_metadata=target_metadata,
|
||||
literal_binds=True,
|
||||
dialect_opts={'paramstyle': 'named'},
|
||||
transaction_per_migration=True,
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
|
|
|
|||
|
|
@ -38,6 +38,9 @@ def _parse_chat_messages(chat_id, user_id, chat_data, now):
|
|||
except Exception:
|
||||
return []
|
||||
|
||||
if not isinstance(chat_data, dict):
|
||||
return []
|
||||
|
||||
history = chat_data.get('history', {})
|
||||
if not isinstance(history, dict):
|
||||
return []
|
||||
|
|
@ -367,23 +370,34 @@ def _upgrade_postgresql() -> None:
|
|||
last_id = rows[-1][0]
|
||||
total_chats += len(rows)
|
||||
|
||||
# Parse all messages from this page of chats
|
||||
# Parse and flush incrementally to avoid memory spikes from
|
||||
# message-dense chats (some chats can have thousands of messages).
|
||||
messages_batch = []
|
||||
for chat_row in rows:
|
||||
messages_batch.extend(
|
||||
_parse_chat_messages(chat_row[0], chat_row[1], chat_row[2], now)
|
||||
)
|
||||
# Flush when batch is full
|
||||
while len(messages_batch) >= BATCH_SIZE:
|
||||
batch = messages_batch[:BATCH_SIZE]
|
||||
messages_batch = messages_batch[BATCH_SIZE:]
|
||||
inserted, failed = _flush_batch_pg(
|
||||
conn, chat_message_table, batch
|
||||
)
|
||||
total_inserted += inserted
|
||||
total_failed += failed
|
||||
|
||||
# Insert in sub-batches with ON CONFLICT DO NOTHING
|
||||
for i in range(0, len(messages_batch), BATCH_SIZE):
|
||||
batch = messages_batch[i : i + BATCH_SIZE]
|
||||
inserted, failed = _flush_batch_pg(conn, chat_message_table, batch)
|
||||
# Flush remaining messages from this page
|
||||
if messages_batch:
|
||||
inserted, failed = _flush_batch_pg(
|
||||
conn, chat_message_table, messages_batch
|
||||
)
|
||||
total_inserted += inserted
|
||||
total_failed += failed
|
||||
|
||||
conn.commit()
|
||||
|
||||
if total_inserted % 50000 < max(len(messages_batch), 1):
|
||||
if total_inserted % 50000 < BATCH_SIZE:
|
||||
log.info(
|
||||
f'Migration progress: {total_chats} chats processed,'
|
||||
f' {total_inserted} messages inserted...'
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue