remove s3 bucket polling when waiting for transformation results (#587)

* remove s3 bucket polling when waiting for transformation results

Author: MattShirley

servicex/models.py

@@ -230,6 +230,16 @@ class TransformedResults(DocStringBaseModel):
     """URL for looking up logs on the ServiceX server"""
 
 
+class ServiceXInfo(DocStringBaseModel):
+    r"""
+    Model for ServiceX Info properties
+    """
+
+    app_version: str = Field(alias="app-version")
+    code_gen_image: dict[str, str] = Field(alias="code-gen-image")
+    capabilities: list[str] = Field(default_factory=list)
+
+
 class DatasetFile(BaseModel):
     """
     Model for a file in a cached dataset

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
@@ -342,13 +343,17 @@ 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,
             )
         )
 
         try:
             signed_urls = []
             downloaded_files = []
+
             download_result = await download_files_task
             if signed_urls_only:
                 signed_urls = download_result
@@ -522,6 +527,7 @@ async def download_files(
         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 = []
@@ -555,40 +561,70 @@ async def get_signed_url(
             if progress:
                 progress.advance(task_id=download_progress, task_type="Download")
 
+        later_than = datetime.datetime.min.replace(tzinfo=datetime.timezone.utc)
+
+        use_local_polling = (
+            "poll_local_transformation_results"
+            in await self.servicex.get_servicex_capabilities()
+        )
+
+        if not use_local_polling:
+            logger.warning(
+                "ServiceX is using legacy S3 bucket polling. Future versions of the "
+                "ServiceX client will not support this method. Please update your "
+                "ServiceX server to the latest version."
+            )
+
         while True:
             if not cached_record:
                 await asyncio.sleep(self.minio_polling_interval)
             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()
+                    if use_local_polling:
+                        files = await self.servicex.get_transformation_results(
+                            self.current_status.request_id, later_than
+                        )
+                    else:
+                        files = await self.minio.list_bucket()
+
                     for file in files:
-                        if file.filename not in files_seen:
+                        filename = file.filename
+
+                        if filename != "" and filename not in files_seen:
                             if signed_urls_only:
                                 download_tasks.append(
                                     loop.create_task(
                                         get_signed_url(
                                             self.minio,
-                                            file.filename,
+                                            filename,
                                             progress,
                                             download_progress,
                                         )
                                     )
                                 )
                             else:
+                                if use_local_polling:
+                                    expected_size = file.total_bytes
+                                else:
+                                    expected_size = file.size
                                 download_tasks.append(
                                     loop.create_task(
                                         download_file(
                                             self.minio,
-                                            file.filename,
+                                            filename,
                                             progress,
                                             download_progress,
                                             shorten_filename=self.configuration.shortened_downloaded_filename,  # NOQA: E501
-                                            expected_size=file.size,
+                                            expected_size=expected_size,
                                         )
                                     )
                                 )  # NOQA 501
-                            files_seen.add(file.filename)
+                            files_seen.add(filename)
+
+                            if use_local_polling:
+                                if file.created_at > later_than:
+                                    later_than = file.created_at
 
             # 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,7 +27,9 @@
 # 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 dataclasses import dataclass
 
 import httpx
 from httpx import AsyncClient, Response
@@ -41,13 +43,25 @@
     retry_if_not_exception_type,
 )
 
-from servicex.models import TransformRequest, TransformStatus, CachedDataset
+from servicex.models import (
+    TransformRequest,
+    TransformStatus,
+    CachedDataset,
+    ServiceXInfo,
+)
 
 
 class AuthorizationError(BaseException):
     pass
 
 
+@dataclass
+class ServiceXFile:
+    created_at: datetime.datetime
+    filename: str
+    total_bytes: int
+
+
 async def _extract_message(r: Response):
     try:
         o = r.json()
@@ -63,6 +77,9 @@ def __init__(self, url: str, refresh_token: Optional[str] = None):
         self.refresh_token = refresh_token
         self.token = None
 
+        # interact with _servicex_info via get_servicex_info
+        self._servicex_info: Optional[ServiceXInfo] = None
+
     async def _get_token(self):
         url = f"{self.url}/token/refresh"
         headers = {"Authorization": f"Bearer {self.refresh_token}"}
@@ -120,6 +137,31 @@ async def _get_authorization(self, force_reauth: bool = False) -> Dict[str, str]
                 await self._get_token()
             return {"Authorization": f"Bearer {self.token}"}
 
+    async def get_servicex_info(self) -> ServiceXInfo:
+        if self._servicex_info:
+            return self._servicex_info
+
+        headers = await self._get_authorization()
+        retry_options = Retry(total=3, backoff_factor=10)
+        async with AsyncClient(transport=RetryTransport(retry=retry_options)) as client:
+            r = await client.get(url=f"{self.url}/servicex", headers=headers)
+            if r.status_code == 401:
+                raise AuthorizationError(
+                    f"Not authorized to access serviceX at {self.url}"
+                )
+            elif r.status_code > 400:
+                error_message = await _extract_message(r)
+                raise RuntimeError(
+                    "ServiceX WebAPI Error during transformation "
+                    f"submission: {r.status_code} - {error_message}"
+                )
+            servicex_info = r.json()
+            self._servicex_info = ServiceXInfo(**servicex_info)
+            return self._servicex_info
+
+    async def get_servicex_capabilities(self) -> List[str]:
+        return (await self.get_servicex_info()).capabilities
+
     async def get_transforms(self) -> List[TransformStatus]:
         headers = await self._get_authorization()
         retry_options = Retry(total=3, backoff_factor=10)
@@ -232,6 +274,48 @@ async def delete_transform(self, transform_id=None):
                 msg = await _extract_message(r)
                 raise RuntimeError(f"Failed to delete transform {transform_id} - {msg}")
 
+    async def get_transformation_results(
+        self, request_id: str, later_than: Optional[datetime.datetime] = None
+    ):
+        if (
+            "poll_local_transformation_results"
+            not in await self.get_servicex_capabilities()
+        ):
+            raise ValueError("ServiceX capabilities not found")
+
+        headers = await self._get_authorization()
+        url = self.url + f"/servicex/transformation/{request_id}/results"
+        params = {}
+        if later_than:
+            params["later_than"] = later_than.isoformat()
+
+        async with AsyncClient() as session:
+            r = await session.get(headers=headers, url=url, params=params)
+            if r.status_code == 403:
+                raise AuthorizationError(
+                    f"Not authorized to access serviceX at {self.url}"
+                )
+
+            if r.status_code == 404:
+                raise ValueError(f"Request {request_id} not found")
+
+            if r.status_code != 200:
+                msg = await _extract_message(r)
+                raise RuntimeError(f"Failed with message: {msg}")
+
+            data = r.json()
+            response = list()
+            for result in data.get("results", []):
+                file = ServiceXFile(
+                    filename=result["s3-object-name"],
+                    created_at=datetime.datetime.fromisoformat(
+                        result["created_at"]
+                    ).replace(tzinfo=datetime.timezone.utc),
+                    total_bytes=result["total-bytes"],
+                )
+                response.append(file)
+            return response
+
     async def cancel_transform(self, transform_id=None):
         headers = await self._get_authorization()
         path_template = f"/servicex/transformation/{transform_id}/cancel"

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
@@ -44,6 +45,8 @@
 )
 from rich.progress import Progress
 
+from servicex.servicex_adapter import ServiceXFile
+
 
 @pytest.mark.asyncio
 async def test_as_signed_urls_happy(transformed_result):
@@ -124,12 +127,27 @@ async def test_download_files(python_dataset):
     minio_mock = AsyncMock()
     config = Configuration(cache_path="temp_dir", api_endpoints=[])
     python_dataset.configuration = config
+    python_dataset.servicex = AsyncMock()
+    python_dataset.servicex.get_servicex_capabilities = AsyncMock(
+        return_value=["poll_local_transformation_results"]
+    )
+
+    python_dataset.servicex.get_transformation_results = AsyncMock()
+    python_dataset.servicex.get_transformation_results.return_value = [
+        ServiceXFile(
+            filename="file1.txt",
+            created_at=datetime.datetime.now(datetime.timezone.utc),
+            total_bytes=100,
+        ),
+        ServiceXFile(
+            filename="file2.txt",
+            created_at=datetime.datetime.now(datetime.timezone.utc),
+            total_bytes=100,
+        ),
+    ]
+
     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
@@ -154,12 +172,27 @@ async def test_download_files_with_signed_urls(python_dataset):
     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_servicex_capabilities = AsyncMock(
+        return_value=["poll_local_transformation_results"]
+    )
+
+    python_dataset.servicex.get_transformation_results = AsyncMock()
+    python_dataset.servicex.get_transformation_results.return_value = [
+        ServiceXFile(
+            filename="file1.txt",
+            created_at=datetime.datetime.now(datetime.timezone.utc),
+            total_bytes=100,
+        ),
+        ServiceXFile(
+            filename="file2.txt",
+            created_at=datetime.datetime.now(datetime.timezone.utc),
+            total_bytes=100,
+        ),
+    ]
+
     python_dataset.minio_polling_interval = 0
     python_dataset.minio = minio_mock
     python_dataset.current_status = Mock(status="Complete", files_completed=2)