add begin_at, update tests

Author: MattShirley

servicex/query_core.py

@@ -27,6 +27,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 from __future__ import annotations
 
+import datetime
 import abc
 import asyncio
 from abc import ABC
@@ -318,6 +319,7 @@ def transform_complete(task: Task):
             else None
         )
 
+        begin_at = datetime.datetime.now(datetime.UTC)
         if not cached_record:
 
             if self.cache.is_transform_request_submitted(sx_request_hash):
@@ -342,13 +344,14 @@ def transform_complete(task: Task):
 
         download_files_task = loop.create_task(
             self.download_files(
-                signed_urls_only, expandable_progress, download_progress, cached_record
+                signed_urls_only, expandable_progress, download_progress, cached_record, begin_at
             )
         )
 
         try:
             signed_urls = []
             downloaded_files = []
+
             download_result = await download_files_task
             if signed_urls_only:
                 signed_urls = download_result
@@ -517,6 +520,7 @@ async def download_files(
         progress: ExpandableProgress,
         download_progress: TaskID,
         cached_record: Optional[TransformedResults],
+        begin_at: datetime.datetime,
     ) -> List[str]:
         """
         Task to monitor the list of files in the transform output's bucket. Any new files
@@ -558,20 +562,25 @@ async def get_signed_url(
             if self.minio:
                 # if self.minio exists, self.current_status will too
                 if self.current_status.files_completed > len(files_seen):
+                    new_begin_at = datetime.datetime.now(datetime.UTC)
                     files = await self.servicex.get_transformation_results(
-                        self.current_status.request_id
+                        self.current_status.request_id,
+                        begin_at
                     )
+                    begin_at = new_begin_at
 
                     for file in files:
                         if "file-path" not in file:
                             continue
 
                         file_path = file["file-path"].replace("/", ":")
+
                         if file_path not in files_seen:
                             if signed_urls_only:
                                 download_tasks.append(
                                     loop.create_task(
                                         get_signed_url(
+                                            self.minio,
                                             file_path,
                                             progress,
                                             download_progress,

servicex/servicex_adapter.py

@@ -27,6 +27,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 import os
 import time
+import datetime
 from typing import Optional, Dict, List
 
 from aiohttp import ClientSession
@@ -228,12 +229,16 @@ async def delete_transform(self, transform_id=None):
                         f"Failed to delete transform {transform_id} - {msg}"
                     )
 
-    async def get_transformation_results(self, request_id: str):
+    async def get_transformation_results(self, request_id: str, begin_at: datetime.datetime):
         headers = await self._get_authorization()
-        url = self.url + f"/servicex/internal/transformation/{request_id}/results"
+        url = self.url + f"/servicex/transformation/{request_id}/results"
+
+        params = {}
+        if begin_at:
+            params["begin_at"] = begin_at.isoformat()
 
         async with ClientSession() as session:
-            async with session.get(headers=headers, url=url) as r:
+            async with session.get(headers=headers, url=url, params=params) as r:
                 if r.status == 403:
                     raise AuthorizationError(
                         f"Not authorized to access serviceX at {self.url}"
@@ -246,7 +251,8 @@ async def get_transformation_results(self, request_id: str):
                     msg = await _extract_message(r)
                     raise RuntimeError(f"Failed with message: {msg}")
 
-                return (await r.json())["results"]
+                data = await r.json()
+                return data.get("results")
 
     async def cancel_transform(self, transform_id=None):
         headers = await self._get_authorization()

tests/test_servicex_dataset.py

@@ -204,8 +204,18 @@ def cache_transform(record: TransformedResults):
 @pytest.mark.asyncio
 async def test_submit(mocker):
     servicex = AsyncMock()
+
     servicex.submit_transform = AsyncMock()
     servicex.submit_transform.return_value = {"request_id": '123-456-789"'}
+
+    servicex.get_transformation_results = AsyncMock(side_effect=[
+        [{"file-path": file1.filename}],
+        [
+         {"file-path": file1.filename},
+         {"file-path": file2.filename},
+        ],
+    ])
+
     servicex.get_transform_status = AsyncMock()
     servicex.get_transform_status.side_effect = [
         transform_status1,
@@ -214,7 +224,6 @@ async def test_submit(mocker):
     ]
 
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1, file2]])
     mock_minio.download_file = AsyncMock(
         side_effect=lambda a, _, shorten_filename: PurePath(a)
     )
@@ -235,6 +244,7 @@ async def test_submit(mocker):
         config=Configuration(api_endpoints=[]),
     )
     datasource.query_string_generator = FuncADLQuery_Uproot().FromTree("nominal")
+
     with ExpandableProgress(display_progress=False) as progress:
         datasource.result_format = ResultFormat.parquet
         result = await datasource.submit_and_download(
@@ -245,11 +255,18 @@ async def test_submit(mocker):
     mock_cache.cache_transform.assert_called_once()
 
 
+
 @pytest.mark.asyncio
 async def test_submit_partial_success(mocker):
     servicex = AsyncMock()
     servicex.submit_transform = AsyncMock()
     servicex.submit_transform.return_value = {"request_id": '123-456-789"'}
+
+    servicex.get_transformation_results = AsyncMock(side_effect=[
+        [{"file-path": file1.filename}],
+        [{"file-path": file1.filename}],
+    ])
+
     servicex.get_transform_status = AsyncMock()
     servicex.get_transform_status.side_effect = [
         transform_status1,
@@ -258,7 +275,6 @@ async def test_submit_partial_success(mocker):
     ]
 
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1]])
     mock_minio.download_file = AsyncMock(
         side_effect=lambda a, _, shorten_filename: PurePath(a)
     )
@@ -295,13 +311,16 @@ async def test_use_of_cache(mocker):
     servicex = AsyncMock()
     servicex.submit_transform = AsyncMock()
     servicex.submit_transform.return_value = {"request_id": '123-456-789"'}
+    servicex.get_transformation_results = AsyncMock(return_value=[
+        {"file-path": file1.filename},
+        {"file-path": file2.filename}
+    ])
     servicex.get_transform_status = AsyncMock()
     servicex.get_transform_status.side_effect = [
         transform_status1,
         transform_status3,
     ]
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(return_value=[file1, file2])
     mock_minio.download_file = AsyncMock(
         side_effect=lambda a, _, shorten_filename: PurePath(a)
     )
@@ -336,7 +355,7 @@ async def test_use_of_cache(mocker):
         # second round, should hit the cache (and not call the sx_adapter, minio, or update_record)
         with ExpandableProgress(display_progress=False) as progress:
             servicex2 = AsyncMock()
-            mock_minio.list_bucket.reset_mock()
+            servicex.get_transformation_results.reset_mock()
             mock_minio.get_signed_url.reset_mock()
             datasource2 = Query(
                 dataset_identifier=did,
@@ -354,14 +373,14 @@ async def test_use_of_cache(mocker):
                 signed_urls_only=True, expandable_progress=progress
             )
             servicex2.assert_not_awaited()
-            mock_minio.list_bucket.assert_not_awaited()
+            servicex.get_transformation_results.assert_not_awaited()
             mock_minio.get_signed_url.assert_not_awaited()
         upd.assert_not_called()
         assert result1 == result2
         upd.reset_mock()
         servicex.get_transform_status.reset_mock(side_effect=True)
         servicex.get_transform_status.return_value = transform_status3
-        mock_minio.list_bucket.reset_mock(side_effect=True)
+        servicex.get_transformation_results.reset_mock(side_effect=True)
         # third round, should hit the cache and download files (and call update_record)
         with ExpandableProgress(display_progress=False) as progress:
             await datasource.submit_and_download(
@@ -371,14 +390,14 @@ async def test_use_of_cache(mocker):
         assert mock_minio.download_file.await_count == 2
         upd.assert_called_once()
         # fourth round, should hit the cache (and nothing else)
-        mock_minio.list_bucket.reset_mock()
+        servicex.get_transformation_results.reset_mock()
         mock_minio.download_file.reset_mock()
         with ExpandableProgress(display_progress=False) as progress:
             await datasource.submit_and_download(
                 signed_urls_only=False, expandable_progress=progress
             )
         servicex.assert_not_awaited()
-        mock_minio.list_bucket.assert_not_awaited()
+        servicex.get_transformation_results.assert_not_awaited()
         mock_minio.download_file.assert_not_awaited()
         upd.assert_called_once()
         cache.close()
@@ -396,7 +415,6 @@ async def test_submit_cancel(mocker):
     ]
 
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1]])
     mock_minio.download_file = AsyncMock(
         side_effect=lambda a, _, shorten_filename: PurePath(a)
     )
@@ -437,7 +455,6 @@ async def test_submit_fatal(mocker):
     ]
 
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1]])
     mock_minio.download_file = AsyncMock(
         side_effect=lambda a, _, shorten_filename: PurePath(a)
     )
@@ -482,7 +499,6 @@ async def test_submit_generic(mocker, codegen_list):
     ]
 
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1, file2]])
     mock_minio.download_file = AsyncMock()
 
     mock_cache = mocker.MagicMock(QueryCache)
@@ -531,7 +547,6 @@ async def test_submit_cancelled(mocker, codegen_list):
     sx.get_transform_status.side_effect = [transform_status4]
 
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1, file2]])
     mock_minio.download_file = AsyncMock()
 
     mock_cache = mocker.MagicMock(QueryCache)
@@ -601,10 +616,12 @@ async def test_use_of_ignore_cache(mocker, servicex):
             transform_status3,
         ]
     )
-
+    servicex.get_transformation_results = AsyncMock(return_value=[
+        {"file-path": file1.filename},
+        {"file-path": file2.filename}
+    ])
     # Prepare Minio
     mock_minio = AsyncMock()
-    mock_minio.list_bucket = AsyncMock(return_value=[file1, file2])
     mock_minio.get_signed_url = AsyncMock(side_effect=["http://file1", "http://file2"])
     mocker.patch("servicex.minio_adapter.MinioAdapter", return_value=mock_minio)
     did = FileListDataset("/foo/bar/baz.root")
@@ -674,13 +691,13 @@ async def test_use_of_ignore_cache(mocker, servicex):
             transform_status1,
             transform_status3,
         ]
-        mock_minio.list_bucket.reset_mock()
+        servicex.get_transformation_results.reset_mock()
         mock_minio.download_file.reset_mock()
         with ExpandableProgress(display_progress=False) as progress:
             res = await datasource_without_ignore_cache.submit_and_download(
                 signed_urls_only=True, expandable_progress=progress
             )  # noqa
-        mock_minio.list_bucket.assert_not_awaited()
+        servicex.get_transformation_results.assert_not_awaited()
         mock_minio.download_file.assert_not_awaited()
         assert len(res.signed_url_list) == 2
         cache.close()