@@ -21,6 +21,7 @@ use libsql_hrana::proto::{Batch, BatchResult, Col, Stmt, StmtResult};
21
21
use std:: collections:: VecDeque ;
22
22
use std:: future:: Future ;
23
23
use std:: pin:: Pin ;
24
+ use std:: sync:: atomic:: Ordering ;
24
25
use std:: sync:: Arc ;
25
26
use std:: task:: { Context , Poll } ;
26
27
@@ -122,7 +123,7 @@ where
122
123
123
124
impl < T > Statement < T >
124
125
where
125
- T : HttpSend ,
126
+ T : HttpSend + Send + Sync + ' static ,
126
127
{
127
128
pub ( crate ) fn new ( stream : HranaStream < T > , sql : String , want_rows : bool ) -> crate :: Result < Self > {
128
129
// in SQLite when a multiple statements are glued together into one string, only the first one is
@@ -170,20 +171,20 @@ where
170
171
pub ( crate ) async fn query_raw (
171
172
& mut self ,
172
173
params : & Params ,
173
- ) -> crate :: Result < HranaRows < T :: Stream > > {
174
+ ) -> crate :: Result < HranaRows < T :: Stream , T > > {
174
175
let mut stmt = self . inner . clone ( ) ;
175
176
bind_params ( params. clone ( ) , & mut stmt) ;
176
177
177
178
let cursor = self . stream . cursor ( Batch :: single ( stmt) ) . await ?;
178
- let rows = HranaRows :: from_cursor ( cursor) . await ?;
179
+ let rows = HranaRows :: from_cursor ( cursor, self . stream . clone ( ) ) . await ?;
179
180
180
181
Ok ( rows)
181
182
}
182
183
}
183
184
184
185
impl < T > Statement < T >
185
186
where
186
- T : HttpSend ,
187
+ T : HttpSend + Send + Sync + ' static ,
187
188
<T as HttpSend >:: Stream : Send + Sync + ' static ,
188
189
{
189
190
pub async fn query ( & mut self , params : & Params ) -> crate :: Result < super :: Rows > {
@@ -192,28 +193,37 @@ where
192
193
}
193
194
}
194
195
195
- pub struct HranaRows < S > {
196
+ pub struct HranaRows < S , T : HttpSend > {
196
197
cursor_step : OwnedCursorStep < S > ,
197
198
column_types : Option < Vec < ValueType > > ,
199
+ stream : HranaStream < T > ,
198
200
}
199
201
200
- impl < S > HranaRows < S >
202
+ impl < S , T > HranaRows < S , T >
201
203
where
204
+ T : HttpSend + Send + Sync + ' static ,
202
205
S : Stream < Item = std:: io:: Result < Bytes > > + Unpin ,
203
206
{
204
- async fn from_cursor ( cursor : Cursor < S > ) -> Result < Self > {
207
+ async fn from_cursor ( cursor : Cursor < S > , stream : HranaStream < T > ) -> Result < Self > {
205
208
let cursor_step = cursor. next_step_owned ( ) . await ?;
206
209
Ok ( HranaRows {
207
210
cursor_step,
208
211
column_types : None ,
212
+ stream,
209
213
} )
210
214
}
211
215
212
216
pub async fn next ( & mut self ) -> crate :: Result < Option < super :: Row > > {
213
217
let row = match self . cursor_step . next ( ) . await {
214
218
Some ( Ok ( row) ) => row,
215
219
Some ( Err ( e) ) => return Err ( crate :: Error :: Hrana ( Box :: new ( e) ) ) ,
216
- None => return Ok ( None ) ,
220
+ None => {
221
+ self . stream
222
+ . inner
223
+ . affected_row_count
224
+ . store ( self . cursor_step . affected_rows ( ) . into ( ) , Ordering :: SeqCst ) ;
225
+ return Ok ( None ) ;
226
+ }
217
227
} ;
218
228
219
229
if self . column_types . is_none ( ) {
@@ -254,17 +264,19 @@ where
254
264
}
255
265
256
266
#[ async_trait:: async_trait]
257
- impl < S > RowsInner for HranaRows < S >
267
+ impl < S , T > RowsInner for HranaRows < S , T >
258
268
where
269
+ T : HttpSend + Send + Sync + ' static ,
259
270
S : Stream < Item = std:: io:: Result < Bytes > > + Send + Sync + Unpin ,
260
271
{
261
272
async fn next ( & mut self ) -> crate :: Result < Option < super :: Row > > {
262
273
self . next ( ) . await
263
274
}
264
275
}
265
276
266
- impl < S > ColumnsInner for HranaRows < S >
277
+ impl < S , T > ColumnsInner for HranaRows < S , T >
267
278
where
279
+ T : HttpSend + Send + Sync + ' static ,
268
280
S : Stream < Item = std:: io:: Result < Bytes > > + Send + Sync + Unpin ,
269
281
{
270
282
fn column_count ( & self ) -> i32 {
0 commit comments