106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
import argparse
|
|
import asyncio
|
|
from collections import namedtuple
|
|
from multiprocessing import Process
|
|
|
|
import apsw
|
|
from elasticsearch import AsyncElasticsearch
|
|
from elasticsearch.helpers import async_bulk
|
|
|
|
from lbry.wallet.server.db.elastic_search import extract_doc, SearchIndex
|
|
|
|
INDEX = 'claims'
|
|
|
|
|
|
async def get_all(db, shard_num, shards_total):
|
|
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)
|
|
total = db.execute(f"select count(*) as total from claim where height % {shards_total} = {shard_num};").fetchone()[0]
|
|
for num, claim in enumerate(db.execute(f"""
|
|
SELECT claimtrie.claim_hash as is_controlling,
|
|
claimtrie.last_take_over_height,
|
|
(select group_concat(tag, ',,') from tag where tag.claim_hash in (claim.claim_hash, claim.reposted_claim_hash)) as tags,
|
|
(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)
|
|
WHERE claim.height % {shards_total} = {shard_num}
|
|
""")):
|
|
claim = dict(claim._asdict())
|
|
claim['censor_type'] = 0
|
|
claim['censoring_channel_hash'] = None
|
|
claim['tags'] = claim['tags'].split(',,') if claim['tags'] else []
|
|
claim['languages'] = claim['languages'].split(' ') if claim['languages'] else []
|
|
if num % 10_000 == 0:
|
|
print(num, total)
|
|
yield extract_doc(claim, INDEX)
|
|
|
|
|
|
async def consume(producer):
|
|
es = AsyncElasticsearch()
|
|
try:
|
|
await async_bulk(es, producer, request_timeout=120)
|
|
await es.indices.refresh(index=INDEX)
|
|
finally:
|
|
await es.close()
|
|
|
|
|
|
async def make_es_index():
|
|
es = AsyncElasticsearch()
|
|
try:
|
|
if await es.indices.exists(index=INDEX):
|
|
print("already synced ES")
|
|
return 1
|
|
index = SearchIndex('')
|
|
await index.start()
|
|
await index.stop()
|
|
return 0
|
|
finally:
|
|
await es.close()
|
|
|
|
|
|
async def run(args, shard):
|
|
db = apsw.Connection(args.db_path, flags=apsw.SQLITE_OPEN_READONLY | apsw.SQLITE_OPEN_URI)
|
|
db.cursor().execute('pragma journal_mode=wal;')
|
|
db.cursor().execute('pragma temp_store=memory;')
|
|
index = SearchIndex('')
|
|
await index.start()
|
|
await index.stop()
|
|
|
|
producer = get_all(db.cursor(), shard, args.clients)
|
|
await asyncio.gather(*(consume(producer) for _ in range(min(8, args.clients))))
|
|
|
|
|
|
def __run(args, shard):
|
|
asyncio.run(run(args, shard))
|
|
|
|
|
|
def __make_index():
|
|
return asyncio.run(make_es_index())
|
|
|
|
|
|
def run_elastic_sync():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("db_path", type=str)
|
|
parser.add_argument("-c", "--clients", type=int, default=16)
|
|
args = parser.parse_args()
|
|
processes = []
|
|
|
|
init_proc = Process(target=__make_index, args=())
|
|
init_proc.start()
|
|
init_proc.join()
|
|
exitcode = init_proc.exitcode
|
|
init_proc.close()
|
|
if exitcode:
|
|
print("ES is already initialized")
|
|
return
|
|
print("bulk-loading ES")
|
|
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()
|