First go at the callback to fetch minio async

Author: gordonwatts

servicex/minio_adaptor.py

@@ -143,7 +143,7 @@ def from_best(self, transaction_info: Optional[Dict[str, str]] = None) -> MinioA
             if all(k in transaction_info for k in keys):
                 logging.getLogger(__name__).debug('Using the request-specific minio_adaptor')
                 return MinioAdaptor(transaction_info['minio-endpoint'],
-                                    transaction_info['minio-secured'],
+                                    bool(transaction_info['minio-secured']),
                                     transaction_info['minio-access-key'],
                                     transaction_info['minio-secret-key'])
         if self._config_adaptor is not None:

servicex/servicex.py

@@ -11,6 +11,7 @@
 import aiohttp
 import backoff
 from backoff import on_exception
+from numpy.lib.function_base import select
 
 from servicex.servicex_config import ServiceXConfigAdaptor
 
@@ -166,13 +167,27 @@ async def get_data_awkward_async(self, selection_query: str):
         return self._converter.combine_awkward(await self._data_return(
             selection_query, lambda f: self._converter.convert_to_awkward(f)))
 
-    # async def get_data_minio_async(self, selection_query: str):
-    #     '''Async iterator return of minio buckets/files for a query. The results are made
-    #     available as soon as they are returned from ServiceX.
+    async def get_data_rootfiles_minio_async(self, selection_query: str) \
+            -> AsyncIterator[Dict[str, str]]:
+        '''Returns, as an async iterator, each of the files from the minio bucket,
+        as the files are added there.
 
-    #     Args:
-    #         selection_query (str): The query string
-    #     '''
+        Args:
+            selection_query (str): The ServiceX Selection
+        '''
+        async for f_info in self._get_minio_buckets(select, 'root-files'):
+            yield f_info
+
+    async def get_data_parquet_minio_async(self, selection_query: str) \
+            -> AsyncIterator[Dict[str, str]]:
+        '''Returns, as an async iterator, each of the files from the minio bucket,
+        as the files are added there.
+
+        Args:
+            selection_query (str): The ServiceX Selection
+        '''
+        async for f_info in self._get_minio_buckets(select, 'parquet'):
+            yield f_info
 
     async def _file_return(self, selection_query: str, data_format: str):
         '''
@@ -198,6 +213,66 @@ async def convert_to_file(f: Path) -> Path:
 
         return await self._data_return(selection_query, convert_to_file, data_format)
 
+    @on_exception(backoff.constant, ServiceXUnknownRequestID, interval=0.1, max_tries=3)
+    async def _get_minio_buckets(self, selection_query: str, data_format: str) \
+            -> AsyncIterator[Dict[str, str]]:
+        '''Get a list of files back for a request
+
+        Args:
+            selection_query (str): The selection query we are to do
+            data_format (str): The requested file format
+
+        Yields:
+            AsyncIterator[Dict[str, str]]: A tuple of the minio bucket and file in that bucket.
+                                           The dict will have entries for:
+                                             bucket: The minio bucket name
+                                             file: the completed file in the bucket
+        '''
+        query = self._build_json_query(selection_query, data_format)
+
+        async with aiohttp.ClientSession() as client:
+
+            # Get a request id - which might be cached, but if not, submit it.
+            request_id = await self._get_request_id(client, query)
+
+            # Get the minio adaptor we are going to use for downloading.
+            minio_adaptor = self._minio_adaptor \
+                .from_best(self._cache.lookup_query_status(request_id))
+
+            # Look up the cache, and then fetch an iterator going thorugh the results
+            # from either servicex or the cache, depending.
+            try:
+                notifier = self._create_notifier()
+                minio_files = self._get_minio_bucket_files_from_servicex(request_id, client,
+                                                                         minio_adaptor, notifier)
+
+                # Reflect the files back up a level.
+                async for r in minio_files:
+                    yield {
+                        'bucket': request_id,
+                        'file': r,
+                    }
+
+                # Cache the final status
+                await self._update_query_status(client, request_id)
+
+            except ServiceXUnknownRequestID as e:
+                self._cache.remove_query(query)
+                raise ServiceXException('Expected the ServiceX backend to know about query '
+                                        f'{request_id}. It did not. Cleared local cache. '
+                                        'Please resubmit to trigger a new query.') from e
+
+            except ServiceXFatalTransformException as e:
+                transform_status = await self._servicex_adaptor.get_query_status(client,
+                                                                                 request_id)
+                raise ServiceXFatalTransformException(
+                    f'ServiceX Fatal Error: {transform_status["failure-info"]}') from e
+
+            except ServiceXFailedFileTransform as e:
+                self._cache.remove_query(query)
+                await self._servicex_adaptor.dump_query_errors(client, request_id)
+                raise ServiceXException(f'Failed to transform all files in {request_id}') from e
+
     @on_exception(backoff.constant, ServiceXUnknownRequestID, interval=0.1, max_tries=3)
     async def _data_return(self, selection_query: str,
                            converter: Callable[[Path], Awaitable[Any]],
@@ -415,7 +490,8 @@ async def _get_files_from_servicex(self, request_id: str,
     async def _get_minio_bucket_files_from_servicex(self, request_id: str,
                                                     client: aiohttp.ClientSession,
                                                     minio_adaptor: MinioAdaptor,
-                                                    notifier: _status_update_wrapper):
+                                                    notifier: _status_update_wrapper) \
+            -> AsyncIterator[str]:
         '''Create an async stream of `minio` bucket/filenames from a request id.
 
         Args:

tests/test_servicex.py

@@ -118,7 +118,7 @@ def do_context():
 
 
 @pytest.mark.asyncio
-async def test_good_run_root_files(mocker):
+async def test_minio_back(mocker):
     'Get a root file with a single file'
     mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
     mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
@@ -390,6 +390,152 @@ async def test_good_run_single_ds_2file_awkward(mocker, good_awkward_file_data):
     assert len(good_awkward_file_data.combine_awkward.call_args[0][0]) == 2
 
 
+@pytest.mark.asyncio
+async def test_good_run_root_files_from_minio(mocker):
+    'Get a root file pulling back minio info as it arrives'
+    mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    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,
+                            local_log=mock_logger,
+                            data_convert_adaptor=data_adaptor)
+    lst = []
+    async for f_info in ds.get_data_rootfiles_minio_async('(valid qastle string)'):
+        lst.append(f_info)
+
+    assert len(lst) == 1
+    assert lst[0]['bucket'] == '123-456'
+    assert lst[0]['file'] == 'one_minio_entry'
+
+    assert mock_servicex_adaptor.query_json['result-format'] == 'root-file'
+
+
+@pytest.mark.asyncio
+async def test_bad_request_id_run_root_files_from_minio(mocker):
+    'Using the minio interface - the request_id is not known'
+    mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
+    transform_status = mocker.MagicMock(side_effect=ServiceXUnknownRequestID('boom'))
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456",
+                                                mock_transform_status=transform_status)
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    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,
+                            local_log=mock_logger,
+                            data_convert_adaptor=data_adaptor)
+
+    with pytest.raises(ServiceXException) as e:
+        lst = []
+        async for f_info in ds.get_data_rootfiles_minio_async('(valid qastle string)'):
+            lst.append(f_info)
+
+    assert 'to know about' in str(e.value)
+
+
+@pytest.mark.asyncio
+async def test_bad_transform_run_root_files_from_minio(mocker):
+    'Using the async minio interface - fail to transform (like bad DID)'
+    mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
+    fatal_transform_status = {
+        "request_id": "24e59fa2-e1d7-4831-8c7e-82b2efc7c658",
+        "did": "mc15_13TeV:mc15_13TeV.361106.PowhegPythia8EvtGen_AZNLOCTEQ6L1_Zee.merge.DAOD_STDM3.e3601_s2576_s2132_r6630_r6264_p2363_tid05630052_0000",  # NOQA
+        "columns": "Electrons.pt(), Electrons.eta(), Electrons.phi(), Electrons.e(), Muons.pt(), Muons.eta(), Muons.phi(), Muons.e()",  # NOQA
+        "selection": None,
+        "tree-name": None,
+        "image": "sslhep/servicex_func_adl_xaod_transformer:130_reset_cwd",
+        "chunk-size": 7000,
+        "workers": 1,
+        "result-destination": "kafka",
+        "result-format": "arrow",
+        "kafka-broker": "servicex-kafka-1.slateci.net:19092",
+        "workflow-name": "straight_transform",
+        "generated-code-cm": None,
+        "status": "Fatal",
+        "failure-info": "DID Not found mc15_13TeV:mc15_13TeV.361106.PowhegPythia8EvtGen_AZNLOCTEQ6L1_Zee.merge.DAOD_STDM3.e3601_s2576_s2132_r6630_r6264_p2363_tid05630052_0000"  # NOQA
+    }
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456",
+        mock_transform_status=mocker.MagicMock(side_effect=ServiceXFatalTransformException('DID was BAD')),  # NOQA
+        mock_transform_query_status=mocker.MagicMock(return_value=fatal_transform_status))
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    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,
+                            local_log=mock_logger,
+                            data_convert_adaptor=data_adaptor)
+
+    with pytest.raises(ServiceXFatalTransformException) as e:
+        lst = []
+        async for f_info in ds.get_data_rootfiles_minio_async('(valid qastle string)'):
+            lst.append(f_info)
+
+    assert 'Fatal Error' in str(e.value)
+
+
+@pytest.mark.asyncio
+async def test_bad_file_transform_run_root_files_from_minio(mocker):
+    'Using the async minio interface, some files will fail to translate.'
+    mock_cache = build_cache_mock(mocker)
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    mock_transform_status = mocker.Mock(return_value=(0, 1, 1))
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456", mock_transform_status)
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry', 'two_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,
+                            local_log=mock_logger,
+                            data_convert_adaptor=data_adaptor)
+
+    with pytest.raises(ServiceXException) as e:
+        lst = []
+        async for f_info in ds.get_data_rootfiles_minio_async('(valid qastle string)'):
+            lst.append(f_info)
+
+    assert 'Failed to transform all files' in str(e.value)
+
+
+@pytest.mark.asyncio
+async def test_good_run_parquet_files_from_minio(mocker):
+    'Get a parquet file pulling back minio info as it arrives'
+    mock_cache = build_cache_mock(mocker, data_file_return="/foo/bar.root")
+    mock_servicex_adaptor = MockServiceXAdaptor(mocker, "123-456")
+    mock_minio_adaptor = MockMinioAdaptor(mocker, files=['one_minio_entry'])
+    mock_logger = mocker.MagicMock(spec=log_adaptor)
+    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,
+                            local_log=mock_logger,
+                            data_convert_adaptor=data_adaptor)
+    lst = []
+    async for f_info in ds.get_data_parquet_minio_async('(valid qastle string)'):
+        lst.append(f_info)
+
+    assert len(lst) == 1
+    assert lst[0]['bucket'] == '123-456'
+    assert lst[0]['file'] == 'one_minio_entry'
+
+    assert mock_servicex_adaptor.query_json['result-format'] == 'parquet'
+
+
 @pytest.mark.asyncio
 async def test_status_exception(mocker):
     'Make sure status error - like transform not found - is reported all the way to the top'