Skip to content

Commit f699792

Browse files
dhermeslukesneeringer
dhermes
authored and
lukesneeringer
committed
Sending x-goog-api-client header for Datastore gRPC implementation. (#3096)
1 parent 732e29f commit f699792

File tree

3 files changed

+41
-40
lines changed

3 files changed

+41
-40
lines changed

‎packages/google-cloud-datastore/google/cloud/datastore/_gax.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,24 @@
1717

1818
import contextlib
1919

20+
from google.cloud.proto.datastore.v1 import datastore_pb2_grpc
21+
from google.gax.utils import metrics
2022
from grpc import StatusCode
2123

2224
from google.cloud._helpers import make_insecure_stub
2325
from google.cloud._helpers import make_secure_stub
2426
from google.cloud import exceptions
2527

26-
from google.cloud.proto.datastore.v1 import datastore_pb2_grpc
28+
from google.cloud.datastore import __version__
2729

2830

31+
_METRICS_HEADERS = (
32+
('gccl', __version__),
33+
)
34+
_HEADER_STR = metrics.stringify(metrics.fill(_METRICS_HEADERS))
35+
_GRPC_EXTRA_OPTIONS = (
36+
('x-goog-api-client', _HEADER_STR),
37+
)
2938
_GRPC_ERROR_MAPPING = {
3039
StatusCode.UNKNOWN: exceptions.InternalServerError,
3140
StatusCode.INVALID_ARGUMENT: exceptions.BadRequest,
@@ -85,10 +94,10 @@ class _DatastoreAPIOverGRPC(object):
8594

8695
def __init__(self, connection, secure):
8796
if secure:
88-
self._stub = make_secure_stub(connection.credentials,
89-
connection.USER_AGENT,
90-
datastore_pb2_grpc.DatastoreStub,
91-
connection.host)
97+
self._stub = make_secure_stub(
98+
connection.credentials, connection.USER_AGENT,
99+
datastore_pb2_grpc.DatastoreStub, connection.host,
100+
extra_options=_GRPC_EXTRA_OPTIONS)
92101
else:
93102
self._stub = make_insecure_stub(datastore_pb2_grpc.DatastoreStub,
94103
connection.host)

‎packages/google-cloud-datastore/setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151

5252
REQUIREMENTS = [
5353
'google-cloud-core >= 0.23.1, < 0.24dev',
54-
'grpcio >= 1.0.2, < 2.0dev',
54+
'google-gax>=0.15.7, <0.16dev',
5555
'gapic-google-cloud-datastore-v1 >= 0.15.0, < 0.16dev',
5656
]
5757

‎packages/google-cloud-datastore/unit_tests/test__gax.py‎

Lines changed: 26 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ def _get_target_class():
9797

9898
return _DatastoreAPIOverGRPC
9999

100-
def _make_one(self, stub, connection=None, secure=True, mock_args=None):
100+
def _make_one(self, stub, connection=None, secure=True):
101101
if connection is None:
102102
connection = mock.Mock(
103103
credentials=object(),
@@ -106,27 +106,21 @@ def _make_one(self, stub, connection=None, secure=True, mock_args=None):
106106
spec=['credentials', 'host', 'USER_AGENT'],
107107
)
108108

109-
if mock_args is None:
110-
mock_args = []
111-
112-
def mock_make_stub(*args):
113-
mock_args.append(args)
114-
return stub
115-
116109
if secure:
117110
patch = mock.patch(
118111
'google.cloud.datastore._gax.make_secure_stub',
119-
new=mock_make_stub)
112+
return_value=stub)
120113
else:
121114
patch = mock.patch(
122115
'google.cloud.datastore._gax.make_insecure_stub',
123-
new=mock_make_stub)
116+
return_value=stub)
124117

125-
with patch:
126-
return self._get_target_class()(connection, secure)
118+
with patch as make_stub_mock:
119+
api_obj = self._get_target_class()(connection, secure)
120+
return api_obj, make_stub_mock
127121

128122
def test_constructor(self):
129-
from google.cloud.proto.datastore.v1 import datastore_pb2_grpc
123+
import google.cloud.datastore._gax as MUT
130124

131125
conn = mock.Mock(
132126
credentials=object(),
@@ -136,17 +130,17 @@ def test_constructor(self):
136130
)
137131

138132
stub = _GRPCStub()
139-
mock_args = []
140-
datastore_api = self._make_one(stub, connection=conn,
141-
mock_args=mock_args)
142-
self.assertIs(datastore_api._stub, stub)
133+
datastore_api, make_stub_mock = self._make_one(
134+
stub, connection=conn)
143135

144-
self.assertEqual(mock_args, [(
136+
self.assertIs(datastore_api._stub, stub)
137+
make_stub_mock.assert_called_once_with(
145138
conn.credentials,
146139
conn.USER_AGENT,
147-
datastore_pb2_grpc.DatastoreStub,
140+
MUT.datastore_pb2_grpc.DatastoreStub,
148141
conn.host,
149-
)])
142+
extra_options=MUT._GRPC_EXTRA_OPTIONS,
143+
)
150144

151145
def test_constructor_insecure(self):
152146
from google.cloud.proto.datastore.v1 import datastore_pb2_grpc
@@ -158,21 +152,19 @@ def test_constructor_insecure(self):
158152
)
159153

160154
stub = _GRPCStub()
161-
mock_args = []
162-
datastore_api = self._make_one(stub, connection=conn,
163-
secure=False,
164-
mock_args=mock_args)
165-
self.assertIs(datastore_api._stub, stub)
155+
datastore_api, make_stub_mock = self._make_one(
156+
stub, connection=conn, secure=False)
166157

167-
self.assertEqual(mock_args, [(
158+
self.assertIs(datastore_api._stub, stub)
159+
make_stub_mock.assert_called_once_with(
168160
datastore_pb2_grpc.DatastoreStub,
169161
conn.host,
170-
)])
162+
)
171163

172164
def test_lookup(self):
173165
return_val = object()
174166
stub = _GRPCStub(return_val)
175-
datastore_api = self._make_one(stub=stub)
167+
datastore_api, _ = self._make_one(stub=stub)
176168

177169
request_pb = mock.Mock(project_id=None, spec=['project_id'])
178170
project = 'PROJECT'
@@ -185,7 +177,7 @@ def test_lookup(self):
185177
def test_run_query(self):
186178
return_val = object()
187179
stub = _GRPCStub(return_val)
188-
datastore_api = self._make_one(stub=stub)
180+
datastore_api, _ = self._make_one(stub=stub)
189181

190182
request_pb = mock.Mock(project_id=None, spec=['project_id'])
191183
project = 'PROJECT'
@@ -197,7 +189,7 @@ def test_run_query(self):
197189

198190
def _run_query_failure_helper(self, exc, err_class):
199191
stub = _GRPCStub(side_effect=exc)
200-
datastore_api = self._make_one(stub=stub)
192+
datastore_api, _ = self._make_one(stub=stub)
201193

202194
request_pb = mock.Mock(project_id=None, spec=['project_id'])
203195
project = 'PROJECT'
@@ -225,7 +217,7 @@ def test_run_query_invalid_argument(self):
225217
def test_begin_transaction(self):
226218
return_val = object()
227219
stub = _GRPCStub(return_val)
228-
datastore_api = self._make_one(stub=stub)
220+
datastore_api, _ = self._make_one(stub=stub)
229221

230222
request_pb = mock.Mock(project_id=None, spec=['project_id'])
231223
project = 'PROJECT'
@@ -239,7 +231,7 @@ def test_begin_transaction(self):
239231
def test_commit_success(self):
240232
return_val = object()
241233
stub = _GRPCStub(return_val)
242-
datastore_api = self._make_one(stub=stub)
234+
datastore_api, _ = self._make_one(stub=stub)
243235

244236
request_pb = mock.Mock(project_id=None, spec=['project_id'])
245237
project = 'PROJECT'
@@ -252,7 +244,7 @@ def test_commit_success(self):
252244
def test_rollback(self):
253245
return_val = object()
254246
stub = _GRPCStub(return_val)
255-
datastore_api = self._make_one(stub=stub)
247+
datastore_api, _ = self._make_one(stub=stub)
256248

257249
request_pb = mock.Mock(project_id=None, spec=['project_id'])
258250
project = 'PROJECT'
@@ -265,7 +257,7 @@ def test_rollback(self):
265257
def test_allocate_ids(self):
266258
return_val = object()
267259
stub = _GRPCStub(return_val)
268-
datastore_api = self._make_one(stub=stub)
260+
datastore_api, _ = self._make_one(stub=stub)
269261

270262
request_pb = mock.Mock(project_id=None, spec=['project_id'])
271263
project = 'PROJECT'

0 commit comments

Comments
 (0)