github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/io/gcp/datastore/v1new/helper.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """
    19  Cloud Datastore client and test functions.
    20  
    21  For internal use only; no backwards-compatibility guarantees.
    22  """
    23  
    24  # pytype: skip-file
    25  
    26  import os
    27  import uuid
    28  from typing import List
    29  from typing import Union
    30  
    31  from google.api_core import exceptions
    32  from google.api_core.gapic_v1 import client_info
    33  from google.cloud import environment_vars
    34  from google.cloud.datastore import __version__
    35  from google.cloud.datastore import client
    36  
    37  from apache_beam.io.gcp.datastore.v1new import types
    38  from apache_beam.version import __version__ as beam_version
    39  from cachetools.func import ttl_cache
    40  
    41  # https://cloud.google.com/datastore/docs/concepts/errors#error_codes
    42  _RETRYABLE_DATASTORE_ERRORS = (
    43      exceptions.Aborted,
    44      exceptions.DeadlineExceeded,
    45      exceptions.InternalServerError,
    46      exceptions.ServiceUnavailable,
    47  )
    48  
    49  
    50  @ttl_cache(maxsize=128, ttl=3600)
    51  def get_client(project, namespace):
    52    """Returns a Cloud Datastore client."""
    53    _client_info = client_info.ClientInfo(
    54        client_library_version=__version__,
    55        gapic_version=__version__,
    56        user_agent=f'beam-python-sdk/{beam_version}')
    57    _client = client.Client(
    58        project=project, namespace=namespace, client_info=_client_info)
    59    # Avoid overwriting user setting. BEAM-7608
    60    if not os.environ.get(environment_vars.GCD_HOST, None):
    61      _client.base_url = 'https://batch-datastore.googleapis.com'  # BEAM-1387
    62    return _client
    63  
    64  
    65  def retry_on_rpc_error(exception):
    66    """A retry filter for Cloud Datastore RPCErrors."""
    67    return isinstance(exception, _RETRYABLE_DATASTORE_ERRORS)
    68  
    69  
    70  def create_entities(count, id_or_name=False):
    71    """Creates a list of entities with random keys."""
    72    if id_or_name:
    73      ids_or_names = [
    74          uuid.uuid4().int & ((1 << 63) - 1) for _ in range(count)
    75      ]  # type: List[Union[str, int]]
    76    else:
    77      ids_or_names = [str(uuid.uuid4()) for _ in range(count)]
    78  
    79    keys = [types.Key(['EntityKind', x], project='project') for x in ids_or_names]
    80    return [types.Entity(key) for key in keys]
    81  
    82  
    83  def create_client_entities(count, id_or_name=False):
    84    """Creates a list of client-style entities with random keys."""
    85    return [
    86        entity.to_client_entity()
    87        for entity in create_entities(count, id_or_name=id_or_name)
    88    ]