sakhteman/core/utils/te/engines/Mongo.py

import datetime
from typing import Union, List
from uuid import UUID

from .AbstractEngine import AbstractEngineClass
from ..utils.Models import FilterItem, FilterScope


class MongoEngine(AbstractEngineClass):
    rows = []
    stack = []
    filters = {}

    async def exe(self):
        _f = self.filters = self.set_filters()
        tmp = self.query_class.find(_f).sort(**self.set_sort())
        self.rows = tmp.limit(
            self.input.get_limit()).skip(self.input.skip())
        return [self.model(**document) async for document in self.rows], await self.count()

    def query_generator(self):
        if not self.input.has_filter():
            return []
        filters = self.input.filters()
        query = self.generate_scope_filter(filters, general=True)
        return query

    def set_filters(self):
        filters = []
        query = self.query_generator()
        if query:
            filters.append(query)
        if self.input.has_text_search():
            items = []
            _search_cols = self.input.get_search_cols() if (
                    self.input.get_search_cols() and len(self.input.get_search_cols())) else self.search_cols
            if _search_cols:
                for col in _search_cols:
                    items.append({col: {'$regex': self.input.get_text_search()}})
                filters.append({"$or": items})
        if self.input.extra_filter:
            filters.extend(self.input.extra_filter)
        if self.soft_deleted:
            filters.append({'deleted_at': {'$exists': False}})
        if not filters:
            return {}
        self.filters = {'$and': filters}
        return self.filters

    def set_sort(self):
        if self.input.o.field:
            return {'key_or_list': self.input.o.field,
                    'direction': 1 if self.input.o.direction == "asc" else -1}
        elif self.default_sort_by:
            return self.default_sort_by
        else:
            return {'key_or_list': 'created_at', 'direction': -1}

    @staticmethod
    def filter_item_generator(item: FilterItem) -> dict:
        if MongoEngine.is_valid_uuid(str(item.col)):
            return {
                'data': {
                    "$elemMatch": {
                        'field_id': UUID(item.col),
                        'value': {
                            **MongoEngine.get_expression_filter(item.op, MongoEngine.cast_value(item.val, item.type), )}
                    }
                }
            }
        return {item.col: {
            **MongoEngine.get_expression_filter(col=item.col, op=item.op,
                                                val=MongoEngine.cast_value(item.val, item.type), )}}

    @staticmethod
    def is_valid_uuid(uuid_to_test, version=4):
        try:
            uuid_obj = UUID(str(uuid_to_test), version=version)
        except ValueError:
            return False
        return str(uuid_obj) == str(uuid_to_test)


    @staticmethod
    def cast_value(val: Union[str, List], type: str) -> any:
        if type in ['NUMBER', 'STACK', 'FLOAT']:
            if isinstance(val, List):
                return [float(x) for x in val]
            return float(val)
        if type in ['DATE', 'DATETIME']:
            if isinstance(val, List):
                return [datetime.datetime.fromisoformat(str(x)) for x in val]
            return datetime.datetime.fromisoformat(val)
        if type == 'UUID':
            if isinstance(val, List):
                return [UUID(str(x)) for x in val]
            return UUID(str(val))
        return val


    @staticmethod
    def get_expression_filter(col=None, op: str = 'eq', val: str = '') -> dict:
        _op = MongoEngine.get_operator(op)
        if _op == "sw":
            return {'$regex': f'^{val}'}
        elif _op == "ew":
            return {'$regex': f'{val}$'}
        elif _op == "n_ew":
            return {'$regex': f'(?<!{val})$'}
        elif _op == "n_sw":
            return {'$regex': f'^(?!{val})'}
        elif _op == "n_like":
            return {'$not': {'$regex': val}}
        elif _op == "bet":
            if isinstance(val, str):
                return {'$gte': val.split('-')[0], '$lte': val.split('-')[1]}
            return {'$gte': val[0], '$lte': val[1]}
        elif _op == "n_bet":
            if isinstance(val, str):
                return {'$gte': val.split('-')[0], '$lte': val.split('-')[1]}
            return {'$not': {'$gte': val[0], '$lte': val[1]}}
        elif _op == "exi":
            return {'$exists': True}
        elif _op == "n_exi":
            return {'$exists': False}
        elif _op == "in":
            return {'$exists': True}
        elif _op == "not_in":
            return {'$exists': False}
        return {_op: val}

    @staticmethod
    def get_operator(op: str) -> str:
        ops = {
            "eq": "$eq",
            "n_eq": "$ne",
            "gt": "$gt",
            "ge": "$ge",
            "lt": "$lt",
            "le": "$le",
            "like": "$regex",
            "n_like": "n_like",
            "sw": "sw",
            "ew": "ew",
            'n_ew': 'n_ew',
            'n_sw': 'n_sw',
            'exi': 'exi',
            "n_exi": "n_exi",
            'lte': '$lte',
            'gte': '$gte',
            'in': 'in',
            'n_in': 'not_in',
            'bet': 'bet',
            'n_bet': 'n_bet',
        }
        return ops[op] if op in ops else '$eq'

    async def count(self):
        return await self.query_class.count_documents(self.filters)

    @staticmethod
    def generate_scope_filter(filters: list, general=False):
        if len(filters) == 1:
            _o = filters[0]
            if isinstance(_o, FilterItem):
                return MongoEngine.filter_item_generator(filters[0])
            elif isinstance(_o, FilterScope):
                return MongoEngine.generate_scope_filter(_o.items)
            else:
                return _o
        else:
            _a = filters.pop(0)
            _a = (MongoEngine.filter_item_generator(_a) if isinstance(_a,
                                                                      FilterItem) else _a) if not general else (
                MongoEngine.generate_scope_filter(_a.items) if isinstance(_a, FilterScope) else _a)
            _b = filters.pop(0)
            op = '$' + (_b.cr if not general else _b.op)
            _b = (MongoEngine.filter_item_generator(_b) if isinstance(_b,
                                                                      FilterItem) else _b) if not general else (
                MongoEngine.generate_scope_filter(_b.items) if isinstance(_b, FilterScope) else _b)
            filters = [{op: [
                _a, _b
            ]}] + filters
            return MongoEngine.generate_scope_filter(filters, general)

    def set_joins(self):
        pass