iframe-proxy

chandra-siri · 2026-01-26T18:59:32Z

feat: add context manager to mrd

gemini-code-assist · 2026-01-26T18:59:47Z

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 introduces an asynchronous context manager for the AsyncMultiRangeDownloader class, enhancing its usability and ensuring proper resource management in asynchronous operations. By enabling the use of async with statements, the PR simplifies the lifecycle of download operations, automatically handling the opening and closing of the underlying gRPC connection, which leads to more robust and readable code.

Highlights

Asynchronous Context Manager Implementation: The AsyncMultiRangeDownloader class now implements the asynchronous context manager protocol by adding __aenter__ and __aexit__ methods. This allows the class to be used with async with statements, automatically handling the opening and closing of the underlying bidi-gRPC connection.
Simplified Resource Management in Tests: System tests have been updated to utilize the new async with syntax for AsyncMultiRangeDownloader, replacing explicit open() and close() calls. This demonstrates the intended usage and simplifies resource lifecycle management.
Unit Test Coverage for Context Manager: A new unit test, test_async_context_manager_calls_open_and_close, has been added to verify that the __aenter__ method correctly calls open() and __aexit__ correctly calls close() when the AsyncMultiRangeDownloader is used as an asynchronous context manager.

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.

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 convenient async context manager for AsyncMultiRangeDownloader, which is a great improvement for ensuring resources are properly managed. The implementation is correct and is accompanied by a new unit test and an update to a system test.

One suggestion for further improvement is to refactor the other system tests in tests/system/test_zonal.py that use AsyncMultiRangeDownloader to also adopt the new async with syntax. This would improve consistency and readability across the test suite.

I've also added a comment with suggestions to enhance the unit tests for the new context manager by covering some important edge cases.

gemini-code-assist · 2026-01-26T19:00:53Z

+    async def test_async_context_manager_calls_open_and_close(
+        self, mock_grpc_client, mock_close, mock_open
+    ):
+        # Arrange
+        mrd = AsyncMultiRangeDownloader(
+            mock_grpc_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME
+        )
+
+        # To simulate the behavior of open and close changing the stream state
+        async def open_side_effect():
+            mrd._is_stream_open = True
+
+        async def close_side_effect():
+            mrd._is_stream_open = False
+
+        mock_open.side_effect = open_side_effect
+        mock_close.side_effect = close_side_effect
+        mrd._is_stream_open = False
+
+        # Act
+        async with mrd as downloader:
+            # Assert
+            mock_open.assert_called_once()
+            assert downloader == mrd
+            assert mrd.is_stream_open
+
+        mock_close.assert_called_once()
+        assert not mrd.is_stream_open


This is a good test for the happy path of the async context manager. To make it more robust, it would be beneficial to also test the behavior in edge cases.

Specifically, consider adding tests for:

When open() fails (e.g., raises an exception). In this scenario, close() should not be called.

When an exception is raised from within the async with block. In this case, close() should still be called to ensure cleanup.

Here are some examples of how you could structure these tests:

Test for open() failure:

@mock.patch("...") @pytest.mark.asyncio async def test_context_manager_no_close_on_open_failure(self, mock_grpc_client, mock_close, mock_open): mock_open.side_effect = ValueError("Failed to open") mrd = AsyncMultiRangeDownloader( mock_grpc_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME ) with pytest.raises(ValueError, match="Failed to open"): async with mrd: pytest.fail("This block should not be executed.") mock_open.assert_called_once() mock_close.assert_not_called()

Test for exception within the block:

@mock.patch("...") @pytest.mark.asyncio async def test_context_manager_closes_on_exception(self, mock_grpc_client, mock_close, mock_open): mrd = AsyncMultiRangeDownloader( mock_grpc_client, _TEST_BUCKET_NAME, _TEST_OBJECT_NAME ) # set up side effects for open/close as in the existing test async def open_side_effect(): mrd._is_stream_open = True mock_open.side_effect = open_side_effect with pytest.raises(RuntimeError, match="Oops"): async with mrd: raise RuntimeError("Oops") mock_open.assert_called_once() mock_close.assert_called_once()

PR created by the Librarian CLI to initialize a release. Merging this PR will auto trigger a release. Librarian Version: v1.0.2-0.20251119154421-36c3e21ad3ac Language Image: us-central1-docker.pkg.dev/cloud-sdk-librarian-prod/images-prod/python-librarian-generator@sha256:8e2c32496077054105bd06c54a59d6a6694287bc053588e24debe6da6920ad91 <details><summary>google-cloud-storage: 3.9.0</summary> ## [3.9.0](v3.8.0...v3.9.0) (2026-02-02) ### Features * update generation for MRD (#1730) ([08bc708](08bc7082)) * add get_object method for async grpc client (#1735) ([0e5ec29](0e5ec29b)) * Add micro-benchmarks for reads comparing standard (regional) vs rapid (zonal) buckets. (#1697) ([1917649](1917649f)) * Add support for opening via `write_handle` and fix `write_handle` type (#1715) ([2bc15fa](2bc15fa5)) * add samples for appendable objects writes and reads ([2e1a1eb](2e1a1eb5)) * add samples for appendable objects writes and reads (#1705) ([2e1a1eb](2e1a1eb5)) * add context manager to mrd (#1724) ([5ac2808](5ac2808a)) * Move Zonal Buckets features of `_experimental` (#1728) ([74c9ecc](74c9ecc5)) * add default user agent for grpc (#1726) ([7b31946](7b319469)) * expose finalized_time in blob.py applicable for GET_OBJECT in ZB (#1719) ([8e21a7f](8e21a7fe)) * expose `DELETE_OBJECT` in `AsyncGrpcClient` (#1718) ([c8dd7a0](c8dd7a0b)) * send `user_agent` to grpc channel (#1712) ([cdb2486](cdb2486b)) * integrate writes strategy and appendable object writer (#1695) ([dbd162b](dbd162b3)) * Add micro-benchmarks for writes comparing standard (regional) vs rapid (zonal) buckets. (#1707) ([dbe9d8b](dbe9d8b8)) * add support for `generation=0` to avoid overwriting existing objects and add `is_stream_open` support (#1709) ([ea0f5bf](ea0f5bf8)) * add support for `generation=0` to prevent overwriting existing objects ([ea0f5bf](ea0f5bf8)) * add `is_stream_open` property to AsyncAppendableObjectWriter for stream status check ([ea0f5bf](ea0f5bf8)) ### Bug Fixes * receive eof while closing reads stream (#1733) ([2ef6339](2ef63396)) * update write handle on every recv() (#1716) ([5d9fafe](5d9fafe1)) * implement requests_done method to signal end of requests in async streams. Gracefully close streams. (#1700) ([6c16079](6c160794)) * implement requests_done method to signal end of requests in async streams. Gracefully close streams. ([6c16079](6c160794)) * instance grpc client once per process in benchmarks (#1725) ([721ea2d](721ea2dd)) * Fix formatting in setup.py dependencies list (#1713) ([cc4831d](cc4831d7)) * Change contructors of MRD and AAOW AsyncGrpcClient.grpc_client to AsyncGrpcClient (#1727) ([e730bf5](e730bf50)) </details>

feat: add context manager to mrd

3531eea

product-auto-label Bot added size: m Pull request size is medium. api: storage Issues related to the googleapis/python-storage API. labels Jan 26, 2026

gemini-code-assist Bot reviewed Jan 26, 2026

View reviewed changes

chandra-siri assigned Pulkit0110 Jan 27, 2026

chandra-siri marked this pull request as ready for review January 27, 2026 13:20

chandra-siri requested review from a team January 27, 2026 13:20

chandra-siri requested a review from a team as a code owner January 27, 2026 13:20

Pulkit0110 approved these changes Jan 28, 2026

View reviewed changes

Merge branch 'main' into context_manager

48c3d50

chandra-siri merged commit 5ac2808 into main Jan 28, 2026
18 of 19 checks passed

chandra-siri deleted the context_manager branch January 28, 2026 10:06

release-please Bot mentioned this pull request Jan 28, 2026

chore(main): release 4.0.0 #1708

Closed

This was referenced Jan 29, 2026

chore: librarian release pull request: 20260129T115903Z #1731

Closed

chore: librarian release pull request: 20260202T123858Z #1736

Merged

Sunbelt Computer Software

PL/B Language Development and Support

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: add context manager to mrd#1724

feat: add context manager to mrd#1724
chandra-siri merged 2 commits into
mainfrom
context_manager

chandra-siri commented Jan 26, 2026

Uh oh!

gemini-code-assist Bot commented Jan 26, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jan 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

Conversation

chandra-siri commented Jan 26, 2026

Uh oh!

gemini-code-assist Bot commented Jan 26, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants