sakhteman/core/utils/te/__init__.py

from typing import List

from starlette.responses import FileResponse

from .engines import MsSqlEngine, MongoEngine, PgSqlEngine, AbstractEngine
from .utils.Models import TableEngineInput
from .utils.Models import WorkingMode, JoinModel
from .utils.Models import TableEngineExport


class NewTableEngine:
    def __init__(self, query_class, model, inputs, mode, engine: AbstractEngine, extra_filter,
                 default_sort_by, joins, search_cols, options=None, soft_deleted=False, unselect={}, return_query=False):
        self.model = model
        self.query_class = query_class
        self.mode = mode
        self.joins = False
        self.raw_input = inputs
        if isinstance(inputs, dict):
            self.input = TableEngineInput(**inputs if inputs else {})
        else:
            self.input = TableEngineInput(**inputs.dict() if inputs else {})

        self.input.extra_filter = extra_filter
        self.default_sort_by = default_sort_by
        self.joins = joins
        self.soft_deleted = soft_deleted
        self.options = options
        self.search_cols = search_cols
        self.unselect = unselect
        self.return_query = return_query
        self.engine = engine(self)

    @staticmethod
    def mongo(query_class,
              model,
              inputs,
              mode: WorkingMode = WorkingMode.Development,
              extra_filter=None,
              soft_deleted=False,
              default_sort_by=None, joins=None, search_cols=[] ,unselect={},  return_query=False):
        return NewTableEngine(query_class, model, inputs, mode, engine=MongoEngine, extra_filter=extra_filter,
                              default_sort_by=default_sort_by, joins=joins, search_cols=search_cols,
                              soft_deleted=soft_deleted,unselect=unselect,  return_query=return_query)

    @staticmethod
    def mssql(query_class, model, inputs, mode: WorkingMode = WorkingMode.Development, extra_filter=None,
              default_sort_by=None, joins=None, search_cols=[], options=None):
        return NewTableEngine(query_class, model, inputs, mode, engine=MsSqlEngine, extra_filter=extra_filter,
                              default_sort_by=default_sort_by, joins=joins, search_cols=search_cols, options=options)

    @staticmethod
    def pgsql(query_class, model, inputs, mode: WorkingMode = WorkingMode.Development, extra_filter=None,
              default_sort_by=None, joins=None, search_cols=None, options=None):
        return NewTableEngine(query_class, model, inputs, mode, engine=PgSqlEngine, extra_filter=extra_filter,
                              default_sort_by=default_sort_by, joins=joins, search_cols=search_cols, options=options)

    async def exe(self):
        return await self.engine.exe()

    async def to_dict(self):
        if self.return_query:
            query = await self.exe()
            return query
        rows, count = await self.exe()
        self.input.extra_filter = None
        del self.engine
        self.set_export_file()
        return TableEngineExport(
            rows=rows,
            page=self.input.virtual_page(),
            total=count,
            te=self.raw_input
        )

    def add_extra_filter(self, filters):
        self.input = filters
        return self

    def joins(self, joins: List[JoinModel]):
        self.joins = joins
        return self

    def set_custom_entities(self, entities=None):
        self.engine.custom_entities = entities
        return self

    def set_export_file(self):
        try:
            self.raw_input['e']['d'] = self.input.e.d
            return self
        except Exception as e:
            return False