fix: Fix bad requests when retrying with streams in multi-part requests

nisystemlink/clients/artifact/_artifact_client.py

@@ -10,13 +10,21 @@
     post,
     response_handler,
 )
+from nisystemlink.clients.core._uplink._multipart_retry import (
+    retryable_multipart_request,
+)
 from nisystemlink.clients.core.helpers._iterator_file_like import IteratorFileLike
 from requests.models import Response
-from uplink import Part, Path
+from uplink import Part, Path, retry
 
 from . import models
 
 
+def _iter_content_filelike_wrapper(response: Response) -> IteratorFileLike:
+    return IteratorFileLike(response.iter_content(chunk_size=4096))
+
+
+@retry(when=retry.when.status(429), stop=retry.stop.after_attempt(5))
 class ArtifactClient(BaseClient):
     def __init__(self, configuration: core.HttpConfiguration | None = None):
         """Initialize an instance.
@@ -35,9 +43,10 @@ def __init__(self, configuration: core.HttpConfiguration | None = None):
 
         super().__init__(configuration, base_path="/ninbartifact/v1/")
 
-    @post("artifacts")
+    @retryable_multipart_request()
+    @post("artifacts", args=[Part("workspace"), Part("artifact")])
     def __upload_artifact(
-        self, workspace: Part, artifact: Part
+        self, workspace: str, artifact: BinaryIO
     ) -> models.UploadArtifactResponse:
         """Uploads an artifact  using multipart/form-data headers to send the file payload in the HTTP body.
 
@@ -49,6 +58,7 @@ def __upload_artifact(
             UploadArtifactResponse: The response containing the artifact ID.
 
         """
+        ...
 
     def upload_artifact(
         self, workspace: str, artifact: BinaryIO
@@ -70,9 +80,6 @@ def upload_artifact(
 
         return response
 
-    def _iter_content_filelike_wrapper(response: Response) -> IteratorFileLike:
-        return IteratorFileLike(response.iter_content(chunk_size=4096))
-
     @response_handler(_iter_content_filelike_wrapper)
     @get("artifacts/{id}")
     def download_artifact(self, id: Path) -> IteratorFileLike:

nisystemlink/clients/core/_uplink/_multipart_retry.py

@@ -0,0 +1,39 @@
+from typing import Any, Callable, cast, TypeVar
+
+from uplink import decorators
+from uplink.clients.io.interfaces import RequestTemplate
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _rewind_retryable_part(part: object) -> None:
+    if hasattr(part, "seek"):
+        cast(Any, part).seek(0)
+        return
+
+    if isinstance(part, tuple):
+        for item in part:
+            if hasattr(item, "seek"):
+                cast(Any, item).seek(0)
+
+
+class _RetryableMultipartRequestTemplate(RequestTemplate):
+    def before_request(self, request: tuple[str, str, dict[str, Any]]) -> None:
+        _, _, extras = request
+        for part in extras.get("files", {}).values():
+            _rewind_retryable_part(part)
+        return None
+
+
+class _RetryableMultipartRequest(decorators.MethodAnnotation):
+    def modify_request(self, request_builder: Any) -> None:
+        request_builder.add_request_template(_RetryableMultipartRequestTemplate())
+
+
+def retryable_multipart_request() -> Callable[[F], F]:
+    """Create a method decorator that rewinds multipart parts before each send."""
+
+    def decorator(func: F) -> F:
+        return _RetryableMultipartRequest()(func)  # type: ignore[return-value]
+
+    return decorator

nisystemlink/clients/feeds/_feeds_client.py

@@ -5,6 +5,9 @@
 from nisystemlink.clients import core
 from nisystemlink.clients.core._uplink._base_client import BaseClient
 from nisystemlink.clients.core._uplink._methods import delete, get, post
+from nisystemlink.clients.core._uplink._multipart_retry import (
+    retryable_multipart_request,
+)
 from uplink import Part, Path, Query, retry
 
 from . import models
@@ -90,6 +93,7 @@ def query_feeds(
 
         return response
 
+    @retryable_multipart_request()
     @post(
         "feeds/{feedId}/packages",
         args=[Path(name="feedId"), Part(), Query(name="shouldOverwrite")],

nisystemlink/clients/file/_file_client.py

@@ -14,6 +14,9 @@
     post,
     response_handler,
 )
+from nisystemlink.clients.core._uplink._multipart_retry import (
+    retryable_multipart_request,
+)
 from nisystemlink.clients.core.helpers import IteratorFileLike
 from requests.models import Response
 from uplink import Body, Field, params, Part, Path, Query, retry
@@ -331,6 +334,7 @@ def download_file(self, id: str) -> IteratorFileLike:
         """
 
     @response_handler(_file_uri_response_handler)
+    @retryable_multipart_request()
     @post("service-groups/Default/upload-files")
     def __upload_file(
         self,
@@ -431,6 +435,7 @@ def start_upload_session(
         ],
     )
     @response_handler(lambda response: None)
+    @retryable_multipart_request()
     def append_to_upload_session(
         self,
         session_id: str,

nisystemlink/clients/notebook/_notebook_client.py

@@ -1,5 +1,5 @@
 import io
-from typing import List
+from typing import BinaryIO, List
 
 from nisystemlink.clients import core
 from nisystemlink.clients.core._api_error import ApiError
@@ -14,6 +14,9 @@
     put,
     response_handler,
 )
+from nisystemlink.clients.core._uplink._multipart_retry import (
+    retryable_multipart_request,
+)
 from nisystemlink.clients.core.helpers._iterator_file_like import IteratorFileLike
 from uplink import Part, Path, retry
 
@@ -54,12 +57,16 @@ def get_notebook(self, id: str) -> models.NotebookMetadata:
         """
         ...
 
-    @put("ninotebook/v1/notebook/{id}")
+    @retryable_multipart_request()
+    @put(
+        "ninotebook/v1/notebook/{id}",
+        args=[Path("id"), Part("metadata"), Part("content")],
+    )
     def __update_notebook(
         self,
-        id: Path,
-        metadata: Part = None,
-        content: Part = None,
+        id: str,
+        metadata: io.BytesIO | None = None,
+        content: BinaryIO | None = None,
     ) -> models.NotebookMetadata:
         """Updates a notebook metadata by ID.
 
@@ -81,7 +88,7 @@ def update_notebook(
         self,
         id: str,
         metadata: models.NotebookMetadata | None = None,
-        content: io.BufferedReader | None = None,
+        content: BinaryIO | None = None,
     ) -> models.NotebookMetadata:
         """Updates a notebook metadata by ID.
 
@@ -121,11 +128,15 @@ def delete_notebook(self, id: str) -> None:
         """
         ...
 
-    @post("ninotebook/v1/notebook")
+    @retryable_multipart_request()
+    @post(
+        "ninotebook/v1/notebook",
+        args=[Part("metadata"), Part("content")],
+    )
     def __create_notebook(
         self,
-        metadata: Part,
-        content: Part,
+        metadata: io.BytesIO,
+        content: BinaryIO,
     ) -> models.NotebookMetadata:
         """Creates a new notebook.
 
@@ -145,7 +156,7 @@ def __create_notebook(
     def create_notebook(
         self,
         metadata: models.NotebookMetadata,
-        content: io.BufferedReader,
+        content: BinaryIO,
     ) -> models.NotebookMetadata:
         """Creates a new notebook.
 

tests/integration/artifact/test_artifact.py

@@ -2,11 +2,18 @@
 from typing import List
 
 import pytest
+import responses
 from nisystemlink.clients.artifact import ArtifactClient
 from nisystemlink.clients.artifact.models._upload_artifact_response import (
     UploadArtifactResponse,
 )
 from nisystemlink.clients.core._http_configuration import HttpConfiguration
+from responses import PassthroughResponse
+from responses.registries import OrderedRegistry
+from uplink.clients.io import blocking_strategy as uplink_blocking_strategy
+
+BASE_URL = "https://test-api.lifecyclesolutions.ni.com"
+DEFAULT_WORKSPACE = "2300760d-38c4-48a1-9acb-800260812337"
 
 
 @pytest.fixture(scope="class")
@@ -23,7 +30,7 @@ def create_artifact(client: ArtifactClient):
     def _create_artifact(
         content: bytes = b"test content",
         cleanup: bool = True,
-        workspace: str = "2300760d-38c4-48a1-9acb-800260812337",
+        workspace: str = DEFAULT_WORKSPACE,
     ):
         # Used the main-test default workspace since the client for creating a workspace has not been added yet
         artifact_stream = io.BytesIO(content)
@@ -51,6 +58,29 @@ def test__upload_artifact__artifact_uploaded(
         assert upload_response is not None
         assert upload_response.id is not None
 
+    def test__upload_artifact_after_rate_limit_retry__artifact_uploaded(
+        self, create_artifact, monkeypatch: pytest.MonkeyPatch
+    ):
+        monkeypatch.setattr(uplink_blocking_strategy.time, "sleep", lambda _: None)
+
+        with responses.RequestsMock(registry=OrderedRegistry) as request_mock:
+            request_mock.add(
+                responses.POST,
+                f"{BASE_URL}/ninbartifact/v1/artifacts",
+                status=429,
+            )
+            request_mock.add(
+                PassthroughResponse(
+                    responses.POST,
+                    f"{BASE_URL}/ninbartifact/v1/artifacts",
+                )
+            )
+
+            upload_response: UploadArtifactResponse = create_artifact()
+
+        assert upload_response is not None
+        assert upload_response.id is not None
+
     def test__download_artifact__artifact_downloaded(
         self, client: ArtifactClient, create_artifact
     ):

tests/integration/feeds/test_feeds_client.py

@@ -6,10 +6,15 @@
 from typing import BinaryIO, Callable
 
 import pytest
+import responses
 from nisystemlink.clients.core import ApiException
 from nisystemlink.clients.feeds import FeedsClient
 from nisystemlink.clients.feeds.models import CreateFeedRequest, Platform
+from responses import PassthroughResponse
+from responses.registries import OrderedRegistry
+from uplink.clients.io import blocking_strategy as uplink_blocking_strategy
 
+BASE_URL = "https://test-api.lifecyclesolutions.ni.com"
 FEED_DESCRIPTION = "Sample feed for uploading packages"
 PACKAGE_PATH = str(
     Path(__file__).parent.resolve()
@@ -236,6 +241,49 @@ def test__upload_package_content__invalid_feed_id_raises(
                 feed_id=invalid_id,
             )
 
+    def test__upload_package_content_after_rate_limit_retry__upload_package_content_succeeds(
+        self,
+        client: FeedsClient,
+        create_feed: Callable,
+        create_feed_request: Callable,
+        get_feed_name: Callable,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        create_feed_request_body = create_feed_request(
+            feed_name=get_feed_name(),
+            description=FEED_DESCRIPTION,
+            platform=Platform.WINDOWS,
+        )
+        create_feed_resp = create_feed(create_feed_request_body)
+        assert create_feed_resp.id is not None
+
+        response = None
+
+        monkeypatch.setattr(uplink_blocking_strategy.time, "sleep", lambda _: None)
+
+        with responses.RequestsMock(registry=OrderedRegistry) as request_mock:
+            request_mock.add(
+                responses.POST,
+                f"{BASE_URL}/nifeed/v1/feeds/{create_feed_resp.id}/packages",
+                status=429,
+            )
+            request_mock.add(
+                PassthroughResponse(
+                    responses.POST,
+                    f"{BASE_URL}/nifeed/v1/feeds/{create_feed_resp.id}/packages",
+                )
+            )
+
+            with open(PACKAGE_PATH, "rb") as package:
+                response = client.upload_package_content(
+                    feed_id=create_feed_resp.id,
+                    package=package,
+                    overwrite=True,
+                )
+
+        assert response is not None
+        assert response.id is not None
+
     def test__delete_windows_feed__succeeds(
         self,
         client: FeedsClient,