Skip to content

feat(zonal): pass Preconditions to AsyncAppendableObjectWriter #1131

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

Open
yuxin00j wants to merge 3 commits into
base: main
Choose a base branch
from
+84 −44
Open

feat(zonal): pass Preconditions to AsyncAppendableObjectWriter #1131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
318a94d
feat(zonal): pass Preconditions to AsyncAppendableObjectWriter
yuxin00j Oct 29, 2026
b6f55c7
refactor(zonal): import Preconditions with an ImportError fallback
yuxin00j Nov 3, 2026
9d517c1
Merge remote-tracking branch 'upstream/main' into zonal-write-precond…
yuxin00j Nov 3, 2026
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
4 changes: 1 addition & 3 deletions gcsfs/tests/test_core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1745,9 +1745,7 @@ def test_write_if_generation_match(gcs):
try:
gcs.pipe_file(fn, b"v2", if_generation_match=gen)
except NotImplementedError:
pytest.skip(
"zonal CAS needs a google-cloud-storage release with if_generation_match"
)
pytest.skip("zonal CAS needs a google-cloud-storage release with Preconditions")

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

test_zonal_file_open_write_mode_with_if_generation_match needs the same update. Its real-GCS branch still says "SDK releases whose writer lacks the parameter" and uses pytest.raises((NotImplementedError, GenerationMismatchError)). With 3.12.0 it reports PASSED without exercising CAS, whereas this test is SKIPPED. Suggestion: skip on NotImplementedError with this same message, expect only GenerationMismatchError, and add a positive write with if_generation_match=<current gen>.

assert gcs.cat(fn) == b"v2"
with pytest.raises(GenerationMismatchError):
gcs.pipe_file(fn, b"v3", if_generation_match=gen)
Expand Down
79 changes: 58 additions & 21 deletions gcsfs/tests/test_zb_hns_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import asyncio
import collections
import concurrent.futures
import dataclasses
import logging
from typing import Optional
from unittest import mock

import pytest
Expand Down Expand Up @@ -99,18 +101,35 @@ async def test_init_aaow():
assert result is mock_writer_instance


@dataclasses.dataclass

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

nit: Because the fake is always patched in, CI never builds the real Preconditions, even after a release ships it. Consider _Preconditions = zb_hns_utils.Preconditions or _FakePreconditions. These tests pass unchanged against the real class from #18641.

class _FakePreconditions:
if_generation_match: Optional[int] = None


@pytest.mark.asyncio
@pytest.mark.parametrize(
"if_generation_match, expected_kwargs",
[(0, {"generation": 0}), (7, {"generation": None, "if_generation_match": 7})],
[
(0, {"generation": 0}),
(
7,
{
"generation": None,
"preconditions": _FakePreconditions(if_generation_match=7),
},
),
],
)
async def test_init_aaow_if_generation_match(if_generation_match, expected_kwargs):

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

This covers init_aaow, but nothing mocked checks that callers pass the precondition through. When I replaced all four if_generation_match=_resolve_generation_match(mode, if_generation_match) arguments in extended_gcsfs.py with None, the mocked zonal and zb_hns_utils suites still passed, because test_zonal.py only asserts if_generation_match=None. Could you add cases to test_pipe_file_zonal, test_put_file_zonal, test_simple_upload_zonal and test_initiate_upload_zonal checking that mode="create" reaches init_aaow as 0 and if_generation_match="7" as 7?

mock_writer_instance = mock.AsyncMock()
with mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
new_callable=mock.Mock,
return_value=mock_writer_instance,
) as mock_writer_class:
with (
mock.patch.object(zb_hns_utils, "Preconditions", _FakePreconditions),
mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
new_callable=mock.Mock,
return_value=mock_writer_instance,
) as mock_writer_class,
):
await zb_hns_utils.init_aaow(
mock_grpc_client,
bucket_name,
Expand Down Expand Up @@ -149,26 +168,41 @@ async def test_init_aaow_rejects_generation_with_if_generation_match(

@pytest.mark.asyncio
async def test_init_aaow_if_generation_match_unsupported_sdk():
with mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
side_effect=TypeError("unexpected keyword argument 'if_generation_match'"),
with (
mock.patch.object(zb_hns_utils, "Preconditions", None),
mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
new_callable=mock.Mock,
) as mock_writer_class,
):
with pytest.raises(NotImplementedError):
await zb_hns_utils.init_aaow(
mock_grpc_client, bucket_name, object_name, if_generation_match=7
)
mock_writer_class.assert_not_called()


@pytest.mark.asyncio
async def test_init_aaow_unrelated_type_error_propagates():
with mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
side_effect=TypeError("boom"),
async def test_init_aaow_if_generation_match_zero_without_preconditions():
mock_writer_instance = mock.AsyncMock()
with (
mock.patch.object(zb_hns_utils, "Preconditions", None),
mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
new_callable=mock.Mock,
return_value=mock_writer_instance,
) as mock_writer_class,
):
with pytest.raises(TypeError, match="boom"):
await zb_hns_utils.init_aaow(
mock_grpc_client, bucket_name, object_name, if_generation_match=7
)
await zb_hns_utils.init_aaow(
mock_grpc_client, bucket_name, object_name, if_generation_match=0
)
mock_writer_class.assert_called_once_with(
client=mock_grpc_client,
bucket_name=bucket_name,
object_name=object_name,
generation=0,
writer_options={},
)


@pytest.mark.asyncio
Expand All @@ -179,10 +213,13 @@ async def test_init_aaow_unrelated_type_error_propagates():
async def test_init_aaow_failed_precondition(if_generation_match, expected):
mock_writer_instance = mock.AsyncMock()
mock_writer_instance.open.side_effect = FailedPrecondition("exists")
with mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
new_callable=mock.Mock,
return_value=mock_writer_instance,
with (
mock.patch.object(zb_hns_utils, "Preconditions", _FakePreconditions),
mock.patch(
"gcsfs.zb_hns_utils.AsyncAppendableObjectWriter",
new_callable=mock.Mock,
return_value=mock_writer_instance,
),
):
with pytest.raises(expected) as excinfo:

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

optional: for the 7 case, also assert excinfo.value.generation == 7 (as test_pipe_file_precondition_error does) and excinfo.value.__cause__ is mock_writer_instance.open.side_effect.

await zb_hns_utils.init_aaow(
Expand Down
45 changes: 25 additions & 20 deletions gcsfs/zb_hns_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,13 @@

from gcsfs.retry import GenerationMismatchError

try:
from google.cloud.storage.asyncio.async_appendable_object_writer import (
Preconditions,
)
except ImportError: # releases before googleapis/google-cloud-python#18641
Preconditions = None

MRD_MAX_RANGES = 1000 # MRD supports up to 1000 ranges per request
try:
DEFAULT_CONCURRENCY = int(os.environ.get("DEFAULT_GCSFS_CONCURRENCY", "4"))
Expand Down Expand Up @@ -175,7 +182,7 @@ async def init_aaow(

``if_generation_match=0`` creates the object only if it does not exist;
a positive value overwrites only that generation (requires a
google-cloud-storage release whose writer accepts ``if_generation_match``).
google-cloud-storage release whose writer accepts ``preconditions``).
A failed precondition raises ``FileExistsError`` for ``0`` and
``GenerationMismatchError`` otherwise; without ``if_generation_match`` the
SDK's ``FailedPrecondition`` propagates unchanged.
Expand All @@ -187,32 +194,30 @@ async def init_aaow(
writer_options["FLUSH_INTERVAL_BYTES"] = flush_interval_bytes
kwargs = {}
if generation and if_generation_match is not None:

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

nit: generation=0 skips this check, so fs.open(p, "wb", generation=0, if_generation_match=7) raises NotImplementedError on older SDKs, or the SDK's own ValueError on #18641. if generation is not None matches how the SDK checks it. Could you also add generation=0 to test_init_aaow_rejects_generation_with_if_generation_match?

# generation > 0 appends to an existing object (AppendObjectSpec), which
# carries no precondition, so the SDK would silently ignore it.
# generation > 0 appends to an existing object, so a generation
# precondition cannot apply. Check here because if_generation_match=0
# is mapped to generation=0 below, which would hide the conflict.
raise ValueError(
"if_generation_match cannot be combined with an explicit generation"
)
if if_generation_match == 0:
# Every SDK release maps generation=0 to WriteObjectSpec(if_generation_match=0).
generation = 0
elif if_generation_match is not None:
kwargs["if_generation_match"] = if_generation_match
try:
writer = AsyncAppendableObjectWriter(
client=grpc_client,
bucket_name=bucket_name,
object_name=object_name,
generation=generation,
writer_options=writer_options,
**kwargs,
)
except TypeError as e:
if "if_generation_match" not in str(e):
raise
raise NotImplementedError(
"if_generation_match > 0 on zonal buckets requires a google-cloud-storage "
"release whose AsyncAppendableObjectWriter accepts if_generation_match"
) from e
if Preconditions is None:
raise NotImplementedError(
"if_generation_match > 0 on zonal buckets requires a google-cloud-storage "
"release whose AsyncAppendableObjectWriter accepts Preconditions"
)
kwargs["preconditions"] = Preconditions(if_generation_match=if_generation_match)
writer = AsyncAppendableObjectWriter(
client=grpc_client,
bucket_name=bucket_name,
object_name=object_name,
generation=generation,

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Pre-existing, fine as a follow-up: ZonalFile passes its generation here in every mode, so open(path, "wb", generation=G) appends to G instead of overwriting. On a real zonal bucket this produced v1-v2 (or "already finalized" if G was finalized), while GCSFile ignores generation for writes and produced v2. Passing generation only in "a" mode would make the two backends consistent.

writer_options=writer_options,
**kwargs,
)
try:
await writer.open()

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Pre-existing, fine as a follow-up: opening the writer replaces the object right away. ExtendedGcsFileSystem._put_file calls init_aaow before open(lpath), and _pipe_file calls it before memoryview(data). On a real zonal bucket, put_file("/missing", existing) and pipe_file(existing, "str") each truncate the remote object to 0 bytes before raising; the standard backend keeps the object. Opening or validating the input before init_aaow fixes both.

except FailedPrecondition as e:

@zhixiangli zhixiangli Nov 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Blocking. The errors raised here (FileExistsError, GenerationMismatchError, and the NotImplementedError above) reach ZonalFile callers, and ZonalFile retries them on close. ZonalFile._ensure_aaow doesn't set self.closed = True when init_aaow raises, unlike fsspec's AbstractBufferedFile.flush() when _initiate_upload() fails. So __exit__ → close() → flush(force=True) → _ensure_aaow() calls init_aaow a second time.

On a real zonal bucket, with both google-cloud-storage 3.12.0 and googleapis/google-cloud-python#18641:

  • with fs.open(existing, "xb") as f: f.write(b"d") makes 2 BidiWriteObject opens, and the FileExistsError has a second FileExistsError as its __context__. A stale if_generation_match does the same with GenerationMismatchError.
  • If the object is deleted between the failed write() and __exit__, the retry succeeds and leaves a 0-byte object even though the write raised.
  • Without a with block, the deferred close on GC makes the extra RPC and logs ERROR gcsfs: deferred close ... failed.
def _ensure_aaow(self):
    if self.aaow is None:
        try:
            self.aaow = asyn.sync(self.gcsfs.loop, self._init_aaow, self.bucket,
                                  self.key, self.generation, self.flush_interval_bytes)
        except Exception:
            self.closed = True  # same as AbstractBufferedFile.flush()
            raise

With this change each case makes 1 RPC and no object is created, and the existing zonal unit suites still pass. Test idea: set init_aaow.side_effect to each of the three errors, then assert init_aaow.await_count == 1 and f.closed.

A related follow-up on the same close path (pre-existing, not blocking): ZonalFile._close_impl doesn't guard super()._close_impl(). If the closing flush fails (e.g. the object is deleted while open, which returns a 500), closed becomes True and the write stream is never closed. Wrapping that call in try/except and then calling close_aaow(..., finalize_on_close=False) fixes it.

Expand Down
Loading