Merge pull request #12 from qaspen-python/feature/add_pipeline_queries

Added pipeline functionality for transactions

Author: chandr-andr

Cargo.lock

@@ -27,3 +27,4 @@ chrono = "0.4.33"
 chrono-tz = "0.8.5"
 uuid = { version = "1.7.0", features = ["v4"] }
 serde_json = "1.0.113"
+futures-util = "0.3.30"

Cargo.toml

@@ -355,6 +355,61 @@ async def main() -> None:
     first_row_result = first_row.result()  # This will be a dict.
 ```
 
+### Transaction pipelining
+When you have a lot of independent queries and want to execute them concurrently, you can use `pipeline`.
+Pipelining can improve performance in use cases in which multiple,
+independent queries need to be executed.
+In a traditional workflow,
+each query is sent to the server after the previous query completes.
+In contrast, pipelining allows the client to send all of the queries to the server up front,
+minimizing time spent by one side waiting for the other to finish sending data:
+```
+           Sequential                              Pipelined
+| Client         | Server          |    | Client         | Server          |
+|----------------|-----------------|    |----------------|-----------------|
+| send query 1   |                 |    | send query 1   |                 |
+|                | process query 1 |    | send query 2   | process query 1 |
+| receive rows 1 |                 |    | send query 3   | process query 2 |
+| send query 2   |                 |    | receive rows 1 | process query 3 |
+|                | process query 2 |    | receive rows 2 |                 |
+| receive rows 2 |                 |    | receive rows 3 |                 |
+| send query 3   |                 |
+|                | process query 3 |
+| receive rows 3 |                 |
+```
+Read more: https://docs.rs/tokio-postgres/latest/tokio_postgres/#pipelining
+
+Let's see some code:
+```python
+import asyncio
+
+from psqlpy import PSQLPool, QueryResult
+
+
+async def main() -> None:
+    db_pool = PSQLPool()
+    await db_pool.startup()
+
+    transaction = await db_pool.transaction()
+
+    results: list[QueryResult] = await transaction.pipeline(
+        queries=[
+            (
+                "SELECT username FROM users WHERE id = $1",
+                [100],
+            ),
+            (
+                "SELECT some_data FROM profiles",
+                None,
+            ),
+            (
+                "INSERT INTO users (username, id) VALUES ($1, $2)",
+                ["PSQLPy", 1],
+            ),
+        ]
+    )
+```
+
 ### Transaction ROLLBACK TO SAVEPOINT
 
 You can rollback your transaction to the specified savepoint, but before it you must create it.

README.md

@@ -458,6 +458,66 @@ class Transaction:
             # This way transaction begins and commits by itself.
         ```
         """
+    async def pipeline(
+        self,
+        queries: list[tuple[str, list[Any] | None]],
+    ) -> list[QueryResult]:
+        """Execute queries in pipeline.
+
+        Pipelining can improve performance in use cases in which multiple,
+        independent queries need to be executed.
+        In a traditional workflow,
+        each query is sent to the server after the previous query completes.
+        In contrast, pipelining allows the client to send all of the
+        queries to the server up front, minimizing time spent
+        by one side waiting for the other to finish sending data:
+        ```
+                            Sequential                              Pipelined
+        | Client         | Server          |    | Client         | Server          |
+        |----------------|-----------------|    |----------------|-----------------|
+        | send query 1   |                 |    | send query 1   |                 |
+        |                | process query 1 |    | send query 2   | process query 1 |
+        | receive rows 1 |                 |    | send query 3   | process query 2 |
+        | send query 2   |                 |    | receive rows 1 | process query 3 |
+        |                | process query 2 |    | receive rows 2 |                 |
+        | receive rows 2 |                 |    | receive rows 3 |                 |
+        | send query 3   |                 |
+        |                | process query 3 |
+        | receive rows 3 |                 |
+        ```
+        Read more: https://docs.rs/tokio-postgres/latest/tokio_postgres/#pipelining
+        ### Example:
+        ```python
+        import asyncio
+
+        from psqlpy import PSQLPool, QueryResult
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await db_pool.startup()
+
+            transaction = await db_pool.transaction()
+
+            results: list[QueryResult] = await transaction.pipeline(
+                queries=[
+                    (
+                        "SELECT username FROM users WHERE id = $1",
+                        [100],
+                    ),
+                    (
+                        "SELECT some_data FROM profiles",
+                        None,
+                    ),
+                    (
+                        "INSERT INTO users (username, id) VALUES ($1, $2)",
+                        ["PSQLPy", 1],
+                    ),
+                ]
+            )
+
+        ```
+        """  # noqa: E501
     async def savepoint(self: Self, savepoint_name: str) -> None:
         """Create new savepoint.
 

python/psqlpy/_internal/__init__.pyi

@@ -9,9 +9,10 @@ use crate::{
     value_converter::{convert_parameters, PythonDTO},
 };
 use deadpool_postgres::Object;
+use futures_util::future;
 use pyo3::{
     pyclass, pymethods,
-    types::{PyList, PyString},
+    types::{PyList, PyString, PyTuple},
     Py, PyAny, PyErr, PyObject, PyRef, PyRefMut, Python,
 };
 use std::{collections::HashSet, sync::Arc, vec};
@@ -163,6 +164,7 @@ impl RustTransaction {
 
         Ok(())
     }
+
     /// Fetch single row from query.
     ///
     /// Method doesn't acquire lock on any structure fields.
@@ -213,6 +215,29 @@ impl RustTransaction {
 
         Ok(PSQLDriverSinglePyQueryResult::new(result))
     }
+
+    /// Run many queries as pipeline.
+    ///
+    /// It can boost up querying speed.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if can't join futures or cannot execute
+    /// any of queries.
+    pub async fn inner_pipeline(
+        &self,
+        queries: Vec<(String, Vec<PythonDTO>)>,
+    ) -> RustPSQLDriverPyResult<Vec<PSQLDriverPyQueryResult>> {
+        let mut futures = vec![];
+        for (querystring, params) in queries {
+            let execute_future = self.inner_execute(querystring, params);
+            futures.push(execute_future);
+        }
+
+        let b = future::try_join_all(futures).await?;
+        Ok(b)
+    }
+
     /// Start transaction
     /// Set up isolation level if specified
     /// Set up deferable if specified
@@ -748,6 +773,48 @@ impl Transaction {
         })
     }
 
+    /// Execute querystrings with parameters and return all results.
+    ///
+    /// Create pipeline of queries.
+    ///
+    /// # Errors
+    ///
+    /// May return Err Result if:
+    /// 1) Cannot convert python parameters
+    /// 2) Cannot execute any of querystring.
+    pub fn pipeline<'a>(
+        &'a self,
+        py: Python<'a>,
+        queries: Option<&'a PyList>,
+    ) -> RustPSQLDriverPyResult<&'a PyAny> {
+        let mut processed_queries: Vec<(String, Vec<PythonDTO>)> = vec![];
+        if let Some(queries) = queries {
+            for single_query in queries {
+                let query_tuple = single_query.downcast::<PyTuple>().map_err(|err| {
+                    RustPSQLDriverError::PyToRustValueConversionError(format!(
+                        "Cannot cast to tuple: {err}",
+                    ))
+                })?;
+                let querystring = query_tuple.get_item(0)?.extract::<String>()?;
+                match query_tuple.get_item(1) {
+                    Ok(params) => {
+                        processed_queries.push((querystring, convert_parameters(params)?));
+                    }
+                    Err(_) => {
+                        processed_queries.push((querystring, vec![]));
+                    }
+                }
+            }
+        }
+
+        let transaction_arc = self.transaction.clone();
+
+        rustengine_future(py, async move {
+            let transaction_guard = transaction_arc.read().await;
+            transaction_guard.inner_pipeline(processed_queries).await
+        })
+    }
+
     /// Start the transaction.
     ///
     /// # Errors