def process(self, query, *unused_args, **unused_kwargs):
if query.namespace is None:
query.namespace = ''
_client = helper.get_client(query.project, query.namespace)
client_query = query._to_client_query(_client)
# Create request count metric
resource = resource_identifiers.DatastoreNamespace(
query.project, query.namespace)
labels = {
monitoring_infos.SERVICE_LABEL: 'Datastore',
monitoring_infos.METHOD_LABEL: 'BatchDatastoreRead',
monitoring_infos.RESOURCE_LABEL: resource,
monitoring_infos.DATASTORE_NAMESPACE_LABEL: query.namespace,
monitoring_infos.DATASTORE_PROJECT_ID_LABEL: query.project,
monitoring_infos.STATUS_LABEL: 'ok'
}
service_call_metric = ServiceCallMetric(
request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN,
base_labels=labels)
try:
for client_entity in client_query.fetch(query.limit):
yield types.Entity.from_client_entity(client_entity)
service_call_metric.call('ok')
except (ClientError, GoogleAPICallError) as e:
# e.code.value contains the numeric http status code.
service_call_metric.call(e.code.value)
except HttpError as e:
service_call_metric.call(e)
def verify_read_call_metric(self, project_id, namespace, status, count):
"""Check if a metric was recorded for the Datastore IO read API call."""
process_wide_monitoring_infos = list(
MetricsEnvironment.process_wide_container(
).to_runner_api_monitoring_infos(None).values())
resource = resource_identifiers.DatastoreNamespace(
project_id, namespace)
labels = {
monitoring_infos.SERVICE_LABEL: 'Datastore',
monitoring_infos.METHOD_LABEL: 'BatchDatastoreRead',
monitoring_infos.RESOURCE_LABEL: resource,
monitoring_infos.DATASTORE_NAMESPACE_LABEL: namespace,
monitoring_infos.DATASTORE_PROJECT_ID_LABEL: project_id,
monitoring_infos.STATUS_LABEL: status
}
expected_mi = monitoring_infos.int64_counter(
monitoring_infos.API_REQUEST_COUNT_URN, count, labels=labels)
expected_mi.ClearField("start_time")
found = False
for actual_mi in process_wide_monitoring_infos:
actual_mi.ClearField("start_time")
if expected_mi == actual_mi:
found = True
break
self.assertTrue(
found, "Did not find read call metric with status: %s" % status)
def write_mutations(self,
throttler,
rpc_stats_callback,
throttle_delay=1):
"""Writes a batch of mutations to Cloud Datastore.
If a commit fails, it will be retried up to 5 times. All mutations in the
batch will be committed again, even if the commit was partially
successful. If the retry limit is exceeded, the last exception from
Cloud Datastore will be raised.
Assumes that the Datastore client library does not perform any retries on
commits. It has not been determined how such retries would interact with
the retries and throttler used here.
See ``google.cloud.datastore_v1.gapic.datastore_client_config`` for
retry config.
Args:
rpc_stats_callback: a function to call with arguments `successes` and
`failures` and `throttled_secs`; this is called to record successful
and failed RPCs to Datastore and time spent waiting for throttling.
throttler: (``apache_beam.io.gcp.datastore.v1new.adaptive_throttler.
AdaptiveThrottler``)
Throttler instance used to select requests to be throttled.
throttle_delay: (:class:`float`) time in seconds to sleep when
throttled.
Returns:
(int) The latency of the successful RPC in milliseconds.
"""
# Client-side throttling.
while throttler.throttle_request(time.time() * 1000):
_LOGGER.info(
"Delaying request for %ds due to previous failures",
throttle_delay)
time.sleep(throttle_delay)
rpc_stats_callback(throttled_secs=throttle_delay)
if self._batch is None:
# this will only happen when we re-try previously failed batch
self._batch = self._client.batch()
self._batch.begin()
for element in self._batch_elements:
self.add_to_batch(element)
# Create request count metric
resource = resource_identifiers.DatastoreNamespace(
self._project, "")
labels = {
monitoring_infos.SERVICE_LABEL: 'Datastore',
monitoring_infos.METHOD_LABEL: 'BatchDatastoreWrite',
monitoring_infos.RESOURCE_LABEL: resource,
monitoring_infos.DATASTORE_NAMESPACE_LABEL: "",
monitoring_infos.DATASTORE_PROJECT_ID_LABEL: self._project,
monitoring_infos.STATUS_LABEL: 'ok'
}
service_call_metric = ServiceCallMetric(
request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN,
base_labels=labels)
try:
start_time = time.time()
self._batch.commit()
end_time = time.time()
service_call_metric.call('ok')
rpc_stats_callback(successes=1)
throttler.successful_request(start_time * 1000)
commit_time_ms = int((end_time - start_time) * 1000)
return commit_time_ms
except (ClientError, GoogleAPICallError) as e:
self._batch = None
# e.code.value contains the numeric http status code.
service_call_metric.call(e.code.value)
rpc_stats_callback(errors=1)
raise
except HttpError as e:
service_call_metric.call(e)
rpc_stats_callback(errors=1)
raise