81
81
tracing:: trace!( "baton not found - skipping finalize for {:?}" , req) ;
82
82
return Ok ( false ) ;
83
83
}
84
- let ( resp, is_autocommit) = client. finalize ( req) . await ?;
84
+ let ( resp, is_autocommit) = client. finalize ( req) . await . map_err ( |e| match e {
85
+ HranaError :: Api ( e) if e == "Received an invalid baton" && !self . is_autocommit ( ) => {
86
+ HranaError :: Api ( "Received an invalid baton: Transaction rolled back" . into ( ) )
87
+ }
88
+ e => e,
89
+ } ) ?;
85
90
self . inner
86
91
. is_autocommit
87
92
. store ( is_autocommit, Ordering :: SeqCst ) ;
94
99
( 0 , 0 )
95
100
} ;
96
101
97
- self . inner . total_changes . fetch_add ( affected_row_count, Ordering :: SeqCst ) ;
102
+ self . inner
103
+ . total_changes
104
+ . fetch_add ( affected_row_count, Ordering :: SeqCst ) ;
98
105
self . inner
99
106
. affected_row_count
100
107
. store ( affected_row_count, Ordering :: SeqCst ) ;
@@ -132,7 +139,15 @@ where
132
139
StreamRequest :: GetAutocommit ( GetAutocommitStreamReq { } ) ,
133
140
StreamRequest :: Close ( CloseStreamReq { } ) ,
134
141
] )
135
- . await ?;
142
+ . await
143
+ . map_err ( |e| match e {
144
+ HranaError :: Api ( e)
145
+ if e == "Received an invalid baton" && !self . is_autocommit ( ) =>
146
+ {
147
+ HranaError :: Api ( "Received an invalid baton: Transaction rolled back" . into ( ) )
148
+ }
149
+ e => e,
150
+ } ) ?;
136
151
client. reset ( ) ;
137
152
( resp, get_autocommit)
138
153
} else {
@@ -142,7 +157,15 @@ where
142
157
StreamRequest :: Batch ( BatchStreamReq { batch } ) ,
143
158
StreamRequest :: GetAutocommit ( GetAutocommitStreamReq { } ) ,
144
159
] )
145
- . await ?;
160
+ . await
161
+ . map_err ( |e| match e {
162
+ HranaError :: Api ( e)
163
+ if e == "Received an invalid baton" && !self . is_autocommit ( ) =>
164
+ {
165
+ HranaError :: Api ( "Received an invalid baton: Transaction rolled back" . into ( ) )
166
+ }
167
+ e => e,
168
+ } ) ?;
146
169
( resp, get_autocommit)
147
170
} ;
148
171
drop ( client) ;
0 commit comments