Fixed resolution of MQM factories (#117)

* Fixed resolution of MQM factories * Style fix
datalens-tech · Nov 23, 2023 · 9a97256 · 9a97256
1 parent efac0a9
commit 9a97256
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 19 deletions.
diff --git a/lib/dl_api_lib/dl_api_lib/query/registry.py b/lib/dl_api_lib/dl_api_lib/query/registry.py
@@ -1,6 +1,7 @@
 from typing import (
     Collection,
     Optional,
+    Sequence,
     Type,
 )
 
@@ -87,25 +88,42 @@ class MQMFactorySettingItem:
 _MQM_FACTORY_REGISTRY: dict[MQMFactoryKey, Type[MultiQueryMutatorFactoryBase]] = {}
 
 
+def _get_default_mqm_factory_cls() -> Type[MultiQueryMutatorFactoryBase]:
+    return DefaultMultiQueryMutatorFactory
+
+
 def get_multi_query_mutator_factory(
     query_proc_mode: QueryProcessingMode,
     backend_type: SourceBackendType,
     dialect: DialectCombo,
     result_schema: ResultSchema,
-) -> Optional[MultiQueryMutatorFactoryBase]:
-    factory_cls = _MQM_FACTORY_REGISTRY.get(
-        # First try with exact dialect
+) -> MultiQueryMutatorFactoryBase:
+    prioritized_keys = (
+        # First try with exact dialect and mode (exact match)
         MQMFactoryKey(query_proc_mode=query_proc_mode, backend_type=backend_type, dialect=dialect),
-        _MQM_FACTORY_REGISTRY.get(
-            # Then try without the dialect, just the backend
-            MQMFactoryKey(query_proc_mode=query_proc_mode, backend_type=backend_type, dialect=None),
-            DefaultMultiQueryMutatorFactory,  # If still nothing, then use the default
-        ),
+        # Now the fallbacks begin...
+        # Try without the dialect (all dialects within backend), just the backend and mode
+        MQMFactoryKey(query_proc_mode=query_proc_mode, backend_type=backend_type, dialect=None),
+        # Fall back to `basic` mode (but still within the backend type)
+        # First try with the specific dialect
+        MQMFactoryKey(query_proc_mode=QueryProcessingMode.basic, backend_type=backend_type, dialect=dialect),
+        # If still nothing, try without specifying the dialect (all dialects within backend)
+        MQMFactoryKey(query_proc_mode=QueryProcessingMode.basic, backend_type=backend_type, dialect=None),
     )
 
+    # Now iterate over all of these combinations IN THAT VERY ORDER(!)
+    factory_cls: Optional[Type[MultiQueryMutatorFactoryBase]] = None
+    for key in prioritized_keys:
+        factory_cls = _MQM_FACTORY_REGISTRY.get(key)
+        if factory_cls is not None:
+            break  # found something
+
     if factory_cls is None:
-        return None
+        # Not found for any of the combinations
+        # Use the ultimate default
+        factory_cls = _get_default_mqm_factory_cls()
 
+    assert factory_cls is not None
     return factory_cls(result_schema=result_schema)
 
 

diff --git a/lib/dl_api_lib/dl_api_lib/service_registry/multi_query_mutator_factory.py b/lib/dl_api_lib/dl_api_lib/service_registry/multi_query_mutator_factory.py
@@ -1,4 +1,5 @@
 import abc
+import logging
 from typing import Sequence
 
 import attr
@@ -14,6 +15,9 @@
 from dl_query_processing.multi_query.mutators.base import MultiQueryMutatorBase
 
 
+LOGGER = logging.getLogger(__name__)
+
+
 class SRMultiQueryMutatorFactory(abc.ABC):
     @abc.abstractmethod
     def get_mqm_factory(
@@ -51,14 +55,10 @@ def get_mqm_factory(
             dialect=dialect,
             result_schema=dataset.result_schema,
         )
-        if factory is None:
-            # Try again for the basic mode
-            factory = get_multi_query_mutator_factory(
-                query_proc_mode=QueryProcessingMode.basic,
-                backend_type=backend_type,
-                dialect=dialect,
-                result_schema=dataset.result_schema,
-            )
-
-        assert factory is not None
+        LOGGER.info(
+            f"Resolved MQM factory for backend_type {backend_type.name} "
+            f"and dialect {dialect.common_name_and_version} "
+            f"in {self._query_proc_mode.name} mode "
+            f"to {type(factory).__name__}"
+        )
         return factory
diff --git a/lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/connector.py b/lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/connector.py
@@ -11,6 +11,7 @@
 )
 from dl_api_lib.query.registry import MQMFactorySettingItem
 from dl_constants.enums import QueryProcessingMode
+from dl_query_processing.multi_query.factory import NoCompengMultiQueryMutatorFactory
 
 from dl_connector_bitrix_gds.api.api_schema.connection import BitrixGDSConnectionSchema
 from dl_connector_bitrix_gds.api.connection_form.form_config import BitrixGDSConnectionFormFactory
@@ -47,6 +48,10 @@ class BitrixGDSApiConnector(ApiConnector):
             query_proc_mode=QueryProcessingMode.basic,
             factory_cls=BitrixGDSMultiQueryMutatorFactory,
         ),
+        MQMFactorySettingItem(
+            query_proc_mode=QueryProcessingMode.no_compeng,
+            factory_cls=NoCompengMultiQueryMutatorFactory,
+        ),
     )
     connection_definitions = (BitrixGDSApiConnectionDefinition,)
     source_definitions = (BitrixGDSApiSourceDefinition,)