43 lines
1.1 KiB
Python
43 lines
1.1 KiB
Python
import asyncio
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import selectinload
|
|
|
|
from ..core.db import SessionLocal
|
|
from ..models.video_model import Video
|
|
from .documents import video_to_document
|
|
from .search_client import client, videos_index
|
|
|
|
BATCH = 1000
|
|
|
|
async def full_reindex():
|
|
async with SessionLocal() as db:
|
|
stmt = (
|
|
select(Video)
|
|
.options(selectinload(Video.channel), selectinload(Video.playlist))
|
|
.order_by(Video.id)
|
|
)
|
|
result = await db.stream(stmt)
|
|
|
|
batch: list[dict] = []
|
|
total = 0
|
|
async for video in result.scalars():
|
|
batch.append(video_to_document(video))
|
|
if len(batch) >= BATCH:
|
|
await _push(batch)
|
|
total += len(batch)
|
|
print(f"Indexed {total}...")
|
|
batch = []
|
|
if batch:
|
|
await _push(batch)
|
|
total += len(batch)
|
|
|
|
print(f"Done. Indexed {total} videos.")
|
|
|
|
async def _push(docs: list[dict]):
|
|
task = videos_index().add_documents(docs)
|
|
|
|
client.wait_for_task(task.task_uid)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(full_reindex()) |