Backend type argument

Author: gordonwatts

servicex/servicex.py

@@ -175,6 +175,7 @@ def __init__(
         self,
         dataset: DatasetType,
         backend_name: Optional[str] = None,
+        backend_type: Optional[str] = None,
         codegen: Optional[str] = None,
         image: Optional[str] = None,
         max_workers: int = 20,
@@ -202,8 +203,13 @@ def __init__(
                                         will default to xaod, unless you have any endpoint listed
                                         in your servicex file. It will default to best match there,
                                         or fail if a name has been given.
+            backend_type                The type of backend. Overrides the `type` in the `yaml`
+                                        config file.
             codegen                     The type of code generator passed to the backend to
                                         generate the code that powers the requested transform.
+                                        Don't use unless you know what you are doing - sometimes
+                                        the return filetype is also needed. Better to use
+                                        `backend_type`!
             image                       Name of transformer image to use to transform the data. If
                                         left as default, `None`, then the default image for the
                                         ServiceX backend will be used.
@@ -276,7 +282,9 @@ def __init__(
 
         if not servicex_adaptor:
             # Given servicex adaptor is none, this should be ok. Fixes type checkers
-            end_point, token = config.get_servicex_adaptor_config(backend_name)
+            end_point, token = config.get_servicex_adaptor_config(
+                backend_name, backend_type=backend_type
+            )
             servicex_adaptor = ServiceXAdaptor(end_point, token)
         self._servicex_adaptor = servicex_adaptor
 
@@ -296,7 +304,11 @@ def __init__(
             else default_client_session
         )
 
-        self._return_types = [config.get_default_returned_datatype(backend_name)]
+        self._return_types = [
+            config.get_default_returned_datatype(
+                backend_name, backend_type=backend_type
+            )
+        ]
         self._converter = (
             data_convert_adaptor
             if data_convert_adaptor is not None
@@ -306,7 +318,9 @@ def __init__(
         # TODO - this should not be in the ABC backend since we have to have some intelligence
         # in setting it, it turns out.
         if self._codegen is None:
-            self._codegen = config.get_backend_info(backend_name, "codegen")
+            self._codegen = config.get_backend_info(
+                backend_name, "codegen", backend_type=backend_type
+            )
 
     def first_supported_datatype(
         self, datatypes: Union[List[str], str]

servicex/servicex_config.py

@@ -37,7 +37,9 @@ def settings(self) -> ConfigView:
         """
         return self._settings
 
-    def get_default_returned_datatype(self, backend_name: Optional[str]) -> str:
+    def get_default_returned_datatype(
+        self, backend_name: Optional[str], backend_type: Optional[str] = None
+    ) -> str:
         """Return the default return data type, given the backend is a certain type.
 
         Args:
@@ -47,7 +49,9 @@ def get_default_returned_datatype(self, backend_name: Optional[str]) -> str:
             str: The backend datatype, like `root` or `parquet`.
         """
         # Check to see if we know about the backend info
-        r = self.get_backend_info(backend_name, "return_data")
+        r = self.get_backend_info(
+            backend_name, "return_data", backend_type=backend_type
+        )
 
         if r is None:
             raise ServiceXException(
@@ -57,22 +61,27 @@ def get_default_returned_datatype(self, backend_name: Optional[str]) -> str:
 
         return r
 
-    def get_backend_info(self, backend_name: Optional[str], key: str) -> Optional[str]:
+    def get_backend_info(
+        self, backend_name: Optional[str], key: str, backend_type: Optional[str] = None
+    ) -> Optional[str]:
         """Find an item in the backend info, searching first for the backend
         name/type and then the defaults with a given type.
 
         Args:
             backend_name (str): Backend name
             key (str): The key for the info we are after
+            backend_type (str): The backend type, if we are overriding normal lookup
 
         Returns:
             Optional[str]: Return a string for the info we are after, or return None if we can't
                            find it.
         """
-        config = self._get_backend_info(backend_name)
+        config = self._get_backend_info(backend_name, backend_type)
         return config[key] if key in config else None
 
-    def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
+    def _get_backend_info(
+        self, backend_name: Optional[str], backend_type: Optional[str]
+    ) -> Dict[str, str]:
         """Returns all the info for a backend name/type.
 
         Search algorithm is non-trivial:
@@ -89,6 +98,8 @@ def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
 
         Args:
             backend_name (str): Name or type of the api end point we are going to look up.
+            backend_type (str): Override the backend type, if we are looking for a specific
+                                type behavior.
 
         Returns:
             Dict[str, str]: Attributes for this backend's configuration
@@ -159,8 +170,14 @@ def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
             )
 
         # Now, extract the type and see if we can figure out any defaults from the
-        # `backend_types` info.
-        type_lookup = config["type"] if "type" in config else backend_name
+        # `backend_types` info. Skip this if we have a type we are passed in.
+        type_lookup = (
+            backend_type
+            if backend_type is not None
+            else config["type"]
+            if "type" in config
+            else backend_name
+        )
         if type_lookup is None:
             return config
 
@@ -181,7 +198,9 @@ def _get_backend_info(self, backend_name: Optional[str]) -> Dict[str, str]:
         return config
 
     def get_servicex_adaptor_config(
-        self, backend_name: Optional[str] = None
+        self,
+        backend_name: Optional[str] = None,
+        backend_type: Optional[str] = None,
     ) -> Tuple[str, Optional[str]]:
         """Return the servicex (endpoint, token) from a given backend configuration.
 
@@ -193,7 +212,7 @@ def get_servicex_adaptor_config(
             Tuple[str, str]: The tuple of info to create a `ServiceXAdaptor`: end point,
             token (optionally).
         """
-        config = self._get_backend_info(backend_name)
+        config = self._get_backend_info(backend_name, backend_type)
 
         endpoint = config["endpoint"]
         token = config["token"] if "token" in config else None

tests/test_servicex.py

@@ -1324,6 +1324,30 @@ async def test_codegen_override(mocker, good_awkward_file_data):
     assert called["codegen"] == "good_codegen"
 
 
+@pytest.mark.asyncio
+async def test_codegen_backend_type(mocker, good_awkward_file_data):
+    mock_cache = build_cache_mock(mocker)
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=["one_minio_entry"])
+    data_adaptor = mocker.MagicMock(spec=DataConverterAdaptor)
+
+    ds = fe.ServiceXDataset(
+        "localds://mc16_tev:13",
+        servicex_adaptor=mock_servicex_adaptor,  # type: ignore
+        minio_adaptor=mock_minio_adaptor,  # type: ignore
+        cache_adaptor=mock_cache,
+        data_convert_adaptor=data_adaptor,
+        local_log=mock_logger,
+        max_workers=50,
+        backend_type="uproot",
+    )
+    await ds.get_data_rootfiles_async("(valid qastle string)")
+
+    called = mock_servicex_adaptor.query_json
+    assert called["codegen"] == "uproot"
+
+
 @pytest.mark.asyncio
 async def test_codegen_default_by_backend(mocker, good_awkward_file_data):
     mock_cache = build_cache_mock(mocker)

tests/test_servicex_config.py

@@ -63,6 +63,36 @@ def test_returned_datatype_from_endpoint():
     assert x.get_default_returned_datatype("forkit") == "spoons"
 
 
+def test_backend_info():
+    c = ConfigSettings("servicex", "servicex")
+    c.clear()
+    c["backend_types"] = [{"type": "forkit", "silly": "spoon"}]
+    c["api_endpoints"] = [{"type": "forkit", "return_data": "spoons"}]
+
+    x = ServiceXConfigAdaptor(c)
+    assert x.get_backend_info("forkit", "silly") == "spoon"
+
+
+def test_backend_info_by_name():
+    c = ConfigSettings("servicex", "servicex")
+    c.clear()
+    c["backend_types"] = [{"type": "forkit", "silly": "spoon"}]
+    c["api_endpoints"] = [{"name": "dude", "type": "forkit", "return_data": "spoons"}]
+
+    x = ServiceXConfigAdaptor(c)
+    assert x.get_backend_info("dude", "silly") == "spoon"
+
+
+def test_backend_info_specified_type():
+    c = ConfigSettings("servicex", "servicex")
+    c.clear()
+    c["backend_types"] = [{"type": "forkit", "silly": "spoon"}]
+    c["api_endpoints"] = [{"name": "dude", "type": "leftfoot", "return_data": "spoons"}]
+
+    x = ServiceXConfigAdaptor(c)
+    assert x.get_backend_info("dude", "silly", backend_type="forkit") == "spoon"
+
+
 def test_default_config_has_default_return_datatype():
     "Test default settings - default_returned_datatype"
     c = ConfigSettings("servicex", "servicex")
@@ -80,6 +110,16 @@ def test_default_config_has_backend_types():
     assert count > 0
 
 
+def test_default_config_has_codegen():
+    c = ConfigSettings("servicex", "servicex")
+    assert c["backend_types"].exists()
+    count = 0
+    for info in c["backend_types"]:
+        count += 1
+        assert info["codegen"].exists()
+    assert count > 0
+
+
 def test_sx_adaptor_settings(caplog):
     from confuse import Configuration