Source code for apache_beam.io.gcp.datastore.v1new.helper

#
# 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.api_core.gapic_v1 import client_info
from google.cloud import environment_vars
from google.cloud.datastore import __version__
from google.cloud.datastore import client

from apache_beam.io.gcp.datastore.v1new import types
from apache_beam.version import __version__ as beam_version
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_info = client_info.ClientInfo( client_library_version=__version__, gapic_version=__version__, user_agent=f'beam-python-sdk/{beam_version}') _client = client.Client( project=project, namespace=namespace, client_info=_client_info) # 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: List[Union[str, int]] = [ uuid.uuid4().int & ((1 << 63) - 1) for _ in range(count) ] 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) ]