diff --git a/lib/dl_formula/dl_formula/core/datatype.py b/lib/dl_formula/dl_formula/core/datatype.py index 4650481e9..a59f10760 100644 --- a/lib/dl_formula/dl_formula/core/datatype.py +++ b/lib/dl_formula/dl_formula/core/datatype.py @@ -7,6 +7,7 @@ ) from typing import ( FrozenSet, + Hashable, Optional, ) @@ -101,6 +102,9 @@ class DataTypeParams: timezone: Optional[str] = attr.ib(default=None) # Other possible cases: decimal precision, datetime sub-second precision, nullable, enum values. + def as_primitive(self) -> tuple[Optional[Hashable]]: + return (self.timezone,) + _AUTOCAST_FROM_TYPES = OrderedDict( ( diff --git a/lib/dl_query_processing/dl_query_processing/compilation/primitives.py b/lib/dl_query_processing/dl_query_processing/compilation/primitives.py index 0fa5b2fab..b63420225 100644 --- a/lib/dl_query_processing/dl_query_processing/compilation/primitives.py +++ b/lib/dl_query_processing/dl_query_processing/compilation/primitives.py @@ -22,13 +22,12 @@ OrderDirection, ) import dl_formula.core.nodes as formula_nodes -from dl_query_processing.compilation.query_meta import QueryMetaInfo +from dl_query_processing.compilation.query_meta import QueryMetaInfo, QueryElementExtract from dl_query_processing.enums import ( ExecutionLevel, QueryPart, ) - _COMPILED_FLA_TV = TypeVar("_COMPILED_FLA_TV", bound="CompiledFormulaInfo") @@ -45,6 +44,17 @@ class CompiledFormulaInfo: alias: Optional[str] = attr.ib() avatar_ids: Set[str] = attr.ib(factory=set) original_field_id: Optional[str] = attr.ib(default=None) + + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.formula_obj.extract, + self.alias, + frozenset(self.avatar_ids), + self.original_field_id, + ), + ) @property def not_none_alias(self) -> str: @@ -96,6 +106,17 @@ class CompiledOrderByFormulaInfo(CompiledFormulaInfo): # noqa direction: OrderDirection = attr.ib(kw_only=True) + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.formula_obj.extract, + self.alias, + frozenset(self.avatar_ids), + self.original_field_id, + self.direction, + ), + ) @attr.s(slots=True, frozen=True) class CompiledJoinOnFormulaInfo(CompiledFormulaInfo): # noqa @@ -109,6 +130,20 @@ class CompiledJoinOnFormulaInfo(CompiledFormulaInfo): # noqa right_id: str = attr.ib(kw_only=True) join_type: JoinType = attr.ib(kw_only=True) + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.formula_obj.extract, + self.alias, + frozenset(self.avatar_ids), + self.original_field_id, + self.left_id, + self.right_id, + self.join_type.name, + ), + ) + _FROM_OBJ_TV = TypeVar("_FROM_OBJ_TV", bound="FromObject") @@ -121,6 +156,15 @@ class FromColumn: def clone(self, **updates: Any) -> FromColumn: return attr.evolve(self, **updates) + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.id, + self.name, + ), + ) + @attr.s(frozen=True) class FromObject: @@ -128,6 +172,16 @@ class FromObject: alias: str = attr.ib(kw_only=True) columns: tuple[FromColumn, ...] = attr.ib(kw_only=True) + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.id, + self.alias, + tuple(col.extract for col in self.columns), + ), + ) + def clone(self: _FROM_OBJ_TV, **updates: Any) -> _FROM_OBJ_TV: return attr.evolve(self, **updates) @@ -137,6 +191,15 @@ class JoinedFromObject: root_from_id: Optional[str] = attr.ib(kw_only=True, default=None) froms: Sequence[FromObject] = attr.ib(kw_only=True, default=()) + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.root_from_id, + tuple(from_obj.extract for from_obj in self.froms), + ), + ) + def iter_ids(self) -> Iterable[str]: return (from_obj.id for from_obj in self.froms) @@ -175,6 +238,24 @@ class CompiledQuery: offset: Optional[int] = attr.ib(kw_only=True, default=None) meta: QueryMetaInfo = attr.ib(kw_only=True, factory=QueryMetaInfo) + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.id, + self.level_type.name, + tuple(formula.extract for formula in self.select), + tuple(formula.extract for formula in self.group_by), + tuple(formula.extract for formula in self.filters), + tuple(formula.extract for formula in self.order_by), + tuple(formula.extract for formula in self.join_on), + self.joined_from.extract, + self.limit, + self.offset, + self.meta.extract, + ), + ) + def get_complexity(self) -> int: return sum(formula.complexity for formula in self.all_formulas) @@ -262,6 +343,14 @@ class SubqueryFromObject(FromObject): @attr.s(frozen=True) class CompiledMultiQueryBase(abc.ABC): + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + tuple(query.extract for query in sorted(self.iter_queries(), key=lambda query: query.id)), + ), + ) + @abc.abstractmethod def iter_queries(self) -> Iterable[CompiledQuery]: raise NotImplementedError diff --git a/lib/dl_query_processing/dl_query_processing/compilation/query_meta.py b/lib/dl_query_processing/dl_query_processing/compilation/query_meta.py index 437d6127b..7699e19aa 100644 --- a/lib/dl_query_processing/dl_query_processing/compilation/query_meta.py +++ b/lib/dl_query_processing/dl_query_processing/compilation/query_meta.py @@ -4,6 +4,8 @@ Any, Optional, TypeVar, + Hashable, + NamedTuple, ) import attr @@ -15,6 +17,10 @@ ) +class QueryElementExtract(NamedTuple): + values: tuple[Optional[Hashable], ...] + + _QUERY_META_TV = TypeVar("_QUERY_META_TV", bound="QueryMetaInfo") @@ -28,6 +34,20 @@ class QueryMetaInfo: subquery_limit: int = attr.ib(kw_only=True, default=None) empty_query_mode: EmptyQueryMode = attr.ib(kw_only=True, default=EmptyQueryMode.error) + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + self.query_type.name, + tuple(self.phantom_select_ids), + tuple(self.field_order) if self.field_order is not None else None, + self.row_count_hard_limit, + self.from_subquery, + self.subquery_limit, + self.empty_query_mode.name, + ), + ) + @property def result_field_ids(self) -> list[str]: assert self.field_order is not None diff --git a/lib/dl_query_processing/dl_query_processing/translation/flat_translator.py b/lib/dl_query_processing/dl_query_processing/translation/flat_translator.py index 9cd321f33..8d4b38965 100644 --- a/lib/dl_query_processing/dl_query_processing/translation/flat_translator.py +++ b/lib/dl_query_processing/dl_query_processing/translation/flat_translator.py @@ -251,6 +251,7 @@ def translate_flat_query( offset=compiled_flat_query.offset, column_list=column_list, meta=translated_meta, + extract=compiled_flat_query.extract, ) def get_collected_stats(self) -> TranslationStats: diff --git a/lib/dl_query_processing/dl_query_processing/translation/primitives.py b/lib/dl_query_processing/dl_query_processing/translation/primitives.py index 8f3272643..f43aee685 100644 --- a/lib/dl_query_processing/dl_query_processing/translation/primitives.py +++ b/lib/dl_query_processing/dl_query_processing/translation/primitives.py @@ -32,7 +32,7 @@ DataTypeParams, ) from dl_query_processing.compilation.primitives import FromObject -from dl_query_processing.compilation.query_meta import QueryMetaInfo +from dl_query_processing.compilation.query_meta import QueryMetaInfo, QueryElementExtract from dl_query_processing.enums import ExecutionLevel @@ -42,6 +42,15 @@ class DetailedType(NamedTuple): # TODO: native_type: Optional[GenericNativeType] = None formula_data_type: Optional[DataType] = None formula_data_type_params: Optional[DataTypeParams] = None + + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract(values=( + self.field_id, + self.data_type.name, + self.formula_data_type.name if self.formula_data_type is not None else None, + self.formula_data_type_params.as_primitive() if self.formula_data_type is not None else None, + )) _META_TV = TypeVar("_META_TV", bound="TranslatedQueryMetaInfo") @@ -51,6 +60,16 @@ class DetailedType(NamedTuple): class TranslatedQueryMetaInfo(QueryMetaInfo): detailed_types: Optional[list[Optional[DetailedType]]] = attr.ib(kw_only=True, factory=list) # type: ignore + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract(values=( + *super().extract, + ( + dt.extract if dt is not None else None + for dt in self.detailed_types + ) if self.detailed_types is not None else None, + )) + @classmethod def from_comp_meta( cls: Type[_META_TV], @@ -107,6 +126,8 @@ class TranslatedFlatQuery: column_list: list[SchemaColumn] = attr.ib(kw_only=True) meta: TranslatedQueryMetaInfo = attr.ib(kw_only=True, factory=TranslatedQueryMetaInfo) + extract: QueryElementExtract = attr.ib(kw_only=True) + def is_empty(self) -> bool: return not self.select @@ -116,6 +137,14 @@ def is_empty(self) -> bool: @attr.s(frozen=True) class TranslatedMultiQueryBase(abc.ABC): + @property + def extract(self) -> QueryElementExtract: + return QueryElementExtract( + values=( + tuple(query.extract for query in sorted(self.iter_queries(), key=lambda query: query.id)), + ), + ) + @abc.abstractmethod def iter_queries(self) -> Iterable[TranslatedFlatQuery]: raise NotImplementedError