mirror of
https://github.com/flibusta-apps/meilie_updater.git
synced 2025-12-06 15:15:37 +01:00
111 lines
3.1 KiB
Python
111 lines
3.1 KiB
Python
import asyncio
|
|
import concurrent.futures
|
|
|
|
from arq.connections import ArqRedis
|
|
import asyncpg
|
|
from meilisearch import Client
|
|
|
|
from core.config import env_config
|
|
|
|
|
|
def get_meilisearch_client() -> Client:
|
|
return Client(url=env_config.MEILI_HOST, api_key=env_config.MEILI_MASTER_KEY)
|
|
|
|
|
|
async def get_postgres_pool() -> asyncpg.Pool:
|
|
return await asyncpg.create_pool( # type: ignore
|
|
database=env_config.POSTGRES_DB_NAME,
|
|
host=env_config.POSTGRES_HOST,
|
|
port=env_config.POSTGRES_PORT,
|
|
user=env_config.POSTGRES_USER,
|
|
password=env_config.POSTGRES_PASSWORD,
|
|
)
|
|
|
|
|
|
async def update_books(ctx) -> bool:
|
|
loop = asyncio.get_event_loop()
|
|
|
|
meili = get_meilisearch_client()
|
|
index = meili.index("books")
|
|
|
|
postgres_pool = await get_postgres_pool()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as pool:
|
|
count = await postgres_pool.fetchval("SELECT count(*) FROM books;")
|
|
|
|
for offset in range(0, count, 4096):
|
|
rows = await postgres_pool.fetch(
|
|
"SELECT id, title, is_deleted FROM books"
|
|
f" ORDER BY id LIMIT 4096 OFFSET {offset}"
|
|
)
|
|
|
|
documents = [dict(row) for row in rows]
|
|
|
|
await loop.run_in_executor(pool, index.add_documents, documents)
|
|
|
|
index.update_searchable_attributes(["title"])
|
|
index.update_filterable_attributes(["is_deleted"])
|
|
|
|
return True
|
|
|
|
|
|
async def update_authors(ctx) -> bool:
|
|
loop = asyncio.get_event_loop()
|
|
|
|
meili = get_meilisearch_client()
|
|
index = meili.index("authors")
|
|
|
|
postgres_pool = await get_postgres_pool()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as pool:
|
|
count = await postgres_pool.fetchval("SELECT count(*) FROM authors;")
|
|
|
|
for offset in range(0, count, 4096):
|
|
rows = await postgres_pool.fetch(
|
|
"SELECT id, first_name, last_name, middle_name FROM authors"
|
|
f" ORDER BY id LIMIT 4096 OFFSET {offset}"
|
|
)
|
|
|
|
documents = [dict(row) for row in rows]
|
|
|
|
await loop.run_in_executor(pool, index.add_documents, documents)
|
|
|
|
index.update_searchable_attributes(["first_name", "last_name", "middle_name"])
|
|
|
|
return True
|
|
|
|
|
|
async def update_sequences(ctx) -> bool:
|
|
loop = asyncio.get_event_loop()
|
|
|
|
meili = get_meilisearch_client()
|
|
index = meili.index("sequences")
|
|
|
|
postgres_pool = await get_postgres_pool()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor() as pool:
|
|
count = await postgres_pool.fetchval("SELECT count(*) FROM sequences;")
|
|
|
|
for offset in range(0, count, 4096):
|
|
rows = await postgres_pool.fetch(
|
|
f"SELECT id, name FROM sequences ORDER BY id LIMIT 4096 OFFSET {offset}"
|
|
)
|
|
|
|
documents = [dict(row) for row in rows]
|
|
|
|
await loop.run_in_executor(pool, index.add_documents, documents)
|
|
|
|
index.update_searchable_attributes(["name"])
|
|
|
|
return True
|
|
|
|
|
|
async def update(ctx) -> bool:
|
|
arq_pool: ArqRedis = ctx["arc_pool"]
|
|
|
|
await arq_pool.enqueue_job("update_books")
|
|
await arq_pool.enqueue_job("update_authors")
|
|
await arq_pool.enqueue_job("update_sequences")
|
|
|
|
return True
|