From 7c0fa47a5a450328bf57a4dbfe68fa3f96044266 Mon Sep 17 00:00:00 2001
From: AsafMah <asafmahlev@microsoft.com>
Date: Thu, 30 Jun 2022 15:10:57 +0300
Subject: [PATCH] Added back consuming

---
 azure-kusto-data/examples/query.rs       |  2 +-
 azure-kusto-data/src/operations/query.rs | 48 ++++++++++++++++--------
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/azure-kusto-data/examples/query.rs b/azure-kusto-data/examples/query.rs
index 5c4e1ba..8356bc7 100644
--- a/azure-kusto-data/examples/query.rs
+++ b/azure-kusto-data/examples/query.rs
@@ -80,7 +80,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     }
 
     // Print the primary tables
-    let primary_results = response.primary_results().collect::<Vec<_>>();
+    let primary_results = response.into_primary_results().collect::<Vec<_>>();
     println!("primary results: {:#?}", primary_results);
 
     println!("Querying {} with streaming client", args.query);
diff --git a/azure-kusto-data/src/operations/query.rs b/azure-kusto-data/src/operations/query.rs
index 4c9f6c4..da0151e 100644
--- a/azure-kusto-data/src/operations/query.rs
+++ b/azure-kusto-data/src/operations/query.rs
@@ -190,12 +190,12 @@ impl std::convert::TryFrom<KustoResponse> for KustoResponseDataSetV1 {
     }
 }
 
-struct KustoResponseDataSetV2TableIterator<'a, T: Iterator<Item = &'a V2QueryResult>> {
+struct KustoResponseDataSetV2TableIterator<T: Iterator<Item = V2QueryResult>> {
     tables: T,
     finished: bool,
 }
 
-impl<'a, T: Iterator<Item = &'a V2QueryResult>> KustoResponseDataSetV2TableIterator<'a, T> {
+impl<T: Iterator<Item = V2QueryResult>> KustoResponseDataSetV2TableIterator<T> {
     fn new(tables: T) -> Self {
         Self {
             tables,
@@ -204,9 +204,7 @@ impl<'a, T: Iterator<Item = &'a V2QueryResult>> KustoResponseDataSetV2TableItera
     }
 }
 
-impl<'a, T: Iterator<Item = &'a V2QueryResult>> Iterator
-    for KustoResponseDataSetV2TableIterator<'a, T>
-{
+impl<T: Iterator<Item = V2QueryResult>> Iterator for KustoResponseDataSetV2TableIterator<T> {
     type Item = DataTable;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -219,7 +217,7 @@ impl<'a, T: Iterator<Item = &'a V2QueryResult>> Iterator
         });
 
         if let Some(V2QueryResult::DataTable(t)) = next_table {
-            return Some(t.clone());
+            return Some(t);
         }
 
         let mut table = DataTable {
@@ -232,9 +230,9 @@ impl<'a, T: Iterator<Item = &'a V2QueryResult>> Iterator
 
         if let Some(V2QueryResult::TableHeader(header)) = next_table {
             table.table_id = header.table_id;
-            table.table_kind = header.table_kind.clone();
-            table.table_name = header.table_name.clone();
-            table.columns = header.columns.clone();
+            table.table_kind = header.table_kind;
+            table.table_name = header.table_name;
+            table.columns = header.columns;
         } else {
             self.finished = true;
             return None;
@@ -247,8 +245,8 @@ impl<'a, T: Iterator<Item = &'a V2QueryResult>> Iterator
                 V2QueryResult::TableFragment(fragment) => {
                     assert_eq!(fragment.table_id, table.table_id);
                     match fragment.table_fragment_type {
-                        TableFragmentType::DataAppend => table.rows.extend(fragment.rows.clone()),
-                        TableFragmentType::DataReplace => table.rows = fragment.rows.clone(),
+                        TableFragmentType::DataAppend => table.rows.extend(fragment.rows),
+                        TableFragmentType::DataReplace => table.rows = fragment.rows,
                     };
                 }
                 V2QueryResult::TableProgress(progress) => {
@@ -256,7 +254,10 @@ impl<'a, T: Iterator<Item = &'a V2QueryResult>> Iterator
                 }
                 V2QueryResult::TableCompletion(completion) => {
                     assert_eq!(completion.table_id, table.table_id);
-                    assert_eq!(completion.row_count, table.rows.len() as i32);
+                    assert_eq!(
+                        completion.row_count,
+                        TryInto::<i32>::try_into(table.rows.len()).expect("Row count overflow")
+                    );
                     finished_table = true;
                     break;
                 }
@@ -278,20 +279,35 @@ impl KustoResponseDataSetV2 {
         self.tables.len()
     }
 
-    pub fn parsed_data_tables<'a>(&'a self) -> impl Iterator<Item = DataTable> + 'a {
-        KustoResponseDataSetV2TableIterator::new(self.tables.iter())
+    pub fn parsed_data_tables(&self) -> impl Iterator<Item = DataTable> + '_ {
+        KustoResponseDataSetV2TableIterator::new(self.tables.iter().cloned())
     }
 
     /// Consumes the response into an iterator over all PrimaryResult tables within the response dataset
-    pub fn primary_results<'a>(&'a self) -> impl Iterator<Item = DataTable> + 'a {
+    pub fn primary_results(&self) -> impl Iterator<Item = DataTable> + '_ {
         self.parsed_data_tables()
             .filter(|t| t.table_kind == TableKind::PrimaryResult)
     }
 
     #[cfg(feature = "arrow")]
-    pub fn record_batches<'a>(&'a self) -> impl Iterator<Item = Result<RecordBatch>> + 'a {
+    pub fn record_batches(&self) -> impl Iterator<Item = Result<RecordBatch>> + '_ {
         self.primary_results().map(convert_table)
     }
+
+    pub fn into_parsed_data_tables(self) -> impl Iterator<Item = DataTable> {
+        KustoResponseDataSetV2TableIterator::new(self.tables.into_iter())
+    }
+
+    /// Consumes the response into an iterator over all PrimaryResult tables within the response dataset
+    pub fn into_primary_results(self) -> impl Iterator<Item = DataTable> {
+        self.into_parsed_data_tables()
+            .filter(|t| t.table_kind == TableKind::PrimaryResult)
+    }
+
+    #[cfg(feature = "arrow")]
+    pub fn into_record_batches(self) -> impl Iterator<Item = Result<RecordBatch>> {
+        self.into_primary_results().map(convert_table)
+    }
 }
 
 #[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]