Merge pull request #173 from UncoderIO/gis-8193

mapping flow changes, render unmapped fields to comment

Author: alexvolha

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

@@ -10,19 +10,14 @@ class BasePlatformException(BaseException):
 
 
 class StrictPlatformException(BasePlatformException):
-    field_name: str = None
-
-    def __init__(
-        self, platform_name: str, field_name: str, mapping: Optional[str] = None, detected_fields: Optional[list] = None
-    ):
+    def __init__(self, platform_name: str, fields: list[str], mapping: Optional[str] = None):
         message = (
             f"Platform {platform_name} has strict mapping. "
-            f"Source fields: {', '.join(detected_fields) if detected_fields else field_name} has no mapping."
+            f"Source fields: {', '.join(fields)} have no mapping."
             f" Mapping file: {mapping}."
             if mapping
             else ""
         )
-        self.field_name = field_name
         super().__init__(message)
 
 

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

@@ -94,17 +94,16 @@ def set_functions_manager(self, manager: PlatformFunctionsManager) -> FunctionRe
     def render(self, function: Function, source_mapping: SourceMapping) -> str:
         raise NotImplementedError
 
-    @staticmethod
-    def map_field(field: Union[Alias, Field], source_mapping: SourceMapping) -> str:
+    def map_field(self, field: Union[Alias, Field], source_mapping: SourceMapping) -> str:
         if isinstance(field, Alias):
             return field.name
 
-        generic_field_name = field.get_generic_field_name(source_mapping.source_id)
-        mapped_field = source_mapping.fields_mapping.get_platform_field_name(generic_field_name=generic_field_name)
-        if isinstance(mapped_field, list):
-            mapped_field = mapped_field[0]
+        if isinstance(field, Field):
+            mappings = self.manager.platform_functions.platform_query_render.mappings
+            mapped_fields = mappings.map_field(field, source_mapping)
+            return mapped_fields[0]
 
-        return mapped_field if mapped_field else field.source_name
+        raise NotSupportedFunctionException
 
 
 class PlatformFunctionsManager:

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

@@ -1,10 +1,16 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Optional, TypeVar
+from typing import TYPE_CHECKING, Optional, TypeVar
 
+from app.translator.core.exceptions.core import StrictPlatformException
+from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.mappings.utils.load_from_files import LoaderFileMappings
 
+if TYPE_CHECKING:
+    from app.translator.core.models.query_tokens.field import Field
+
+
 DEFAULT_MAPPING_NAME = "default"
 
 
@@ -85,12 +91,16 @@ def __init__(
 
 
 class BasePlatformMappings:
+    details: PlatformDetails = None
+
+    is_strict_mapping: bool = False
     skip_load_default_mappings: bool = True
     extend_default_mapping_with_all_fields: bool = False
 
-    def __init__(self, platform_dir: str):
+    def __init__(self, platform_dir: str, platform_details: PlatformDetails):
         self._loader = LoaderFileMappings()
         self._platform_dir = platform_dir
+        self.details = platform_details
         self._source_mappings = self.prepare_mapping()
 
     def update_default_source_mapping(self, default_mapping: SourceMapping, fields_mapping: FieldsMapping) -> None:
@@ -148,6 +158,32 @@ def get_source_mapping(self, source_id: str) -> Optional[SourceMapping]:
     def default_mapping(self) -> SourceMapping:
         return self._source_mappings[DEFAULT_MAPPING_NAME]
 
+    def check_fields_mapping_existence(self, field_tokens: list[Field], source_mapping: SourceMapping) -> list[str]:
+        unmapped = []
+        for field in field_tokens:
+            generic_field_name = field.get_generic_field_name(source_mapping.source_id)
+            mapped_field = source_mapping.fields_mapping.get_platform_field_name(generic_field_name=generic_field_name)
+            if not mapped_field and field.source_name not in unmapped:
+                unmapped.append(field.source_name)
+
+        if self.is_strict_mapping and unmapped:
+            raise StrictPlatformException(
+                platform_name=self.details.name, fields=unmapped, mapping=source_mapping.source_id
+            )
+
+        return unmapped
+
+    @staticmethod
+    def map_field(field: Field, source_mapping: SourceMapping) -> list[str]:
+        generic_field_name = field.get_generic_field_name(source_mapping.source_id)
+        # field can be mapped to corresponding platform field name or list of platform field names
+        mapped_field = source_mapping.fields_mapping.get_platform_field_name(generic_field_name=generic_field_name)
+
+        if isinstance(mapped_field, str):
+            mapped_field = [mapped_field]
+
+        return mapped_field if mapped_field else [generic_field_name] if generic_field_name else [field.source_name]
+
 
 class BaseCommonPlatformMappings(ABC, BasePlatformMappings):
     def prepare_mapping(self) -> dict[str, SourceMapping]:

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

@@ -184,6 +184,7 @@ class QueryRender(ABC):
     details: PlatformDetails = None
     is_single_line_comment: bool = False
     unsupported_functions_text = "Unsupported functions were excluded from the result query:"
+    unmapped_fields_text = "Unmapped fields: "
 
     platform_functions: PlatformFunctions = None
 
@@ -206,6 +207,11 @@ def wrap_with_not_supported_functions(self, query: str, not_supported_functions:
 
         return query
 
+    def wrap_with_unmapped_fields(self, query: str, fields: Optional[list[str]]) -> str:
+        if fields:
+            return query + "\n\n" + self.wrap_with_comment(f"{self.unmapped_fields_text}{', '.join(fields)}")
+        return query
+
     def wrap_with_comment(self, value: str) -> str:
         return f"{self.comment_symbol} {value}"
 
@@ -216,7 +222,6 @@ def generate(self, query_container: Union[RawQueryContainer, TokenizedQueryConta
 
 class PlatformQueryRender(QueryRender):
     mappings: BasePlatformMappings = None
-    is_strict_mapping: bool = False
 
     or_token = "or"
     and_token = "and"
@@ -247,22 +252,10 @@ def generate_prefix(self, log_source_signature: Optional[LogSourceSignature], fu
     def generate_functions(self, functions: list[Function], source_mapping: SourceMapping) -> RenderedFunctions:
         return self.platform_functions.render(functions, source_mapping)
 
-    def map_field(self, field: Field, source_mapping: SourceMapping) -> list[str]:
-        generic_field_name = field.get_generic_field_name(source_mapping.source_id)
-        # field can be mapped to corresponding platform field name or list of platform field names
-        mapped_field = source_mapping.fields_mapping.get_platform_field_name(generic_field_name=generic_field_name)
-        if not mapped_field and self.is_strict_mapping:
-            raise StrictPlatformException(field_name=field.source_name, platform_name=self.details.name)
-
-        if isinstance(mapped_field, str):
-            mapped_field = [mapped_field]
-
-        return mapped_field if mapped_field else [generic_field_name] if generic_field_name else [field.source_name]
-
     def map_predefined_field(self, predefined_field: PredefinedField) -> str:
         if not (mapped_predefined_field_name := self.predefined_fields_map.get(predefined_field.name)):
-            if self.is_strict_mapping:
-                raise StrictPlatformException(field_name=predefined_field.name, platform_name=self.details.name)
+            if self.mappings.is_strict_mapping:
+                raise StrictPlatformException(platform_name=self.details.name, fields=[predefined_field.name])
 
             return predefined_field.name
 
@@ -275,7 +268,7 @@ def apply_token(self, token: QUERY_TOKEN_TYPE, source_mapping: SourceMapping) ->
             elif token.predefined_field:
                 mapped_fields = [self.map_predefined_field(token.predefined_field)]
             else:
-                mapped_fields = self.map_field(token.field, source_mapping)
+                mapped_fields = self.mappings.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)
@@ -285,9 +278,13 @@ def apply_token(self, token: QUERY_TOKEN_TYPE, source_mapping: SourceMapping) ->
             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)
+            mapped_fields_left = (
+                [alias_left.name] if alias_left else self.mappings.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)
+            mapped_fields_right = (
+                [alias_right.name] if alias_right else self.mappings.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(
                 [
@@ -311,14 +308,9 @@ def apply_token(self, token: QUERY_TOKEN_TYPE, source_mapping: SourceMapping) ->
 
     def generate_query(self, tokens: list[QUERY_TOKEN_TYPE], source_mapping: SourceMapping) -> str:
         result_values = []
-        unmapped_fields = set()
         for token in tokens:
-            try:
-                result_values.append(self.apply_token(token=token, source_mapping=source_mapping))
-            except StrictPlatformException as err:
-                unmapped_fields.add(err.field_name)
-        if unmapped_fields:
-            raise StrictPlatformException(self.details.name, "", source_mapping.source_id, sorted(unmapped_fields))
+            result_values.append(self.apply_token(token=token, source_mapping=source_mapping))
+
         return "".join(result_values)
 
     def wrap_with_meta_info(self, query: str, meta_info: Optional[MetaInfoContainer]) -> str:
@@ -351,11 +343,13 @@ def finalize_query(
         meta_info: Optional[MetaInfoContainer] = None,
         source_mapping: Optional[SourceMapping] = None,  # noqa: ARG002
         not_supported_functions: Optional[list] = None,
+        unmapped_fields: Optional[list[str]] = None,
         *args,  # noqa: ARG002
         **kwargs,  # noqa: ARG002
     ) -> str:
         query = self._join_query_parts(prefix, query, functions)
         query = self.wrap_with_meta_info(query, meta_info)
+        query = self.wrap_with_unmapped_fields(query, unmapped_fields)
         return self.wrap_with_not_supported_functions(query, not_supported_functions)
 
     @staticmethod
@@ -417,8 +411,10 @@ def generate_raw_log_fields(self, fields: list[Field], source_mapping: SourceMap
                 mapped_field = source_mapping.fields_mapping.get_platform_field_name(
                     generic_field_name=generic_field_name
                 )
-            if not mapped_field and self.is_strict_mapping:
-                raise StrictPlatformException(field_name=field.source_name, platform_name=self.details.name)
+            if not mapped_field and self.mappings.is_strict_mapping:
+                raise StrictPlatformException(
+                    platform_name=self.details.name, fields=[field.source_name], mapping=source_mapping.source_id
+                )
             if prefix_list := self.process_raw_log_field_prefix(field=mapped_field, source_mapping=source_mapping):
                 for prefix in prefix_list:
                     if prefix not in defined_raw_log_fields:
@@ -428,6 +424,9 @@ def generate_raw_log_fields(self, fields: list[Field], source_mapping: SourceMap
     def _generate_from_tokenized_query_container_by_source_mapping(
         self, query_container: TokenizedQueryContainer, source_mapping: SourceMapping
     ) -> str:
+        unmapped_fields = self.mappings.check_fields_mapping_existence(
+            query_container.meta_info.query_fields, source_mapping
+        )
         rendered_functions = self.generate_functions(query_container.functions.functions, source_mapping)
         prefix = self.generate_prefix(source_mapping.log_source_signature, rendered_functions.rendered_prefix)
 
@@ -443,6 +442,7 @@ def _generate_from_tokenized_query_container_by_source_mapping(
             query=query,
             functions=rendered_functions.rendered,
             not_supported_functions=not_supported_functions,
+            unmapped_fields=unmapped_fields,
             meta_info=query_container.meta_info,
             source_mapping=source_mapping,
         )

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

@@ -9,4 +9,4 @@
     "alt_platform_name": "OCSF",
 }
 
-athena_details = PlatformDetails(**ATHENA_QUERY_DETAILS)
+athena_query_details = PlatformDetails(**ATHENA_QUERY_DETAILS)

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

@@ -1,6 +1,7 @@
 from typing import Optional
 
 from app.translator.core.mapping import DEFAULT_MAPPING_NAME, BasePlatformMappings, LogSourceSignature, SourceMapping
+from app.translator.platforms.athena.const import athena_query_details
 
 
 class AthenaLogSourceSignature(LogSourceSignature):
@@ -40,4 +41,4 @@ def get_suitable_source_mappings(self, field_names: list[str], table: Optional[s
         return suitable_source_mappings
 
 
-athena_mappings = AthenaMappings(platform_dir="athena")
+athena_query_mappings = AthenaMappings(platform_dir="athena", platform_details=athena_query_details)

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

@@ -18,14 +18,14 @@
 
 from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.managers import parser_manager
-from app.translator.platforms.athena.const import athena_details
-from app.translator.platforms.athena.mapping import AthenaMappings, athena_mappings
+from app.translator.platforms.athena.const import athena_query_details
+from app.translator.platforms.athena.mapping import AthenaMappings, athena_query_mappings
 from app.translator.platforms.base.sql.parsers.sql import SqlQueryParser
 
 
 @parser_manager.register_supported_by_roota
 class AthenaQueryParser(SqlQueryParser):
-    details: PlatformDetails = athena_details
-    mappings: AthenaMappings = athena_mappings
+    details: PlatformDetails = athena_query_details
+    mappings: AthenaMappings = athena_query_mappings
     query_delimiter_pattern = r"\sFROM\s\S*\sWHERE\s"
     table_pattern = r"\sFROM\s(?P<table>[a-zA-Z\.\-\*]+)\sWHERE\s"

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

@@ -19,19 +19,19 @@
 
 from app.translator.core.models.platform_details import PlatformDetails
 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.athena.const import athena_query_details
+from app.translator.platforms.athena.mapping import AthenaMappings, athena_query_mappings
 from app.translator.platforms.base.sql.renders.sql import SqlFieldValueRender, SqlQueryRender
 
 
 class AthenaFieldValueRender(SqlFieldValueRender):
-    details: PlatformDetails = athena_details
+    details: PlatformDetails = athena_query_details
 
 
 @render_manager.register
 class AthenaQueryRender(SqlQueryRender):
-    details: PlatformDetails = athena_details
-    mappings: AthenaMappings = athena_mappings
+    details: PlatformDetails = athena_query_details
+    mappings: AthenaMappings = athena_query_mappings
 
     or_token = "OR"
 

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

@@ -20,13 +20,13 @@
 from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.core.render_cti import RenderCTI
 from app.translator.managers import render_cti_manager
-from app.translator.platforms.athena.const import athena_details
+from app.translator.platforms.athena.const import athena_query_details
 from app.translator.platforms.athena.mappings.athena_cti import DEFAULT_ATHENA_MAPPING
 
 
 @render_cti_manager.register
 class AthenaCTI(RenderCTI):
-    details: PlatformDetails = athena_details
+    details: PlatformDetails = athena_query_details
 
     field_value_template: str = "{key} = '{value}'"
     or_operator: str = " OR "

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

@@ -90,6 +90,3 @@ def get_suitable_source_mappings(
             suitable_source_mappings = [self._source_mappings[DEFAULT_MAPPING_NAME]]
 
         return suitable_source_mappings
-
-
-aql_mappings = AQLMappings(platform_dir="qradar")