8000 feat: named database support (#398) · googleapis/python-datastore@ef4fce3 · GitHub
[go: up one dir, main page]
Skip to content

Commit ef4fce3

Browse files
rwhoggMariatta
rwhogg
authored and
Mariatta
committed
feat: named database support (#398)
* feat: Add named database support * test: Use named db in system tests
1 parent b31a944 commit ef4fce3

File tree

25 files changed

+667
-93
lines changed
  • conftest.py
  • test_put.py
  • utils
  • unit

    • 25 files changed

    +667
    -93
    lines changed

    ‎CONTRIBUTING.rst

    Lines changed: 1 addition & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -174,6 +174,7 @@ Running System Tests
    174174

    175175
    - You'll also need stored data in your dataset. To populate this data, run::
    176176

    177+
    $ export SYSTEM_TESTS_DATABASE=system-tests-named-db
    177178
    $ python tests/system/utils/populate_datastore.py
    178179

    179180
    - If you make a mistake during development (i.e. a failing test that

    ‎google/cloud/datastore/__init__.py

    Lines changed: 3 additions & 3 deletions
    Original file line numberDiff line numberDiff line change
    @@ -34,9 +34,9 @@
    3434
    The main concepts with this API are:
    3535
    3636
    - :class:`~google.cloud.datastore.client.Client`
    37-
    which represents a project (string) and namespace (string) bundled with
    38-
    a connection and has convenience methods for constructing objects with that
    39-
    project / namespace.
    37+
    which represents a project (string), database (string), and namespace
    38+
    (string) bundled with a connection and has convenience methods for
    39+
    constructing objects with that project/database/namespace.
    4040
    4141
    - :class:`~google.cloud.datastore.entity.Entity`
    4242
    which represents a single entity in the datastore

    ‎google/cloud/datastore/aggregation.py

    Lines changed: 16 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -22,6 +22,7 @@
    2222
    from google.cloud.datastore_v1.types import query as query_pb2
    2323
    from google.cloud.datastore import helpers
    2424
    from google.cloud.datastore.query import _pb_from_query
    25+
    from google.cloud.datastore.constants import DEFAULT_DATABASE
    2526

    2627

    2728
    _NOT_FINISHED = query_pb2.QueryResultBatch.MoreResultsType.NOT_FINISHED
    @@ -123,6 +124,18 @@ def project(self):
    123124
    """
    124125
    return self._nested_query._project or self._client.project
    125126

    127+
    @property
    128+
    def database(self):
    129+
    """Get the database for this AggregationQuery.
    130+
    :rtype: str
    131+
    :returns: The database for the query.
    132+
    """
    133+
    if self._nested_query._database or (
    134+
    self._nested_query._database == DEFAULT_DATABASE
    135+
    ):
    136+
    return self._nested_query._database
    137+
    return self._client.database
    138+
    126139
    @property
    127140
    def namespace(self):
    128141
    """The nested query's namespace
    @@ -376,6 +389,7 @@ def _next_page(self):
    376389

    377390
    partition_id = entity_pb2.PartitionId(
    378391
    project_id=self._aggregation_query.project,
    392+
    database_id=self._aggregation_query.database,
    379393
    namespace_id=self._aggregation_query.namespace,
    380394
    )
    381395

    @@ -390,6 +404,7 @@ def _next_page(self):
    390404
    response_pb = self.client._datastore_api.run_aggregation_query(
    391405
    request={
    392406
    "project_id": self._aggregation_query.project,
    407+
    "database_id": self._aggregation_query.database,
    393408
    "partition_id": partition_id,
    394409
    "read_options": read_options,
    395410
    "aggregation_query": query_pb,
    @@ -409,6 +424,7 @@ def _next_page(self):
    409424
    response_pb = self.client._datastore_api.run_aggregation_query(
    410425
    request={
    411426
    "project_id": self._aggregation_query.project,
    427+
    "database_id": self._aggregation_query.database,
    412428
    "partition_id": partition_id,
    413429
    "read_options": read_options,
    414430
    "aggregation_query": query_pb,

    ‎google/cloud/datastore/batch.py

    Lines changed: 23 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -23,6 +23,7 @@
    2323

    2424
    from google.cloud.datastore import helpers
    2525
    from google.cloud.datastore_v1.types import datastore as _datastore_pb2
    26+
    from google.cloud.datastore.constants import DEFAULT_DATABASE
    2627

    2728

    2829
    class Batch(object):
    @@ -122,6 +123,15 @@ def project(self):
    122123
    """
    123124
    return self._client.project
    124125

    126+
    @property
    127+
    def database(self):
    128+
    """Getter for database in which the batch will run.
    129+
    130+
    :rtype: :class:`str`
    131+
    :returns: The database in which the batch will run.
    132+
    """
    133+
    return self._client.database
    134+
    125135
    @property
    126136
    def namespace(self):
    127137
    """Getter for namespace in which the batch will run.
    @@ -218,6 +228,12 @@ def put(self, entity):
    218228
    if self.project != entity.key.project:
    219229
    raise ValueError("Key must be from same project as batch")
    220230

    231+
    entity_key_database = entity.key.database
    232+
    if entity_key_database is None:
    233+
    entity_key_database = DEFAULT_DATABASE
    234+
    if self.database != entity_key_database:
    235+
    raise ValueError("Key must be from same database as batch")
    236+
    221237
    if entity.key.is_partial:
    222238
    entity_pb = self._add_partial_key_entity_pb()
    223239
    self._partial_key_entities.append(entity)
    @@ -245,6 +261,12 @@ def delete(self, key):
    245261
    if self.project != key.project:
    246262
    raise ValueError("Key must be from same project as batch")
    247263

    264+
    key_db = key.database
    265+
    if key_db is None:
    266+
    key_db = DEFAULT_DATABASE
    267+
    if self.database != key_db:
    268+
    raise ValueError("Key must be from same database as batch")
    269+
    248270
    key_pb = key.to_protobuf()
    249271
    self._add_delete_key_pb()._pb.CopyFrom(key_pb._pb)
    250272

    @@ -284,6 +306,7 @@ def _commit(self, retry, timeout):
    284306
    commit_response_pb = self._client._datastore_api.commit(
    285307
    request={
    286308
    "project_id": self.project,
    309+
    "database_id": self.database,
    287310
    "mode": mode,
    288311
    "transaction": self._id,
    289312
    "mutations": self._mutations,

    ‎google/cloud/datastore/client.py

    Lines changed: 43 additions & 6 deletions
    Original file line numberDiff line numberDiff line change
    @@ -25,6 +25,7 @@
    2525
    from google.cloud.datastore import helpers
    2626
    from google.cloud.datastore._http import HTTPDatastoreAPI
    2727
    from google.cloud.datastore.batch import Batch
    28+
    from google.cloud.datastore.constants import DEFAULT_DATABASE
    2829
    from google.cloud.datastore.entity import Entity
    2930
    from google.cloud.datastore.key import Key
    3031
    from google.cloud.datastore.query import Query
    @@ -126,6 +127,7 @@ def _extended_lookup(
    126127
    retry=None,
    127128
    timeout=None,
    128129
    read_time=None,
    130+
    database=DEFAULT_DATABASE,
    129131
    ):
    130132
    """Repeat lookup until all keys found (unless stop requested).
    131133
    @@ -179,6 +181,10 @@ def _extended_lookup(
    179181
    ``eventual==True`` or ``transaction_id``.
    180182
    This feature is in private preview.
    181183
    184+
    :type database: str
    185+
    :param database:
    186+
    (Optional) Database from which to fetch data. Defaults to the (default) database.
    187+
    182188
    :rtype: list of :class:`.entity_pb2.Entity`
    183189
    :returns: The requested entities.
    184190
    :raises: :class:`ValueError` if missing / deferred are not null or
    @@ -201,6 +207,7 @@ def _extended_lookup(
    201207
    lookup_response = datastore_api.lookup(
    202208
    request={
    203209
    "project_id": project,
    210+
    "database_id": database,
    204211
    "keys": key_pbs,
    205212
    "read_options": read_options,
    206213
    },
    @@ -276,6 +283,9 @@ class Client(ClientWithProject):
    276283
    environment variable.
    277284
    This parameter should be considered private, and could
    278285
    change in the future.
    286+
    287+
    :type database: str
    28 179B 8+
    :param database: (Optional) database to pass to proxied API methods.
    279289
    """
    280290

    281291
    SCOPE = ("https://www.googleapis.com/auth/datastore",)
    @@ -290,6 +300,8 @@ def __init__(
    290300
    client_options=None,
    291301
    _http=None,
    292302
    _use_grpc=None,
    303+
    *,
    304+
    database=DEFAULT_DATABASE,
    293305
    ):
    294306
    emulator_host = os.getenv(DATASTORE_EMULATOR_HOST)
    295307

    @@ -306,6 +318,7 @@ def __init__(
    306318
    client_options=client_options,
    307319
    _http=_http,
    308320
    )
    321+
    self.database = database
    309322
    self.namespace = namespace
    310323
    self._client_info = client_info
    311324
    self._client_options = client_options
    @@ -549,6 +562,7 @@ def get_multi(
    549562
    entity_pbs = _extended_lookup(
    550563
    datastore_api=self._datastore_api,
    551564
    project=self.project,
    565+
    database=self.database,
    552566
    key_pbs=[key.to_protobuf() for key in keys],
    553567
    eventual=eventual,
    554568
    missing=missing,
    @@ -740,7 +754,11 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None):
    740754
    kwargs = _make_retry_timeout_kwargs(retry, timeout)
    741755

    742756
    response_pb = self._datastore_api.allocate_ids(
    743-
    request={"project_id": incomplete_key.project, "keys": incomplete_key_pbs},
    757+
    request={
    758+
    "project_id": incomplete_key.project,
    759+
    "database_id": incomplete_key.database,
    760+
    "keys": incomplete_key_pbs,
    761+
    },
    744762
    **kwargs,
    745763
    )
    746764
    allocated_ids = [
    @@ -753,11 +771,14 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None):
    753771
    def key(self, *path_args, **kwargs):
    754772
    """Proxy to :class:`google.cloud.datastore.key.Key`.
    755773
    756-
    Passes our ``project``.
    774+
    Passes our ``project`` and our ``database``.
    757775
    """
    758776
    if "project" in kwargs:
    759777
    raise TypeError("Cannot pass project")
    760778
    kwargs["project"] = self.project
    779+
    if "database" in kwargs:
    780+
    raise TypeError("Cannot pass database")
    781+
    kwargs["database"] = self.database
    761782
    if "namespace" not in kwargs:
    762783
    kwargs["namespace"] = self.namespace
    763784
    return Key(*path_args, **kwargs)
    @@ -780,7 +801,7 @@ def transaction(self, **kwargs):
    780801
    def query(self, **kwargs):
    781802
    """Proxy to :class:`google.cloud.datastore.query.Query`.
    782803
    783-
    Passes our ``project``.
    804+
    Passes our ``project`` and our ``database``.
    784805
    785806
    Using query to search a datastore:
    786807
    @@ -834,7 +855,10 @@ def do_something_with(entity):
    834855
    raise TypeError("Cannot pass client")
    835856
    if "project" in kwargs:
    836857
    raise TypeError("Cannot pass project")
    858+
    if "database" in kwargs:
    859+
    raise TypeError("Cannot pass database")
    837860
    kwargs["project"] = self.project
    861+
    kwargs["database"] = self.database
    838862
    if "namespace" not in kwargs:
    839863
    kwargs["namespace"] = self.namespace
    840864
    return Query(self, **kwargs)
    @@ -963,18 +987,26 @@ def reserve_ids_sequential(self, complete_key, num_ids, retry=None, timeout=None
    963987
    key_class = type(complete_key)
    964988
    namespace = complete_key._namespace
    965989
    project = complete_key._project
    990+
    database = complete_key._database
    966991
    flat_path = list(complete_key._flat_path[:-1])
    967992
    start_id = complete_key._flat_path[-1]
    968993

    969994
    key_pbs = []
    970995
    for id in range(start_id, start_id + num_ids):
    971996
    path = flat_path + [id]
    972-
    key = key_class(*path, project=project, namespace=namespace)
    997+
    key = key_class(
    998+
    *path, project=project, database=database, namespace=namespace
    999+
    )
    9731000
    key_pbs.append(key.to_protobuf())
    9741001

    9751002
    kwargs = _make_retry_timeout_kwargs(retry, timeout)
    9761003
    self._datastore_api.reserve_ids(
    977-
    request={"project_id": complete_key.project, "keys": key_pbs}, **kwargs
    1004+
    request={
    1005+
    "project_id": complete_key.project,
    1006+
    "database_id": complete_key.database,
    1007+
    "keys": key_pbs,
    1008+
    },
    1009+
    **kwargs,
    9781010
    )
    9791011
    return None
    9801012

    @@ -1021,7 +1053,12 @@ def reserve_ids_multi(self, complete_keys, retry=None, timeout=None):
    10211053
    kwargs = _make_retry_timeout_kwargs(retry, timeout)
    10221054
    key_pbs = [key.to_protobuf() for key in complete_keys]
    10231055
    self._datastore_api.reserve_ids(
    1024-
    request={"project_id": complete_keys[0].project, "keys": key_pbs}, **kwargs
    1056+
    request={
    1057+
    "project_id": complete_keys[0].project,
    1058+
    "database_id": complete_keys[0].database,
    1059+
    "keys": key_pbs,
    1060+
    },
    1061+
    **kwargs,
    10251062
    )
    10261063

    10271064
    return None

    ‎google/cloud/datastore/constants.py

    Lines changed: 4 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -0,0 +1,4 @@
    1+
    """Constants for Datastore."""
    2+
    3+
    DEFAULT_DATABASE = ""
    4+
    """Datastore default database."""

    ‎google/cloud/datastore/helpers.py

    Lines changed: 5 additions & 1 deletion
    Original file line numberDiff line numberDiff line change
    @@ -29,6 +29,7 @@
    2929
    from google.cloud.datastore_v1.types import entity as entity_pb2
    3030
    from google.cloud.datastore.entity import Entity
    3131
    from google.cloud.datastore.key import Key
    32+
    from google.cloud.datastore.constants import DEFAULT_DATABASE
    3233
    from google.protobuf import timestamp_pb2
    3334

    3435

    @@ -300,11 +301,14 @@ def key_from_protobuf(pb):
    300301
    project = None
    301302
    if pb.partition_id.project_id: # Simple field (string)
    302303
    project = pb.partition_id.project_id
    304+
    database = DEFAULT_DATABASE
    305+
    if pb.partition_id.database_id: # Simple field (string)
    306+
    database = pb.partition_id.database_id
    303307
    namespace = None
    304308
    if pb.partition_id.namespace_id: # Simple field (string)
    305309
    namespace = pb.partition_id.namespace_id
    306310

    307-
    return Key(*path_args, namespace=namespace, project=project)
    311+
    return Key(*path_args, namespace=namespace, project=project, database=database)
    308312

    309313

    310314
    def _pb_attr_value(val):

    0 commit comments

    Comments
     (0)
    0