1+ import contextlib
12import os
23from datetime import datetime
34from unittest import mock
1112from sqlalchemy import text
1213
1314import sentry_sdk
14- from sentry_sdk import capture_message , start_transaction
1515from sentry_sdk .consts import DEFAULT_MAX_VALUE_LENGTH , SPANDATA
1616from sentry_sdk .integrations .sqlalchemy import SqlalchemyIntegration
1717from sentry_sdk .serializer import MAX_EVENT_BYTES
@@ -54,7 +54,7 @@ class Address(Base):
5454
5555 assert session .query (Person ).first () == bob
5656
57- capture_message ("hi" )
57+ sentry_sdk . capture_message ("hi" )
5858
5959 (event ,) = events
6060
@@ -111,7 +111,7 @@ class Address(Base):
111111 Session = sessionmaker (bind = engine ) # noqa: N806
112112 session = Session ()
113113
114- with start_transaction (name = "test_transaction" , sampled = True ):
114+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
115115 with session .begin_nested ():
116116 session .query (Person ).first ()
117117
@@ -135,7 +135,7 @@ class Address(Base):
135135 assert (
136136 render_span_tree (event )
137137 == """\
138- - op=null : description=null
138+ - op="test_transaction" : description=null
139139 - op="db": description="SAVEPOINT sa_savepoint_1"
140140 - op="db": description="SELECT person.id AS person_id, person.name AS person_name \\ nFROM person\\ n LIMIT ? OFFSET ?"
141141 - op="db": description="RELEASE SAVEPOINT sa_savepoint_1"
@@ -185,7 +185,7 @@ class Address(Base):
185185 Session = sessionmaker (bind = engine ) # noqa: N806
186186 session = Session ()
187187
188- with start_transaction (name = "test_transaction" , sampled = True ):
188+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
189189 with session .begin_nested ():
190190 session .query (Person ).first ()
191191
@@ -217,7 +217,7 @@ def test_long_sql_query_preserved(sentry_init, capture_events):
217217 engine = create_engine (
218218 "sqlite:///:memory:" , connect_args = {"check_same_thread" : False }
219219 )
220- with start_transaction (name = "test" ):
220+ with sentry_sdk . start_span (name = "test" ):
221221 with engine .connect () as con :
222222 con .execute (text (" UNION " .join ("SELECT {}" .format (i ) for i in range (100 ))))
223223
@@ -246,7 +246,7 @@ def processor(event, hint):
246246 engine = create_engine (
247247 "sqlite:///:memory:" , connect_args = {"check_same_thread" : False }
248248 )
249- with start_transaction (name = "test" ):
249+ with sentry_sdk . start_span (name = "test" ):
250250 with engine .connect () as con :
251251 for _ in range (1500 ):
252252 con .execute (
@@ -306,7 +306,7 @@ def test_query_source_disabled(sentry_init, capture_events):
306306
307307 events = capture_events ()
308308
309- with start_transaction (name = "test_transaction" , sampled = True ):
309+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
310310 Base = declarative_base () # noqa: N806
311311
312312 class Person (Base ):
@@ -358,7 +358,7 @@ def test_query_source_enabled(sentry_init, capture_events, enable_db_query_sourc
358358
359359 events = capture_events ()
360360
361- with start_transaction (name = "test_transaction" , sampled = True ):
361+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
362362 Base = declarative_base () # noqa: N806
363363
364364 class Person (Base ):
@@ -405,7 +405,7 @@ def test_query_source(sentry_init, capture_events):
405405 )
406406 events = capture_events ()
407407
408- with start_transaction (name = "test_transaction" , sampled = True ):
408+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
409409 Base = declarative_base () # noqa: N806
410410
411411 class Person (Base ):
@@ -475,7 +475,7 @@ def test_query_source_with_module_in_search_path(sentry_init, capture_events):
475475 query_first_model_from_session ,
476476 )
477477
478- with start_transaction (name = "test_transaction" , sampled = True ):
478+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
479479 Base = declarative_base () # noqa: N806
480480
481481 class Person (Base ):
@@ -533,7 +533,7 @@ def test_no_query_source_if_duration_too_short(sentry_init, capture_events):
533533 )
534534 events = capture_events ()
535535
536- with start_transaction (name = "test_transaction" , sampled = True ):
536+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
537537 Base = declarative_base () # noqa: N806
538538
539539 class Person (Base ):
@@ -601,7 +601,7 @@ def test_query_source_if_duration_over_threshold(sentry_init, capture_events):
601601 )
602602 events = capture_events ()
603603
604- with start_transaction (name = "test_transaction" , sampled = True ):
604+ with sentry_sdk . start_span (name = "test_transaction" , sampled = True ):
605605 Base = declarative_base () # noqa: N806
606606
607607 class Person (Base ):
@@ -620,21 +620,15 @@ class Person(Base):
620620 bob = Person (name = "Bob" )
621621 session .add (bob )
622622
623- class fake_record_sql_queries : # noqa: N801
624- def __init__ ( self , * args , ** kwargs ):
625- with freeze_time (datetime (2024 , 1 , 1 , microsecond = 0 )):
626- with record_sql_queries (* args , ** kwargs ) as span :
627- self . span = span
628- freezer = freeze_time ( datetime ( 2024 , 1 , 1 , microsecond = 99999 ) )
629- freezer . start ()
623+ @ contextlib . contextmanager
624+ def fake_record_sql_queries ( * args , ** kwargs ): # noqa: N801
625+ with freeze_time (datetime (2024 , 1 , 1 , second = 0 )):
626+ with record_sql_queries (* args , ** kwargs ) as span :
627+ freezer = freeze_time ( datetime ( 2024 , 1 , 1 , second = 1 ))
628+ freezer . start ( )
629+ yield span
630630
631- freezer .stop ()
632-
633- def __enter__ (self ):
634- return self .span
635-
636- def __exit__ (self , type , value , traceback ):
637- pass
631+ freezer .stop ()
638632
639633 with mock .patch (
640634 "sentry_sdk.integrations.sqlalchemy.record_sql_queries" ,
@@ -687,7 +681,7 @@ def test_span_origin(sentry_init, capture_events):
687681 engine = create_engine (
688682 "sqlite:///:memory:" , connect_args = {"check_same_thread" : False }
689683 )
690- with start_transaction (name = "foo" ):
684+ with sentry_sdk . start_span (name = "foo" ):
691685 with engine .connect () as con :
692686 con .execute (text ("SELECT 0" ))
693687
0 commit comments