fix: store affected rows for HranaRows

Author: levydsa

libsql/src/hrana/connection.rs

@@ -28,7 +28,7 @@ where
 
 impl<T> HttpConnection<T>
 where
-    T: HttpSend,
+    T: HttpSend + Send + Sync + 'static,
 {
     pub fn new(url: String, token: String, inner: T) -> Self {
         // The `libsql://` protocol is an alias for `https://`.

libsql/src/hrana/mod.rs

@@ -21,6 +21,7 @@ use libsql_hrana::proto::{Batch, BatchResult, Col, Stmt, StmtResult};
 use std::collections::VecDeque;
 use std::future::Future;
 use std::pin::Pin;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
@@ -122,7 +123,7 @@ where
 
 impl<T> Statement<T>
 where
-    T: HttpSend,
+    T: HttpSend + Send + Sync + 'static,
 {
     pub(crate) fn new(stream: HranaStream<T>, sql: String, want_rows: bool) -> crate::Result<Self> {
         // in SQLite when a multiple statements are glued together into one string, only the first one is
@@ -170,20 +171,20 @@ where
     pub(crate) async fn query_raw(
         &mut self,
         params: &Params,
-    ) -> crate::Result<HranaRows<T::Stream>> {
+    ) -> crate::Result<HranaRows<T::Stream, T>> {
         let mut stmt = self.inner.clone();
         bind_params(params.clone(), &mut stmt);
 
         let cursor = self.stream.cursor(Batch::single(stmt)).await?;
-        let rows = HranaRows::from_cursor(cursor).await?;
+        let rows = HranaRows::from_cursor(cursor, self.stream.clone()).await?;
 
         Ok(rows)
     }
 }
 
 impl<T> Statement<T>
 where
-    T: HttpSend,
+    T: HttpSend + Send + Sync + 'static,
     <T as HttpSend>::Stream: Send + Sync + 'static,
 {
     pub async fn query(&mut self, params: &Params) -> crate::Result<super::Rows> {
@@ -192,28 +193,37 @@ where
     }
 }
 
-pub struct HranaRows<S> {
+pub struct HranaRows<S, T: HttpSend> {
     cursor_step: OwnedCursorStep<S>,
     column_types: Option<Vec<ValueType>>,
+    stream: HranaStream<T>,
 }
 
-impl<S> HranaRows<S>
+impl<S, T> HranaRows<S, T>
 where
+    T: HttpSend + Send + Sync + 'static,
     S: Stream<Item = std::io::Result<Bytes>> + Unpin,
 {
-    async fn from_cursor(cursor: Cursor<S>) -> Result<Self> {
+    async fn from_cursor(cursor: Cursor<S>, stream: HranaStream<T>) -> Result<Self> {
         let cursor_step = cursor.next_step_owned().await?;
         Ok(HranaRows {
             cursor_step,
             column_types: None,
+            stream,
         })
     }
 
     pub async fn next(&mut self) -> crate::Result<Option<super::Row>> {
         let row = match self.cursor_step.next().await {
             Some(Ok(row)) => row,
             Some(Err(e)) => return Err(crate::Error::Hrana(Box::new(e))),
-            None => return Ok(None),
+            None => {
+                self.stream
+                    .inner
+                    .affected_row_count
+                    .store(self.cursor_step.affected_rows().into(), Ordering::SeqCst);
+                return Ok(None);
+            }
         };
 
         if self.column_types.is_none() {
@@ -254,17 +264,19 @@ where
 }
 
 #[async_trait::async_trait]
-impl<S> RowsInner for HranaRows<S>
+impl<S, T> RowsInner for HranaRows<S, T>
 where
+    T: HttpSend + Send + Sync + 'static,
     S: Stream<Item = std::io::Result<Bytes>> + Send + Sync + Unpin,
 {
     async fn next(&mut self) -> crate::Result<Option<super::Row>> {
         self.next().await
     }
 }
 
-impl<S> ColumnsInner for HranaRows<S>
+impl<S, T> ColumnsInner for HranaRows<S, T>
 where
+    T: HttpSend + Send + Sync + 'static,
     S: Stream<Item = std::io::Result<Bytes>> + Send + Sync + Unpin,
 {
     fn column_count(&self) -> i32 {

libsql/src/hrana/stream.rs

@@ -2,12 +2,13 @@ use crate::hrana::cursor::{Cursor, CursorReq};
 use crate::hrana::proto::{Batch, BatchResult, DescribeResult, Stmt, StmtResult};
 use crate::hrana::{CursorResponseError, HranaError, HttpSend, Result};
 use bytes::{Bytes, BytesMut};
-use futures::Stream;
+use futures::{Stream, StreamExt};
 use libsql_hrana::proto::{
     BatchStreamReq, CloseSqlStreamReq, CloseStreamReq, CloseStreamResp, DescribeStreamReq,
     GetAutocommitStreamReq, PipelineReqBody, PipelineRespBody, SequenceStreamReq,
     StoreSqlStreamReq, StreamRequest, StreamResponse, StreamResult,
 };
+use std::pin::pin;
 use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering};
 use std::sync::Arc;
 use tokio::sync::Mutex;
@@ -30,7 +31,7 @@ pub struct HranaStream<T>
 where
     T: HttpSend,
 {
-    inner: Arc<Inner<T>>,
+    pub inner: Arc<Inner<T>>,
 }
 
 impl<T> Clone for HranaStream<T>
@@ -94,7 +95,9 @@ where
             (0, 0)
         };
 
-        self.inner.total_changes.fetch_add(affected_row_count, Ordering::SeqCst);
+        self.inner
+            .total_changes
+            .fetch_add(affected_row_count, Ordering::SeqCst);
         self.inner
             .affected_row_count
             .store(affected_row_count, Ordering::SeqCst);
@@ -279,11 +282,11 @@ where
 }
 
 #[derive(Debug)]
-struct Inner<T>
+pub struct Inner<T>
 where
     T: HttpSend,
 {
-    affected_row_count: AtomicU64,
+    pub affected_row_count: AtomicU64,
     total_changes: AtomicU64,
     last_insert_rowid: AtomicI64,
     is_autocommit: AtomicBool,