def _handle_response(self, response: SubscribeResponse):
if "messages" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid subsequent response on the subscribe stream."
)
)
return
# Workaround for incredibly slow proto-plus-python accesses
messages = list(response.messages.messages._pb)
self._outstanding_flow_control.on_messages(messages)
for message in messages:
if (
self._last_received_offset is not None
and message.cursor.offset <= self._last_received_offset
):
self._connection.fail(
FailedPrecondition(
"Received an invalid out of order message from the server. Message is {}, previous last received is {}.".format(
message.cursor.offset, self._last_received_offset
)
)
)
return
self._last_received_offset = message.cursor.offset
# queue is unbounded.
self._message_queue.put_nowait(messages)
async def reinitialize(self, connection: Connection[SubscribeRequest,
SubscribeResponse]):
self._reinitializing = True
await self._stop_loopers()
await connection.write(SubscribeRequest(initial=self._initial))
response = await connection.read()
if "initial" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the subscribe stream."
))
return
if self._last_received_offset is not None:
# Perform a seek to get the next message after the one we received.
await connection.write(
SubscribeRequest(seek=SeekRequest(cursor=Cursor(
offset=self._last_received_offset + 1))))
seek_response = await connection.read()
if "seek" not in seek_response:
self._connection.fail(
FailedPrecondition(
"Received an invalid seek response on the subscribe stream."
))
return
tokens = self._outstanding_flow_control.request_for_restart()
if tokens is not None:
await connection.write(SubscribeRequest(flow_control=tokens))
self._reinitializing = False
self._start_loopers()
async def read(self) -> Response:
if self._response_it is None:
self.fail(FailedPrecondition("GapicConnection not initialized."))
raise self.error()
try:
response_it = cast(AsyncIterator[Response], self._response_it)
return await self.await_unless_failed(response_it.__anext__())
except StopAsyncIteration:
self.fail(FailedPrecondition("Server sent unprompted half close."))
except GoogleAPICallError as e:
self.fail(e)
raise self.error()
def _handle_response(self, response: StreamingCommitCursorResponse):
if "commit" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid subsequent response on the commit stream."
))
if response.commit.acknowledged_commits > len(
self._outstanding_commits):
self._connection.fail(
FailedPrecondition(
"Received a commit response on the stream with no outstanding commits."
))
for _ in range(response.commit.acknowledged_commits):
batch = self._outstanding_commits.pop(0)
for item in batch:
item.response_future.set_result(None)
async def reinitialize(
self, connection: Connection[SubscribeRequest, SubscribeResponse]
):
initial = deepcopy(self._base_initial)
if self._last_received_offset is not None:
initial.initial_location = SeekRequest(
cursor=Cursor(offset=self._last_received_offset + 1)
)
else:
initial.initial_location = SeekRequest(
named_target=SeekRequest.NamedTarget.COMMITTED_CURSOR
)
await connection.write(SubscribeRequest(initial=initial))
response = await connection.read()
if "initial" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the subscribe stream."
)
)
return
tokens = self._outstanding_flow_control.request_for_restart()
if tokens is not None:
await connection.write(SubscribeRequest(flow_control=tokens))
self._start_loopers()
async def reinitialize(
self,
connection: Connection[StreamingCommitCursorRequest,
StreamingCommitCursorResponse],
last_error: Optional[GoogleAPICallError],
):
await self._stop_loopers()
await connection.write(
StreamingCommitCursorRequest(initial=self._initial))
response = await connection.read()
if "initial" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the publish stream."
))
if self._outstanding_commits:
# Roll up outstanding commits
rollup: List[WorkItem[Cursor, None]] = []
for batch in self._outstanding_commits:
for item in batch:
rollup.append(item)
self._outstanding_commits = [rollup]
req = StreamingCommitCursorRequest()
req.commit.cursor = rollup[-1].request
await connection.write(req)
self._start_loopers()
async def test_ack_failure(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue)
async with subscriber:
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
read: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
read.ack()
await ack_called_queue.get()
ack_set_tracker.ack.assert_has_calls([call(1)])
await ack_result_queue.put(FailedPrecondition("Bad ack"))
async def sleep_forever():
await asyncio.sleep(float("inf"))
underlying.read.side_effect = sleep_forever
with pytest.raises(FailedPrecondition):
await subscriber.read()
async def test_iterator(
default_subscriber,
subscriber_factory,
multiplexed_client: AsyncSubscriberClientInterface,
):
read_queues = wire_queues(default_subscriber.read)
subscription = SubscriptionPath(1, CloudZone.parse("us-central1-a"), "abc")
message = Message(PubsubMessage(message_id="1")._pb, "", 0, None)
async with multiplexed_client:
iterator = await multiplexed_client.subscribe(
subscription, DISABLED_FLOW_CONTROL
)
subscriber_factory.assert_has_calls(
[call(subscription, None, DISABLED_FLOW_CONTROL)]
)
read_fut_1 = asyncio.ensure_future(iterator.__anext__())
assert not read_fut_1.done()
await read_queues.called.get()
default_subscriber.read.assert_has_calls([call()])
await read_queues.results.put(message)
assert await read_fut_1 is message
read_fut_2 = asyncio.ensure_future(iterator.__anext__())
assert not read_fut_2.done()
await read_queues.called.get()
default_subscriber.read.assert_has_calls([call(), call()])
await read_queues.results.put(FailedPrecondition(""))
with pytest.raises(FailedPrecondition):
await read_fut_2
default_subscriber.__aexit__.assert_called_once()
def _handle_response(self, response: PublishResponse):
if "message_response" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid subsequent response on the publish stream."
))
if not self._outstanding_writes:
self._connection.fail(
FailedPrecondition(
"Received an publish response on the stream with no outstanding publishes."
))
next_offset: Cursor = response.message_response.start_cursor.offset
batch: List[WorkItem[PubSubMessage]] = self._outstanding_writes.pop(0)
for item in batch:
item.response_future.set_result(Cursor(offset=next_offset))
next_offset += 1
def _handle_response(self, response: StreamingCommitCursorResponse):
if "commit" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid subsequent response on the commit stream."
))
if response.commit.acknowledged_commits > len(
self._outstanding_commits):
self._connection.fail(
FailedPrecondition(
"Received a commit response on the stream with no outstanding commits."
))
for _ in range(response.commit.acknowledged_commits):
self._outstanding_commits.pop(0)
if len(self._outstanding_commits) == 0:
self._empty.set()
async def test_assigner_failure(subscriber, assigner, subscriber_factory):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
await assign_queues.called.get()
await assign_queues.results.put(FailedPrecondition("bad assign"))
with pytest.raises(FailedPrecondition):
await subscriber.read()
def track(self, offset: int):
if len(self._receipts) > 0:
last = self._receipts[0]
if last >= offset:
raise FailedPrecondition(
f"Tried to track message {offset} which is before last tracked message {last}."
)
self._receipts.append(offset)
async def create_or_fail(self, key: _Key,
factory: _OpenedClientFactory) -> _Client:
async with self._lock:
if key in self._live_clients:
raise FailedPrecondition(
f"Cannot create two clients with the same key. {_Key}")
self._live_clients[key] = await factory()
return self._live_clients[key]
async def read(self) -> Response:
try:
return await self.await_unless_failed(
self._response_it.__anext__())
except StopAsyncIteration:
self.fail(FailedPrecondition("Server sent unprompted half close."))
except GoogleAPICallError as e:
self.fail(e)
raise self.error()
def test_access_destroyed_secret_version(env, secrets):
version = '2'
key = '[SUPER SECRET KEY]'
expected_response = FailedPrecondition('[projects/75936188733/secrets/vibe-cdc/versions/2] is in DESTROYED state')
crypto = Crypto(
env=env['env'],
key=key
)
client = secrets.get_client(expected_response)
secret = crypto.access_secret_version(version=version, client=client)
if secret is not None:
assert False
async def test_subscriber_failure(subscriber, assigner, subscriber_factory):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
await assign_queues.called.get()
sub1 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub1_queues = wire_queues(sub1.read)
subscriber_factory.return_value = sub1
await assign_queues.results.put({Partition(1)})
await sub1_queues.called.get()
await sub1_queues.results.put(FailedPrecondition("sub failed"))
with pytest.raises(FailedPrecondition):
await subscriber.read()
async def test_track_failure(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
async with subscriber:
ack_set_tracker.track.side_effect = FailedPrecondition("Bad track")
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
with pytest.raises(FailedPrecondition):
await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
def test_failed(subscriber: SubscriberImpl, async_subscriber, close_callback):
error = FailedPrecondition("bad read")
async_subscriber.read.side_effect = error
close_called = concurrent.futures.Future()
close_callback.side_effect = lambda manager, err: close_called.set_result(
None)
subscriber.add_close_callback(close_callback)
subscriber.__enter__()
async_subscriber.__aenter__.assert_called_once()
close_called.result()
async_subscriber.__aexit__.assert_called_once()
close_callback.assert_called_once_with(subscriber, error)
async def _receive_loop(self):
while True:
response = await self._connection.read()
if self._outstanding_assignment or not self._new_assignment.empty(
):
self._connection.fail(
FailedPrecondition(
"Received a duplicate assignment on the stream while one was outstanding."
))
return
self._outstanding_assignment = True
partitions = set()
for partition in response.partitions:
partitions.add(Partition(partition))
self._new_assignment.put_nowait(partitions)
async def _handle_queue_message(
self,
message: Union[requests.AckRequest, requests.DropRequest,
requests.ModAckRequest, requests.NackRequest, ],
):
if isinstance(message, requests.DropRequest) or isinstance(
message, requests.ModAckRequest):
self.fail(
FailedPrecondition(
"Called internal method of google.cloud.pubsub_v1.subscriber.message.Message "
f"Pub/Sub Lite does not support: {message}"))
elif isinstance(message, requests.AckRequest):
await self._handle_ack(message)
else:
self._handle_nack(message)
async def reinitialize(
self,
connection: Connection[StreamingCommitCursorRequest,
StreamingCommitCursorResponse],
):
await connection.write(
StreamingCommitCursorRequest(initial=self._initial))
response = await connection.read()
if "initial" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the publish stream."
))
if self._next_to_commit is None:
if self._outstanding_commits:
self._next_to_commit = self._outstanding_commits[-1]
self._outstanding_commits = []
self._start_loopers()
async def reinitialize(
self,
connection: Connection[PublishRequest, PublishResponse],
):
await connection.write(PublishRequest(initial_request=self._initial))
response = await connection.read()
if "initial_response" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the publish stream."
))
for batch in self._outstanding_writes:
aggregate = PublishRequest()
aggregate.message_publish_request.messages = [
item.request for item in batch
]
await connection.write(aggregate)
self._start_loopers()
async def test_nack_failure(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
nack_handler,
):
async with subscriber:
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
read: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
nack_handler.on_nack.side_effect = FailedPrecondition("Bad nack")
read.nack()
async def sleep_forever():
await asyncio.sleep(float("inf"))
underlying.read.side_effect = sleep_forever
with pytest.raises(FailedPrecondition):
await subscriber.read()
def create_task(self, queue, task):
task.name = task.name or "%s/tasks/%s" % (
queue, int(datetime.now().timestamp())
)
if task.app_engine_http_request.relative_uri:
# Set a default http_method
task.app_engine_http_request.http_method = (
task.app_engine_http_request.http_method or target_pb2.HttpMethod.Value("POST")
)
elif task.http_request.url:
# Set a default http_method
task.http_request.http_method = (
task.http_request.http_method or target_pb2.HttpMethod.Value("POST")
)
if queue not in self._queue_tasks:
raise FailedPrecondition("Queue does not exist.")
self._queue_tasks[queue].append(task)
logger.info("[TASKS] Created task %s", task.name)
return task
async def reinitialize(
self,
connection: Connection[SubscribeRequest, SubscribeResponse],
last_error: Optional[GoogleAPICallError],
):
self._reinitializing = True
await self._stop_loopers()
if last_error and is_reset_signal(last_error):
# Discard undelivered messages and refill flow control tokens.
while not self._message_queue.empty():
msg = self._message_queue.get_nowait()
self._outstanding_flow_control.add(
FlowControlRequest(
allowed_messages=1,
allowed_bytes=msg.size_bytes,
))
await self._reset_handler.handle_reset()
self._last_received_offset = None
initial = deepcopy(self._base_initial)
if self._last_received_offset is not None:
initial.initial_location = SeekRequest(cursor=Cursor(
offset=self._last_received_offset + 1))
else:
initial.initial_location = SeekRequest(
named_target=SeekRequest.NamedTarget.COMMITTED_CURSOR)
await connection.write(SubscribeRequest(initial=initial))
response = await connection.read()
if "initial" not in response:
self._connection.fail(
FailedPrecondition(
"Received an invalid initial response on the subscribe stream."
))
return
tokens = self._outstanding_flow_control.request_for_restart()
if tokens is not None:
await connection.write(SubscribeRequest(flow_control=tokens))
self._reinitializing = False
self._start_loopers()
def __enter__(self):
if self._started:
raise FailedPrecondition("__enter__ called twice.")
self._started = True
return self
def on_nack(self, message: PubsubMessage, ack: Callable[[], None]):
raise FailedPrecondition(
"You may not nack messages by default when using a PubSub Lite client. See NackHandler for how to customize"
" this.")
def require_started(self):
if not self._started:
raise FailedPrecondition("__enter__ has never been called.")
async def test_failed_transform(subscriber, underlying, transformer):
async with subscriber:
transformer.transform.side_effect = FailedPrecondition("Bad message")
underlying.read.return_value = SequencedMessage()
with pytest.raises(FailedPrecondition):
await subscriber.read()
