Added close method for cursor and scroll option for cursor

Author: chandr-andr

README.md

@@ -3,7 +3,7 @@
 Driver for PostgreSQL written fully in Rust and exposed to Python.  
 *Normal documentation is in development.*
 
-## Installation
+# Installation
 
 You can install package with `pip` or `poetry`.
 
@@ -53,7 +53,7 @@ async def main() -> None:
 ```
 Please take into account that each new execute gets new connection from connection pool. 
 
-## Query parameters
+# Query parameters
 You can pass parameters into queries.  
 Parameters can be passed in any `execute` method as the second parameter, it must be a list.  
 Any placeholder must be marked with `$< num>`.
@@ -65,7 +65,7 @@ Any placeholder must be marked with `$< num>`.
     )
 ```
 
-## Connection
+# Connection
 You can work with connection instead of DatabasePool.
 ```python
 from typing import Any
@@ -96,7 +96,7 @@ async def main() -> None:
     # rust does it instead.
 ```
 
-## Transactions
+# Transactions
 Of course it's possible to use transactions with this driver.  
 It's as easy as possible and sometimes it copies common functionality from PsycoPG and AsyncPG.
 
@@ -254,11 +254,18 @@ async def main() -> None:
     await transaction.commit()
 ```
 
-## Cursors
+# Cursors
 Library supports PostgreSQL cursors.
 
 Cursors can be created only in transaction. In addition, cursor supports async iteration.
 
+### Cursor parameters
+In process of cursor creation you can specify some configuration parameters.
+- `querystring`: query for the cursor. Required.
+- `parameters`: parameters for the query. Not Required.
+- `fetch_number`: number of records per fetch if cursor is used as an async iterator. If you are using `.fetch()` method you can pass different fetch number. Not required. Default - 10.
+- `scroll`: set `SCROLL` if True or `NO SCROLL` if False. Not required. By default - `None`.
+
 ```python
 from typing import Any
 
@@ -283,9 +290,16 @@ async def main() -> None:
         fetch_number=100,
     )
 
+    # You can manually fetch results from cursor
+    results: list[dict[str, Any]] = await cursor.fetch(fetch_number=8)
+
+    # Or you can use it as an async iterator.
     async for fetched_result in cursor:
         print(fetched_result.result())
 
+    # If you want to close cursor, please do it manually.
+    await cursor.close()
+    
     await transaction.commit()
 ```
 

python/psqlpy/_internal/__init__.pyi

@@ -49,6 +49,12 @@ class Cursor:
         result as `QueryResult`.
         """
 
+    async def close(self: Self) -> None:
+        """Close the cursor.
+
+        Execute CLOSE command for the cursor.
+        """
+    
     def __aiter__(self: Self) -> Self:
         ...
 
@@ -247,6 +253,7 @@ class Transaction:
         querystring: str,
         parameters: List[Any] | None = None,
         fetch_number: int | None = None,
+        scroll: bool | None = None,
     ) -> Cursor:
         """Create new cursor object.
 
@@ -256,6 +263,7 @@ class Transaction:
         - `querystring`: querystring to execute.
         - `parameters`: list of parameters to pass in the query.
         - `fetch_number`: how many rows need to fetch.
+        - `scroll`: SCROLL or NO SCROLL cursor.
 
         ### Returns:
         new initialized cursor.

python/psqlpy/_internal/exceptions.pyi

@@ -40,3 +40,7 @@ class DBPoolConfigurationError(RustPSQLDriverPyBaseError):
 
 class UUIDValueConvertError(RustPSQLDriverPyBaseError):
     """Error if it's impossible to convert py string UUID into rust UUID."""
+
+
+class CursorError(RustPSQLDriverPyBaseError):
+    """Error if something goes wrong with the cursor."""

python/psqlpy/exceptions.py

@@ -6,6 +6,7 @@
     DBTransactionError,
     DBPoolConfigurationError,
     UUIDValueConvertError,
+    CursorError,
 )
 
 __all__ = [
@@ -16,4 +17,5 @@
     "DBTransactionError",
     "DBPoolConfigurationError",
     "UUIDValueConvertError",
+    "CursorError",
 ]

src/driver/cursor.rs

@@ -12,6 +12,7 @@ pub struct Cursor {
     db_client: Arc<tokio::sync::RwLock<Object>>,
     cursor_name: String,
     fetch_number: usize,
+    closed: Arc<tokio::sync::RwLock<bool>>,
 }
 
 impl Cursor {
@@ -24,6 +25,7 @@ impl Cursor {
             db_client,
             cursor_name,
             fetch_number,
+            closed: Arc::new(tokio::sync::RwLock::new(false)),
         };
     }
 }
@@ -82,4 +84,39 @@ impl Cursor {
 
         Ok(Some(future?.into()))
     }
+
+    /// Close cursor.
+    ///
+    /// # Errors
+    /// May return Err Result if cannot execute CLOSE command
+    pub fn close<'a>(&'a self, py: Python<'a>) -> RustPSQLDriverPyResult<&PyAny> {
+        let db_client_arc = self.db_client.clone();
+        let cursor_name = self.cursor_name.clone();
+        let closed = self.closed.clone();
+
+        rustengine_future(py, async move {
+            let is_closed = {
+                let closed_read = closed.write().await;
+                *closed_read
+            };
+            if is_closed {
+                return Err(
+                    crate::exceptions::rust_errors::RustPSQLDriverError::DBCursorError(
+                        "Cursor is already closed".into(),
+                    ),
+                );
+            }
+
+            let db_client_guard = db_client_arc.read().await;
+
+            db_client_guard
+                .batch_execute(format!("CLOSE {cursor_name}").as_str())
+                .await?;
+
+            let mut closed_write = closed.write().await;
+            *closed_write = true;
+
+            Ok(())
+        })
+    }
 }

src/driver/transaction.rs

@@ -427,6 +427,7 @@ impl RustTransaction {
         querystring: String,
         parameters: Vec<PythonDTO>,
         fetch_number: usize,
+        scroll: Option<bool>,
     ) -> RustPSQLDriverPyResult<Cursor> {
         let db_client_arc = self.db_client.clone();
         let db_client_arc2 = self.db_client.clone();
@@ -437,12 +438,22 @@ impl RustTransaction {
             vec_parameters.push(param);
         }
 
+        let mut cursor_init_query = "DECLARE".to_string();
+        cursor_init_query.push_str(format!(" cur{}", self.cursor_num).as_str());
+
+        if let Some(scroll) = scroll {
+            if scroll == true {
+                cursor_init_query.push_str(" SCROLL");
+            } else {
+                cursor_init_query.push_str(" NO SCROLL");
+            }
+        }
+
+        cursor_init_query.push_str(format!(" CURSOR FOR {querystring}").as_str());
+
         let cursor_name = format!("cur{}", self.cursor_num);
         db_client_guard
-            .execute(
-                &format!("DECLARE {} CURSOR FOR {querystring}", cursor_name),
-                &vec_parameters.into_boxed_slice(),
-            )
+            .execute(&cursor_init_query, &vec_parameters.into_boxed_slice())
             .await?;
 
         self.cursor_num = self.cursor_num + 1;
@@ -682,6 +693,7 @@ impl Transaction {
         querystring: String,
         parameters: Option<&'a PyAny>,
         fetch_number: Option<usize>,
+        scroll: Option<bool>,
     ) -> RustPSQLDriverPyResult<&PyAny> {
         let transaction_arc = self.transaction.clone();
         let mut params: Vec<PythonDTO> = vec![];
@@ -692,7 +704,7 @@ impl Transaction {
         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))
+                .inner_cursor(querystring, params, fetch_number.unwrap_or(10), scroll)
                 .await?)
         })
     }

src/exceptions/python_errors.rs

@@ -33,6 +33,8 @@ create_exception!(
     RustPSQLDriverPyBaseError
 );
 
+create_exception!(psqlpy.exceptions, CursorError, RustPSQLDriverPyBaseError);
+
 pub fn python_exceptions_module(py: Python<'_>, pymod: &PyModule) -> PyResult<()> {
     pymod.add(
         "RustPSQLDriverPyBaseError",
@@ -56,5 +58,6 @@ pub fn python_exceptions_module(py: Python<'_>, pymod: &PyModule) -> PyResult<()
         "UUIDValueConvertError",
         py.get_type::<UUIDValueConvertError>(),
     )?;
+    pymod.add("CursorError", py.get_type::<CursorError>())?;
     Ok(())
 }

src/exceptions/rust_errors.rs

@@ -5,7 +5,7 @@ use crate::exceptions::python_errors::{
     RustPSQLDriverPyBaseError, RustToPyValueMappingError,
 };
 
-use super::python_errors::UUIDValueConvertError;
+use super::python_errors::{CursorError, UUIDValueConvertError};
 
 pub type RustPSQLDriverPyResult<T> = Result<T, RustPSQLDriverError>;
 
@@ -21,6 +21,8 @@ pub enum RustPSQLDriverError {
     DataBaseTransactionError(String),
     #[error("Configuration database pool error: {0}")]
     DataBasePoolConfigurationError(String),
+    #[error("Cursor error: {0}")]
+    DBCursorError(String),
 
     #[error("Python exception: {0}.")]
     PyError(#[from] pyo3::PyErr),
@@ -62,6 +64,7 @@ impl From<RustPSQLDriverError> for pyo3::PyErr {
                 DBPoolConfigurationError::new_err((error_desc,))
             }
             RustPSQLDriverError::UUIDConvertError(_) => UUIDValueConvertError::new_err(error_desc),
+            RustPSQLDriverError::DBCursorError(_) => CursorError::new_err(error_desc),
         }
     }
 }