remove s3 bucket polling when waiting for transformation results

pyproject.toml

@@ -128,3 +128,10 @@ dynamic_context = "test_function"
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
+
+[tool.hatch.envs.test]
+features = ["test"]
+
+[tool.hatch.envs.test.scripts]
+test = "pytest {args}"
+cov = "pytest --cov=servicex {args}"

servicex/models.py

@@ -237,7 +237,7 @@ class ServiceXInfo(DocStringBaseModel):
 
     app_version: str = Field(alias="app-version")
     code_gen_image: dict[str, str] = Field(alias="code-gen-image")
-    capabilities: list[str]
+    capabilities: list[str] = Field(default_factory=list)
 
 
 class DatasetFile(BaseModel):

servicex/query_core.py

@@ -567,6 +567,14 @@ async def get_signed_url(
             "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)
@@ -597,7 +605,7 @@ async def get_signed_url(
                                 )
                             else:
                                 if use_local_polling:
-                                    expected_size = None
+                                    expected_size = file.total_bytes
                                 else:
                                     expected_size = file.size
                                 download_tasks.append(

servicex/servicex_adapter.py

@@ -59,9 +59,10 @@ class AuthorizationError(BaseException):
 class ServiceXFile:
     created_at: datetime.datetime
     filename: str
+    total_bytes: int
 
 
-async def _extract_message(r: ClientResponse):
+async def _extract_message(r: Response):
     try:
         o = r.json()
         error_message = o.get("message", str(r))
@@ -141,20 +142,20 @@ async def get_servicex_info(self) -> ServiceXInfo:
             return self._servicex_info
 
         headers = await self._get_authorization()
-        retry_options = ExponentialRetry(attempts=3, start_timeout=10)
-        async with RetryClient(retry_options=retry_options) as client:
-            async with client.get(url=f"{self.url}/servicex", headers=headers) as r:
-                if r.status == 401:
-                    raise AuthorizationError(
-                        f"Not authorized to access serviceX at {self.url}"
-                    )
-                elif r.status > 400:
-                    error_message = await _extract_message(r)
-                    raise RuntimeError(
-                        "ServiceX WebAPI Error during transformation "
-                        f"submission: {r.status} - {error_message}"
-                    )
-                servicex_info = await r.json()
+        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
 
@@ -288,50 +289,49 @@ async def get_transformation_results(
         if later_than:
             params["later_than"] = later_than.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()
-                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),
-                    )
-                    response.append(file)
-                return response
+        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"
         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}"
-                    )
-                elif r.status == 404:
-                    raise ValueError(f"Transform {transform_id} not found")
-                elif r.status != 200:
-                    msg = await _extract_message(r)
-                    raise RuntimeError(
-                        f"Failed to cancel transform {transform_id} - {msg}"
-                    )
+        async with AsyncClient() as session:
+            r = await session.get(headers=headers, url=url)
+            if r.status_code == 403:
+                raise AuthorizationError(
+                    f"Not authorized to access serviceX at {self.url}"
+                )
+            elif r.status_code == 404:
+                raise ValueError(f"Transform {transform_id} not found")
+            elif r.status_code != 200:
+                msg = await _extract_message(r)
+                raise RuntimeError(f"Failed to cancel transform {transform_id} - {msg}")
 
     async def submit_transform(self, transform_request: TransformRequest) -> str:
         headers = await self._get_authorization()

tests/test_dataset.py

@@ -137,10 +137,12 @@ async def test_download_files(python_dataset):
         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,
         ),
     ]
 
@@ -182,10 +184,12 @@ async def test_download_files_with_signed_urls(python_dataset):
         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,
         ),
     ]