LongArticleTaskServer/app/infra/external/elastic_search.py

import ssl

from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk


class AsyncElasticSearchClient:
    def __init__(self, index_):
        self.password = "nkvvASQuQ0XUGRq5OLvm"
        self.hosts = ["https://192.168.205.85:9200", "https://192.168.205.85:9300"]
        self.ctx = ssl.create_default_context(
            cafile="app/core/config/cert/es_certs.crt"
        )
        self.es = AsyncElasticsearch(
            self.hosts, basic_auth=("elastic", self.password), ssl_context=self.ctx
        )
        self.index_name = index_

    async def create_index(self, settings, mappings):
        exists = await self.es.indices.exists(index=self.index_name)
        if exists:
            await self.es.indices.delete(index=self.index_name)

        try:
            await self.es.indices.create(
                index=self.index_name, settings=settings, mappings=mappings
            )
            print("Index created successfully")
        except Exception as e:
            print("fail to create index, reason:", e)

    async def get_max_article_id(self):
        response = await self.es.search(
            index=self.index_name,
            size=1,
            sort="article_id:desc",
            _source=["article_id"],
        )
        return response["hits"]["hits"][0]["_source"]["article_id"]

    async def search(self, search_keys, size=10):
        query = {
            "query": {"match": {"title": search_keys}},
            "_source": ["article_id", "title"],
            "size": size,
        }
        resp = await self.es.search(index=self.index_name, body=query)
        return [i["_source"] for i in resp["hits"]["hits"]]

    async def bulk_insert(self, docs):
        await async_bulk(self.es, docs)

    async def close(self):
        await self.es.close()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.es.close()