Skip to content

Commit bae395d

Browse files
odeke-emodeke-em
authored
chore(x-goog-spanner-request-id): add x_goog_spanner_request_id as an attribute to OTEL spans (#1378)
* chore(x-goog-spanner-request-id): add x_goog_spanner_request_id as an attribute to OTEL spans This change is effectively 3/3 of the work to complete x-goog-spanner-request-id propagation. While here instrumented batch_write as well to send over the header too. Updates #1261 * Remove debug printf
1 parent de322f8 commit bae395d

File tree

16 files changed

+416
-123
lines changed

16 files changed

+416
-123
lines changed

‎google/cloud/spanner_v1/batch.py‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
_metadata_with_prefix,
2727
_metadata_with_leader_aware_routing,
2828
_merge_Transaction_Options,
29+
AtomicCounter,
2930
)
3031
from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
3132
from google.cloud.spanner_v1 import RequestOptions
@@ -248,7 +249,7 @@ def commit(
248249
trace_attributes,
249250
observability_options=observability_options,
250251
metadata=metadata,
251-
), MetricsCapture():
252+
) as span, MetricsCapture():
252253

253254
def wrapped_method(*args, **kwargs):
254255
method = functools.partial(
@@ -261,6 +262,7 @@ def wrapped_method(*args, **kwargs):
261262
getattr(database, "_next_nth_request", 0),
262263
1,
263264
metadata,
265+
span,
264266
),
265267
)
266268
return method(*args, **kwargs)
@@ -384,14 +386,25 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
384386
trace_attributes,
385387
observability_options=observability_options,
386388
metadata=metadata,
387-
), MetricsCapture():
388-
method = functools.partial(
389-
api.batch_write,
390-
request=request,
391-
metadata=metadata,
392-
)
389+
) as span, MetricsCapture():
390+
attempt = AtomicCounter(0)
391+
nth_request = getattr(database, "_next_nth_request", 0)
392+
393+
def wrapped_method(*args, **kwargs):
394+
method = functools.partial(
395+
api.batch_write,
396+
request=request,
397+
metadata=database.metadata_with_request_id(
398+
nth_request,
399+
attempt.increment(),
400+
metadata,
401+
span,
402+
),
403+
)
404+
return method(*args, **kwargs)
405+
393406
response = _retry(
394-
method,
407+
wrapped_method,
395408
allowed_exceptions={
396409
InternalServerError: _check_rst_stream_error,
397410
},

‎google/cloud/spanner_v1/database.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -462,13 +462,19 @@ def spanner_api(self):
462462

463463
return self._spanner_api
464464

465-
def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
465+
def metadata_with_request_id(
466+
self, nth_request, nth_attempt, prior_metadata=[], span=None
467+
):
468+
if span is None:
469+
span = get_current_span()
470+
466471
return _metadata_with_request_id(
467472
self._nth_client_id,
468473
self._channel_id,
469474
nth_request,
470475
nth_attempt,
471476
prior_metadata,
477+
span,
472478
)
473479

474480
def __eq__(self, other):
@@ -762,6 +768,7 @@ def execute_pdml():
762768
self._next_nth_request,
763769
1,
764770
metadata,
771+
span,
765772
),
766773
)
767774

‎google/cloud/spanner_v1/pool.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,10 @@ def bind(self, database):
257257
resp = api.batch_create_sessions(
258258
request=request,
259259
metadata=database.metadata_with_request_id(
260-
database._next_nth_request, 1, metadata
260+
database._next_nth_request,
261+
1,
262+
metadata,
263+
span,
261264
),
262265
)
263266

@@ -564,7 +567,10 @@ def bind(self, database):
564567
resp = api.batch_create_sessions(
565568
request=request,
566569
metadata=database.metadata_with_request_id(
567-
database._next_nth_request, 1, metadata
570+
database._next_nth_request,
571+
1,
572+
metadata,
573+
span,
568574
),
569575
)
570576

‎google/cloud/spanner_v1/request_id_header.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,32 @@ def generate_rand_uint64():
3333

3434

3535
REQ_RAND_PROCESS_ID = generate_rand_uint64()
36+
X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR = "x_goog_spanner_request_id"
3637

3738

38-
def with_request_id(client_id, channel_id, nth_request, attempt, other_metadata=[]):
39+
def with_request_id(
40+
client_id, channel_id, nth_request, attempt, other_metadata=[], span=None
41+
):
3942
req_id = f"{REQ_ID_VERSION}.{REQ_RAND_PROCESS_ID}.{client_id}.{channel_id}.{nth_request}.{attempt}"
4043
all_metadata = (other_metadata or []).copy()
4144
all_metadata.append((REQ_ID_HEADER_KEY, req_id))
45+
46+
if span is not None:
47+
span.set_attribute(X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR, req_id)
48+
4249
return all_metadata
50+
51+
52+
def parse_request_id(request_id_str):
53+
splits = request_id_str.split(".")
54+
version, rand_process_id, client_id, channel_id, nth_request, nth_attempt = list(
55+
map(lambda v: int(v), splits)
56+
)
57+
return (
58+
version,
59+
rand_process_id,
60+
client_id,
61+
channel_id,
62+
nth_request,
63+
nth_attempt,
64+
)

‎google/cloud/spanner_v1/session.py‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -167,11 +167,14 @@ def create(self):
167167
self._labels,
168168
observability_options=observability_options,
169169
metadata=metadata,
170-
), MetricsCapture():
170+
) as span, MetricsCapture():
171171
session_pb = api.create_session(
172172
request=request,
173173
metadata=self._database.metadata_with_request_id(
174-
self._database._next_nth_request, 1, metadata
174+
self._database._next_nth_request,
175+
1,
176+
metadata,
177+
span,
175178
),
176179
)
177180
self._session_id = session_pb.name.split("/")[-1]
@@ -218,7 +221,10 @@ def exists(self):
218221
api.get_session(
219222
name=self.name,
220223
metadata=database.metadata_with_request_id(
221-
database._next_nth_request, 1, metadata
224+
database._next_nth_request,
225+
1,
226+
metadata,
227+
span,
222228
),
223229
)
224230
if span:
@@ -263,11 +269,14 @@ def delete(self):
263269
},
264270
observability_options=observability_options,
265271
metadata=metadata,
266-
), MetricsCapture():
272+
) as span, MetricsCapture():
267273
api.delete_session(
268274
name=self.name,
269275
metadata=database.metadata_with_request_id(
270-
database._next_nth_request, 1, metadata
276+
database._next_nth_request,
277+
1,
278+
metadata,
279+
span,
271280
),
272281
)
273282

‎google/cloud/spanner_v1/snapshot.py‎

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -104,11 +104,14 @@ def _restart_on_unavailable(
104104
attributes,
105105
observability_options=observability_options,
106106
metadata=metadata,
107-
), MetricsCapture():
107+
) as span, MetricsCapture():
108108
iterator = method(
109109
request=request,
110110
metadata=request_id_manager.metadata_with_request_id(
111-
nth_request, attempt, metadata
111+
nth_request,
112+
attempt,
113+
metadata,
114+
span,
112115
),
113116
)
114117
for item in iterator:
@@ -133,7 +136,7 @@ def _restart_on_unavailable(
133136
attributes,
134137
observability_options=observability_options,
135138
metadata=metadata,
136-
), MetricsCapture():
139+
) as span, MetricsCapture():
137140
request.resume_token = resume_token
138141
if transaction is not None:
139142
transaction_selector = transaction._make_txn_selector()
@@ -142,7 +145,10 @@ def _restart_on_unavailable(
142145
iterator = method(
143146
request=request,
144147
metadata=request_id_manager.metadata_with_request_id(
145-
nth_request, attempt, metadata
148+
nth_request,
149+
attempt,
150+
metadata,
151+
span,
146152
),
147153
)
148154
continue
@@ -160,7 +166,7 @@ def _restart_on_unavailable(
160166
attributes,
161167
observability_options=observability_options,
162168
metadata=metadata,
163-
), MetricsCapture():
169+
) as span, MetricsCapture():
164170
request.resume_token = resume_token
165171
if transaction is not None:
166172
transaction_selector = transaction._make_txn_selector()
@@ -169,7 +175,10 @@ def _restart_on_unavailable(
169175
iterator = method(
170176
request=request,
171177
metadata=request_id_manager.metadata_with_request_id(
172-
nth_request, attempt, metadata
178+
nth_request,
179+
attempt,
180+
metadata,
181+
span,
173182
),
174183
)
175184
continue
@@ -745,13 +754,16 @@ def partition_read(
745754
extra_attributes=trace_attributes,
746755
observability_options=getattr(database, "observability_options", None),
747756
metadata=metadata,
748-
), MetricsCapture():
757+
) as span, MetricsCapture():
749758
nth_request = getattr(database, "_next_nth_request", 0)
750759
attempt = AtomicCounter()
751760

752761
def attempt_tracking_method():
753762
all_metadata = database.metadata_with_request_id(
754-
nth_request, attempt.increment(), metadata
763+
nth_request,
764+
attempt.increment(),
765+
metadata,
766+
span,
755767
)
756768
method = functools.partial(
757769
api.partition_read,
@@ -858,13 +870,16 @@ def partition_query(
858870
trace_attributes,
859871
observability_options=getattr(database, "observability_options", None),
860872
metadata=metadata,
861-
), MetricsCapture():
873+
) as span, MetricsCapture():
862874
nth_request = getattr(database, "_next_nth_request", 0)
863875
attempt = AtomicCounter()
864876

865877
def attempt_tracking_method():
866878
all_metadata = database.metadata_with_request_id(
867-
nth_request, attempt.increment(), metadata
879+
nth_request,
880+
attempt.increment(),
881+
metadata,
882+
span,
868883
)
869884
method = functools.partial(
870885
api.partition_query,
@@ -1014,13 +1029,16 @@ def begin(self):
10141029
self._session,
10151030
observability_options=getattr(database, "observability_options", None),
10161031
metadata=metadata,
1017-
), MetricsCapture():
1032+
) as span, MetricsCapture():
10181033
nth_request = getattr(database, "_next_nth_request", 0)
10191034
attempt = AtomicCounter()
10201035

10211036
def attempt_tracking_method():
10221037
all_metadata = database.metadata_with_request_id(
1023-
nth_request, attempt.increment(), metadata
1038+
nth_request,
1039+
attempt.increment(),
1040+
metadata,
1041+
span,
10241042
)
10251043
method = functools.partial(
10261044
api.begin_transaction,

‎google/cloud/spanner_v1/testing/interceptors.py‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from grpc_interceptor import ClientInterceptor
1919
from google.api_core.exceptions import Aborted
20+
from google.cloud.spanner_v1.request_id_header import parse_request_id
2021

2122

2223
class MethodCountInterceptor(ClientInterceptor):
@@ -119,18 +120,3 @@ def stream_request_ids(self):
119120
def reset(self):
120121
self._stream_req_segments.clear()
121122
self._unary_req_segments.clear()
122-
123-
124-
def parse_request_id(request_id_str):
125-
splits = request_id_str.split(".")
126-
version, rand_process_id, client_id, channel_id, nth_request, nth_attempt = list(
127-
map(lambda v: int(v), splits)
128-
)
129-
return (
130-
version,
131-
rand_process_id,
132-
client_id,
133-
channel_id,
134-
nth_request,
135-
nth_attempt,
136-
)

‎google/cloud/spanner_v1/transaction.py‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,10 @@ def wrapped_method(*args, **kwargs):
191191
session=self._session.name,
192192
options=txn_options,
193193
metadata=database.metadata_with_request_id(
194-
nth_request, attempt.increment(), metadata
194+
nth_request,
195+
attempt.increment(),
196+
metadata,
197+
span,
195198
),
196199
)
197200
return method(*args, **kwargs)
@@ -232,7 +235,7 @@ def rollback(self):
232235
self._session,
233236
observability_options=observability_options,
234237
metadata=metadata,
235-
), MetricsCapture():
238+
) as span, MetricsCapture():
236239
attempt = AtomicCounter(0)
237240
nth_request = database._next_nth_request
238241

@@ -243,7 +246,10 @@ def wrapped_method(*args, **kwargs):
243246
session=self._session.name,
244247
transaction_id=self._transaction_id,
245248
metadata=database.metadata_with_request_id(
246-
nth_request, attempt.value, metadata
249+
nth_request,
250+
attempt.value,
251+
metadata,
252+
span,
247253
),
248254
)
249255
return method(*args, **kwargs)
@@ -334,7 +340,10 @@ def wrapped_method(*args, **kwargs):
334340
api.commit,
335341
request=request,
336342
metadata=database.metadata_with_request_id(
337-
nth_request, attempt.value, metadata
343+
nth_request,
344+
attempt.value,
345+
metadata,
346+
span,
338347
),
339348
)
340349
return method(*args, **kwargs)

0 commit comments

Comments
 (0)