remove s3 bucket polling when waiting for transformation results

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(tz=datetime.timezone.utc)
         if not cached_record:
 
             if self.cache.is_transform_request_submitted(sx_request_hash):
@@ -342,13 +344,18 @@ 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,11 +524,13 @@ 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
         will be downloaded.
         """
+
         files_seen = set()
         result_uris = []
         download_tasks = []
@@ -557,15 +566,22 @@ 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):
-                    files = await self.minio.list_bucket()
+                    new_begin_at = datetime.datetime.now(tz=datetime.timezone.utc)
+                    files = await self.servicex.get_transformation_results(
+                        self.current_status.request_id, begin_at
+                    )
+                    begin_at = new_begin_at
+
                     for file in files:
-                        if file.filename not in files_seen:
+                        file_path = file.get("file-path", "").replace("/", ":")
+
+                        if file_path != "" and file_path not in files_seen:
                             if signed_urls_only:
                                 download_tasks.append(
                                     loop.create_task(
                                         get_signed_url(
                                             self.minio,
-                                            file.filename,
+                                            file_path,
                                             progress,
                                             download_progress,
                                         )
@@ -576,14 +592,14 @@ async def get_signed_url(
                                     loop.create_task(
                                         download_file(
                                             self.minio,
-                                            file.filename,
+                                            file_path,
                                             progress,
                                             download_progress,
                                             shorten_filename=self.configuration.shortened_downloaded_filename,  # NOQA: E501
                                         )
                                     )
                                 )  # NOQA 501
-                            files_seen.add(file.filename)
+                            files_seen.add(file_path)
 
             # Once the transform is complete and all files are seen we can stop polling.
             # Also, if we are just downloading or signing urls for a previous transform

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,14 +229,39 @@ 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, begin_at: datetime.datetime
+    ):
+        headers = await self._get_authorization()
+        url = self.url + f"/servicex/transformation/{request_id}/results"
+        params = {
+            "begin_at": begin_at.isoformat(),
+        }
+
+        async with ClientSession() as session:
+            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}"
+                    )
+
+                if r.status == 404:
+                    raise ValueError(f"Request {request_id} not found")
+
+                if r.status != 200:
+                    msg = await _extract_message(r)
+                    raise RuntimeError(f"Failed with message: {msg}")
+
+                data = await r.json()
+                return data.get("results")
+
     async def cancel_transform(self, transform_id=None):
         headers = await self._get_authorization()
         path_template = f"/servicex/transformation/{transform_id}/cancel"
         url = self.url + path_template.format(transform_id=transform_id)
 
         async with ClientSession() as session:
             async with session.get(headers=headers, url=url) as r:
-
                 if r.status == 403:
                     raise AuthorizationError(
                         f"Not authorized to access serviceX at {self.url}"

tests/test_dataset.py

@@ -28,6 +28,7 @@
 import pytest
 import tempfile
 import os
+import datetime
 
 from unittest.mock import AsyncMock, Mock, patch
 from servicex.dataset_identifier import FileListDataset
@@ -120,16 +121,21 @@ async def test_as_files_cached(transformed_result, python_dataset):
 @pytest.mark.asyncio
 async def test_download_files(python_dataset):
     signed_urls_only = False
+    begin_at = datetime.datetime.now(tz=datetime.timezone.utc)
     download_progress = "download_task_id"
     minio_mock = AsyncMock()
     config = Configuration(cache_path="temp_dir", api_endpoints=[])
     python_dataset.configuration = config
+    python_dataset.servicex = AsyncMock()
+    python_dataset.servicex.get_transformation_results = AsyncMock(
+        side_effect=[
+            [{"file-path": "file1.txt"}],
+            [{"file-path": "file2.txt"}],
+        ]
+    )
+
     minio_mock.download_file.return_value = Path("/path/to/downloaded_file")
     minio_mock.get_signed_url.return_value = Path("http://example.com/signed_url")
-    minio_mock.list_bucket.return_value = [
-        Mock(filename="file1.txt"),
-        Mock(filename="file2.txt"),
-    ]
 
     progress_mock = Mock()
     python_dataset.minio_polling_interval = 0
@@ -138,7 +144,7 @@ async def test_download_files(python_dataset):
     python_dataset.configuration.shortened_downloaded_filename = False
 
     result_uris = await python_dataset.download_files(
-        signed_urls_only, progress_mock, download_progress, None
+        signed_urls_only, progress_mock, download_progress, None, begin_at
     )
     minio_mock.download_file.assert_awaited()
     minio_mock.get_signed_url.assert_not_awaited()
@@ -148,25 +154,29 @@ async def test_download_files(python_dataset):
 @pytest.mark.asyncio
 async def test_download_files_with_signed_urls(python_dataset):
     signed_urls_only = True
+    begin_at = datetime.datetime.now(tz=datetime.timezone.utc)
     download_progress = "download_task_id"
     minio_mock = AsyncMock()
     config = Configuration(cache_path="temp_dir", api_endpoints=[])
     python_dataset.configuration = config
     minio_mock.download_file.return_value = "/path/to/downloaded_file"
     minio_mock.get_signed_url.return_value = "http://example.com/signed_url"
-    minio_mock.list_bucket.return_value = [
-        Mock(filename="file1.txt"),
-        Mock(filename="file2.txt"),
-    ]
     progress_mock = Mock()
 
+    python_dataset.servicex = AsyncMock()
+    python_dataset.servicex.get_transformation_results = AsyncMock(
+        side_effect=[
+            [{"file-path": "file1.txt"}],
+            [{"file-path": "file2.txt"}],
+        ]
+    )
     python_dataset.minio_polling_interval = 0
     python_dataset.minio = minio_mock
     python_dataset.current_status = Mock(status="Complete", files_completed=2)
     python_dataset.configuration.shortened_downloaded_filename = False
 
     result_uris = await python_dataset.download_files(
-        signed_urls_only, progress_mock, download_progress, None
+        signed_urls_only, progress_mock, download_progress, None, begin_at
     )
     minio_mock.download_file.assert_not_called()
     minio_mock.get_signed_url.assert_called()

tests/test_servicex_adapter.py

@@ -28,6 +28,7 @@
 import os
 import tempfile
 import time
+import datetime
 from unittest.mock import patch
 
 import httpx
@@ -502,3 +503,84 @@ async def test_get_authorization(servicex):
         with patch("google.auth.jwt.decode", return_value={"exp": time.time() - 90}):
             r = await servicex._get_authorization()
             get_token.assert_called_once()
+
+
+@pytest.mark.asyncio
+@patch("servicex.servicex_adapter.ClientSession.get")
+async def test_get_transformation_results_success(get_transformation_results, servicex):
+    get_transformation_results.return_value.__aenter__.return_value.status = 200
+
+    request_id = "123-45-6789"
+    now = datetime.datetime.now(datetime.timezone.utc)
+    await servicex.get_transformation_results(request_id, now)
+
+    get_transformation_results.assert_called_with(
+        url=f"https://servicex.org/servicex/transformation/{request_id}/results",
+        headers={},
+        params={
+            "begin_at": now.isoformat(),
+        },
+    )
+
+
+@pytest.mark.asyncio
+@patch("servicex.servicex_adapter.ClientSession.get")
+async def test_get_transformation_results_not_found(
+    get_transformation_results, servicex
+):
+    get_transformation_results.return_value.__aenter__.return_value.status = 404
+    request_id = "123-45-6789"
+    now = datetime.datetime.now(datetime.timezone.utc)
+
+    with pytest.raises(ValueError):
+        await servicex.get_transformation_results(request_id, now)
+
+    get_transformation_results.assert_called_with(
+        url=f"https://servicex.org/servicex/transformation/{request_id}/results",
+        headers={},
+        params={
+            "begin_at": now.isoformat(),
+        },
+    )
+
+
+@pytest.mark.asyncio
+@patch("servicex.servicex_adapter.ClientSession.get")
+async def test_get_transformation_results_not_authorized(
+    get_transformation_results, servicex
+):
+    get_transformation_results.return_value.__aenter__.return_value.status = 403
+    request_id = "123-45-6789"
+    now = datetime.datetime.now(datetime.timezone.utc)
+
+    with pytest.raises(AuthorizationError):
+        await servicex.get_transformation_results(request_id, now)
+
+    get_transformation_results.assert_called_with(
+        url=f"https://servicex.org/servicex/transformation/{request_id}/results",
+        headers={},
+        params={
+            "begin_at": now.isoformat(),
+        },
+    )
+
+
+@pytest.mark.asyncio
+@patch("servicex.servicex_adapter.ClientSession.get")
+async def test_get_transformation_results_server_error(
+    get_transformation_results, servicex
+):
+    get_transformation_results.return_value.__aenter__.return_value.status = 500
+    request_id = "123-45-6789"
+    now = datetime.datetime.now(datetime.timezone.utc)
+
+    with pytest.raises(RuntimeError):
+        await servicex.get_transformation_results(request_id, now)
+
+    get_transformation_results.assert_called_with(
+        url=f"https://servicex.org/servicex/transformation/{request_id}/results",
+        headers={},
+        params={
+            "begin_at": now.isoformat(),
+        },
+    )