Skip to content

Commit 9242220

Browse files
odeke-emodeke-em
committed
Add trace span helper add_span_event to check if a span is non-None before setting event
1 parent 6f400f9 commit 9242220

File tree

6 files changed

+123
-113
lines changed

6 files changed

+123
-113
lines changed

‎google/cloud/spanner_v1/_opentelemetry_tracing.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,3 +113,8 @@ def get_current_span():
113113
if not HAS_OPENTELEMETRY_INSTALLED:
114114
return None
115115
return trace.get_current_span()
116+
117+
118+
def add_span_event(span, commentary, event_attributes=None):
119+
if span:
120+
span.add_event(commentary, event_attributes)

‎google/cloud/spanner_v1/database.py

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,10 @@
6767
SpannerGrpcTransport,
6868
)
6969
from google.cloud.spanner_v1.table import Table
70-
from google.cloud.spanner_v1._opentelemetry_tracing import get_current_span
70+
from google.cloud.spanner_v1._opentelemetry_tracing import (
71+
add_span_event,
72+
get_current_span,
73+
)
7174

7275

7376
SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data"
@@ -1167,8 +1170,7 @@ def __enter__(self):
11671170
"""Begin ``with`` block."""
11681171
current_span = get_current_span()
11691172
session = self._session = self._database._pool.get()
1170-
if current_span:
1171-
current_span.add_event("Using session", {"id": session.session_id})
1173+
add_span_event(current_span, "Using session", {"id": session.session_id})
11721174
batch = self._batch = Batch(session)
11731175
if self._request_options.transaction_tag:
11741176
batch.transaction_tag = self._request_options.transaction_tag
@@ -1192,10 +1194,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
11921194
)
11931195
self._database._pool.put(self._session)
11941196
current_span = get_current_span()
1195-
if current_span:
1196-
current_span.add_event(
1197-
"Returned session to pool", {"id": self._session.session_id}
1198-
)
1197+
add_span_event(
1198+
current_span,
1199+
"Returned session to pool",
1200+
{"id": self._session.session_id},
1201+
)
11991202

12001203

12011204
class MutationGroupsCheckout(object):

‎google/cloud/spanner_v1/pool.py

Lines changed: 80 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
_metadata_with_leader_aware_routing,
2727
)
2828
from google.cloud.spanner_v1._opentelemetry_tracing import (
29+
add_span_event,
2930
get_current_span,
3031
trace_call,
3132
)
@@ -214,27 +215,27 @@ def bind(self, database):
214215
session_template=Session(creator_role=self.database_role),
215216
)
216217

217-
if requested_session_count > 0 and span:
218-
span.add_event(
218+
if requested_session_count > 0:
219+
add_span_event(
220+
span,
219221
f"Requesting {requested_session_count} sessions",
220222
span_event_attributes,
221223
)
222224

223225
if self._sessions.full():
224-
if span:
225-
span.add_event(
226-
"Session pool is already full", span_event_attributes
227-
)
226+
add_span_event(
227+
span, "Session pool is already full", span_event_attributes
228+
)
228229
return
229230

230231
returned_session_count = 0
231232
while not self._sessions.full():
232233
request.session_count = requested_session_count - self._sessions.qsize()
233-
if span:
234-
span.add_event(
235-
f"Creating {request.session_count} sessions",
236-
span_event_attributes,
237-
)
234+
add_span_event(
235+
span,
236+
f"Creating {request.session_count} sessions",
237+
span_event_attributes,
238+
)
238239
resp = api.batch_create_sessions(
239240
request=request,
240241
metadata=metadata,
@@ -245,11 +246,11 @@ def bind(self, database):
245246
self._sessions.put(session)
246247
returned_session_count += 1
247248

248-
if span:
249-
span.add_event(
250-
f"Requested for {requested_session_count} sessions, returned {returned_session_count}",
251-
span_event_attributes,
252-
)
249+
add_span_event(
250+
span,
251+
f"Requested for {requested_session_count} sessions, returned {returned_session_count}",
252+
span_event_attributes,
253+
)
253254

254255
def get(self, timeout=None):
255256
"""Check a session out from the pool.
@@ -268,35 +269,34 @@ def get(self, timeout=None):
268269
start_time = time.time()
269270
current_span = get_current_span()
270271
span_event_attributes = {"kind": type(self).__name__}
271-
if current_span:
272-
current_span.add_event("Acquiring session", span_event_attributes)
272+
add_span_event(current_span, "Acquiring session", span_event_attributes)
273273

274274
session = None
275275
try:
276-
if current_span:
277-
current_span.add_event(
278-
"Waiting for a session to become available", span_event_attributes
279-
)
276+
add_span_event(
277+
current_span,
278+
"Waiting for a session to become available",
279+
span_event_attributes,
280+
)
280281
session = self._sessions.get(block=True, timeout=timeout)
281282
except queue.Empty as e:
282-
if current_span:
283-
current_span.add_event("No session available", span_event_attributes)
283+
add_span_event(current_span, "No session available", span_event_attributes)
284284
raise e
285285

286286
span_event_attributes["session.id"] = session._session_id
287287

288288
if not session.exists():
289-
if current_span:
290-
current_span.add_event(
291-
"Session is not valid, recreating it", span_event_attributes
292-
)
289+
add_span_event(
290+
current_span,
291+
"Session is not valid, recreating it",
292+
span_event_attributes,
293+
)
293294
session = self._database.session()
294295
session.create()
295296
span_event_attributes["session.id"] = session._session_id
296297

297298
span_event_attributes["time.elapsed"] = time.time() - start_time
298-
if current_span:
299-
current_span.add_event("Acquired session", span_event_attributes)
299+
add_span_event(current_span, "Acquired session", span_event_attributes)
300300
return session
301301

302302
def put(self, session):
@@ -371,28 +371,30 @@ def get(self):
371371
"""
372372
current_span = get_current_span()
373373
span_event_attributes = {"kind": type(self).__name__}
374-
if current_span:
375-
current_span.add_event("Acquiring session", span_event_attributes)
374+
add_span_event(current_span, "Acquiring session", span_event_attributes)
376375

377376
try:
378-
if current_span:
379-
current_span.add_event(
380-
"Waiting for a session to become available", span_event_attributes
381-
)
377+
add_span_event(
378+
current_span,
379+
"Waiting for a session to become available",
380+
span_event_attributes,
381+
)
382382
session = self._sessions.get_nowait()
383383
except queue.Empty:
384-
if current_span:
385-
current_span.add_event(
386-
"No session available. Creating session", span_event_attributes
387-
)
384+
add_span_event(
385+
current_span,
386+
"No session available. Creating session",
387+
span_event_attributes,
388+
)
388389
session = self._new_session()
389390
session.create()
390391
else:
391392
if not session.exists():
392-
if current_span:
393-
current_span.add_event(
394-
"Session is not valid, recreating it", span_event_attributes
395-
)
393+
add_span_event(
394+
current_span,
395+
"Session is not valid, recreating it",
396+
span_event_attributes,
397+
)
396398
session = self._new_session()
397399
session.create()
398400
return session
@@ -504,23 +506,25 @@ def bind(self, database):
504506

505507
requested_session_count = request.session_count
506508
current_span = get_current_span()
507-
if current_span:
508-
current_span.add_event(
509-
f"Requesting {requested_session_count} sessions", span_event_attributes
510-
)
509+
add_span_event(
510+
current_span,
511+
f"Requesting {requested_session_count} sessions",
512+
span_event_attributes,
513+
)
511514

512515
if created_session_count >= self.size:
513-
if current_span:
514-
current_span.add_event(
515-
"Created no new sessions as sessionPool is full",
516-
span_event_attributes,
517-
)
516+
add_span_event(
517+
current_span,
518+
"Created no new sessions as sessionPool is full",
519+
span_event_attributes,
520+
)
518521
return
519522

520-
if current_span:
521-
current_span.add_event(
522-
f"Creating {request.session_count} sessions", span_event_attributes
523-
)
523+
add_span_event(
524+
current_span,
525+
f"Creating {request.session_count} sessions",
526+
span_event_attributes,
527+
)
524528

525529
returned_session_count = 0
526530
while created_session_count < self.size:
@@ -536,11 +540,11 @@ def bind(self, database):
536540

537541
created_session_count += len(resp.session)
538542

539-
if current_span:
540-
current_span.add_event(
541-
f"Requested for {requested_session_count} sessions, return {returned_session_count}",
542-
span_event_attributes,
543-
)
543+
add_span_event(
544+
current_span,
545+
f"Requested for {requested_session_count} sessions, return {returned_session_count}",
546+
span_event_attributes,
547+
)
544548

545549
def get(self, timeout=None):
546550
"""Check a session out from the pool.
@@ -559,18 +563,18 @@ def get(self, timeout=None):
559563
start_time = time.time()
560564
span_event_attributes = {"kind": type(self).__name__}
561565
current_span = get_current_span()
562-
if current_span:
563-
current_span.add_event(
564-
"Waiting for a session to become available", span_event_attributes
565-
)
566+
add_span_event(
567+
current_span,
568+
"Waiting for a session to become available",
569+
span_event_attributes,
570+
)
566571

567572
ping_after = None
568573
session = None
569574
try:
570575
ping_after, session = self._sessions.get(block=True, timeout=timeout)
571576
except queue.Empty as e:
572-
if current_span:
573-
current_span.add_event("No session available", span_event_attributes)
577+
add_span_event(current_span, "No session available", span_event_attributes)
574578
raise e
575579

576580
if _NOW() > ping_after:
@@ -581,15 +585,14 @@ def get(self, timeout=None):
581585
session = self._new_session()
582586
session.create()
583587

584-
if current_span:
585-
span_event_attributes.update(
586-
{
587-
"time.elapsed": time.time() - start_time,
588-
"session.id": session._session_id,
589-
"kind": "pinging_pool",
590-
}
591-
)
592-
current_span.add_event("Acquired session", span_event_attributes)
588+
span_event_attributes.update(
589+
{
590+
"time.elapsed": time.time() - start_time,
591+
"session.id": session._session_id,
592+
"kind": "pinging_pool",
593+
}
594+
)
595+
add_span_event(current_span, "Acquired session", span_event_attributes)
593596
return session
594597

595598
def put(self, session):

‎google/cloud/spanner_v1/session.py

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
_metadata_with_leader_aware_routing,
3232
)
3333
from google.cloud.spanner_v1._opentelemetry_tracing import (
34+
add_span_event,
3435
get_current_span,
3536
trace_call,
3637
)
@@ -128,8 +129,7 @@ def create(self):
128129
:raises ValueError: if :attr:`session_id` is already set.
129130
"""
130131
current_span = get_current_span()
131-
if current_span:
132-
current_span.add_event("Creating Session")
132+
add_span_event(current_span, "Creating Session")
133133

134134
if self._session_id is not None:
135135
raise ValueError("Session ID already set by back-end")
@@ -173,13 +173,14 @@ def exists(self):
173173
"""
174174
current_span = get_current_span()
175175
if self._session_id is None:
176-
current_span.add_event("Checking if Session failed due to unset session_id")
176+
add_span_event(
177+
current_span, "Checking if Session failed due to unset session_id"
178+
)
177179
return False
178180

179-
if current_span:
180-
current_span.add_event(
181-
"Checking if Session exists", {"session.id": self._session_id}
182-
)
181+
add_span_event(
182+
current_span, "Checking if Session exists", {"session.id": self._session_id}
183+
)
183184

184185
api = self._database.spanner_api
185186
metadata = _metadata_with_prefix(self._database.name)
@@ -216,14 +217,14 @@ def delete(self):
216217
"""
217218
current_span = get_current_span()
218219
if self._session_id is None:
219-
if current_span:
220-
current_span.add_event(
221-
"Deleting Session failed due to unset session_id"
222-
)
220+
add_span_event(
221+
current_span, "Deleting Session failed due to unset session_id"
222+
)
223223
raise ValueError("Session ID not set by back-end")
224224

225-
if current_span:
226-
current_span.add_event("Deleting Session", {"session.id": self._session_id})
225+
add_span_event(
226+
current_span, "Deleting Session", {"session.id": self._session_id}
227+
)
227228

228229
api = self._database.spanner_api
229230
metadata = _metadata_with_prefix(self._database.name)

‎google/cloud/spanner_v1/snapshot.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -538,7 +538,7 @@ def _get_streamed_result_set(
538538
iterator = _restart_on_unavailable(
539539
restart,
540540
request,
541-
"CloudSpanner.execute_sql",
541+
"CloudSpanner.ReadWriteTransaction",
542542
self._session,
543543
trace_attributes,
544544
transaction=self,

0 commit comments

Comments
 (0)