|
| 1 | +# Copyright 2014 Google Inc. All rights reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""Methods for interacting with Google Cloud Datastore. |
| 16 | +
|
| 17 | +Allows interacting with the datastore via user-friendly Key, Entity and |
| 18 | +Query objects rather than via protobufs. |
| 19 | +""" |
| 20 | + |
| 21 | +from gcloud.datastore import _implicit_environ |
| 22 | +from gcloud.datastore import helpers |
| 23 | + |
| 24 | + |
| 25 | +def _require_dataset_id(dataset_id=None): |
| 26 | + """Infer a dataset ID from the environment, if not passed explicitly. |
| 27 | +
|
| 28 | + :type dataset_id: string |
| 29 | + :param dataset_id: Optional. |
| 30 | +
|
| 31 | + :rtype: string |
| 32 | + :returns: A dataset ID based on the current environment. |
| 33 | + :raises: :class:`EnvironmentError` if ``dataset_id`` is ``None``, |
| 34 | + and cannot be inferred from the environment. |
| 35 | + """ |
| 36 | + if dataset_id is None: |
| 37 | + if _implicit_environ.DATASET_ID is None: |
| 38 | + raise EnvironmentError('Dataset ID could not be inferred.') |
| 39 | + dataset_id = _implicit_environ.DATASET_ID |
| 40 | + return dataset_id |
| 41 | + |
| 42 | + |
| 43 | +def _require_connection(connection=None): |
| 44 | + """Infer a connection from the environment, if not passed explicitly. |
| 45 | +
|
| 46 | + :type connection: :class:`gcloud.datastore.connection.Connection` |
| 47 | + :param connection: Optional. |
| 48 | +
|
| 49 | + :rtype: :class:`gcloud.datastore.connection.Connection` |
| 50 | + :returns: A connection based on the current environment. |
| 51 | + :raises: :class:`EnvironmentError` if ``connection`` is ``None``, and |
| 52 | + cannot be inferred from the environment. |
| 53 | + """ |
| 54 | + if connection is None: |
| 55 | + if _implicit_environ.CONNECTION is None: |
| 56 | + raise EnvironmentError('Connection could not be inferred.') |
| 57 | + connection = _implicit_environ.CONNECTION |
| 58 | + return connection |
| 59 | + |
| 60 | + |
| 61 | +def get_entities(keys, missing=None, deferred=None, connection=None): |
| 62 | + """Retrieves entities, along with their attributes. |
| 63 | +
|
| 64 | + :type keys: list of :class:`gcloud.datastore.key.Key` |
| 65 | + :param keys: The name of the item to retrieve. |
| 66 | +
|
| 67 | + :type missing: an empty list or None. |
| 68 | + :param missing: If a list is passed, the key-only entities returned |
| 69 | + by the backend as "missing" will be copied into it. |
| 70 | + Use only as a keyword param. |
| 71 | +
|
| 72 | + :type deferred: an empty list or None. |
| 73 | + :param deferred: If a list is passed, the keys returned |
| 74 | + by the backend as "deferred" will be copied into it. |
| 75 | + Use only as a keyword param. |
| 76 | +
|
| 77 | + :type connection: :class:`gcloud.datastore.connection.Connection` |
| 78 | + :param connection: Optional. The connection used to connect to datastore. |
| 79 | +
|
| 80 | + :rtype: list of :class:`gcloud.datastore.entity.Entity` |
| 81 | + :returns: The requested entities. |
| 82 | + :raises: :class:`ValueError` if the key dataset IDs don't agree. |
| 83 | + """ |
| 84 | + if not keys: |
| 85 | + return [] |
| 86 | + |
| 87 | + connection = _require_connection(connection) |
| 88 | + dataset_id = keys[0].dataset_id |
| 89 | + # Rather than creating a list or set of all dataset IDs, we iterate |
| 90 | + # and check. We could allow the backend to check this for us if IDs |
| 91 | + # with no prefix worked (GoogleCloudPlatform/google-cloud-datastore#59) |
| 92 | + # or if we made sure that a prefix s~ or e~ was on each key. |
| 93 | + for key in keys[1:]: |
| 94 | + if key.dataset_id != dataset_id: |
| 95 | + raise ValueError('All keys in get_entities must be from the ' |
| 96 | + 'same dataset.') |
| 97 | + |
| 98 | + entity_pbs = connection.lookup( |
| 99 | + dataset_id=dataset_id, |
| 100 | + key_pbs=[k.to_protobuf() for k in keys], |
| 101 | + missing=missing, deferred=deferred, |
| 102 | + ) |
| 103 | + |
| 104 | + if missing is not None: |
| 105 | + missing[:] = [ |
| 106 | + helpers.entity_from_protobuf(missed_pb) |
| 107 | + for missed_pb in missing] |
| 108 | + |
| 109 | + if deferred is not None: |
| 110 | + deferred[:] = [ |
| 111 | + helpers.key_from_protobuf(deferred_pb) |
| 112 | + for deferred_pb in deferred] |
| 113 | + |
| 114 | + entities = [] |
| 115 | + for entity_pb in entity_pbs: |
| 116 | + entities.append(helpers.entity_from_protobuf(entity_pb)) |
| 117 | + |
| 118 | + return entities |
| 119 | + |
| 120 | + |
| 121 | +def allocate_ids(incomplete_key, num_ids, connection=None): |
| 122 | + """Allocates a list of IDs from a partial key. |
| 123 | +
|
| 124 | + :type incomplete_key: A :class:`gcloud.datastore.key.Key` |
| 125 | + :param incomplete_key: Partial key to use as base for allocated IDs. |
| 126 | +
|
| 127 | + :type num_ids: integer |
| 128 | + :param num_ids: The number of IDs to allocate. |
| 129 | +
|
| 130 | + :type connection: :class:`gcloud.datastore.connection.Connection` |
| 131 | + :param connection: Optional. The connection used to connect to datastore. |
| 132 | +
|
| 133 | + :rtype: list of :class:`gcloud.datastore.key.Key` |
| 134 | + :returns: The (complete) keys allocated with ``incomplete_key`` as root. |
| 135 | + :raises: :class:`ValueError` if ``incomplete_key`` is not a partial key. |
| 136 | + """ |
| 137 | + connection = _require_connection(connection) |
| 138 | + |
| 139 | + if not incomplete_key.is_partial: |
| 140 | + raise ValueError(('Key is not partial.', incomplete_key)) |
| 141 | + |
| 142 | + incomplete_key_pb = incomplete_key.to_protobuf() |
| 143 | + incomplete_key_pbs = [incomplete_key_pb] * num_ids |
| 144 | + |
| 145 | + allocated_key_pbs = connection.allocate_ids(incomplete_key.dataset_id, |
| 146 | + incomplete_key_pbs) |
| 147 | + allocated_ids = [allocated_key_pb.path_element[-1].id |
| 148 | + for allocated_key_pb in allocated_key_pbs] |
| 149 | + return [incomplete_key.completed_key(allocated_id) |
| 150 | + for allocated_id in allocated_ids] |
0 commit comments