elasticsearch.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
# @Version     : Python 3.11.4
# @Software    : Sublime Text 4
# @Author      : StudentCWZ
# @Email       : StudentCWZ@outlook.com
# @Date        : 2023/10/28 12:19
# @File        : elasticsearch.py
# @Description : A utility class to interact with Elasticsearch.
"""

import datetime

from application.extensions.init_elasticsearch import es


import datetime

from application.extensions.init_elasticsearch import es


class ElasticsearchUtil:
    @classmethod
    def dsl(cls, _start: str, _end: str, size=5000) -> dict:
        """
        Constructs a DSL (Domain Specific Language) query for Elasticsearch.

        Args:
            _start (str): Start time.
            _end (str): End time.
            size (int, optional): Number of data results. Defaults to 5000.

        Returns:
            dict: The DSL query.
        """
        _dsl = {
            "size": size,
            "query": {
                "bool": {
                    "must": [{
                        "match_phrase": {
                            "tag": {
                                "query": "global"
                            }
                        }
                    },
                        {
                        "range": {
                            "@timestamp": {
                                "gte": (datetime.datetime.strptime(_start, "%Y-%m-%d %H:%M:%S") -
                                        datetime.timedelta(hours=8)).strftime("%Y-%m-%dT%H:%M:%SZ"),
                                "lte": (datetime.datetime.strptime(_end, "%Y-%m-%d %H:%M:%S") -
                                        datetime.timedelta(hours=8)).strftime("%Y-%m-%dT%H:%M:%SZ")
                            }
                        }
                    }
                    ],
                    "filter": [{
                        "match_all": {}
                    }],
                    "should": [],
                    "must_not": []
                }
            }
        }
        return _dsl

    @classmethod
    def search(cls, _index: str, _dsl: dict, _scroll="5m") -> dict:
        """
        Searches data in Elasticsearch.

        Args:
            _index (str): Index name.
            _dsl (dict): DSL query.
            _scroll (str, optional): Scroll time. Defaults to "5m".

        Returns:
            dict: Data after the search.
        """
        return es.search(index=_index, scroll=_scroll, body=_dsl)

    @classmethod
    def scroll_search(cls, _id,  _scroll="5m") -> dict:
        """
        Searches data in Elasticsearch using scroll.

        Args:
            _id (str): Scroll ID.
            _scroll (str, optional): Scroll time. Defaults to "5m".

        Returns:
            dict: Data after the search by scroll.
        """
        return es.scroll(scroll_id=_id, scroll=_scroll, request_timeout=30)

    @classmethod
    def insert_data(cls, index: str, doc_type: str, data: dict) -> None:
        """
        Inserts data into Elasticsearch.

        Args:
            index (str): Index name.
            doc_type (str): Document type.
            data (dict): Data to be inserted.

        Returns:
            None
        """
        es.index(index=index, doc_type=doc_type, body=data)