@@ -934,25 +934,19 @@ class Person(Base):
934934
935935 class fake_record_sql_queries : # noqa: N801
936936 def __init__ (self , * args , ** kwargs ):
937- with record_sql_queries_supporting_streaming (
937+ self . _ctx_mgr = record_sql_queries_supporting_streaming (
938938 * args , ** kwargs
939- ) as span :
940- self .span = span
941-
942- if span_streaming :
943- self .span ._start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
944- self .span ._end_timestamp = datetime (
945- 2024 , 1 , 1 , microsecond = 99999
946- )
947- else :
948- self .span .start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
949- self .span .timestamp = datetime (2024 , 1 , 1 , microsecond = 99999 )
939+ )
950940
951941 def __enter__ (self ):
942+ self .span = self ._ctx_mgr .__enter__ ()
943+ self .span ._start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
944+ self .span ._end_timestamp = datetime (2024 , 1 , 1 , microsecond = 99999 )
952945 return self .span
953946
954947 def __exit__ (self , type , value , traceback ):
955- pass
948+ self .span ._end_timestamp = None
949+ self ._ctx_mgr .__exit__ (type , value , traceback )
956950
957951 with mock .patch (
958952 "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming" ,
@@ -1000,25 +994,18 @@ class Person(Base):
1000994
1001995 class fake_record_sql_queries : # noqa: N801
1002996 def __init__ (self , * args , ** kwargs ):
1003- with record_sql_queries_supporting_streaming (
997+ self . _ctx_mgr = record_sql_queries_supporting_streaming (
1004998 * args , ** kwargs
1005- ) as span :
1006- self .span = span
1007-
1008- if span_streaming :
1009- self .span ._start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
1010- self .span ._end_timestamp = datetime (
1011- 2024 , 1 , 1 , microsecond = 99999
1012- )
1013- else :
1014- self .span .start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
1015- self .span .timestamp = datetime (2024 , 1 , 1 , microsecond = 99999 )
999+ )
10161000
10171001 def __enter__ (self ):
1002+ self .span = self ._ctx_mgr .__enter__ ()
1003+ self .span .start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
10181004 return self .span
10191005
10201006 def __exit__ (self , type , value , traceback ):
1021- pass
1007+ self ._ctx_mgr .__exit__ (type , value , traceback )
1008+ self .span .timestamp = datetime (2024 , 1 , 1 , microsecond = 99999 )
10221009
10231010 with mock .patch (
10241011 "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming" ,
@@ -1159,19 +1146,18 @@ class Person(Base):
11591146
11601147 class fake_record_sql_queries : # noqa: N801
11611148 def __init__ (self , * args , ** kwargs ):
1162- with record_sql_queries_supporting_streaming (
1149+ self . _ctx_mgr = record_sql_queries_supporting_streaming (
11631150 * args , ** kwargs
1164- ) as span :
1165- self .span = span
1166-
1167- self .span .start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
1168- self .span .timestamp = datetime (2024 , 1 , 1 , microsecond = 101000 )
1151+ )
11691152
11701153 def __enter__ (self ):
1154+ self .span = self ._ctx_mgr .__enter__ ()
1155+ self .span .start_timestamp = datetime (2024 , 1 , 1 , microsecond = 0 )
11711156 return self .span
11721157
11731158 def __exit__ (self , type , value , traceback ):
1174- pass
1159+ self ._ctx_mgr .__exit__ (type , value , traceback )
1160+ self .span .timestamp = datetime (2024 , 1 , 1 , microsecond = 101000 )
11751161
11761162 with mock .patch (
11771163 "sentry_sdk.integrations.sqlalchemy.record_sql_queries_supporting_streaming" ,
0 commit comments