Minor changes to the model for streaming

Fixing it in its final "way".

Author: gordonwatts

.vscode/settings.json

@@ -1,7 +1,5 @@
 {
     "python.pythonPath": ".venv\\Scripts\\python.exe",
-    "pyright.openFilesOnly": false,
-    "pyright.useLibraryCodeForTypes": true,
     "python.analysis.logLevel": "Information",
     "python.analysis.memory.keepLibraryAst": true,
     "python.linting.flake8Enabled": true,
@@ -82,6 +80,7 @@
     ],
     "python.analysis.typeCheckingMode": "basic",
     "python.testing.pytestArgs": [
+        "tests",
         "--no-cov"
-    ]
+    ],
 }

servicex/__init__.py

@@ -1,4 +1,4 @@
-from .servicex import ServiceXDataset  # NOQA
+from .servicex import ServiceXDataset, StreamInfoUrl  # NOQA
 from .utils import (  # NOQA
     ServiceXException,
     ServiceXUnknownRequestID,

servicex/servicex.py

@@ -3,6 +3,7 @@
 import functools
 import logging
 import time
+from dataclasses import dataclass
 from datetime import timedelta
 from pathlib import Path
 from typing import (Any, AsyncIterator, Awaitable, Callable, Dict, List,
@@ -30,6 +31,15 @@
                     stream_status_updates, stream_unique_updates_only)
 
 
+@dataclass
+class StreamInfoUrl:
+    '''Contains information on accessing ServiceX data via a url
+    '''
+    url: str
+    file: str
+    bucket: str
+
+
 class ServiceXDataset(ServiceXABC):
     '''
     Used to access an instance of ServiceX at an end point on the internet. Support convieration
@@ -166,19 +176,20 @@ 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_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.
+    async def get_data_rootfiles_url_stream(self, selection_query: str) \
+            -> AsyncIterator[StreamInfoUrl]:
+        '''Returns, as an async iterator, each completed batch of work from ServiceX.
+        The data that comes back includes a `url` that can be accessed to download the
+        data.
 
         Args:
             selection_query (str): The ServiceX Selection
         '''
         async for f_info in self._get_minio_buckets(selection_query, 'root-files'):
             yield f_info
 
-    async def get_data_parquet_minio_async(self, selection_query: str) \
-            -> AsyncIterator[Dict[str, str]]:
+    async def get_data_parquet_minio_stream(self, selection_query: str) \
+            -> AsyncIterator[StreamInfoUrl]:
         '''Returns, as an async iterator, each of the files from the minio bucket,
         as the files are added there.
 
@@ -214,7 +225,7 @@ async def convert_to_file(f: Path) -> Path:
 
     @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]]:
+            -> AsyncIterator[StreamInfoUrl]:
         '''Get a list of files back for a request
 
         Args:
@@ -247,11 +258,7 @@ async def _get_minio_buckets(self, selection_query: str, data_format: str) \
 
                 # Reflect the files back up a level.
                 async for r in minio_files:
-                    yield {
-                        'bucket': request_id,
-                        'file': r,
-                        'url': minio_adaptor.get_access_url(request_id, r),
-                    }
+                    yield StreamInfoUrl(minio_adaptor.get_access_url(request_id, r), r, request_id)
 
                 # Cache the final status
                 await self._update_query_status(client, request_id)

tests/test_servicex.py

@@ -1,6 +1,7 @@
 import asyncio
 from contextlib import contextmanager
 from pathlib import Path
+from servicex.servicex import StreamInfoUrl
 from servicex.cache import Cache
 
 from confuse.core import Configuration
@@ -405,14 +406,14 @@ async def test_good_run_root_files_from_minio(mocker):
                             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)
+    lst = [f async for f in ds.get_data_rootfiles_url_stream('(valid qastle string)')]
 
     assert len(lst) == 1
-    assert lst[0]['bucket'] == '123-456'
-    assert lst[0]['file'] == 'one_minio_entry'
-    assert lst[0]['url'] == 'http://the.url.com'
+    r = lst[0]
+    assert isinstance(r, StreamInfoUrl)
+    assert r.bucket == '123-456'
+    assert r.file == 'one_minio_entry'
+    assert r.url == 'http://the.url.com'
 
     assert mock_servicex_adaptor.query_json['result-format'] == 'root-file'
     assert mock_minio_adaptor.access_called_with == ('123-456', 'one_minio_entry')
@@ -438,7 +439,7 @@ async def test_bad_request_id_run_root_files_from_minio(mocker):
 
     with pytest.raises(ServiceXException) as e:
         lst = []
-        async for f_info in ds.get_data_rootfiles_minio_async('(valid qastle string)'):
+        async for f_info in ds.get_data_rootfiles_url_stream('(valid qastle string)'):
             lst.append(f_info)
 
     assert 'to know about' in str(e.value)
@@ -481,7 +482,7 @@ async def test_bad_transform_run_root_files_from_minio(mocker):
 
     with pytest.raises(ServiceXFatalTransformException) as e:
         lst = []
-        async for f_info in ds.get_data_rootfiles_minio_async('(valid qastle string)'):
+        async for f_info in ds.get_data_rootfiles_url_stream('(valid qastle string)'):
             lst.append(f_info)
 
     assert 'Fatal Error' in str(e.value)
@@ -506,7 +507,7 @@ async def test_bad_file_transform_run_root_files_from_minio(mocker):
 
     with pytest.raises(ServiceXException) as e:
         lst = []
-        async for f_info in ds.get_data_rootfiles_minio_async('(valid qastle string)'):
+        async for f_info in ds.get_data_rootfiles_url_stream('(valid qastle string)'):
             lst.append(f_info)
 
     assert 'Failed to transform all files' in str(e.value)
@@ -528,12 +529,12 @@ async def test_good_run_parquet_files_from_minio(mocker):
                             local_log=mock_logger,
                             data_convert_adaptor=data_adaptor)
     lst = []
-    async for f_info in ds.get_data_parquet_minio_async('(valid qastle string)'):
+    async for f_info in ds.get_data_parquet_minio_stream('(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 lst[0].bucket == '123-456'
+    assert lst[0].file == 'one_minio_entry'
 
     assert mock_servicex_adaptor.query_json['result-format'] == 'parquet'