Merge pull request #183 from UncoderIO/gis-7887

alexvolha · web-flow · commit 94e16aa92bd4 · 2024-08-01T10:31:33.000+03:00
update field tokens collection
diff --git a/uncoder-core/app/translator/core/models/functions/base.py b/uncoder-core/app/translator/core/models/functions/base.py
@@ -3,7 +3,7 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Optional, Union
 
-from app.translator.core.models.query_tokens.field import Alias, Field
+from app.translator.core.models.query_tokens.field import Alias, BaseFieldsGetter, Field
 from app.translator.core.models.query_tokens.field_field import FieldField
 from app.translator.core.models.query_tokens.field_value import FieldValue
 from app.translator.core.models.query_tokens.identifier import Identifier
@@ -14,14 +14,25 @@
 
 
 @dataclass
-class Function:
+class Function(BaseFieldsGetter):
     name: str = None
     args: list[
         Union[Alias, Field, FieldField, FieldValue, FunctionValue, Keyword, Function, Identifier, int, str, bool]
     ] = field(default_factory=list)
     alias: Optional[Alias] = None
     raw: str = ""
 
+    @property
+    def fields(self) -> list[Field]:
+        fields = []
+        for arg in self.args:
+            if isinstance(arg, Field):
+                fields.append(arg)
+            elif isinstance(arg, (BaseFieldsGetter, Function)):
+                fields.extend(arg.fields)
+
+        return fields
+
 
 @dataclass
 class ParsedFunctions:
diff --git a/uncoder-core/app/translator/core/models/functions/bin.py b/uncoder-core/app/translator/core/models/functions/bin.py
@@ -19,3 +19,7 @@ class BinFunction(Function):
     span: Optional[Span] = None
     field: Optional[Field] = None
     bins: Optional[int] = None
+
+    @property
+    def fields(self) -> list[Field]:
+        return [self.field] if self.field else []
diff --git a/uncoder-core/app/translator/core/models/functions/eval.py b/uncoder-core/app/translator/core/models/functions/eval.py
@@ -17,3 +17,16 @@ class EvalArg:
 class EvalFunction(Function):
     name: str = FunctionType.eval
     args: list[EvalArg] = None
+
+    @property
+    def fields(self) -> list[Field]:
+        fields = []
+        for arg in self.args:
+            if isinstance(arg.field_, Field):
+                fields.append(arg.field_)
+            for el in arg.expression:
+                if isinstance(el, Field):
+                    fields.append(el)
+                if isinstance(el, Function):
+                    fields.extend(el.fields)
+        return fields
diff --git a/uncoder-core/app/translator/core/models/functions/group_by.py b/uncoder-core/app/translator/core/models/functions/group_by.py
@@ -1,9 +1,9 @@
-from dataclasses import Field, dataclass, field
+from dataclasses import dataclass, field
 from typing import Union
 
 from app.translator.core.custom_types.functions import FunctionType
 from app.translator.core.models.functions.base import Function
-from app.translator.core.models.query_tokens.field import Alias, PredefinedField
+from app.translator.core.models.query_tokens.field import Alias, Field, PredefinedField
 
 
 @dataclass
@@ -12,3 +12,16 @@ class GroupByFunction(Function):
     args: list[Function] = field(default_factory=list)
     by_clauses: list[Union[Alias, Field, PredefinedField]] = field(default_factory=list)
     filter_: Function = None
+
+    @property
+    def fields(self) -> list[Field]:
+        fields = []
+        for arg in self.args:
+            fields.extend(arg.fields)
+        for by_clause in self.by_clauses:
+            if isinstance(by_clause, Field):
+                fields.append(by_clause)
+        if self.filter_:
+            fields.extend(self.filter_.fields)
+
+        return fields
diff --git a/uncoder-core/app/translator/core/models/functions/join.py b/uncoder-core/app/translator/core/models/functions/join.py
@@ -4,7 +4,10 @@
 from app.translator.core.custom_types.functions import FunctionType
 from app.translator.core.models.functions.base import Function
 from app.translator.core.models.query_container import TokenizedQueryContainer
-from app.translator.core.models.query_tokens.field import Alias, Field
+from app.translator.core.models.query_tokens.field import Alias, BaseFieldsGetter, Field
+from app.translator.core.models.query_tokens.field_field import FieldField
+from app.translator.core.models.query_tokens.field_value import FieldValue
+from app.translator.core.models.query_tokens.function_value import FunctionValue
 from app.translator.core.models.query_tokens.identifier import Identifier
 from app.translator.tools.custom_enum import CustomEnum
 
@@ -22,5 +25,14 @@ class JoinFunction(Function):
     alias: Alias = None
     type_: str = JoinType.inner
     tokenized_query_container: TokenizedQueryContainer = None
-    condition: list[Union[Alias, Field, Identifier]] = field(default_factory=list)
+    condition: list[Union[FieldField, FieldValue, FunctionValue, Identifier]] = field(default_factory=list)
     preset_log_source_str: str = None
+
+    @property
+    def fields(self) -> list[Field]:
+        fields = []
+        for arg in self.condition:
+            if isinstance(arg, BaseFieldsGetter):
+                fields.extend(arg.fields)
+
+        return fields
diff --git a/uncoder-core/app/translator/core/models/functions/rename.py b/uncoder-core/app/translator/core/models/functions/rename.py
@@ -15,3 +15,11 @@ class RenameArg:
 class RenameFunction(Function):
     name: str = FunctionType.rename
     args: list[RenameArg] = None
+
+    @property
+    def fields(self) -> list[Field]:
+        fields = []
+        for arg in self.args:
+            fields.append(arg.field_)
+
+        return fields
diff --git a/uncoder-core/app/translator/core/models/functions/sort.py b/uncoder-core/app/translator/core/models/functions/sort.py
@@ -24,3 +24,15 @@ class SortLimitFunction(Function):
     name: str = FunctionType.sort_limit
     args: list[SortArg] = None
     limit: str = None
+
+    @property
+    def fields(self) -> list[Field]:
+        fields = []
+        for arg in self.args:
+            if isinstance(arg.field, Field):
+                fields.append(arg.field)
+
+            if arg.function:
+                fields.extend(arg.function.fields)
+
+        return fields
diff --git a/uncoder-core/app/translator/core/models/functions/union.py b/uncoder-core/app/translator/core/models/functions/union.py
@@ -3,10 +3,15 @@
 from app.translator.core.custom_types.functions import FunctionType
 from app.translator.core.models.functions.base import Function
 from app.translator.core.models.query_container import TokenizedQueryContainer
+from app.translator.core.models.query_tokens.field import Field
 
 
 @dataclass
 class UnionFunction(Function):
     name: str = FunctionType.union
     tokenized_query_container: TokenizedQueryContainer = None
     preset_log_source_str: str = None
+
+    @property
+    def fields(self) -> list[Field]:
+        return []
diff --git a/uncoder-core/app/translator/core/models/query_tokens/field.py b/uncoder-core/app/translator/core/models/query_tokens/field.py
@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from typing import Optional
 
 from app.translator.core.mapping import DEFAULT_MAPPING_NAME, SourceMapping
@@ -37,3 +38,10 @@ def set_generic_names_map(self, source_mappings: list[SourceMapping], default_ma
 class PredefinedField:
     def __init__(self, name: str):
         self.name = name
+
+
+class BaseFieldsGetter(ABC):
+    @property
+    @abstractmethod
+    def fields(self) -> list[Field]:
+        raise NotImplementedError("Abstract method")
diff --git a/uncoder-core/app/translator/core/models/query_tokens/field_field.py b/uncoder-core/app/translator/core/models/query_tokens/field_field.py
@@ -1,8 +1,8 @@
-from app.translator.core.models.query_tokens.field import Alias, Field
+from app.translator.core.models.query_tokens.field import Alias, BaseFieldsGetter, Field
 from app.translator.core.models.query_tokens.identifier import Identifier
 
 
-class FieldField:
+class FieldField(BaseFieldsGetter):
     def __init__(
         self,
         source_name_left: str,
@@ -16,3 +16,13 @@ def __init__(
         self.operator = operator
         self.field_right = Field(source_name=source_name_right) if not is_alias_right else None
         self.alias_right = Alias(name=source_name_right) if is_alias_right else None
+
+    @property
+    def fields(self) -> list[Field]:
+        fields = []
+        if self.field_left:
+            fields.append(self.field_left)
+        if self.field_right:
+            fields.append(self.field_right)
+
+        return fields
diff --git a/uncoder-core/app/translator/core/models/query_tokens/field_value.py b/uncoder-core/app/translator/core/models/query_tokens/field_value.py
@@ -1,13 +1,13 @@
 from typing import Union
 
 from app.translator.core.custom_types.tokens import STR_SEARCH_OPERATORS
-from app.translator.core.models.query_tokens.field import Alias, Field, PredefinedField
+from app.translator.core.models.query_tokens.field import Alias, BaseFieldsGetter, Field, PredefinedField
 from app.translator.core.models.query_tokens.identifier import Identifier
 from app.translator.core.models.query_tokens.value import Value
 from app.translator.core.str_value_manager import StrValue
 
 
-class FieldValue(Value):
+class FieldValue(BaseFieldsGetter, Value):
     def __init__(
         self,
         source_name: str,
@@ -33,3 +33,7 @@ def __repr__(self):
             return f"{self.predefined_field.name} {self.operator.token_type} {self.values}"
 
         return f"{self.field.source_name} {self.operator.token_type} {self.values}"
+
+    @property
+    def fields(self) -> list[Field]:
+        return [self.field] if self.field else []
diff --git a/uncoder-core/app/translator/core/models/query_tokens/function_value.py b/uncoder-core/app/translator/core/models/query_tokens/function_value.py
@@ -2,13 +2,18 @@
 
 from app.translator.core.custom_types.tokens import STR_SEARCH_OPERATORS
 from app.translator.core.models.functions.base import Function
+from app.translator.core.models.query_tokens.field import BaseFieldsGetter, Field
 from app.translator.core.models.query_tokens.identifier import Identifier
 from app.translator.core.models.query_tokens.value import Value
 from app.translator.core.str_value_manager import StrValue
 
 
-class FunctionValue(Value):
+class FunctionValue(BaseFieldsGetter, Value):
     def __init__(self, function: Function, operator: Identifier, value: Union[int, str, StrValue, list, tuple]):
         super().__init__(value, cast_to_int=operator.token_type not in STR_SEARCH_OPERATORS)
         self.function = function
         self.operator = operator
+
+    @property
+    def fields(self) -> list[Field]:
+        return self.function.fields
diff --git a/uncoder-core/app/translator/core/parser.py b/uncoder-core/app/translator/core/parser.py
@@ -62,23 +62,17 @@ def get_query_tokens(self, query: str) -> list[QUERY_TOKEN_TYPE]:
             raise TokenizerGeneralException("Can't translate empty query. Please provide more details")
         return self.tokenizer.tokenize(query=query)
 
+    @staticmethod
     def get_field_tokens(
-        self, query_tokens: list[QUERY_TOKEN_TYPE], functions: Optional[list[Function]] = None
+        query_tokens: list[QUERY_TOKEN_TYPE], functions: Optional[list[Function]] = None
     ) -> list[Field]:
         field_tokens = []
         for token in query_tokens:
-            if isinstance(token, FieldValue):
-                field_tokens.append(token.field)
-            elif isinstance(token, FieldField):
-                if token.field_left:
-                    field_tokens.append(token.field_left)
-                if token.field_right:
-                    field_tokens.append(token.field_right)
-            elif isinstance(token, FunctionValue):
-                field_tokens.extend(self.tokenizer.get_field_tokens_from_func_args([token.function]))
+            if isinstance(token, (FieldField, FieldValue, FunctionValue)):
+                field_tokens.extend(token.fields)
 
         if functions:
-            field_tokens.extend(self.tokenizer.get_field_tokens_from_func_args(functions))
+            field_tokens.extend([field for func in functions for field in func.fields])
 
         return field_tokens
 
diff --git a/uncoder-core/app/translator/core/tokenizer.py b/uncoder-core/app/translator/core/tokenizer.py
@@ -31,13 +31,6 @@
 )
 from app.translator.core.functions import PlatformFunctions
 from app.translator.core.mapping import SourceMapping
-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.query_tokens.field import Field
 from app.translator.core.models.query_tokens.field_field import FieldField
 from app.translator.core.models.query_tokens.field_value import FieldValue
@@ -356,43 +349,6 @@ def filter_tokens(
     ) -> list[QUERY_TOKEN_TYPE]:
         return [token for token in tokens if isinstance(token, token_type)]
 
-    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 arg.field_left:
-                    result.append(arg.field_left)
-                if arg.field_right:
-                    result.append(arg.field_right)
-            elif isinstance(arg, FieldValue):
-                if arg.field:
-                    result.append(arg.field)
-            elif isinstance(arg, FunctionValue):
-                result.extend(self.get_field_tokens_from_func_args(args=[arg.function]))
-            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):
-                result.extend(self.get_field_tokens_from_func_args(args=arg.condition))
-            elif isinstance(arg, UnionFunction):
-                continue
-            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):
-                result.append(arg.field)
-            elif isinstance(arg, RenameArg):
-                result.append(arg.field_)
-            elif isinstance(arg, EvalArg):
-                if isinstance(arg.field_, Field):
-                    result.append(arg.field_)
-                result.extend(self.get_field_tokens_from_func_args(args=arg.expression))
-        return result
-
     @staticmethod
     def set_field_tokens_generic_names_map(
         tokens: list[Field], source_mappings: list[SourceMapping], default_mapping: SourceMapping
