#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#
"""
Cloud Datastore client and test functions.
For internal use only; no backwards-compatibility guarantees.
"""
# pytype: skip-file
import os
import uuid
from typing import List
from typing import Union
from google.api_core import exceptions
from google.cloud import environment_vars
from google.cloud.datastore import client
from apache_beam.io.gcp.datastore.v1new import types
from cachetools.func import ttl_cache
# https://cloud.google.com/datastore/docs/concepts/errors#error_codes
_RETRYABLE_DATASTORE_ERRORS = (
exceptions.Aborted,
exceptions.DeadlineExceeded,
exceptions.InternalServerError,
exceptions.ServiceUnavailable,
)
[docs]@ttl_cache(maxsize=128, ttl=3600)
def get_client(project, namespace):
"""Returns a Cloud Datastore client."""
_client = client.Client(project=project, namespace=namespace)
# Avoid overwriting user setting. BEAM-7608
if not os.environ.get(environment_vars.GCD_HOST, None):
_client.base_url = 'https://batch-datastore.googleapis.com' # BEAM-1387
return _client
[docs]def retry_on_rpc_error(exception):
"""A retry filter for Cloud Datastore RPCErrors."""
return isinstance(exception, _RETRYABLE_DATASTORE_ERRORS)
[docs]def create_entities(count, id_or_name=False):
"""Creates a list of entities with random keys."""
if id_or_name:
ids_or_names = [
uuid.uuid4().int & ((1 << 63) - 1) for _ in range(count)
] # type: List[Union[str, int]]
else:
ids_or_names = [str(uuid.uuid4()) for _ in range(count)]
keys = [types.Key(['EntityKind', x], project='project') for x in ids_or_names]
return [types.Entity(key) for key in keys]
[docs]def create_client_entities(count, id_or_name=False):
"""Creates a list of client-style entities with random keys."""
return [
entity.to_client_entity()
for entity in create_entities(count, id_or_name=id_or_name)
]