refac
This commit is contained in:
Timothy Jaeryang Baek committed
1 parent
73421c5b42
commit
060d5da473
1 file changed
+169
-85
+169
-85
@@ -17,10 +17,13 @@ down_revision: Union[str, None] = '856c5b02fb54'
|
|||||||
branch_labels: Union[str, Sequence[str], None] = None
|
branch_labels: Union[str, Sequence[str], None] = None
|
||||||
depends_on: Union[str, Sequence[str], None] = None
|
depends_on: Union[str, Sequence[str], None] = None
|
||||||
|
|
||||||
|
BATCH_SIZE = 150
|
||||||
|
|
||||||
|
|
||||||
def upgrade() -> None:
|
def upgrade() -> None:
|
||||||
conn = op.get_bind()
|
conn = op.get_bind()
|
||||||
columns = [col['name'] for col in sa.inspect(conn).get_columns('chat')]
|
inspector = sa.inspect(conn)
|
||||||
|
columns = [col['name'] for col in inspector.get_columns('chat')]
|
||||||
if 'current_message_id' not in columns:
|
if 'current_message_id' not in columns:
|
||||||
op.add_column('chat', sa.Column('current_message_id', sa.Text(), nullable=True))
|
op.add_column('chat', sa.Column('current_message_id', sa.Text(), nullable=True))
|
||||||
|
|
||||||
@@ -38,100 +41,181 @@ def upgrade() -> None:
|
|||||||
sa.column('created_at', sa.BigInteger()),
|
sa.column('created_at', sa.BigInteger()),
|
||||||
)
|
)
|
||||||
|
|
||||||
messages_by_chat: dict[str, dict[str, dict]] = {}
|
has_chat_message = 'chat_message' in inspector.get_table_names()
|
||||||
for row in conn.execute(
|
result = conn.execute(
|
||||||
sa.select(
|
sa.select(chat.c.id, chat.c.chat, chat.c.current_message_id).execution_options(
|
||||||
chat_message.c.chat_id,
|
yield_per=BATCH_SIZE,
|
||||||
chat_message.c.id,
|
stream_results=True,
|
||||||
chat_message.c.parent_id,
|
|
||||||
chat_message.c.created_at,
|
|
||||||
)
|
)
|
||||||
):
|
)
|
||||||
values = row._mapping
|
|
||||||
chat_id = values['chat_id']
|
|
||||||
prefix = f'{chat_id}-'
|
|
||||||
message_id = values['id']
|
|
||||||
if message_id and message_id.startswith(prefix):
|
|
||||||
message_id = message_id[len(prefix) :]
|
|
||||||
if not message_id:
|
|
||||||
continue
|
|
||||||
parent_id = values['parent_id']
|
|
||||||
if parent_id and parent_id.startswith(prefix):
|
|
||||||
parent_id = parent_id[len(prefix) :]
|
|
||||||
messages_by_chat.setdefault(chat_id, {})[message_id] = {
|
|
||||||
'parent_id': parent_id,
|
|
||||||
'created_at': values['created_at'] or 0,
|
|
||||||
}
|
|
||||||
|
|
||||||
for row in conn.execute(sa.select(chat.c.id, chat.c.chat, chat.c.current_message_id)):
|
while True:
|
||||||
values = row._mapping
|
rows = result.fetchmany(BATCH_SIZE)
|
||||||
chat_id = values['id']
|
if not rows:
|
||||||
prefix = f'{chat_id}-'
|
break
|
||||||
chat_messages = messages_by_chat.get(chat_id, {})
|
|
||||||
|
|
||||||
chat_data = {}
|
batch_chat_ids: list[str] = []
|
||||||
if isinstance(values['chat'], dict):
|
candidates_by_chat: dict[str, list[str]] = {}
|
||||||
chat_data = values['chat']
|
current_by_chat: dict[str, str | None] = {}
|
||||||
elif isinstance(values['chat'], str):
|
json_messages_by_chat: dict[str, dict[str, dict]] = {}
|
||||||
try:
|
|
||||||
parsed = json.loads(values['chat'])
|
|
||||||
chat_data = parsed if isinstance(parsed, dict) else {}
|
|
||||||
except (TypeError, ValueError, json.JSONDecodeError):
|
|
||||||
chat_data = {}
|
|
||||||
|
|
||||||
history = chat_data.get('history') if isinstance(chat_data.get('history'), dict) else {}
|
for row in rows:
|
||||||
messages = history.get('messages') if isinstance(history.get('messages'), dict) else {}
|
values = row._mapping
|
||||||
if not messages and isinstance(chat_data.get('messages'), list):
|
chat_id = values['id']
|
||||||
messages = {
|
prefix = f'{chat_id}-'
|
||||||
message['id']: message
|
batch_chat_ids.append(chat_id)
|
||||||
for message in chat_data['messages']
|
current_by_chat[chat_id] = values['current_message_id']
|
||||||
if isinstance(message, dict) and message.get('id')
|
|
||||||
}
|
|
||||||
json_messages = {
|
|
||||||
message_id: {
|
|
||||||
'parent_id': message.get('parentId') if isinstance(message, dict) else None,
|
|
||||||
'created_at': message.get('timestamp', 0) if isinstance(message, dict) else 0,
|
|
||||||
}
|
|
||||||
for message_id, message in messages.items()
|
|
||||||
}
|
|
||||||
available_messages = chat_messages or json_messages
|
|
||||||
candidates = [
|
|
||||||
values['current_message_id'],
|
|
||||||
history.get('currentId'),
|
|
||||||
chat_data.get('currentId'),
|
|
||||||
chat_data.get('branchPointMessageId'),
|
|
||||||
]
|
|
||||||
|
|
||||||
current_message_id = next(
|
chat_data = {}
|
||||||
(
|
if isinstance(values['chat'], dict):
|
||||||
candidate[len(prefix) :] if candidate.startswith(prefix) else candidate
|
chat_data = values['chat']
|
||||||
|
elif isinstance(values['chat'], str):
|
||||||
|
try:
|
||||||
|
parsed = json.loads(values['chat'])
|
||||||
|
chat_data = parsed if isinstance(parsed, dict) else {}
|
||||||
|
except (TypeError, ValueError, json.JSONDecodeError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
history = chat_data.get('history') if isinstance(chat_data.get('history'), dict) else {}
|
||||||
|
candidates_by_chat[chat_id] = []
|
||||||
|
for candidate in (
|
||||||
|
values['current_message_id'],
|
||||||
|
history.get('currentId'),
|
||||||
|
chat_data.get('currentId'),
|
||||||
|
chat_data.get('branchPointMessageId'),
|
||||||
|
):
|
||||||
|
if not isinstance(candidate, str) or not candidate:
|
||||||
|
continue
|
||||||
|
candidate = candidate[len(prefix) :] if candidate.startswith(prefix) else candidate
|
||||||
|
if candidate not in candidates_by_chat[chat_id]:
|
||||||
|
candidates_by_chat[chat_id].append(candidate)
|
||||||
|
|
||||||
|
messages = history.get('messages') if isinstance(history.get('messages'), dict) else {}
|
||||||
|
if not messages and isinstance(chat_data.get('messages'), list):
|
||||||
|
messages = {
|
||||||
|
message['id']: message
|
||||||
|
for message in chat_data['messages']
|
||||||
|
if isinstance(message, dict) and message.get('id')
|
||||||
|
}
|
||||||
|
if messages:
|
||||||
|
json_messages_by_chat[chat_id] = {
|
||||||
|
message_id: {
|
||||||
|
'parent_id': message.get('parentId') if isinstance(message, dict) else None,
|
||||||
|
'created_at': message.get('timestamp', 0) if isinstance(message, dict) else 0,
|
||||||
|
}
|
||||||
|
for message_id, message in messages.items()
|
||||||
|
}
|
||||||
|
|
||||||
|
resolved: dict[str, str] = {}
|
||||||
|
|
||||||
|
if has_chat_message:
|
||||||
|
candidate_ids = {
|
||||||
|
f'{chat_id}-{candidate}'
|
||||||
|
for chat_id, candidates in candidates_by_chat.items()
|
||||||
for candidate in candidates
|
for candidate in candidates
|
||||||
if candidate
|
|
||||||
and (
|
|
||||||
not available_messages
|
|
||||||
or (candidate[len(prefix) :] if candidate.startswith(prefix) else candidate) in available_messages
|
|
||||||
)
|
|
||||||
),
|
|
||||||
None,
|
|
||||||
)
|
|
||||||
|
|
||||||
if not current_message_id and available_messages:
|
|
||||||
parent_ids = {
|
|
||||||
message['parent_id']
|
|
||||||
for message in available_messages.values()
|
|
||||||
if message.get('parent_id') in available_messages
|
|
||||||
}
|
}
|
||||||
leaf_ids = [message_id for message_id in available_messages if message_id not in parent_ids]
|
if candidate_ids:
|
||||||
current_message_id = max(
|
valid_by_chat: dict[str, set[str]] = {}
|
||||||
leaf_ids or list(available_messages),
|
for row in conn.execute(
|
||||||
key=lambda message_id: available_messages[message_id].get('created_at') or 0,
|
sa.select(chat_message.c.chat_id, chat_message.c.id).where(
|
||||||
)
|
chat_message.c.chat_id.in_(batch_chat_ids),
|
||||||
|
chat_message.c.id.in_(candidate_ids),
|
||||||
|
)
|
||||||
|
):
|
||||||
|
values = row._mapping
|
||||||
|
chat_id = values['chat_id']
|
||||||
|
prefix = f'{chat_id}-'
|
||||||
|
message_id = values['id']
|
||||||
|
if message_id and message_id.startswith(prefix):
|
||||||
|
message_id = message_id[len(prefix) :]
|
||||||
|
if message_id:
|
||||||
|
valid_by_chat.setdefault(chat_id, set()).add(message_id)
|
||||||
|
for chat_id, candidates in candidates_by_chat.items():
|
||||||
|
valid_ids = valid_by_chat.get(chat_id, set())
|
||||||
|
for candidate in candidates:
|
||||||
|
if candidate in valid_ids:
|
||||||
|
resolved[chat_id] = candidate
|
||||||
|
break
|
||||||
|
|
||||||
if current_message_id and current_message_id != values['current_message_id']:
|
unresolved_chat_ids = [chat_id for chat_id in batch_chat_ids if chat_id not in resolved]
|
||||||
|
messages_by_chat: dict[str, dict[str, dict]] = {}
|
||||||
|
if unresolved_chat_ids:
|
||||||
|
for row in conn.execute(
|
||||||
|
sa.select(
|
||||||
|
chat_message.c.chat_id,
|
||||||
|
chat_message.c.id,
|
||||||
|
chat_message.c.parent_id,
|
||||||
|
chat_message.c.created_at,
|
||||||
|
).where(chat_message.c.chat_id.in_(unresolved_chat_ids))
|
||||||
|
):
|
||||||
|
values = row._mapping
|
||||||
|
chat_id = values['chat_id']
|
||||||
|
prefix = f'{chat_id}-'
|
||||||
|
message_id = values['id']
|
||||||
|
if message_id and message_id.startswith(prefix):
|
||||||
|
message_id = message_id[len(prefix) :]
|
||||||
|
if not message_id:
|
||||||
|
continue
|
||||||
|
parent_id = values['parent_id']
|
||||||
|
if parent_id and parent_id.startswith(prefix):
|
||||||
|
parent_id = parent_id[len(prefix) :]
|
||||||
|
messages_by_chat.setdefault(chat_id, {})[message_id] = {
|
||||||
|
'parent_id': parent_id,
|
||||||
|
'created_at': values['created_at'] or 0,
|
||||||
|
}
|
||||||
|
|
||||||
|
for chat_id, messages in messages_by_chat.items():
|
||||||
|
parent_ids = {
|
||||||
|
message['parent_id']
|
||||||
|
for message in messages.values()
|
||||||
|
if message.get('parent_id') in messages
|
||||||
|
}
|
||||||
|
leaf_ids = [message_id for message_id in messages if message_id not in parent_ids]
|
||||||
|
resolved[chat_id] = max(
|
||||||
|
leaf_ids or list(messages),
|
||||||
|
key=lambda message_id: messages[message_id].get('created_at') or 0,
|
||||||
|
)
|
||||||
|
|
||||||
|
for chat_id in batch_chat_ids:
|
||||||
|
if chat_id in resolved:
|
||||||
|
continue
|
||||||
|
|
||||||
|
messages = json_messages_by_chat.get(chat_id, {})
|
||||||
|
valid_candidate = next(
|
||||||
|
(candidate for candidate in candidates_by_chat[chat_id] if candidate in messages),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
if valid_candidate:
|
||||||
|
resolved[chat_id] = valid_candidate
|
||||||
|
elif messages:
|
||||||
|
parent_ids = {
|
||||||
|
message['parent_id']
|
||||||
|
for message in messages.values()
|
||||||
|
if message.get('parent_id') in messages
|
||||||
|
}
|
||||||
|
leaf_ids = [message_id for message_id in messages if message_id not in parent_ids]
|
||||||
|
resolved[chat_id] = max(
|
||||||
|
leaf_ids or list(messages),
|
||||||
|
key=lambda message_id: messages[message_id].get('created_at') or 0,
|
||||||
|
)
|
||||||
|
|
||||||
|
updates = [
|
||||||
|
{'chat_id': chat_id, 'current_message_id': message_id}
|
||||||
|
for chat_id, message_id in resolved.items()
|
||||||
|
if message_id and message_id != current_by_chat.get(chat_id)
|
||||||
|
]
|
||||||
|
if updates:
|
||||||
conn.execute(
|
conn.execute(
|
||||||
sa.update(chat)
|
sa.update(chat)
|
||||||
.where(chat.c.id == chat_id)
|
.where(chat.c.id == sa.bindparam('update_chat_id'))
|
||||||
.values(current_message_id=current_message_id)
|
.values(current_message_id=sa.bindparam('update_current_message_id')),
|
||||||
|
[
|
||||||
|
{
|
||||||
|
'update_chat_id': row['chat_id'],
|
||||||
|
'update_current_message_id': row['current_message_id'],
|
||||||
|
}
|
||||||
|
for row in updates
|
||||||
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user