16
16
17
17
from . import _helpers
18
18
from google .cloud .spanner_v1 import Client
19
+ from google .api_core .exceptions import Aborted
20
+ from google .auth .credentials import AnonymousCredentials
21
+ from google .api_core .exceptions import Aborted
22
+ from google .rpc import code_pb2
19
23
20
24
HAS_OTEL_INSTALLED = False
21
25
@@ -132,18 +136,7 @@ def test_propagation(enable_extended_tracing):
132
136
test_propagation (False )
133
137
134
138
135
- @pytest .mark .skipif (
136
- not _helpers .USE_EMULATOR ,
137
- reason = "Emulator needed to run this tests" ,
138
- )
139
- @pytest .mark .skipif (
140
- not HAS_OTEL_INSTALLED ,
141
- reason = "Tracing requires OpenTelemetry" ,
142
- )
143
- def test_transaction_abort_then_retry_spans ():
144
- from google .auth .credentials import AnonymousCredentials
145
- from google .api_core .exceptions import Aborted
146
- from google .rpc import code_pb2
139
+ def create_db_trace_exporter ():
147
140
from opentelemetry .sdk .trace .export import SimpleSpanProcessor
148
141
from opentelemetry .sdk .trace .export .in_memory_span_exporter import (
149
142
InMemorySpanExporter ,
@@ -160,20 +153,6 @@ def test_transaction_abort_then_retry_spans():
160
153
NODE_COUNT = 5
161
154
LABELS = {"test" : "true" }
162
155
163
- counters = dict (aborted = 0 )
164
-
165
- def select_in_txn (txn ):
166
- results = txn .execute_sql ("SELECT 1" )
167
- for row in results :
168
- _ = row
169
-
170
- if counters ["aborted" ] == 0 :
171
- counters ["aborted" ] = 1
172
- raise Aborted (
173
- "Thrown from ClientInterceptor for testing" ,
174
- errors = [_helpers .FauxCall (code_pb2 .ABORTED )],
175
- )
176
-
177
156
tracer_provider = TracerProvider (sampler = ALWAYS_ON )
178
157
trace_exporter = InMemorySpanExporter ()
179
158
tracer_provider .add_span_processor (SimpleSpanProcessor (trace_exporter ))
@@ -207,22 +186,74 @@ def select_in_txn(txn):
207
186
except Exception :
208
187
pass
209
188
189
+ return db , trace_exporter
190
+
191
+
192
+ @pytest .mark .skipif (
193
+ not _helpers .USE_EMULATOR ,
194
+ reason = "Emulator needed to run this test" ,
195
+ )
196
+ @pytest .mark .skipif (
197
+ not HAS_OTEL_INSTALLED ,
198
+ reason = "Tracing requires OpenTelemetry" ,
199
+ )
200
+ def test_transaction_abort_then_retry_spans ():
201
+ from opentelemetry .trace .status import StatusCode
202
+
203
+ db , trace_exporter = create_db_trace_exporter ()
204
+
205
+ counters = dict (aborted = 0 )
206
+
207
+ def select_in_txn (txn ):
208
+ results = txn .execute_sql ("SELECT 1" )
209
+ for row in results :
210
+ _ = row
211
+
212
+ if counters ["aborted" ] == 0 :
213
+ counters ["aborted" ] = 1
214
+ raise Aborted (
215
+ "Thrown from ClientInterceptor for testing" ,
216
+ errors = [_helpers .FauxCall (code_pb2 .ABORTED )],
217
+ )
218
+
210
219
db .run_in_transaction (select_in_txn )
211
220
221
+ got_statuses , got_events = finished_spans_statuses (trace_exporter )
222
+
223
+ # Check for the series of events
224
+ want_events = [
225
+ ("Acquiring session" , {"kind" : "BurstyPool" }),
226
+ ("Waiting for a session to become available" , {"kind" : "BurstyPool" }),
227
+ ("No sessions available in pool. Creating session" , {"kind" : "BurstyPool" }),
228
+ ("Creating Session" , {}),
229
+ ("Creating Transaction" , {}),
230
+ (
231
+ "Transaction was aborted in user operation, retrying" ,
232
+ {"delay_seconds" : "EPHEMERAL" , "cause" : "EPHEMERAL" , "attempt" : 1 },
233
+ ),
234
+ ("Creating Transaction" , {}),
235
+ ("Starting Commit" , {}),
236
+ ("Commit Done" , {}),
237
+ ]
238
+ assert got_events == want_events
239
+
240
+ # Check for the statues.
241
+ codes = StatusCode
242
+ want_statuses = [
243
+ ("CloudSpanner.Database.run_in_transaction" , codes .OK , None ),
244
+ ("CloudSpanner.CreateSession" , codes .OK , None ),
245
+ ("CloudSpanner.Session.run_in_transaction" , codes .OK , None ),
246
+ ("CloudSpanner.Transaction.execute_streaming_sql" , codes .OK , None ),
247
+ ("CloudSpanner.Transaction.execute_streaming_sql" , codes .OK , None ),
248
+ ("CloudSpanner.Transaction.commit" , codes .OK , None ),
249
+ ]
250
+ assert got_statuses == want_statuses
251
+
252
+
253
+ def finished_spans_statuses (trace_exporter ):
212
254
span_list = trace_exporter .get_finished_spans ()
213
255
# Sort the spans by their start time in the hierarchy.
214
256
span_list = sorted (span_list , key = lambda span : span .start_time )
215
- got_span_names = [span .name for span in span_list ]
216
- want_span_names = [
217
- "CloudSpanner.Database.run_in_transaction" ,
218
- "CloudSpanner.CreateSession" ,
219
- "CloudSpanner.Session.run_in_transaction" ,
220
- "CloudSpanner.Transaction.execute_streaming_sql" ,
221
- "CloudSpanner.Transaction.execute_streaming_sql" ,
222
- "CloudSpanner.Transaction.commit" ,
223
- ]
224
-
225
- assert got_span_names == want_span_names
226
257
227
258
got_events = []
228
259
got_statuses = []
@@ -234,6 +265,7 @@ def select_in_txn(txn):
234
265
got_statuses .append (
235
266
(span .name , span .status .status_code , span .status .description )
236
267
)
268
+
237
269
for event in span .events :
238
270
evt_attributes = event .attributes .copy ()
239
271
for attr_name in imprecise_event_attributes :
@@ -242,30 +274,70 @@ def select_in_txn(txn):
242
274
243
275
got_events .append ((event .name , evt_attributes ))
244
276
277
+ return got_statuses , got_events
278
+
279
+
280
+ @pytest .mark .skipif (
281
+ not _helpers .USE_EMULATOR ,
282
+ reason = "Emulator needed to run this test" ,
283
+ )
284
+ @pytest .mark .skipif (
285
+ not HAS_OTEL_INSTALLED ,
286
+ reason = "Tracing requires OpenTelemetry" ,
287
+ )
288
+ def test_database_partitioned ():
289
+ from opentelemetry .trace .status import StatusCode
290
+
291
+ db , trace_exporter = create_db_trace_exporter ()
292
+
293
+ try :
294
+ db .execute_partitioned_dml ("UPDATE NonExistent SET name = 'foo' WHERE id > 1" )
295
+ except Exception :
296
+ pass
297
+
298
+ got_statuses , got_events = finished_spans_statuses (trace_exporter )
245
299
# Check for the series of events
246
300
want_events = [
247
301
("Acquiring session" , {"kind" : "BurstyPool" }),
248
302
("Waiting for a session to become available" , {"kind" : "BurstyPool" }),
249
303
("No sessions available in pool. Creating session" , {"kind" : "BurstyPool" }),
250
304
("Creating Session" , {}),
305
+ ("Starting BeginTransaction" , {}),
251
306
(
252
- "Transaction was aborted in user operation, retrying" ,
253
- {"delay_seconds" : "EPHEMERAL" , "cause" : "EPHEMERAL" , "attempt" : 1 },
307
+ "exception" ,
308
+ {
309
+ "exception.type" : "google.api_core.exceptions.InvalidArgument" ,
310
+ "exception.message" : "400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
311
+ "exception.stacktrace" : "EPHEMERAL" ,
312
+ "exception.escaped" : "False" ,
313
+ },
314
+ ),
315
+ (
316
+ "exception" ,
317
+ {
318
+ "exception.type" : "google.api_core.exceptions.InvalidArgument" ,
319
+ "exception.message" : "400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
320
+ "exception.stacktrace" : "EPHEMERAL" ,
321
+ "exception.escaped" : "False" ,
322
+ },
254
323
),
255
- ("Starting Commit" , {}),
256
- ("Commit Done" , {}),
257
324
]
258
325
assert got_events == want_events
259
326
260
327
# Check for the statues.
261
328
codes = StatusCode
262
329
want_statuses = [
263
- ("CloudSpanner.Database.run_in_transaction" , codes .OK , None ),
330
+ (
331
+ "CloudSpanner.Database.execute_partitioned_pdml" ,
332
+ codes .ERROR ,
333
+ "InvalidArgument: 400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
334
+ ),
264
335
("CloudSpanner.CreateSession" , codes .OK , None ),
265
- ("CloudSpanner.Session.run_in_transaction" , codes .OK , None ),
266
- ("CloudSpanner.Transaction.execute_streaming_sql" , codes .OK , None ),
267
- ("CloudSpanner.Transaction.execute_streaming_sql" , codes .OK , None ),
268
- ("CloudSpanner.Transaction.commit" , codes .OK , None ),
336
+ (
337
+ "CloudSpanner.ExecuteStreamingSql" ,
338
+ codes .ERROR ,
339
+ "InvalidArgument: 400 Table not found: NonExistent [at 1:8]\n UPDATE NonExistent SET name = 'foo' WHERE id > 1\n ^" ,
340
+ ),
269
341
]
270
342
assert got_statuses == want_statuses
271
343
0 commit comments