lbry-sdk/lbry/wallet/server/db/elasticsearch/sync.py

106 lines
3.5 KiB
Python
Raw Normal View History

2021-01-20 05:41:54 +01:00
import argparse
import asyncio
import logging
2021-02-16 16:52:32 +01:00
import os
2021-01-20 05:41:54 +01:00
from collections import namedtuple
2021-01-27 05:43:06 +01:00
from multiprocessing import Process
2021-01-20 05:41:54 +01:00
import apsw
from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk
from .search import extract_doc, SearchIndex
2021-01-20 05:41:54 +01:00
INDEX = 'claims'
async def get_all(db, shard_num, shards_total, limit=0):
logging.info("shard %d starting", shard_num)
2021-01-20 05:41:54 +01:00
def exec_factory(cursor, statement, bindings):
tpl = namedtuple('row', (d[0] for d in cursor.getdescription()))
cursor.setrowtrace(lambda cursor, row: tpl(*row))
return True
db.setexectrace(exec_factory)
2021-02-02 21:11:13 +01:00
total = db.execute(f"select count(*) as total from claim where height % {shards_total} = {shard_num};").fetchone()[0]
2021-01-20 05:41:54 +01:00
for num, claim in enumerate(db.execute(f"""
SELECT claimtrie.claim_hash as is_controlling,
claimtrie.last_take_over_height,
2021-01-27 02:33:17 +01:00
(select group_concat(tag, ',,') from tag where tag.claim_hash in (claim.claim_hash, claim.reposted_claim_hash)) as tags,
2021-01-20 05:41:54 +01:00
(select group_concat(language, ' ') from language where language.claim_hash in (claim.claim_hash, claim.reposted_claim_hash)) as languages,
claim.*
FROM claim LEFT JOIN claimtrie USING (claim_hash)
2021-02-02 21:11:13 +01:00
WHERE claim.height % {shards_total} = {shard_num}
ORDER BY claim.height desc
2021-01-20 05:41:54 +01:00
""")):
claim = dict(claim._asdict())
claim['censor_type'] = 0
claim['censoring_channel_hash'] = None
2021-01-27 02:33:17 +01:00
claim['tags'] = claim['tags'].split(',,') if claim['tags'] else []
2021-01-20 05:41:54 +01:00
claim['languages'] = claim['languages'].split(' ') if claim['languages'] else []
2021-01-27 02:26:45 +01:00
if num % 10_000 == 0:
logging.info("%d/%d", num, total)
2021-01-20 05:41:54 +01:00
yield extract_doc(claim, INDEX)
if 0 < limit <= num:
break
2021-01-20 05:41:54 +01:00
async def consume(producer):
es = AsyncElasticsearch()
2021-02-12 05:10:30 +01:00
try:
await async_bulk(es, producer, request_timeout=120)
await es.indices.refresh(index=INDEX)
finally:
await es.close()
async def make_es_index():
index = SearchIndex('')
2021-02-12 05:10:30 +01:00
try:
return await index.start()
2021-02-12 05:10:30 +01:00
finally:
index.stop()
2021-01-27 05:43:06 +01:00
async def run(args, shard):
2021-02-16 16:52:32 +01:00
def itsbusy(*_):
logging.info("shard %d: db is busy, retry", shard)
return True
2021-01-20 05:41:54 +01:00
db = apsw.Connection(args.db_path, flags=apsw.SQLITE_OPEN_READONLY | apsw.SQLITE_OPEN_URI)
db.setbusyhandler(itsbusy)
2021-02-02 21:11:13 +01:00
db.cursor().execute('pragma journal_mode=wal;')
db.cursor().execute('pragma temp_store=memory;')
2021-02-12 05:10:30 +01:00
producer = get_all(db.cursor(), shard, args.clients, limit=args.blocks)
2021-01-27 06:56:43 +01:00
await asyncio.gather(*(consume(producer) for _ in range(min(8, args.clients))))
2021-01-27 05:43:06 +01:00
2021-02-12 05:10:30 +01:00
2021-01-27 05:43:06 +01:00
def __run(args, shard):
asyncio.run(run(args, shard))
2021-02-12 05:10:30 +01:00
def run_elastic_sync():
logging.basicConfig(level=logging.INFO)
logging.info('lbry.server starting')
2021-01-27 05:43:06 +01:00
parser = argparse.ArgumentParser()
parser.add_argument("db_path", type=str)
parser.add_argument("-c", "--clients", type=int, default=16)
parser.add_argument("-b", "--blocks", type=int, default=0)
parser.add_argument("-f", "--force", default=False, action='store_true')
2021-01-27 05:43:06 +01:00
args = parser.parse_args()
processes = []
2021-02-12 05:10:30 +01:00
if not args.force and not os.path.exists(args.db_path):
2021-02-16 16:52:32 +01:00
logging.info("DB path doesnt exist")
return
if not args.force and not asyncio.run(make_es_index()):
logging.info("ES is already initialized")
2021-02-12 05:10:30 +01:00
return
2021-01-27 05:43:06 +01:00
for i in range(args.clients):
processes.append(Process(target=__run, args=(args, i)))
processes[-1].start()
for process in processes:
process.join()
process.close()