Repository: scylladb/scylla-rust-driver
Pull request #309 (2021-10-01 14:39:22)

Add a method to query data using multiple paged queries
## Description
This PR adds new methods: `query_all` and `execute_all`.  
They use multiple paged queries to fetch results, and then merge all those results into one.  
`page_size` must be set on the query that is passed to those methods.

Currently we use a single unpaged query to fetch data, but this might cause stalls (#298).
It's better to use many smaller queries and merge results locally.

```rust
Connection::query_all(
        &self,
        query: &Query,
        values: impl ValueList,
    ) -> Result<QueryResult, QueryError>;

Connection::execute_all(
        &self,
        prepared_statement: &PreparedStatement,
        values: impl ValueList,
    ) -> Result<QueryResult, QueryError>;
```

## Fixed Issues
Fixes: #298

## Pre-review checklist

<!--
    Make sure you took care of the issues on the list.
    Put 'x' into those boxes which apply.
    You can also create the PR now and click on all relevant checkboxes.
    See CONTRIBUTING.md for more details.
-->

- [x] I have split my patch into logically separate commits.
- [x] All commit messages clearly explain what they change and why.
- [x] I added relevant tests for new features and bug fixes.
- [x] All commits compile, pass static checks and pass test.
- [x] PR description sums up the changes and reasons why they should be introduced.
- [x] I added appropriate `Fixes:` annotations to PR description.


Jasperav: I thought the `execute_iter` gave back an iterator which lazy loads new page row results when the current page is exhausted, doesn't it do the same thing? It looks more memory efficient as well (if you have millions of rows, `execute_iter` won't care since it lazy loads the results and deallocates old rows). If you have millions of rows and you call this method, changes are you get out of memory exceptions, or am I misunderstanding this method?
psarna: @Jasperav using iterators is absolutely preferred, especially for large datasets. This set of helper functions is mainly for internal use, where we know that the set is small, but we still want to use paging in order to avoid unnecessarily stressing the server - e.g. when fetching topology information.
cvybhu: Ok @psarna @piodul I think it's ready for review
psarna: @Jasperav ongoing. Note however, that this interface is intended to be used mostly for internal purposes, e.g. topology fetching. User data is often too large to gather everything in memory at one go (like these functions do), so it's **strongly** preferred to use our iterator interface. `query_all` and friends should only be used if you know that the result will be small and you don't want to worry about fetching multiple pages manually or advancing an iterator by hand.
Jasperav: @piodul I am using this crate in my toy project and I actually needed `query_iter` on `Session`. It is a query on a small (<20 rows) set. If I don't have this method, I will end up writing my own implementation (probably initializing an empty vec, use the iterator to add rows). 

That is something this method already will do. I think other people sometimes just need all the rows in memory to send back to a client. So IMO this methods should remain on `Session`.
Jasperav: The `try_collect` method bounds are Default and Extend, which are implemented for simple primitives in your example (i32, String) but not for every struct. Forcing clients to implement default and extend for types that will collect an async stream isn't desired IMO.
cvybhu: v2:
* Removed Session::query_all, it's backed up on the [queryall-backup](https://github.com/cvybhu/scylla-rust-driver/tree/queryall-backup) branch in case we wanted it in the future
cvybhu: v2:
* Removed Session::query_all, it's backed up on the [queryall-backup](https://github.com/cvybhu/scylla-rust-driver/tree/queryall-backup) branch in case we wanted it in the future
* Removed new `BadQuery` variant, `Connection::query_all` now uses `ProtocolError` to report bad usage.

Patch:
diff --git a/docs/source/queries/batch.md b/docs/source/queries/batch.md
index 5343fd5a..28bc7cce 100644
--- a/docs/source/queries/batch.md
+++ b/docs/source/queries/batch.md
@@ -20,7 +20,7 @@ let mut batch: Batch = Default::default();
 batch.append_statement("INSERT INTO ks.tab(a, b) VALUES(?, ?)");
 
 // Add a simple query created manually to the batch
-let simple: Query = Query::new("INSERT INTO ks.tab (a, b) VALUES(3, 4)".to_string());
+let simple: Query = Query::new("INSERT INTO ks.tab (a, b) VALUES(3, 4)");
 batch.append_statement(simple);
 
 // Add a prepared query to the batch
diff --git a/docs/source/queries/paged.md b/docs/source/queries/paged.md
index dd419b29..10d91d35 100644
--- a/docs/source/queries/paged.md
+++ b/docs/source/queries/paged.md
@@ -70,7 +70,7 @@ On a `Query`:
 # async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
 use scylla::query::Query;
 
-let mut query: Query = Query::new("SELECT a, b FROM ks.t".to_string());
+let mut query: Query = Query::new("SELECT a, b FROM ks.t");
 query.set_page_size(16);
 
 let _ = session.query_iter(query, &[]).await?; // ...
@@ -110,7 +110,7 @@ On a `Query`:
 # async fn check_only_compiles(session: &Session) -> Result<(), Box<dyn Error>> {
 use scylla::query::Query;
 
-let paged_query = Query::new("SELECT a, b, c FROM ks.t".to_owned()).with_page_size(6);
+let paged_query = Query::new("SELECT a, b, c FROM ks.t").with_page_size(6);
 let res1 = session.query(paged_query.clone(), &[]).await?;
 let res2 = session
     .query_paged(paged_query.clone(), &[], res1.paging_state)
@@ -128,7 +128,7 @@ On a `PreparedStatement`:
 use scylla::query::Query;
 
 let paged_prepared = session
-    .prepare(Query::new("SELECT a, b, c FROM ks.t".to_owned()).with_page_size(7))
+    .prepare(Query::new("SELECT a, b, c FROM ks.t").with_page_size(7))
     .await?;
 let res1 = session.execute(&paged_prepared, &[]).await?;
 let res2 = session
diff --git a/docs/source/queries/prepared.md b/docs/source/queries/prepared.md
index 7f7ebb06..307537e5 100644
--- a/docs/source/queries/prepared.md
+++ b/docs/source/queries/prepared.md
@@ -27,8 +27,11 @@ session.execute(&prepared, (to_insert,)).await?;
 > must be sent as bound values (see [performance section](#performance))
 
 > ***Warning***  
-> Prepared query returns only a single page of results.
-> If number of rows might exceed single page size use a [paged query](paged.md) instead.
+> Don't use `execute` to receive large amounts of data.  
+> By default the query is unpaged and might cause heavy load on the cluster.
+> In such cases set a page size and use a [paged query](paged.md) instead.
+> 
+> When page size is set, `execute` will return only the first page of results.
 
 ### `Session::prepare`
 `Session::prepare` takes query text and prepares the query on all nodes and shards.
diff --git a/docs/source/queries/simple.md b/docs/source/queries/simple.md
index 79b0682f..0bf54d9e 100644
--- a/docs/source/queries/simple.md
+++ b/docs/source/queries/simple.md
@@ -16,8 +16,11 @@ session
 ```
 
 > ***Warning***  
-> Simple query returns only a single page of results.
-> If number of rows might exceed single page size use a [paged query](paged.md) instead.  
+> Don't use simple query to receive large amounts of data.  
+> By default the query is unpaged and might cause heavy load on the cluster.  
+> In such cases set a page size and use [paged query](paged.md) instead.  
+> 
+> When page size is set, `query` will return only the first page of results.
 
 ### First argument - the query
 As the first argument `Session::query` takes anything implementing `Into<Query>`.  
@@ -31,7 +34,7 @@ use scylla::query::Query;
 use scylla::statement::Consistency;
 
 // Create a Query manually to change the Consistency to ONE
-let mut my_query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(?)".to_string());
+let mut my_query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(?)");
 my_query.set_consistency(Consistency::One);
 
 // Insert a value into the table
@@ -86,8 +89,8 @@ if let Some(rows) = session.query("SELECT a FROM ks.tab", &[]).await?.rows {
 # Ok(())
 # }
 ```
-> Simple query returns only a single page of results.
-> If number of rows might exceed single page size use a [paged query](paged.md) instead.  
+> In cases where page size is set, simple query returns only a single page of results.  
+> To receive all pages use a [paged query](paged.md) instead.  
 
 See [Query result](result.md) for more information about handling query results
 
diff --git a/docs/source/retry-policy/default.md b/docs/source/retry-policy/default.md
index c57ac88c..d113521a 100644
--- a/docs/source/retry-policy/default.md
+++ b/docs/source/retry-policy/default.md
@@ -32,7 +32,7 @@ use scylla::query::Query;
 use scylla::transport::retry_policy::DefaultRetryPolicy;
 
 // Create a Query manually and set the retry policy
-let mut my_query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(?)".to_string());
+let mut my_query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(?)");
 my_query.set_retry_policy(Box::new(DefaultRetryPolicy::new()));
 
 // Run the query using this retry policy
diff --git a/docs/source/retry-policy/fallthrough.md b/docs/source/retry-policy/fallthrough.md
index a965d0cb..ca78898e 100644
--- a/docs/source/retry-policy/fallthrough.md
+++ b/docs/source/retry-policy/fallthrough.md
@@ -31,7 +31,7 @@ use scylla::query::Query;
 use scylla::transport::retry_policy::FallthroughRetryPolicy;
 
 // Create a Query manually and set the retry policy
-let mut my_query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(?)".to_string());
+let mut my_query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(?)");
 my_query.set_retry_policy(Box::new(FallthroughRetryPolicy::new()));
 
 // Run the query using this retry policy
diff --git a/docs/source/retry-policy/retry-policy.md b/docs/source/retry-policy/retry-policy.md
index 23f19601..eb220bf9 100644
--- a/docs/source/retry-policy/retry-policy.md
+++ b/docs/source/retry-policy/retry-policy.md
@@ -26,7 +26,7 @@ use scylla::query::Query;
 use scylla::prepared_statement::PreparedStatement;
 
 // Specify that a Query is idempotent
-let mut my_query: Query = Query::new("SELECT a FROM ks.tab".to_string());
+let mut my_query: Query = Query::new("SELECT a FROM ks.tab");
 my_query.set_is_idempotent(true);
 
 
diff --git a/docs/source/tracing/paged.md b/docs/source/tracing/paged.md
index cdbcb995..774f39e7 100644
--- a/docs/source/tracing/paged.md
+++ b/docs/source/tracing/paged.md
@@ -19,7 +19,7 @@ use futures::StreamExt;
 use uuid::Uuid;
 
 // Create a Query manually and enable tracing
-let mut query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(4)".to_string());
+let mut query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(4)");
 query.set_tracing(true);
 
 // Create a paged query iterator and fetch pages
diff --git a/docs/source/tracing/prepare.md b/docs/source/tracing/prepare.md
index a520b02a..2f3850cb 100644
--- a/docs/source/tracing/prepare.md
+++ b/docs/source/tracing/prepare.md
@@ -14,7 +14,7 @@ use scylla::tracing::TracingInfo;
 use uuid::Uuid;
 
 // Prepare the query with tracing enabled
-let mut to_prepare: Query = Query::new("SELECT a FROM ks.tab".to_string());
+let mut to_prepare: Query = Query::new("SELECT a FROM ks.tab");
 to_prepare.set_tracing(true);
 
 let mut prepared: PreparedStatement = session
diff --git a/docs/source/tracing/simple-prepared.md b/docs/source/tracing/simple-prepared.md
index 782c98f2..6eab2e64 100644
--- a/docs/source/tracing/simple-prepared.md
+++ b/docs/source/tracing/simple-prepared.md
@@ -16,7 +16,7 @@ use scylla::tracing::TracingInfo;
 use uuid::Uuid;
 
 // Create a Query manually and enable tracing
-let mut query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(4)".to_string());
+let mut query: Query = Query::new("INSERT INTO ks.tab (a) VALUES(4)");
 query.set_tracing(true);
 
 let res: QueryResult = session.query(query, &[]).await?;
diff --git a/examples/select-paging.rs b/examples/select-paging.rs
index 3de4e64f..ac6e2826 100644
--- a/examples/select-paging.rs
+++ b/examples/select-paging.rs
@@ -40,7 +40,7 @@ async fn main() -> Result<()> {
         println!("a, b, c: {}, {}, {}", a, b, c);
     }
 
-    let paged_query = Query::new("SELECT a, b, c FROM ks.t".to_owned()).with_page_size(6);
+    let paged_query = Query::new("SELECT a, b, c FROM ks.t").with_page_size(6);
     let res1 = session.query(paged_query.clone(), &[]).await?;
     println!(
         "Paging state: {:#?} ({} rows)",
@@ -65,7 +65,7 @@ async fn main() -> Result<()> {
     );
 
     let paged_prepared = session
-        .prepare(Query::new("SELECT a, b, c FROM ks.t".to_owned()).with_page_size(7))
+        .prepare(Query::new("SELECT a, b, c FROM ks.t").with_page_size(7))
         .await?;
     let res4 = session.execute(&paged_prepared, &[]).await?;
     println!(
diff --git a/examples/tracing.rs b/examples/tracing.rs
index f7ff0702..98ba7ead 100644
--- a/examples/tracing.rs
+++ b/examples/tracing.rs
@@ -34,7 +34,7 @@ async fn main() -> Result<()> {
 
     // QUERY
     // Create a simple query and enable tracing for it
-    let mut query: Query = Query::new("SELECT val from ks.tracing_example".to_string());
+    let mut query: Query = Query::new("SELECT val from ks.tracing_example");
     query.set_tracing(true);
     query.set_serial_consistency(Some(SerialConsistency::LocalSerial));
 
diff --git a/scylla/src/statement/query.rs b/scylla/src/statement/query.rs
index 378d219b..eac654c0 100644
--- a/scylla/src/statement/query.rs
+++ b/scylla/src/statement/query.rs
@@ -15,9 +15,9 @@ pub struct Query {
 
 impl Query {
     /// Creates a new `Query` from a CQL query string.
-    pub fn new(contents: String) -> Self {
+    pub fn new(query_text: impl Into<String>) -> Self {
         Self {
-            contents,
+            contents: query_text.into(),
             page_size: None,
             config: Default::default(),
         }
diff --git a/scylla/src/transport/connection.rs b/scylla/src/transport/connection.rs
index 22640744..3cbde348 100644
--- a/scylla/src/transport/connection.rs
+++ b/scylla/src/transport/connection.rs
@@ -102,6 +102,23 @@ impl QueryResult {
             .enumerate()
             .find(|(_id, spec)| spec.name == name)
     }
+
+    /// This function is used to merge results of multiple paged queries into one.  
+    /// other is the result of a new paged query.  
+    /// It is merged with current result kept in self.  
+    pub(crate) fn merge_with_next_page_res(&mut self, other: QueryResult) {
+        if let Some(other_rows) = other.rows {
+            match &mut self.rows {
+                Some(self_rows) => self_rows.extend(other_rows),
+                None => self.rows = Some(other_rows),
+            }
+        };
+
+        self.warnings.extend(other.warnings);
+        self.tracing_id = other.tracing_id;
+        self.paging_state = other.paging_state;
+        self.col_specs = other.col_specs;
+    }
 }
 
 /// Result of Session::batch(). Contains no rows, only some useful information.
@@ -351,6 +368,45 @@ impl Connection {
             .await
     }
 
+    /// Performs query_single_page multiple times to query all available pages
+    pub async fn query_all(
+        &self,
+        query: &Query,
+        values: impl ValueList,
+    ) -> Result<QueryResult, QueryError> {
+        if query.get_page_size().is_none() {
+            // Page size should be set when someone wants to use paging
+            // This is an internal function, we can use ProtocolError even if it's not technically desigen for this
+            return Err(QueryError::ProtocolError(
+                "Called Connection::query_all without page size set!",
+            ));
+        }
+
+        let mut final_result = QueryResult::default();
+
+        let serialized_values = values.serialized()?;
+        let mut paging_state: Option<Bytes> = None;
+
+        loop {
+            // Send next paged query
+            let mut cur_result: QueryResult = self
+                .query(query, &serialized_values, paging_state)
+                .await?
+                .into_query_result()?;
+
+            // Set paging_state for the next query
+            paging_state = cur_result.paging_state.take();
+
+            // Add current query results to the final_result
+            final_result.merge_with_next_page_res(cur_result);
+
+            if paging_state.is_none() {
+                // No more pages to query, we can return the final result
+                return Ok(final_result);
+            }
+        }
+    }
+
     pub async fn execute_single_page(
         &self,
         prepared_statement: &PreparedStatement,
@@ -408,6 +464,43 @@ impl Connection {
         Ok(query_response)
     }
 
+    /// Performs execute_single_page multiple times to fetch all available pages
+    pub async fn execute_all(
+        &self,
+        prepared_statement: &PreparedStatement,
+        values: impl ValueList,
+    ) -> Result<QueryResult, QueryError> {
+        if prepared_statement.get_page_size().is_none() {
+            // This is an internal function, we can use ProtocolError even if it's not technically desigen for this
+            return Err(QueryError::ProtocolError(
+                "Called Connection::execute_all without page size set!",
+            ));
+        }
+
+        let mut final_result = QueryResult::default();
+
+        let serialized_values = values.serialized()?;
+        let mut paging_state: Option<Bytes> = None;
+
+        loop {
+            // Send next paged query
+            let mut cur_result: QueryResult = self
+                .execute_single_page(prepared_statement, &serialized_values, paging_state)
+                .await?;
+
+            // Set paging_state for the next query
+            paging_state = cur_result.paging_state.take();
+
+            // Add current query results to the final_result
+            final_result.merge_with_next_page_res(cur_result);
+
+            if paging_state.is_none() {
+                // No more pages to query, we can return the final result
+                return Ok(final_result);
+            }
+        }
+    }
+
     pub async fn batch(
         &self,
         batch: &Batch,
@@ -1082,3 +1175,126 @@ impl VerifiedKeyspaceName {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::super::errors::QueryError;
+    use super::ConnectionConfig;
+    use crate::query::Query;
+    use crate::IntoTypedRows;
+    use std::net::SocketAddr;
+
+    // Just like resolve_hostname in session.rs
+    async fn resolve_hostname(hostname: &str) -> SocketAddr {
+        match tokio::net::lookup_host(hostname).await {
+            Ok(mut addrs) => addrs.next().unwrap(),
+            Err(_) => {
+                tokio::net::lookup_host((hostname, 9042)) // Port might not be specified, try default
+                    .await
+                    .unwrap()
+                    .next()
+                    .unwrap()
+            }
+        }
+    }
+
+    /// Tests for Connection::query_all and Connection::execute_all
+    /// 1. SELECT from an empty table.
+    /// 2. Create table and insert ints 0..100.
+    ///    Then use query_all and execute_all with page_size set to 7 to select all 100 rows.
+    /// 3. INSERT query_all should have None in result rows.
+    /// 4. Calling query_all with a Query that doesn't have page_size set should result in an error.
+    #[tokio::test]
+    async fn connection_query_all_execute_all_test() {
+        let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+        let addr: SocketAddr = resolve_hostname(&uri).await;
+
+        let (connection, _) = super::open_connection(addr, None, ConnectionConfig::default())
+            .await
+            .unwrap();
+
+        connection.query_single_page("CREATE KEYSPACE IF NOT EXISTS ks WITH REPLICATION = {'class' : 'SimpleStrategy', 'replication_factor' : 1}", &[]).await.unwrap();
+        connection
+            .query_single_page("DROP TABLE IF EXISTS ks.connection_query_all_tab", &[])
+            .await
+            .unwrap();
+        connection
+            .query_single_page(
+                "CREATE TABLE IF NOT EXISTS ks.connection_query_all_tab (p int primary key)",
+                &[],
+            )
+            .await
+            .unwrap();
+
+        // 1. SELECT from an empty table returns query result where rows are Some(Vec::new())
+        let select_query =
+            Query::new("SELECT p FROM ks.connection_query_all_tab").with_page_size(7);
+        let empty_res = connection.query_all(&select_query, &[]).await.unwrap();
+        assert!(empty_res.rows.unwrap().is_empty());
+
+        let mut prepared_select = connection.prepare(&select_query).await.unwrap();
+        prepared_select.set_page_size(7);
+        let empty_res_prepared = connection.execute_all(&prepared_select, &[]).await.unwrap();
+        assert!(empty_res_prepared.rows.unwrap().is_empty());
+
+        // 2. Insert 100 and select using query_all with page_size 7
+        let values: Vec<i32> = (0..100).collect();
+        let mut insert_futures = Vec::new();
+        let insert_query =
+            Query::new("INSERT INTO ks.connection_query_all_tab (p) VALUES (?)").with_page_size(7);
+        for v in &values {
+            insert_futures.push(connection.query_single_page(insert_query.clone(), (v,)));
+        }
+
+        futures::future::try_join_all(insert_futures).await.unwrap();
+
+        let mut results: Vec<i32> = connection
+            .query_all(&select_query, &[])
+            .await
+            .unwrap()
+            .rows
+            .unwrap()
+            .into_typed::<(i32,)>()
+            .map(|r| r.unwrap().0)
+            .collect();
+        results.sort_unstable(); // Clippy recommended to use sort_unstable instead of sort()
+        assert_eq!(results, values);
+
+        let mut results2: Vec<i32> = connection
+            .execute_all(&prepared_select, &[])
+            .await
+            .unwrap()
+            .rows
+            .unwrap()
+            .into_typed::<(i32,)>()
+            .map(|r| r.unwrap().0)
+            .collect();
+        results2.sort_unstable();
+        assert_eq!(results2, values);
+
+        // 3. INSERT query_all should have None in result rows.
+        let insert_res1 = connection.query_all(&insert_query, (0,)).await.unwrap();
+        assert!(insert_res1.rows.is_none());
+
+        let prepared_insert = connection.prepare(&insert_query).await.unwrap();
+        let insert_res2 = connection
+            .execute_all(&prepared_insert, (0,))
+            .await
+            .unwrap();
+        assert!(insert_res2.rows.is_none(),);
+
+        // 4. Calling query_all with a Query that doesn't have page_size set should result in an error.
+        let no_page_size_query = Query::new("SELECT p FROM ks.connection_query_all_tab");
+        let no_page_res = connection.query_all(&no_page_size_query, &[]).await;
+        assert!(matches!(no_page_res, Err(QueryError::ProtocolError(_))));
+
+        let prepared_no_page_size_query = connection.prepare(&no_page_size_query).await.unwrap();
+        let prepared_no_page_res = connection
+            .execute_all(&prepared_no_page_size_query, &[])
+            .await;
+        assert!(matches!(
+            prepared_no_page_res,
+            Err(QueryError::ProtocolError(_))
+        ));
+    }
+}
diff --git a/scylla/src/transport/session.rs b/scylla/src/transport/session.rs
index 8954f93f..21a5d7d6 100644
--- a/scylla/src/transport/session.rs
+++ b/scylla/src/transport/session.rs
@@ -896,15 +896,15 @@ impl Session {
         consistency: Consistency,
     ) -> Result<Option<TracingInfo>, QueryError> {
         // Query system_traces.sessions for TracingInfo
-        let mut traces_session_query =
-            Query::new(crate::tracing::TRACES_SESSION_QUERY_STR.to_string());
+        let mut traces_session_query = Query::new(crate::tracing::TRACES_SESSION_QUERY_STR);
         traces_session_query.config.consistency = consistency;
+        traces_session_query.set_page_size(1024);
 
         // Query system_traces.events for TracingEvents
-        let mut traces_events_query =
-            Query::new(crate::tracing::TRACES_EVENTS_QUERY_STR.to_string());
+        let mut traces_events_query = Query::new(crate::tracing::TRACES_EVENTS_QUERY_STR);
         traces_events_query.config.consistency = Consistency::One;
         traces_events_query.config.consistency = consistency;
+        traces_events_query.set_page_size(1024);
 
         let (traces_session_res, traces_events_res) = tokio::try_join!(
             self.query(traces_session_query, (tracing_id,)),
diff --git a/scylla/src/transport/topology.rs b/scylla/src/transport/topology.rs
index db17423d..1f0f96f5 100644
--- a/scylla/src/transport/topology.rs
+++ b/scylla/src/transport/topology.rs
@@ -5,6 +5,7 @@ use crate::transport::connection_pool::{NodeConnectionPool, PoolConfig, PoolSize
 use crate::transport::errors::QueryError;
 use crate::transport::session::IntoTypedRows;
 
+use crate::statement::query::Query;
 use rand::seq::SliceRandom;
 use rand::thread_rng;
 use std::collections::HashMap;
@@ -209,17 +210,16 @@ async fn query_topology_info(
 }
 
 async fn query_peers(conn: &Connection, connect_port: u16) -> Result<Vec<Peer>, QueryError> {
-    // There shouldn't be more peers than a single page capacity
-    let peers_query = conn.query_single_page(
-        "select peer, data_center, rack, tokens from system.peers",
-        &[],
-    );
-    let local_query = conn.query_single_page(
-        "select rpc_address, data_center, rack, tokens from system.local",
-        &[],
-    );
-
-    let (peers_res, local_res) = tokio::try_join!(peers_query, local_query)?;
+    let mut peers_query = Query::new("select peer, data_center, rack, tokens from system.peers");
+    peers_query.set_page_size(1024);
+    let peers_query_future = conn.query_all(&peers_query, &[]);
+
+    let mut local_query =
+        Query::new("select rpc_address, data_center, rack, tokens from system.local");
+    local_query.set_page_size(1024);
+    let local_query_future = conn.query_all(&local_query, &[]);
+
+    let (peers_res, local_res) = tokio::try_join!(peers_query_future, local_query_future)?;
 
     let peers_rows = peers_res.rows.ok_or(QueryError::ProtocolError(
         "system.peers query response was not Rows",
@@ -269,16 +269,17 @@ async fn query_peers(conn: &Connection, connect_port: u16) -> Result<Vec<Peer>,
 }
 
 async fn query_keyspaces(conn: &Connection) -> Result<HashMap<String, Keyspace>, QueryError> {
-    let rows = conn
-        .query_single_page(
-            "select keyspace_name, replication from system_schema.keyspaces",
-            &[],
-        )
-        .await?
-        .rows
-        .ok_or(QueryError::ProtocolError(
-            "system_schema.keyspaces query response was not Rows",
-        ))?;
+    let mut keyspaces_query =
+        Query::new("select keyspace_name, replication from system_schema.keyspaces");
+    keyspaces_query.set_page_size(1024);
+
+    let rows =
+        conn.query_all(&keyspaces_query, &[])
+            .await?
+            .rows
+            .ok_or(QueryError::ProtocolError(
+                "system_schema.keyspaces query response was not Rows",
+            ))?;
 
     let mut result = HashMap::with_capacity(rows.len());
 


Test patch:
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 3edb4c89..d9b4d22d 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -76,7 +76,7 @@ async fn test_unprepared_statement() {
         ]
     );
     let mut results_from_manual_paging: Vec<Row> = vec![];
-    let query = Query::new("SELECT a, b, c FROM ks.t".to_owned()).with_page_size(1);
+    let query = Query::new("SELECT a, b, c FROM ks.t").with_page_size(1);
     let mut paging_state: Option<Bytes> = None;
     let mut watchdog = 0;
     loop {
@@ -196,7 +196,7 @@ async fn test_prepared_statement() {
         assert_eq!((a, b, c), (17, 16, &String::from("I'm prepared!!!")));
 
         let mut results_from_manual_paging: Vec<Row> = vec![];
-        let query = Query::new("SELECT a, b, c FROM ks.t2".to_owned()).with_page_size(1);
+        let query = Query::new("SELECT a, b, c FROM ks.t2").with_page_size(1);
         let prepared_paged = session.prepare(query).await.unwrap();
         let mut paging_state: Option<Bytes> = None;
         let mut watchdog = 0;
@@ -724,13 +724,13 @@ async fn test_tracing() {
 
 async fn test_tracing_query(session: &Session) {
     // A query without tracing enabled has no tracing uuid in result
-    let untraced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab".to_string());
+    let untraced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab");
     let untraced_query_result: QueryResult = session.query(untraced_query, &[]).await.unwrap();
 
     assert!(untraced_query_result.tracing_id.is_none());
 
     // A query with tracing enabled has a tracing uuid in result
-    let mut traced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab".to_string());
+    let mut traced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab");
     traced_query.config.tracing = true;
 
     let traced_query_result: QueryResult = session.query(traced_query, &[]).await.unwrap();
@@ -777,7 +777,7 @@ async fn test_tracing_prepare(session: &Session) {
     assert!(untraced_prepared.prepare_tracing_ids.is_empty());
 
     // Preparing a statement with tracing enabled has tracing uuids in result
-    let mut to_prepare_traced = Query::new("SELECT * FROM test_tracing_ks.tab".to_string());
+    let mut to_prepare_traced = Query::new("SELECT * FROM test_tracing_ks.tab");
     to_prepare_traced.config.tracing = true;
 
     let traced_prepared = session.prepare(to_prepare_traced).await.unwrap();
@@ -791,7 +791,7 @@ async fn test_tracing_prepare(session: &Session) {
 
 async fn test_get_tracing_info(session: &Session) {
     // A query with tracing enabled has a tracing uuid in result
-    let mut traced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab".to_string());
+    let mut traced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab");
     traced_query.config.tracing = true;
 
     let traced_query_result: QueryResult = session.query(traced_query, &[]).await.unwrap();
@@ -804,7 +804,7 @@ async fn test_get_tracing_info(session: &Session) {
 
 async fn test_tracing_query_iter(session: &Session) {
     // A query without tracing enabled has no tracing ids
-    let untraced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab".to_string());
+    let untraced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab");
 
     let mut untraced_row_iter = session.query_iter(untraced_query, &[]).await.unwrap();
     while let Some(_row) = untraced_row_iter.next().await {
@@ -818,7 +818,7 @@ async fn test_tracing_query_iter(session: &Session) {
     assert!(untraced_typed_row_iter.get_tracing_ids().is_empty());
 
     // A query with tracing enabled has a tracing ids in result
-    let mut traced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab".to_string());
+    let mut traced_query: Query = Query::new("SELECT * FROM test_tracing_ks.tab");
     traced_query.config.tracing = true;
 
     let mut traced_row_iter = session.query_iter(traced_query, &[]).await.unwrap();
@@ -898,8 +898,7 @@ async fn test_tracing_batch(session: &Session) {
 }
 
 async fn assert_in_tracing_table(session: &Session, tracing_uuid: Uuid) {
-    let mut traces_query =
-        Query::new("SELECT * FROM system_traces.sessions WHERE session_id = ?".to_string());
+    let mut traces_query = Query::new("SELECT * FROM system_traces.sessions WHERE session_id = ?");
     traces_query.config.consistency = Consistency::One;
 
     // Tracing info might not be immediately available