fix: update write handle on every recv()

[chandra-siri]

fix: update write_handle on every recv() from write object stream.

[gemini-code-assist]

Summary of Changes

Hello @chandra-siri, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request focuses on improving the robustness and correctness of asynchronous object writing and reading operations, particularly concerning the management of write_handles. It ensures that these handles are consistently updated during stream interactions, preventing potential issues with stale references. Additionally, the PR refines the type annotations for these handles and significantly enhances the testing framework for asynchronous components, making the system more reliable and easier to maintain.

Highlights

• Write Handle Update Logic: Introduced a new utility function update_write_handle_if_exists to ensure that the write_handle attribute of an object is updated after every recv() operation in asynchronous write streams, addressing potential staleness.
• Type Hint Refinement: Updated type hints for handle and write_handle across several asynchronous stream classes (e.g., _AsyncAbstractObjectStream, AsyncAppendableObjectWriter, AsyncMultiRangeDownloader, _AsyncReadObjectStream, _AsyncWriteObjectStream) from bytes to Any or more specific _storage_v2.BidiWriteHandle types for improved clarity and type safety.
• Asynchronous Test Infrastructure Improvement: Refactored system tests in test_zonal.py to use session-scoped event_loop and grpc_clients fixtures, ensuring gRPC clients are instantiated within the correct event loop context and removing redundant @pytest.mark.asyncio decorators from individual test functions.
• Enhanced Error Handling and Testing: Added new unit tests to test_async_write_object_stream.py to cover scenarios like missing persisted_size when opening a stream with a write_handle and explicitly verifying that the write_handle is updated correctly by recv() and requests_done methods.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a fix to update the write_handle on every recv() call, ensuring the most current handle is used. This is achieved by adding a new utility function and applying it across different stream operations. The changes also include updating type hints for handles from bytes to their specific proto types, which improves code clarity and correctness. Additionally, there's a significant and valuable refactoring of the system tests to use session-scoped event loops and gRPC clients, which should improve test performance and stability. I've identified one critical bug in async_write_object_stream.py where an AttributeError could occur when opening a stream with an existing write_handle. A code suggestion to fix this is provided.

I am having trouble creating individual review comments. Click here to see my feedback.

google/cloud/storage/_experimental/asyncio/async_write_object_stream.py (153-181)

When opening a stream with is_open_via_write_handle=True, response.resource is not guaranteed to be populated. Accessing response.resource.generation will raise an AttributeError. The generation number is already known in this case, so it should only be updated from the response when a new stream is created without a handle.

self.generation_number = response.resource.generation should be moved inside the else block that handles the case where is_open_via_write_handle is false.

        if is_open_via_write_handle:
            # Don't use if not response.persisted_size because this will be true
            # if persisted_size==0 (0 is considered "Falsy" in Python)
            if response.persisted_size is None:
                raise ValueError(
                    "Failed to obtain persisted_size after opening the stream via write_handle"
                )
            self.persisted_size = response.persisted_size
        else:
            if not response.resource:
                raise ValueError(
                    "Failed to obtain object resource after opening the stream"
                )
            if not response.resource.generation:
                raise ValueError(
                    "Failed to obtain object generation after opening the stream"
                )
            if not response.resource.size:
                # Appending to a 0 byte appendable object.
                self.persisted_size = 0
            else:
                self.persisted_size = response.resource.size
            self.generation_number = response.resource.generation

        if not response.write_handle:
            raise ValueError("Failed to obtain write_handle after opening the stream")

        self.write_handle = response.write_handle

The merge-base changed after approval.