queues/venv/lib/python3.11/site-packages/pymongo/asynchronous/change_stream.py
Egor Matveev 6c6a549aff
All checks were successful
Deploy Prod / Build (pull_request) Successful in 9s
Deploy Prod / Push (pull_request) Successful in 12s
Deploy Prod / Deploy prod (pull_request) Successful in 10s
fix
2024-12-28 22:48:16 +03:00

499 lines
19 KiB
Python

# Copyright 2017 MongoDB, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you
# may not use this file except in compliance with the License. You
# may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.
"""Watch changes on a collection, a database, or the entire cluster."""
from __future__ import annotations
import copy
from typing import TYPE_CHECKING, Any, Generic, Mapping, Optional, Type, Union
from bson import CodecOptions, _bson_to_dict
from bson.raw_bson import RawBSONDocument
from bson.timestamp import Timestamp
from pymongo import _csot, common
from pymongo.asynchronous.aggregation import (
_AggregationCommand,
_CollectionAggregationCommand,
_DatabaseAggregationCommand,
)
from pymongo.asynchronous.command_cursor import AsyncCommandCursor
from pymongo.collation import validate_collation_or_none
from pymongo.errors import (
ConnectionFailure,
CursorNotFound,
InvalidOperation,
OperationFailure,
PyMongoError,
)
from pymongo.operations import _Op
from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
_IS_SYNC = False
# The change streams spec considers the following server errors from the
# getMore command non-resumable. All other getMore errors are resumable.
_RESUMABLE_GETMORE_ERRORS = frozenset(
[
6, # HostUnreachable
7, # HostNotFound
89, # NetworkTimeout
91, # ShutdownInProgress
189, # PrimarySteppedDown
262, # ExceededTimeLimit
9001, # SocketException
10107, # NotWritablePrimary
11600, # InterruptedAtShutdown
11602, # InterruptedDueToReplStateChange
13435, # NotPrimaryNoSecondaryOk
13436, # NotPrimaryOrSecondary
63, # StaleShardVersion
150, # StaleEpoch
13388, # StaleConfig
234, # RetryChangeStream
133, # FailedToSatisfyReadPreference
]
)
if TYPE_CHECKING:
from pymongo.asynchronous.client_session import AsyncClientSession
from pymongo.asynchronous.collection import AsyncCollection
from pymongo.asynchronous.database import AsyncDatabase
from pymongo.asynchronous.mongo_client import AsyncMongoClient
from pymongo.asynchronous.pool import AsyncConnection
def _resumable(exc: PyMongoError) -> bool:
"""Return True if given a resumable change stream error."""
if isinstance(exc, (ConnectionFailure, CursorNotFound)):
return True
if isinstance(exc, OperationFailure):
if exc._max_wire_version is None:
return False
return (
exc._max_wire_version >= 9 and exc.has_error_label("ResumableChangeStreamError")
) or (exc._max_wire_version < 9 and exc.code in _RESUMABLE_GETMORE_ERRORS)
return False
class AsyncChangeStream(Generic[_DocumentType]):
"""The internal abstract base class for change stream cursors.
Should not be called directly by application developers. Use
:meth:`pymongo.asynchronous.collection.AsyncCollection.watch`,
:meth:`pymongo.asynchronous.database.AsyncDatabase.watch`, or
:meth:`pymongo.asynchronous.mongo_client.AsyncMongoClient.watch` instead.
.. versionadded:: 3.6
.. seealso:: The MongoDB documentation on `changeStreams <https://mongodb.com/docs/manual/changeStreams/>`_.
"""
def __init__(
self,
target: Union[
AsyncMongoClient[_DocumentType],
AsyncDatabase[_DocumentType],
AsyncCollection[_DocumentType],
],
pipeline: Optional[_Pipeline],
full_document: Optional[str],
resume_after: Optional[Mapping[str, Any]],
max_await_time_ms: Optional[int],
batch_size: Optional[int],
collation: Optional[_CollationIn],
start_at_operation_time: Optional[Timestamp],
session: Optional[AsyncClientSession],
start_after: Optional[Mapping[str, Any]],
comment: Optional[Any] = None,
full_document_before_change: Optional[str] = None,
show_expanded_events: Optional[bool] = None,
) -> None:
if pipeline is None:
pipeline = []
pipeline = common.validate_list("pipeline", pipeline)
common.validate_string_or_none("full_document", full_document)
validate_collation_or_none(collation)
common.validate_non_negative_integer_or_none("batchSize", batch_size)
self._decode_custom = False
self._orig_codec_options: CodecOptions[_DocumentType] = target.codec_options
if target.codec_options.type_registry._decoder_map:
self._decode_custom = True
# Keep the type registry so that we support encoding custom types
# in the pipeline.
self._target = target.with_options( # type: ignore
codec_options=target.codec_options.with_options(document_class=RawBSONDocument)
)
else:
self._target = target
self._pipeline = copy.deepcopy(pipeline)
self._full_document = full_document
self._full_document_before_change = full_document_before_change
self._uses_start_after = start_after is not None
self._uses_resume_after = resume_after is not None
self._resume_token = copy.deepcopy(start_after or resume_after)
self._max_await_time_ms = max_await_time_ms
self._batch_size = batch_size
self._collation = collation
self._start_at_operation_time = start_at_operation_time
self._session = session
self._comment = comment
self._closed = False
self._timeout = self._target._timeout
self._show_expanded_events = show_expanded_events
async def _initialize_cursor(self) -> None:
# Initialize cursor.
self._cursor = await self._create_cursor()
@property
def _aggregation_command_class(self) -> Type[_AggregationCommand]:
"""The aggregation command class to be used."""
raise NotImplementedError
@property
def _client(self) -> AsyncMongoClient:
"""The client against which the aggregation commands for
this AsyncChangeStream will be run.
"""
raise NotImplementedError
def _change_stream_options(self) -> dict[str, Any]:
"""Return the options dict for the $changeStream pipeline stage."""
options: dict[str, Any] = {}
if self._full_document is not None:
options["fullDocument"] = self._full_document
if self._full_document_before_change is not None:
options["fullDocumentBeforeChange"] = self._full_document_before_change
resume_token = self.resume_token
if resume_token is not None:
if self._uses_start_after:
options["startAfter"] = resume_token
else:
options["resumeAfter"] = resume_token
elif self._start_at_operation_time is not None:
options["startAtOperationTime"] = self._start_at_operation_time
if self._show_expanded_events:
options["showExpandedEvents"] = self._show_expanded_events
return options
def _command_options(self) -> dict[str, Any]:
"""Return the options dict for the aggregation command."""
options = {}
if self._max_await_time_ms is not None:
options["maxAwaitTimeMS"] = self._max_await_time_ms
if self._batch_size is not None:
options["batchSize"] = self._batch_size
return options
def _aggregation_pipeline(self) -> list[dict[str, Any]]:
"""Return the full aggregation pipeline for this AsyncChangeStream."""
options = self._change_stream_options()
full_pipeline: list = [{"$changeStream": options}]
full_pipeline.extend(self._pipeline)
return full_pipeline
def _process_result(self, result: Mapping[str, Any], conn: AsyncConnection) -> None:
"""Callback that caches the postBatchResumeToken or
startAtOperationTime from a changeStream aggregate command response
containing an empty batch of change documents.
This is implemented as a callback because we need access to the wire
version in order to determine whether to cache this value.
"""
if not result["cursor"]["firstBatch"]:
if "postBatchResumeToken" in result["cursor"]:
self._resume_token = result["cursor"]["postBatchResumeToken"]
elif (
self._start_at_operation_time is None
and self._uses_resume_after is False
and self._uses_start_after is False
and conn.max_wire_version >= 7
):
self._start_at_operation_time = result.get("operationTime")
# PYTHON-2181: informative error on missing operationTime.
if self._start_at_operation_time is None:
raise OperationFailure(
"Expected field 'operationTime' missing from command "
f"response : {result!r}"
)
async def _run_aggregation_cmd(
self, session: Optional[AsyncClientSession], explicit_session: bool
) -> AsyncCommandCursor:
"""Run the full aggregation pipeline for this AsyncChangeStream and return
the corresponding AsyncCommandCursor.
"""
cmd = self._aggregation_command_class(
self._target,
AsyncCommandCursor,
self._aggregation_pipeline(),
self._command_options(),
explicit_session,
result_processor=self._process_result,
comment=self._comment,
)
return await self._client._retryable_read(
cmd.get_cursor,
self._target._read_preference_for(session),
session,
operation=_Op.AGGREGATE,
)
async def _create_cursor(self) -> AsyncCommandCursor:
async with self._client._tmp_session(self._session, close=False) as s:
return await self._run_aggregation_cmd(
session=s, explicit_session=self._session is not None
)
async def _resume(self) -> None:
"""Reestablish this change stream after a resumable error."""
try:
await self._cursor.close()
except PyMongoError:
pass
self._cursor = await self._create_cursor()
async def close(self) -> None:
"""Close this AsyncChangeStream."""
self._closed = True
await self._cursor.close()
def __aiter__(self) -> AsyncChangeStream[_DocumentType]:
return self
@property
def resume_token(self) -> Optional[Mapping[str, Any]]:
"""The cached resume token that will be used to resume after the most
recently returned change.
.. versionadded:: 3.9
"""
return copy.deepcopy(self._resume_token)
@_csot.apply
async def next(self) -> _DocumentType:
"""Advance the cursor.
This method blocks until the next change document is returned or an
unrecoverable error is raised. This method is used when iterating over
all changes in the cursor. For example::
try:
resume_token = None
pipeline = [{'$match': {'operationType': 'insert'}}]
async with await db.collection.watch(pipeline) as stream:
async for insert_change in stream:
print(insert_change)
resume_token = stream.resume_token
except pymongo.errors.PyMongoError:
# The AsyncChangeStream encountered an unrecoverable error or the
# resume attempt failed to recreate the cursor.
if resume_token is None:
# There is no usable resume token because there was a
# failure during AsyncChangeStream initialization.
logging.error('...')
else:
# Use the interrupted AsyncChangeStream's resume token to create
# a new AsyncChangeStream. The new stream will continue from the
# last seen insert change without missing any events.
async with await db.collection.watch(
pipeline, resume_after=resume_token) as stream:
async for insert_change in stream:
print(insert_change)
Raises :exc:`StopIteration` if this AsyncChangeStream is closed.
"""
while self.alive:
doc = await self.try_next()
if doc is not None:
return doc
raise StopAsyncIteration
__anext__ = next
@property
def alive(self) -> bool:
"""Does this cursor have the potential to return more data?
.. note:: Even if :attr:`alive` is ``True``, :meth:`next` can raise
:exc:`StopIteration` and :meth:`try_next` can return ``None``.
.. versionadded:: 3.8
"""
return not self._closed
@_csot.apply
async def try_next(self) -> Optional[_DocumentType]:
"""Advance the cursor without blocking indefinitely.
This method returns the next change document without waiting
indefinitely for the next change. For example::
async with await db.collection.watch() as stream:
while stream.alive:
change = await stream.try_next()
# Note that the AsyncChangeStream's resume token may be updated
# even when no changes are returned.
print("Current resume token: %r" % (stream.resume_token,))
if change is not None:
print("Change document: %r" % (change,))
continue
# We end up here when there are no recent changes.
# Sleep for a while before trying again to avoid flooding
# the server with getMore requests when no changes are
# available.
asyncio.sleep(10)
If no change document is cached locally then this method runs a single
getMore command. If the getMore yields any documents, the next
document is returned, otherwise, if the getMore returns no documents
(because there have been no changes) then ``None`` is returned.
:return: The next change document or ``None`` when no document is available
after running a single getMore or when the cursor is closed.
.. versionadded:: 3.8
"""
if not self._closed and not self._cursor.alive:
await self._resume()
# Attempt to get the next change with at most one getMore and at most
# one resume attempt.
try:
try:
change = await self._cursor._try_next(True)
except PyMongoError as exc:
if not _resumable(exc):
raise
await self._resume()
change = await self._cursor._try_next(False)
except PyMongoError as exc:
# Close the stream after a fatal error.
if not _resumable(exc) and not exc.timeout:
await self.close()
raise
except Exception:
await self.close()
raise
# Check if the cursor was invalidated.
if not self._cursor.alive:
self._closed = True
# If no changes are available.
if change is None:
# We have either iterated over all documents in the cursor,
# OR the most-recently returned batch is empty. In either case,
# update the cached resume token with the postBatchResumeToken if
# one was returned. We also clear the startAtOperationTime.
if self._cursor._post_batch_resume_token is not None:
self._resume_token = self._cursor._post_batch_resume_token
self._start_at_operation_time = None
return change
# Else, changes are available.
try:
resume_token = change["_id"]
except KeyError:
await self.close()
raise InvalidOperation(
"Cannot provide resume functionality when the resume token is missing."
) from None
# If this is the last change document from the current batch, cache the
# postBatchResumeToken.
if not self._cursor._has_next() and self._cursor._post_batch_resume_token:
resume_token = self._cursor._post_batch_resume_token
# Hereafter, don't use startAfter; instead use resumeAfter.
self._uses_start_after = False
self._uses_resume_after = True
# Cache the resume token and clear startAtOperationTime.
self._resume_token = resume_token
self._start_at_operation_time = None
if self._decode_custom:
return _bson_to_dict(change.raw, self._orig_codec_options)
return change
async def __aenter__(self) -> AsyncChangeStream[_DocumentType]:
return self
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
await self.close()
class AsyncCollectionChangeStream(AsyncChangeStream[_DocumentType]):
"""A change stream that watches changes on a single collection.
Should not be called directly by application developers. Use
helper method :meth:`pymongo.asynchronous.collection.AsyncCollection.watch` instead.
.. versionadded:: 3.7
"""
_target: AsyncCollection[_DocumentType]
@property
def _aggregation_command_class(self) -> Type[_CollectionAggregationCommand]:
return _CollectionAggregationCommand
@property
def _client(self) -> AsyncMongoClient[_DocumentType]:
return self._target.database.client
class AsyncDatabaseChangeStream(AsyncChangeStream[_DocumentType]):
"""A change stream that watches changes on all collections in a database.
Should not be called directly by application developers. Use
helper method :meth:`pymongo.asynchronous.database.AsyncDatabase.watch` instead.
.. versionadded:: 3.7
"""
_target: AsyncDatabase[_DocumentType]
@property
def _aggregation_command_class(self) -> Type[_DatabaseAggregationCommand]:
return _DatabaseAggregationCommand
@property
def _client(self) -> AsyncMongoClient[_DocumentType]:
return self._target.client
class AsyncClusterChangeStream(AsyncDatabaseChangeStream[_DocumentType]):
"""A change stream that watches changes on all collections in the cluster.
Should not be called directly by application developers. Use
helper method :meth:`pymongo.asynchronous.mongo_client.AsyncMongoClient.watch` instead.
.. versionadded:: 3.7
"""
def _change_stream_options(self) -> dict[str, Any]:
options = super()._change_stream_options()
options["allChangesForCluster"] = True
return options