Skip to content

Commit e267183

Browse files
author
Chris Rossi
authored
Wire up queries to Datastore. (#44)
This is the bare minimum necessary to run the simplest query possible: all of the entities of a given kind.
1 parent 6441e77 commit e267183

File tree

9 files changed

+826
-330
lines changed

9 files changed

+826
-330
lines changed

‎packages/google-cloud-ndb/noxfile.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ def unit(session):
4545
run_args.extend(
4646
[
4747
"--cov=google.cloud.ndb",
48-
"--cov=tests",
48+
"--cov=tests.unit",
4949
"--cov-config",
5050
get_path(".coveragerc"),
5151
"--cov-report=",

‎packages/google-cloud-ndb/src/google/cloud/ndb/_datastore_query.py‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
# Copyright 2018 Google LLC
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+
# https://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+
"""Translate NDB queries to Datastore calls."""
16+
17+
from google.cloud.datastore_v1.proto import datastore_pb2
18+
from google.cloud.datastore_v1.proto import query_pb2
19+
20+
from google.cloud.ndb import context as context_module
21+
from google.cloud.ndb import _datastore_api
22+
from google.cloud.ndb import model
23+
from google.cloud.ndb import tasklets
24+
25+
MoreResultsType = query_pb2.QueryResultBatch.MoreResultsType
26+
MORE_RESULTS_TYPE_NOT_FINISHED = MoreResultsType.Value("NOT_FINISHED")
27+
ResultType = query_pb2.EntityResult.ResultType
28+
RESULT_TYPE_FULL = ResultType.Value("FULL")
29+
30+
31+
@tasklets.tasklet
32+
def fetch(query):
33+
"""Fetch query results.
34+
35+
Args:
36+
query (query.Query): The query.
37+
38+
Returns:
39+
tasklets.Future: Result is List[model.Model]: The query results.
40+
"""
41+
for name in (
42+
"ancestor",
43+
"filters",
44+
"orders",
45+
"app",
46+
"namespace",
47+
"default_options",
48+
"projection",
49+
"group_by",
50+
):
51+
if getattr(query, name, None):
52+
raise NotImplementedError(
53+
"{} is not yet implemented for queries.".format(name)
54+
)
55+
56+
query_pb = _query_to_protobuf(query)
57+
results = yield _run_query(query_pb)
58+
return [
59+
_process_result(result_type, result) for result_type, result in results
60+
]
61+
62+
63+
def _process_result(result_type, result):
64+
"""Process a single entity result.
65+
66+
Args:
67+
result_type (query_pb2.EntityResult.ResultType): The type of the result
68+
(full entity, projection, or key only).
69+
result (query_pb2.EntityResult): The protocol buffer representation of
70+
the query result.
71+
72+
Returns:
73+
Union[model.Model, key.Key]: The processed result.
74+
"""
75+
if result_type == RESULT_TYPE_FULL:
76+
return model._entity_from_protobuf(result.entity)
77+
78+
raise NotImplementedError(
79+
"Processing for projection and key only entity results is not yet "
80+
"implemented for queries."
81+
)
82+
83+
84+
def _query_to_protobuf(query):
85+
"""Convert an NDB query to a Datastore protocol buffer.
86+
87+
Args:
88+
query (query.Query): The query.
89+
90+
Returns:
91+
query_pb2.Query: The protocol buffer representation of the query.
92+
"""
93+
query_args = {}
94+
if query.kind:
95+
query_args["kind"] = [query_pb2.KindExpression(name=query.kind)]
96+
97+
return query_pb2.Query(**query_args)
98+
99+
100+
@tasklets.tasklet
101+
def _run_query(query_pb):
102+
"""Run a query in Datastore.
103+
104+
Will potentially repeat the query to get all results.
105+
106+
Args:
107+
query_pb (query_pb2.Query): The query protocol buffer representation.
108+
109+
Returns:
110+
tasklets.Future: List[Tuple[query_pb2.EntityResult.ResultType,
111+
query_pb2.EntityResult]]: The raw query results.
112+
"""
113+
client = context_module.get_context().client
114+
results = []
115+
116+
while True:
117+
# See what results we get from the backend
118+
request = datastore_pb2.RunQueryRequest(
119+
project_id=client.project, query=query_pb
120+
)
121+
response = yield _datastore_api.make_call("RunQuery", request)
122+
batch = response.batch
123+
results.extend(
124+
(
125+
(batch.entity_result_type, result)
126+
for result in batch.entity_results
127+
)
128+
)
129+
130+
# Did we get all of them?
131+
if batch.more_results != MORE_RESULTS_TYPE_NOT_FINISHED:
132+
break
133+
134+
# Still some results left to fetch. Update cursors and try again.
135+
query_pb.start_cursor = batch.end_cursor
136+
137+
return results

‎packages/google-cloud-ndb/src/google/cloud/ndb/query.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
"""High-level wrapper for datastore queries."""
1616

17+
from google.cloud.ndb import _datastore_query
1718
from google.cloud.ndb import exceptions
1819
from google.cloud.ndb import model
1920

@@ -1041,6 +1042,43 @@ def _check_properties(self, fixed, **kwargs):
10411042
if modelclass is not None:
10421043
modelclass._check_properties(fixed, **kwargs)
10431044

1045+
def fetch(self, limit=None, **options):
1046+
"""Run a query, fetching results.
1047+
1048+
Args:
1049+
limit (int): Maximum number of results to fetch. data:`None`
1050+
or data:`0` indicates no limit.
1051+
options (Dict[str, Any]): TBD.
1052+
1053+
Returns:
1054+
List([model.Model]): The query results.
1055+
"""
1056+
return self.fetch_async(limit, **options).result()
1057+
1058+
def fetch_async(self, limit=None, **options):
1059+
"""Run a query, asynchronously fetching the results.
1060+
1061+
Args:
1062+
limit (int): Maximum number of results to fetch. data:`None`
1063+
or data:`0` indicates no limit.
1064+
options (Dict[str, Any]): TBD.
1065+
1066+
Returns:
1067+
tasklets.Future: Eventual result will be a List[model.Model] of the
1068+
results.
1069+
"""
1070+
if limit:
1071+
raise NotImplementedError(
1072+
"'limit' is not implemented yet for queries"
1073+
)
1074+
1075+
if options:
1076+
raise NotImplementedError(
1077+
"'options' are not implemented yet for queries"
1078+
)
1079+
1080+
return _datastore_query.fetch(self)
1081+
10441082

10451083
def gql(*args, **kwargs):
10461084
raise NotImplementedError

‎packages/google-cloud-ndb/tests/system/__init__.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright 2018 Google LLC
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+
# https://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+
KIND = "SomeKind"

‎packages/google-cloud-ndb/tests/system/conftest.py‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
import pytest
2+
3+
from google.cloud import datastore
4+
from google.cloud import ndb
5+
6+
from . import KIND
7+
8+
9+
@pytest.fixture(scope="module", autouse=True)
10+
def initial_clean():
11+
# Make sure database is in clean state at beginning of test run
12+
client = datastore.Client()
13+
query = client.query(kind=KIND)
14+
for entity in query.fetch():
15+
client.delete(entity.key)
16+
17+
18+
@pytest.fixture
19+
def ds_client():
20+
client = datastore.Client()
21+
22+
# Make sure we're leaving database as clean as we found it after each test
23+
query = client.query(kind=KIND)
24+
results = list(query.fetch())
25+
assert not results
26+
27+
yield client
28+
29+
results = list(query.fetch())
30+
assert not results
31+
32+
33+
@pytest.fixture
34+
def ds_entity(ds_client, dispose_of):
35+
def make_entity(*key_args, **entity_kwargs):
36+
key = ds_client.key(*key_args)
37+
assert ds_client.get(key) is None
38+
entity = datastore.Entity(key=key)
39+
entity.update(entity_kwargs)
40+
ds_client.put(entity)
41+
dispose_of(key)
42+
43+
return entity
44+
45+
yield make_entity
46+
47+
48+
@pytest.fixture
49+
def dispose_of(ds_client):
50+
ds_keys = []
51+
52+
def delete_entity(ds_key):
53+
ds_keys.append(ds_key)
54+
55+
yield delete_entity
56+
57+
for ds_key in ds_keys:
58+
ds_client.delete(ds_key)
59+
60+
61+
@pytest.fixture
62+
def client_context():
63+
client = ndb.Client()
64+
with client.context():
65+
yield

packages/google-cloud-ndb/tests/system/test_system.py renamed to packages/google-cloud-ndb/tests/system/test_crud.py

Lines changed: 5 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -12,74 +12,18 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
"""
16+
System tests for Create, Update, Delete. (CRUD)
17+
"""
18+
1519
import pytest
1620

1721
import test_utils.system
1822

1923
from google.cloud import datastore
2024
from google.cloud import ndb
2125

22-
23-
KIND = "SomeKind"
24-
25-
26-
@pytest.fixture(scope="module", autouse=True)
27-
def initial_clean():
28-
# Make sure database is in clean state at beginning of test run
29-
client = datastore.Client()
30-
query = client.query(kind=KIND)
31-
for entity in query.fetch():
32-
client.delete(entity.key)
33-
34-
35-
@pytest.fixture
36-
def ds_client():
37-
client = datastore.Client()
38-
39-
# Make sure we're leaving database as clean as we found it after each test
40-
query = client.query(kind=KIND)
41-
results = list(query.fetch())
42-
assert not results
43-
44-
yield client
45-
46-
results = list(query.fetch())
47-
assert not results
48-
49-
50-
@pytest.fixture
51-
def ds_entity(ds_client, dispose_of):
52-
def make_entity(*key_args, **entity_kwargs):
53-
key = ds_client.key(*key_args)
54-
assert ds_client.get(key) is None
55-
entity = datastore.Entity(key=key)
56-
entity.update(entity_kwargs)
57-
ds_client.put(entity)
58-
dispose_of(key)
59-
60-
return entity
61-
62-
yield make_entity
63-
64-
65-
@pytest.fixture
66-
def dispose_of(ds_client):
67-
ds_keys = []
68-
69-
def delete_entity(ds_key):
70-
ds_keys.append(ds_key)
71-
72-
yield delete_entity
73-
74-
for ds_key in ds_keys:
75-
ds_client.delete(ds_key)
76-
77-
78-
@pytest.fixture
79-
def client_context():
80-
client = ndb.Client()
81-
with client.context():
82-
yield
26+
from . import KIND
8327

8428

8529
@pytest.mark.usefixtures("client_context")

0 commit comments

Comments
 (0)