FieldField class, renders improvements

Author: alexvolha

uncoder-core/app/translator/core/const.py

@@ -0,0 +1,6 @@
+from typing import Union
+
+from app.translator.core.models.field import Alias, Field, FieldValue, Keyword
+from app.translator.core.models.identifier import Identifier
+
+TOKEN_TYPE = Union[FieldValue, Keyword, Identifier, Field, Alias]

uncoder-core/app/translator/core/context_vars.py

@@ -1,4 +1,10 @@
 from contextvars import ContextVar
+from typing import Optional
 
 return_only_first_query_ctx_var: ContextVar[bool] = ContextVar("return_only_first_query_ctx_var", default=False)
 """Set to True to return only first query if rendered multiple options"""
+
+wrap_query_with_meta_info_ctx_var: ContextVar[bool] = ContextVar("wrap_query_with_meta_info_ctx_var", default=True)
+"""Set to False not to wrap query with meta info commentary"""
+
+preset_log_source_str_ctx_var: ContextVar[Optional[str]] = ContextVar("preset_log_source_str_ctx_var", default=None)

uncoder-core/app/translator/core/models/field.py

@@ -37,6 +37,22 @@ def set_generic_names_map(self, source_mappings: list[SourceMapping], default_ma
         self.__generic_names_map = generic_names_map
 
 
+class FieldField:
+    def __init__(
+        self,
+        source_name_left: str,
+        operator: Identifier,
+        source_name_right: str,
+        is_alias_left: bool = False,
+        is_alias_right: bool = False,
+    ):
+        self.field_left = Field(source_name=source_name_left)
+        self.alias_left = Alias(name=source_name_left) if is_alias_left else None
+        self.operator = operator
+        self.field_right = Field(source_name=source_name_right)
+        self.alias_right = Alias(name=source_name_right) if is_alias_right else None
+
+
 class FieldValue:
     def __init__(
         self,

uncoder-core/app/translator/core/models/query_container.py

@@ -3,11 +3,11 @@
 from datetime import datetime
 from typing import Optional
 
+from app.translator.core.const import TOKEN_TYPE
 from app.translator.core.custom_types.meta_info import SeverityType
 from app.translator.core.mapping import DEFAULT_MAPPING_NAME
 from app.translator.core.models.field import Field
 from app.translator.core.models.functions.base import ParsedFunctions
-from app.translator.core.tokenizer import TOKEN_TYPE
 
 
 class MetaInfoContainer:

uncoder-core/app/translator/core/parser.py

@@ -20,6 +20,7 @@
 from abc import ABC, abstractmethod
 from typing import Union
 
+from app.translator.core.const import TOKEN_TYPE
 from app.translator.core.exceptions.parser import TokenizerGeneralException
 from app.translator.core.functions import PlatformFunctions
 from app.translator.core.mapping import BasePlatformMappings, SourceMapping
@@ -28,7 +29,7 @@
 from app.translator.core.models.identifier import Identifier
 from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.core.models.query_container import RawQueryContainer, TokenizedQueryContainer
-from app.translator.core.tokenizer import TOKEN_TYPE, QueryTokenizer
+from app.translator.core.tokenizer import QueryTokenizer
 
 
 class QueryParser(ABC):

uncoder-core/app/translator/core/render.py

@@ -16,36 +16,36 @@
 limitations under the License.
 -----------------------------------------------------------------
 """
-
+import itertools
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from typing import ClassVar, Optional, Union
 
 from app.translator.const import DEFAULT_VALUE_TYPE
-from app.translator.core.context_vars import return_only_first_query_ctx_var
+from app.translator.core.const import TOKEN_TYPE
+from app.translator.core.context_vars import return_only_first_query_ctx_var, wrap_query_with_meta_info_ctx_var
 from app.translator.core.custom_types.tokens import LogicalOperatorType, OperatorType
 from app.translator.core.custom_types.values import ValueType
 from app.translator.core.escape_manager import EscapeManager
 from app.translator.core.exceptions.core import NotImplementedException, StrictPlatformException
 from app.translator.core.exceptions.parser import UnsupportedOperatorException
 from app.translator.core.functions import PlatformFunctions
 from app.translator.core.mapping import DEFAULT_MAPPING_NAME, BasePlatformMappings, LogSourceSignature, SourceMapping
-from app.translator.core.models.field import Field, FieldValue, Keyword
+from app.translator.core.models.field import Field, FieldField, FieldValue, Keyword
 from app.translator.core.models.functions.base import Function, RenderedFunctions
 from app.translator.core.models.identifier import Identifier
 from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.core.models.query_container import MetaInfoContainer, RawQueryContainer, TokenizedQueryContainer
 from app.translator.core.str_value_manager import StrValue, StrValueManager
-from app.translator.core.tokenizer import TOKEN_TYPE
 
 
-class BaseQueryFieldValue(ABC):
+class BaseFieldValueRender(ABC):
     details: PlatformDetails = None
     escape_manager: EscapeManager = None
     str_value_manager: StrValueManager = None
 
     def __init__(self, or_token: str):
-        self.field_value: dict[str, Callable[[str, DEFAULT_VALUE_TYPE], str]] = {
+        self.modifiers_map: dict[str, Callable[[str, DEFAULT_VALUE_TYPE], str]] = {
             OperatorType.EQ: self.equal_modifier,
             OperatorType.NOT_EQ: self.not_equal_modifier,
             OperatorType.LT: self.less_modifier,
@@ -155,11 +155,20 @@ def apply_value(self, value: Union[str, int], value_type: str = ValueType.value)
         return self.escape_manager.escape(value, value_type)
 
     def apply_field_value(self, field: str, operator: Identifier, value: DEFAULT_VALUE_TYPE) -> str:
-        if modifier_function := self.field_value.get(operator.token_type):
+        if modifier_function := self.modifiers_map.get(operator.token_type):
             return modifier_function(field, value)
         raise UnsupportedOperatorException(operator.token_type)
 
 
+class BaseFieldFieldRender(ABC):
+    operators_map: ClassVar[dict[str, str]] = {}
+
+    def apply_field_field(self, field_left: str, operator: Identifier, field_right: str) -> str:
+        if mapped_operator := self.operators_map.get(operator.token_type):
+            return f"{field_left} {mapped_operator} {field_right}"
+        raise UnsupportedOperatorException(operator.token_type)
+
+
 class QueryRender(ABC):
     comment_symbol: str = None
     details: PlatformDetails = None
@@ -180,6 +189,13 @@ def render_not_supported_functions(self, not_supported_functions: list) -> str:
         not_supported_functions_str = "\n".join(line_template + func.lstrip() for func in not_supported_functions)
         return "\n\n" + self.wrap_with_comment(f"{self.unsupported_functions_text}\n{not_supported_functions_str}")
 
+    def wrap_with_not_supported_functions(self, query: str, not_supported_functions: Optional[list] = None) -> str:
+        if not_supported_functions and wrap_query_with_meta_info_ctx_var.get():
+            rendered_not_supported = self.render_not_supported_functions(not_supported_functions)
+            return query + rendered_not_supported
+
+        return query
+
     def wrap_with_comment(self, value: str) -> str:
         return f"{self.comment_symbol} {value}"
 
@@ -199,13 +215,14 @@ class PlatformQueryRender(QueryRender):
     group_token = "(%s)"
     query_parts_delimiter = " "
 
-    field_value_map = BaseQueryFieldValue(or_token=or_token)
+    field_field_render = BaseFieldFieldRender()
+    field_value_render = BaseFieldValueRender(or_token=or_token)
 
     raw_log_field_pattern_map: ClassVar[dict[str, str]] = None
 
     def __init__(self):
         super().__init__()
-        self.operator_map = {
+        self.logical_operators_map = {
             LogicalOperatorType.AND: f" {self.and_token} ",
             LogicalOperatorType.OR: f" {self.or_token} ",
             LogicalOperatorType.NOT: f" {self.not_token} ",
@@ -233,31 +250,34 @@ def map_field(self, field: Field, source_mapping: SourceMapping) -> list[str]:
 
     def apply_token(self, token: Union[FieldValue, Keyword, Identifier], source_mapping: SourceMapping) -> str:
         if isinstance(token, FieldValue):
-            if token.alias:
-                field_name = token.alias.name
-            else:
-                mapped_fields = self.map_field(token.field, source_mapping)
-                if len(mapped_fields) > 1:
-                    return self.group_token % self.operator_map[LogicalOperatorType.OR].join(
-                        [
-                            self.field_value_map.apply_field_value(
-                                field=field, operator=token.operator, value=token.value
-                            )
-                            for field in mapped_fields
-                        ]
-                    )
-
-                field_name = mapped_fields[0]
-
-            return self.field_value_map.apply_field_value(field=field_name, operator=token.operator, value=token.value)
-
+            mapped_fields = [token.alias.name] if token.alias else self.map_field(token.field, source_mapping)
+            joined = self.logical_operators_map[LogicalOperatorType.OR].join(
+                [
+                    self.field_value_render.apply_field_value(field=field, operator=token.operator, value=token.value)
+                    for field in mapped_fields
+                ]
+            )
+            return self.group_token % joined if len(mapped_fields) > 1 else joined
+        if isinstance(token, FieldField):
+            alias_left, field_left = token.alias_left, token.field_left
+            mapped_fields_left = [alias_left.name] if alias_left else self.map_field(field_left, source_mapping)
+            alias_right, field_right = token.alias_right, token.field_right
+            mapped_fields_right = [alias_right.name] if alias_right else self.map_field(field_right, source_mapping)
+            cross_paired_fields = list(itertools.product(mapped_fields_left, mapped_fields_right))
+            joined = self.logical_operators_map[LogicalOperatorType.OR].join(
+                [
+                    self.field_field_render.apply_field_field(pair[0], token.operator, pair[1])
+                    for pair in cross_paired_fields
+                ]
+            )
+            return self.group_token % joined if len(cross_paired_fields) > 1 else joined
         if isinstance(token, Function):
             func_render = self.platform_functions.manager.get_in_query_render(token.name)
             return func_render.render(token, source_mapping)
         if isinstance(token, Keyword):
-            return self.field_value_map.apply_field_value(field="", operator=token.operator, value=token.value)
+            return self.field_value_render.apply_field_value(field="", operator=token.operator, value=token.value)
         if token.token_type in LogicalOperatorType:
-            return self.operator_map.get(token.token_type)
+            return self.logical_operators_map.get(token.token_type)
 
         return token.token_type
 
@@ -267,8 +287,8 @@ def generate_query(self, tokens: list[TOKEN_TYPE], source_mapping: SourceMapping
             result_values.append(self.apply_token(token=token, source_mapping=source_mapping))
         return "".join(result_values)
 
-    def wrap_query_with_meta_info(self, meta_info: MetaInfoContainer, query: str) -> str:
-        if meta_info and (meta_info.id or meta_info.title):
+    def wrap_with_meta_info(self, query: str, meta_info: MetaInfoContainer) -> str:
+        if wrap_query_with_meta_info_ctx_var.get() and meta_info and (meta_info.id or meta_info.title):
             meta_info_dict = {
                 "name: ": meta_info.title,
                 "uuid: ": meta_info.id,
@@ -301,11 +321,8 @@ def finalize_query(
         **kwargs,  # noqa: ARG002
     ) -> str:
         query = self._join_query_parts(prefix, query, functions)
-        query = self.wrap_query_with_meta_info(meta_info=meta_info, query=query)
-        if not_supported_functions:
-            rendered_not_supported = self.render_not_supported_functions(not_supported_functions)
-            return query + rendered_not_supported
-        return query
+        query = self.wrap_with_meta_info(query, meta_info)
+        return self.wrap_with_not_supported_functions(query, not_supported_functions)
 
     @staticmethod
     def unique_queries(queries_map: dict[str, str]) -> dict[str, dict[str]]:
@@ -336,7 +353,7 @@ def _get_source_mappings(self, source_mapping_ids: list[str]) -> list[SourceMapp
 
         return source_mappings
 
-    def _generate_from_raw_query_container(self, query_container: RawQueryContainer) -> str:
+    def generate_from_raw_query_container(self, query_container: RawQueryContainer) -> str:
         return self.finalize_query(
             prefix="", query=query_container.query, functions="", meta_info=query_container.meta_info
         )
@@ -374,7 +391,7 @@ def generate_raw_log_fields(self, fields: list[Field], source_mapping: SourceMap
                         defined_raw_log_fields.append(prefix)
         return "\n".join(defined_raw_log_fields)
 
-    def _generate_from_tokenized_query_container(self, query_container: TokenizedQueryContainer) -> str:
+    def generate_from_tokenized_query_container(self, query_container: TokenizedQueryContainer) -> str:
         queries_map = {}
         errors = []
         source_mappings = self._get_source_mappings(query_container.meta_info.source_mapping_ids)
@@ -411,6 +428,6 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue
 
     def generate(self, query_container: Union[RawQueryContainer, TokenizedQueryContainer]) -> str:
         if isinstance(query_container, RawQueryContainer):
-            return self._generate_from_raw_query_container(query_container)
+            return self.generate_from_raw_query_container(query_container)
 
-        return self._generate_from_tokenized_query_container(query_container)
+        return self.generate_from_tokenized_query_container(query_container)

uncoder-core/app/translator/core/tokenizer.py

@@ -20,6 +20,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, ClassVar, Optional, Union
 
+from app.translator.core.const import TOKEN_TYPE
 from app.translator.core.custom_types.tokens import GroupType, LogicalOperatorType, OperatorType
 from app.translator.core.custom_types.values import ValueType
 from app.translator.core.escape_manager import EscapeManager
@@ -29,18 +30,18 @@
     UnsupportedOperatorException,
 )
 from app.translator.core.mapping import SourceMapping
-from app.translator.core.models.field import Field, FieldValue, Keyword
+from app.translator.core.models.field import Field, FieldField, FieldValue, Keyword
 from app.translator.core.models.functions.base import Function
 from app.translator.core.models.functions.eval import EvalArg
 from app.translator.core.models.functions.group_by import GroupByFunction
+from app.translator.core.models.functions.join import JoinFunction
 from app.translator.core.models.functions.rename import RenameArg
 from app.translator.core.models.functions.sort import SortArg
+from app.translator.core.models.functions.union import UnionFunction
 from app.translator.core.models.identifier import Identifier
 from app.translator.core.str_value_manager import StrValue, StrValueManager
 from app.translator.tools.utils import get_match_group
 
-TOKEN_TYPE = Union[FieldValue, Keyword, Identifier, Field]
-
 
 class BaseTokenizer(ABC):
     @abstractmethod
@@ -323,20 +324,31 @@ def filter_tokens(
     ) -> list[TOKEN_TYPE]:
         return [token for token in tokens if isinstance(token, token_type)]
 
-    def get_field_tokens_from_func_args(
+    def get_field_tokens_from_func_args(  # noqa: PLR0912
         self, args: list[Union[Field, FieldValue, Keyword, Identifier, Function, SortArg]]
     ) -> list[Field]:
         result = []
         for arg in args:
             if isinstance(arg, Field):
                 result.append(arg)
+            elif isinstance(arg, FieldField):
+                if not arg.alias_left or arg.alias_left.name != arg.field_left.source_name:
+                    result.append(arg.field_left)
+                if not arg.alias_right or arg.alias_right.name != arg.field_right.source_name:
+                    result.append(arg.field_right)
             elif isinstance(arg, FieldValue):
                 if not arg.alias or arg.alias.name != arg.field.source_name:
                     result.append(arg.field)
             elif isinstance(arg, GroupByFunction):
                 result.extend(self.get_field_tokens_from_func_args(args=arg.args))
                 result.extend(self.get_field_tokens_from_func_args(args=arg.by_clauses))
                 result.extend(self.get_field_tokens_from_func_args(args=[arg.filter_]))
+            elif isinstance(arg, (JoinFunction, UnionFunction)):
+                result.extend(self.get_field_tokens_from_func_args(args=arg.tokenized_query_container.tokens))
+                result.extend(
+                    self.get_field_tokens_from_func_args(args=arg.tokenized_query_container.functions.functions)
+                )
+                result.extend(self.get_field_tokens_from_func_args(args=arg.condition))
             elif isinstance(arg, Function):
                 result.extend(self.get_field_tokens_from_func_args(args=arg.args))
             elif isinstance(arg, SortArg) and isinstance(arg.field, Field):

uncoder-core/app/translator/platforms/athena/renders/athena.py

@@ -21,10 +21,10 @@
 from app.translator.managers import render_manager
 from app.translator.platforms.athena.const import athena_details
 from app.translator.platforms.athena.mapping import AthenaMappings, athena_mappings
-from app.translator.platforms.base.sql.renders.sql import SqlFieldValue, SqlQueryRender
+from app.translator.platforms.base.sql.renders.sql import SqlFieldValueRender, SqlQueryRender
 
 
-class AthenaFieldValue(SqlFieldValue):
+class AthenaFieldValueRender(SqlFieldValueRender):
     details: PlatformDetails = athena_details
 
 
@@ -35,7 +35,7 @@ class AthenaQueryRender(SqlQueryRender):
 
     or_token = "OR"
 
-    field_value_map = AthenaFieldValue(or_token=or_token)
+    field_value_render = AthenaFieldValueRender(or_token=or_token)
     comment_symbol = "--"
     is_single_line_comment = True
 

uncoder-core/app/translator/platforms/base/aql/renders/aql.py

@@ -21,13 +21,13 @@
 
 from app.translator.const import DEFAULT_VALUE_TYPE
 from app.translator.core.custom_types.values import ValueType
-from app.translator.core.render import BaseQueryFieldValue, PlatformQueryRender
+from app.translator.core.render import BaseFieldValueRender, PlatformQueryRender
 from app.translator.core.str_value_manager import StrValue
 from app.translator.platforms.base.aql.mapping import AQLLogSourceSignature, AQLMappings, aql_mappings
 from app.translator.platforms.base.aql.str_value_manager import aql_str_value_manager
 
 
-class AQLFieldValue(BaseQueryFieldValue):
+class AQLFieldValueRender(BaseFieldValueRender):
     str_value_manager = aql_str_value_manager
 
     @staticmethod
@@ -127,8 +127,6 @@ class AQLQueryRender(PlatformQueryRender):
     and_token = "AND"
     not_token = "NOT"
 
-    field_value_map = AQLFieldValue(or_token=or_token)
-
     def generate_prefix(self, log_source_signature: AQLLogSourceSignature, functions_prefix: str = "") -> str:  # noqa: ARG002
         table = str(log_source_signature)
         extra_condition = log_source_signature.extra_condition

uncoder-core/app/translator/platforms/base/lucene/renders/lucene.py

@@ -21,13 +21,13 @@
 
 from app.translator.const import DEFAULT_VALUE_TYPE
 from app.translator.core.custom_types.values import ValueType
-from app.translator.core.render import BaseQueryFieldValue, PlatformQueryRender
+from app.translator.core.render import BaseFieldValueRender, PlatformQueryRender
 from app.translator.core.str_value_manager import StrValue
 from app.translator.platforms.base.lucene.mapping import LuceneLogSourceSignature
 from app.translator.platforms.base.lucene.str_value_manager import lucene_str_value_manager
 
 
-class LuceneFieldValue(BaseQueryFieldValue):
+class LuceneFieldValueRender(BaseFieldValueRender):
     str_value_manager = lucene_str_value_manager
 
     @staticmethod