Merge pull request #150 from UncoderIO/gis-7222

alexvolha · web-flow · commit 780bc1559776 · 2024-06-18T16:31:03.000+03:00
Gis 7222
diff --git a/uncoder-core/app/routers/translate.py b/uncoder-core/app/routers/translate.py
@@ -1,6 +1,7 @@
 from fastapi import APIRouter, Body
 
 from app.models.translation import InfoMessage, OneTranslationData, Platform, TranslatorPlatforms
+from app.translator.core.context_vars import return_only_first_query_ctx_var
 from app.translator.cti_translator import CTITranslator
 from app.translator.translator import Translator
 
@@ -15,7 +16,9 @@ def translate_one(
     source_platform_id: str = Body(..., embed=True),
     target_platform_id: str = Body(..., embed=True),
     text: str = Body(..., embed=True),
+    return_only_first_query: bool = False,
 ) -> OneTranslationData:
+    return_only_first_query_ctx_var.set(return_only_first_query)
     status, data = translator.translate_one(text=text, source=source_platform_id, target=target_platform_id)
     if status:
         return OneTranslationData(status=status, translation=data, target_platform_id=target_platform_id)
@@ -27,8 +30,11 @@ def translate_one(
 @st_router.post("/translate/all", tags=["translator"], description="Generate all translations")
 @st_router.post("/translate/all/", include_in_schema=False)
 def translate_all(
-    source_platform_id: str = Body(..., embed=True), text: str = Body(..., embed=True)
+    source_platform_id: str = Body(..., embed=True),
+    text: str = Body(..., embed=True),
+    return_only_first_query: bool = False,
 ) -> list[OneTranslationData]:
+    return_only_first_query_ctx_var.set(return_only_first_query)
     result = translator.translate_all(text=text, source=source_platform_id)
     translations = []
     for platform_result in result:
diff --git a/uncoder-core/app/translator/platforms/base/aql/const.py b/uncoder-core/app/translator/platforms/base/aql/const.py
@@ -1,5 +1,7 @@
 UTF8_PAYLOAD_PATTERN = r"UTF8\(payload\)"
 NUM_VALUE_PATTERN = r"(?P<num_value>\d+(?:\.\d+)*)"
-SINGLE_QUOTES_VALUE_PATTERN = r"""'(?P<s_q_value>(?:[:a-zA-Z\*0-9=+%#\-\/\\|,;_<>`~".$&^@!?\(\)\{\}\[\]\s]|'')*)'"""
+SINGLE_QUOTES_VALUE_PATTERN = (
+    r"""'(?P<s_q_value>(?:[:a-zA-Zа-яА-Я\*0-9=+%#\-\/\\|,;_<>`~".$&^@!?\(\)\{\}\[\]\s]|'')*)'"""
+)
 TABLE_PATTERN = r"\s+FROM\s+[a-zA-Z.\-*]+"
 TABLE_GROUP_PATTERN = r"\s+FROM\s+(?P<table>[a-zA-Z.\-*]+)"
diff --git a/uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py b/uncoder-core/app/translator/platforms/forti_siem/renders/forti_siem_rule.py
@@ -18,6 +18,7 @@
 from typing import 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.custom_types.meta_info import SeverityType
 from app.translator.core.custom_types.tokens import GroupType, LogicalOperatorType, OperatorType
 from app.translator.core.custom_types.values import ValueType
@@ -273,6 +274,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue
                 source_mapping=source_mapping,
                 fields=mapped_fields_set,
             )
+            if return_only_first_query_ctx_var.get() is True:
+                return finalized_query
             queries_map[source_mapping.source_id] = finalized_query
 
         return self.finalize(queries_map)
diff --git a/uncoder-core/app/translator/platforms/logrhythm_axon/renders/logrhythm_axon_query.py b/uncoder-core/app/translator/platforms/logrhythm_axon/renders/logrhythm_axon_query.py
@@ -19,6 +19,7 @@
 from typing import 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.custom_types.tokens import LogicalOperatorType
 from app.translator.core.custom_types.values import ValueType
 from app.translator.core.exceptions.core import StrictPlatformException
@@ -262,6 +263,8 @@ def _generate_from_tokenized_query_container(self, query_container: TokenizedQue
                 meta_info=query_container.meta_info,
                 source_mapping=source_mapping,
             )
+            if return_only_first_query_ctx_var.get() is True:
+                return finalized_query
             queries_map[source_mapping.source_id] = finalized_query
 
         return self.finalize(queries_map)
