Skip to content

Fix retries for get_timeseries_buckets_proto.sync_detailed #58

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
gabrys merged 2 commits into from
Sep 15, 2025
Merged

Fix retries for get_timeseries_buckets_proto.sync_detailed #58

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
79ed7f8
PY-251 Fix retries for `get_timeseries_buckets_proto.sync_detailed`
gabrys Sep 15, 2025
9d6231e
Test for ProtobufPayload
gabrys Sep 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions src/neptune_query/internal/retrieval/metric_buckets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
# limitations under the License.

from dataclasses import dataclass
from io import BytesIO
from typing import (
Iterable,
Literal,
Expand All @@ -34,13 +33,13 @@
XSteps,
)
from neptune_api.proto.neptune_pb.api.v1.model.series_values_pb2 import ProtoTimeseriesBucketsDTO
from neptune_api.types import File

from ..identifiers import RunAttributeDefinition
from ..logger import get_logger
from ..query_metadata_context import with_neptune_client_metadata
from . import retry
from .search import ContainerType
from .util import ProtobufPayload

logger = get_logger()

Expand Down Expand Up @@ -147,7 +146,7 @@ def fetch_time_series_buckets(
call_api = retry.handle_errors_default(with_neptune_client_metadata(get_timeseries_buckets_proto.sync_detailed))
Copy link
Contributor

@michalsosn michalsosn Sep 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm, I guess

call_api = retry.handle_errors_default(with_neptune_client_metadata(lambda : get_timeseries_buckets_proto.sync_detailed(client=client, body=File(payload=BytesIO(request_object.SerializeToString())))))

response = call_api()

should recreate BytesIO on each retry.

I assume File2 is more like a test of an idea given the name :p

This bug seems like sth that could return in the future, so a test would be nice

response = call_api(
client=client,
body=File(payload=BytesIO(request_object.SerializeToString())),
body=ProtobufPayload(request_object),
)

logger.debug(
Expand Down
15 changes: 15 additions & 0 deletions src/neptune_query/internal/retrieval/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from __future__ import annotations

from dataclasses import dataclass
from io import BytesIO
from typing import (
Any,
Callable,
Expand All @@ -25,7 +26,9 @@
TypeVar,
)

from google.protobuf.message import Message
from neptune_api import AuthenticatedClient
from neptune_api.types import File

T = TypeVar("T")
R = TypeVar("R")
Expand All @@ -50,3 +53,15 @@ def fetch_pages(
page = process_page(data)
yield page
page_params = make_new_page_params(page_params, data)


class ProtobufPayload(File):
"""A version of the neptune_api.types.File class that uses a protobuf message as payload"""

@property
def payload(self) -> BytesIO:
return self.get_payload()

@payload.setter
def payload(self, message: Message) -> None:
self.get_payload = lambda: BytesIO(message.SerializeToString())
Loading