m2-traitement-benchmark/drivers/elastic.py

from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk 
from models.data_thread import Threads
from interfaces.db_testing_interface import DB_Testing

INDEX = "threads" 

class Elastic_Testing(DB_Testing):
    driver_name = "ES"
    def __init__(self) -> None:
        super().__init__()
        self.singles_data = []
        for t in Threads:
            x = t.__dict__.copy()
            self.singles_data.append(x)
        
        self.bulk_data = []
        for t in Threads:
            y = t.__dict__.copy()
            y["_index"] = INDEX
            self.bulk_data.append(y)
    
    def delete_table(self, driver: Elasticsearch):
        if driver.indices.exists(index=INDEX):
            driver.indices.delete(index=INDEX)

    def create_table(self, driver: Elasticsearch):
        self.delete_table(driver)
        driver.indices.create(index=INDEX, body={
            'settings': {
                'index': {
                    'number_of_replicas': 0
                }
            }
        })
    
    def delete_data(self, driver: Elasticsearch):
        if driver.indices.exists(index=INDEX):
            driver.indices.delete(index=INDEX)

    def add_bulk(self, driver: Elasticsearch):
        bulk(driver, self.bulk_data)

    def add_singles(self, driver: Elasticsearch):
        for t in self.singles_data:
            driver.index(index=INDEX, document=t)

    def attach5_mr_mrs(self, driver: Elasticsearch):
        query = {
            "query": {
                "bool": {
                    "must": [
                        {
                            "range": {
                                "attachment_count": {
                                    "gt": 5
                                }
                            }
                        },
                        {
                            "terms": {
                                "subject.keyword": ["Mr", "Mrs"]
                            }
                        }
                    ]
                }
            }
        }
        return f"Got {driver.count(index=INDEX, body=query).body['count']}"

    def search_mails_ends_microsoftcom(self, driver: Elasticsearch):
        query = {
            "query": {
                "wildcard": {
                    "cc_recipients.keyword": "*@microsoft.com"
                }
            }
        }
        return f"Got {driver.count(index=INDEX, body=query).body['count']}"
        
    def search_lorem(self, driver: Elasticsearch):
        query = {
            "query": {
                "match_phrase": {
                    "body": "Nullam sit amet turpis elementum ligula vehicula consequat. Morbi a ipsum. Integer a nibh."
                }
            }
        }
        return f"Got {driver.count(index=INDEX, body=query).body['count']}"
        
    def get_sum_attachment_less_5(self, driver: Elasticsearch):
        query = {
            "size": 0,
            "query": {
                "bool": {
                    "must": [
                        {
                            "range": {
                                "attachment_count": {
                                    "lt": 5
                                }
                            }
                        }
                    ]
                }
            }, 
            "aggs": {
                "attachment_count": {
                    "sum": {
                        "field": "attachment_count"
                    }
                }
            }
        }
        return f"Got {driver.search(index=INDEX, body=query).body['aggregations']['attachment_count']['value']}"