Added cursor first implementation

Author: chandr-andr

.github/workflows/CI.yml

@@ -112,4 +112,4 @@ jobs:
           MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
         with:
           command: upload
-          args: --non-interactive --skip-existing *
+          args: --non-interactive --skip-existing *

Cargo.lock

@@ -25,5 +25,5 @@ byteorder = "1.5.0"
 postgres-protocol = "0.6.6"
 chrono = "0.4.33"
 chrono-tz = "0.8.5"
-uuid = "1.7.0"
+uuid = { version = "1.7.0", features = ["v4"] }
 serde_json = "1.0.113"

Cargo.toml

@@ -21,6 +21,34 @@ class IsolationLevel(Enum):
     Serializable = 4
 
 
+class Cursor:
+    """Represent opened cursor in a transaction.
+    
+    It can be used as an asynchronous iterator.
+    """
+
+    async def fetch(
+        self: Self,
+        fetch_number: int | None = None,
+    ) -> QueryResult:
+        """Fetch next <fetch_number> rows.
+        
+        By default fetches 10 next rows.
+
+        ### Parameters:
+        - `fetch_number`: how many rows need to fetch.
+
+        ### Returns:
+        result as `QueryResult`.
+        """
+    
+    def __aiter__(self: Self) -> Self:
+        ...
+    
+    async def __anext__(self: Self) -> QueryResult:
+        ...
+
+
 class Transaction:
     """Single connection for executing queries.
 
@@ -207,6 +235,50 @@ class Transaction:
         ```
         """
 
+    async def cursor(
+        self: Self,
+        querystring: str,
+        parameters: List[Any] | None = None,
+        fetch_number: int | None = None,
+    ) -> Cursor:
+        """Create new cursor object.
+
+        Cursor can be used as an asynchronous iterator.
+        
+        ### Parameters:
+        - `querystring`: querystring to execute.
+        - `parameters`: list of parameters to pass in the query.
+        - `fetch_number`: how many rows need to fetch.
+
+        ### Returns:
+        new initialized cursor.
+
+        ### Example:
+        ```python
+        import asyncio
+
+        from psql_rust_driver import PSQLPool, QueryResult
+
+
+        async def main() -> None:
+            db_pool = PSQLPool()
+            await db_pool.startup()
+
+            connection = await db_pool.connection()
+            transaction = await connection.transaction()
+
+            cursor = await transaction.cursor(
+                querystring="SELECT * FROM users WHERE username = $1",
+                parameters=["Some_Username"],
+                fetch_number=5,
+            )
+
+            async for fetched_result in cursor:
+                dict_result: List[Dict[Any, Any]] = fetched_result.result()
+                ... # do something with this result.
+        ```
+        """
+    
 
 class Connection:
     """Connection from Database Connection Pool.
@@ -228,6 +300,9 @@ class Connection:
         - `querystring`: querystring to execute.
         - `parameters`: list of parameters to pass in the query.
 
+        ### Returns:
+        query result as `QueryResult`
+
         ### Example:
         ```python
         import asyncio

python/psql_rust_driver/_internal/__init__.pyi

@@ -62,6 +62,7 @@ impl RustConnection {
             is_done: Arc::new(tokio::sync::RwLock::new(false)),
             rollback_savepoint: Arc::new(tokio::sync::RwLock::new(HashSet::new())),
             isolation_level: isolation_level,
+            cursor_num: Default::default(),
         };
 
         Transaction {

src/driver/connection.rs

@@ -136,6 +136,7 @@ impl RustPSQLPool {
             is_done: Arc::new(tokio::sync::RwLock::new(false)),
             rollback_savepoint: Arc::new(tokio::sync::RwLock::new(HashSet::new())),
             isolation_level: isolation_level,
+            cursor_num: Default::default(),
         };
 
         Ok(Transaction {

src/driver/connection_pool.rs

@@ -0,0 +1,85 @@
+use deadpool_postgres::Object;
+use pyo3::{exceptions::PyStopAsyncIteration, pyclass, pymethods, PyAny, PyObject, PyRef, Python};
+use std::sync::Arc;
+
+use crate::{
+    common::rustengine_future, exceptions::rust_errors::RustPSQLDriverPyResult,
+    query_result::PSQLDriverPyQueryResult,
+};
+
+#[pyclass]
+pub struct Cursor {
+    db_client: Arc<tokio::sync::RwLock<Object>>,
+    cursor_name: String,
+    fetch_number: usize,
+}
+
+impl Cursor {
+    pub fn new(
+        db_client: Arc<tokio::sync::RwLock<Object>>,
+        cursor_name: String,
+        fetch_number: usize,
+    ) -> Self {
+        return Cursor {
+            db_client,
+            cursor_name,
+            fetch_number,
+        };
+    }
+}
+
+#[pymethods]
+impl Cursor {
+    pub fn fetch<'a>(
+        &'a self,
+        py: Python<'a>,
+        fetch_number: Option<usize>,
+    ) -> RustPSQLDriverPyResult<&PyAny> {
+        let db_client_arc = self.db_client.clone();
+        let cursor_name = self.cursor_name.clone();
+        let fetch_number = match fetch_number {
+            Some(usize) => usize,
+            None => self.fetch_number.clone(),
+        };
+
+        rustengine_future(py, async move {
+            let db_client_guard = db_client_arc.read().await;
+            let result = db_client_guard
+                .query(
+                    format!("FETCH {fetch_number} FROM {cursor_name}").as_str(),
+                    &[],
+                )
+                .await?;
+            Ok(PSQLDriverPyQueryResult::new(result))
+        })
+    }
+
+    #[must_use]
+    pub fn __aiter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
+        slf
+    }
+
+    pub fn __anext__(&self, py: Python<'_>) -> RustPSQLDriverPyResult<Option<PyObject>> {
+        let db_client_arc = self.db_client.clone();
+        let cursor_name = self.cursor_name.clone();
+        let fetch_number = self.fetch_number.clone();
+
+        let future = rustengine_future(py, async move {
+            let db_client_guard = db_client_arc.read().await;
+            let result = db_client_guard
+                .query(
+                    format!("FETCH {fetch_number} FROM {cursor_name}").as_str(),
+                    &[],
+                )
+                .await?;
+
+            if result.len() == 0 {
+                return Err(PyStopAsyncIteration::new_err("Error").into());
+            };
+
+            Ok(PSQLDriverPyQueryResult::new(result))
+        });
+
+        Ok(Some(future?.into()))
+    }
+}

src/driver/cursor.rs

@@ -1,4 +1,5 @@
 pub mod connection;
 pub mod connection_pool;
+pub mod cursor;
 pub mod transaction;
 pub mod transaction_options;

src/driver/mod.rs

@@ -10,7 +10,7 @@ use crate::{
     value_converter::{convert_parameters, PythonDTO},
 };
 
-use super::transaction_options::IsolationLevel;
+use super::{cursor::Cursor, transaction_options::IsolationLevel};
 
 /// Transaction for internal use only.
 ///
@@ -22,6 +22,7 @@ pub struct RustTransaction {
     pub rollback_savepoint: Arc<tokio::sync::RwLock<HashSet<String>>>,
 
     pub isolation_level: Option<IsolationLevel>,
+    pub cursor_num: usize,
 }
 
 impl RustTransaction {
@@ -393,6 +394,34 @@ impl RustTransaction {
 
         Ok(())
     }
+
+    pub async fn inner_cursor<'a>(
+        &'a mut self,
+        querystring: String,
+        parameters: Vec<PythonDTO>,
+        fetch_number: usize,
+    ) -> RustPSQLDriverPyResult<Cursor> {
+        let db_client_arc = self.db_client.clone();
+        let db_client_arc2 = self.db_client.clone();
+        let db_client_guard = db_client_arc.read().await;
+
+        let mut vec_parameters: Vec<&(dyn ToSql + Sync)> = Vec::with_capacity(parameters.len());
+        for param in parameters.iter() {
+            vec_parameters.push(param);
+        }
+
+        let cursor_name = format!("cur{}", self.cursor_num);
+        db_client_guard
+            .execute(
+                &format!("DECLARE {} CURSOR FOR {querystring}", cursor_name),
+                &vec_parameters.into_boxed_slice(),
+            )
+            .await?;
+
+        self.cursor_num = self.cursor_num + 1;
+
+        Ok(Cursor::new(db_client_arc2, cursor_name, fetch_number))
+    }
 }
 
 #[pyclass()]
@@ -619,4 +648,25 @@ impl Transaction {
             Ok(())
         })
     }
+
+    pub fn cursor<'a>(
+        &'a self,
+        py: Python<'a>,
+        querystring: String,
+        parameters: Option<&'a PyAny>,
+        fetch_number: Option<usize>,
+    ) -> RustPSQLDriverPyResult<&PyAny> {
+        let transaction_arc = self.transaction.clone();
+        let mut params: Vec<PythonDTO> = vec![];
+        if let Some(parameters) = parameters {
+            params = convert_parameters(parameters)?
+        }
+
+        rustengine_future(py, async move {
+            let mut transaction_guard = transaction_arc.write().await;
+            Ok(transaction_guard
+                .inner_cursor(querystring, params, fetch_number.unwrap_or(10))
+                .await?)
+        })
+    }
 }

src/driver/transaction.rs

@@ -15,6 +15,7 @@ use pyo3::{pymodule, types::PyModule, PyResult, Python};
 fn psql_rust_engine(py: Python<'_>, pymod: &PyModule) -> PyResult<()> {
     pymod.add_class::<driver::connection_pool::PSQLPool>()?;
     pymod.add_class::<driver::transaction::Transaction>()?;
+    pymod.add_class::<driver::cursor::Cursor>()?;
     pymod.add_class::<driver::transaction_options::IsolationLevel>()?;
     pymod.add_class::<query_result::PSQLDriverPyQueryResult>()?;
     add_module(py, pymod, "extra_types", extra_types_module)?;