[abhishekacharya200]

Fixes #18768

When a query or aggregation stream failed midway with a retryable error such as UNAVAILABLE, _retry_query_after_exception looked for the default retry policy on transport.run_query (or run_aggregation_query). On the real transports that is the raw stub, which has no _retry, so the stream raised AttributeError: '_UnaryStreamMultiCallable' object has no attribute '_retry' instead of restarting. With retry=None the helper called None._predicate(exc), so callers got an AttributeError instead of the original error. This is the traceback reported in googleapis/python-firestore#596 and airbytehq/airbyte#70837.

This reads the policy from transport._wrapped_methods[transport.run_query] (and run_aggregation_query), the same place the generated client reads it for every call, and raises the error as is when retry is None. I checked the lookup on both the gRPC and the REST transport, and both return the Retry the client was built with. The async query and aggregation have no restart logic, so they are unchanged.

One thing worth knowing: this lets the existing restart path run for the first time. A query restarts after the last snapshot it returned, so nothing is repeated. An aggregation reruns the whole query, which is what the existing unit test expects, so if a stream ever failed after its result had already been delivered, the restarted call would deliver it again. I kept that design as it is, but I'm happy to make aggregations raise instead once they have yielded a result, if you'd prefer.

With the script from the issue (the real client against a local gRPC stand in, no mocks):

main (57366eef90)
stream(), error midway: AttributeError: '_UnaryStreamMultiCallable' object has no attribute '_retry'
count().stream(), error midway: AttributeError: '_UnaryStreamMultiCallable' object has no attribute '_retry'
stream(retry=None): AttributeError: 'NoneType' object has no attribute '_predicate'
count().stream(retry=None): AttributeError: 'NoneType' object has no attribute '_predicate'

this branch (d40eb1b522)
stream(), error midway: ['sf', 'la']
count().stream(), error midway: [[2]]
stream(retry=None): ServiceUnavailable: 503 backend restarted mid stream
count().stream(retry=None): ServiceUnavailable: 503 backend restarted mid stream

The unit test helpers _query_stream_w_retriable_exc_helper and _aggregation_query_stream_w_retriable_exc_helper mocked run_query as a callable that carries _retry, which the real transport never does, and the truthy mock returned by the predicate made every error look retryable. They now model the real transport: a raw stub without _retry, plus _wrapped_methods holding a real Retry. New test_query_stream_w_retriable_exc_w_retry_none and test_aggregation_query_stream_w_retriable_exc_w_retry_none cover retry=None. On main, 6 of these tests fail (the defaults, read time and retry=None cases for both); with this change tests/unit/v1/test_query.py and tests/unit/v1/test_aggregation.py pass (118 passed), and the noxfile's ruff==0.14.14 import and format checks and flake8 report nothing.

• Issue opened for this bug (firestore: streams raise AttributeError instead of restarting after a retryable error midway #18768)
• Tests and linter pass
• Code coverage does not decrease

[gemini-code-assist]

Code Review

This pull request updates the retry logic in aggregation.py and query.py to retrieve the default retry policy from the wrapped methods (_wrapped_methods) on the transport object, aligning with the real gRPC transport behavior. It also adds handling for when retry is explicitly set to None to disable retries, and includes corresponding unit tests to verify this behavior. There are no review comments, so I have no feedback to provide.