Merge pull request #15 from qaspen-python/feature/new_cursor

Rewrote the whole implementation of Cursor. Made some performance changes in Transaction

Author: chandr-andr

README.md

@@ -506,7 +506,7 @@ async def main() -> None:
 
     await transaction.begin()
     # Create new savepoint
-    cursor = await transaction.cursor(
+    cursor = transaction.cursor(
         querystring="SELECT * FROM users WHERE username = $1",
         parameters=["SomeUserName"],
         fetch_number=100,
@@ -525,6 +525,32 @@ async def main() -> None:
     await transaction.commit()
 ```
 
+### Cursor as an async context manager
+```python
+from typing import Any
+
+from psqlpy import PSQLPool, IsolationLevel, QueryResult, Transaction, Cursor
+
+
+db_pool = PSQLPool()
+
+
+async def main() -> None:
+    await db_pool.startup()
+
+    connection = await db_pool.connection()
+    transaction: Transaction
+    cursor: Cursor
+    async with connection.transaction() as transaction:
+        async with transaction.cursor(
+            querystring="SELECT * FROM users WHERE username = $1",
+            parameters=["SomeUserName"],
+            fetch_number=100,
+        ) as cursor:
+            async for fetched_result in cursor:
+                print(fetched_result.result())
+```
+
 ### Cursor operations
 
 Available cursor operations:

python/psqlpy/_internal/__init__.pyi

@@ -160,6 +160,25 @@ class Cursor:
     It can be used as an asynchronous iterator.
     """
 
+    def __aiter__(self: Self) -> Self: ...
+    async def __anext__(self: Self) -> QueryResult: ...
+    async def __aenter__(self: Self) -> Self: ...
+    async def __aexit__(
+        self: Self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> None: ...
+    async def start(self: Self) -> None:
+        """Start the cursor.
+
+        Execute DECLARE command for the cursor.
+        """
+    async def close(self: Self) -> None:
+        """Close the cursor.
+
+        Execute CLOSE command for the cursor.
+        """
     async def fetch(
         self: Self,
         fetch_number: int | None = None,
@@ -267,13 +286,6 @@ class Cursor:
         ### Returns:
         result as `QueryResult`.
         """
-    async def close(self: Self) -> None:
-        """Close the cursor.
-
-        Execute CLOSE command for the cursor.
-        """
-    def __aiter__(self: Self) -> Self: ...
-    async def __anext__(self: Self) -> QueryResult: ...
 
 class Transaction:
     """Single connection for executing queries.
@@ -330,7 +342,8 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
             await transaction.begin()
             query_result: QueryResult = await transaction.execute(
                 "SELECT username FROM users WHERE id = $1",
@@ -339,21 +352,6 @@ class Transaction:
             dict_result: List[Dict[Any, Any]] = query_result.result()
             # You must call commit manually
             await transaction.commit()
-
-        # Or you can transaction as a async context manager
-
-        async def main() -> None:
-            db_pool = PSQLPool()
-            await psqlpy.startup()
-
-            transaction = await db_pool.transaction()
-            async with transaction:
-                query_result: QueryResult = await transaction.execute(
-                    "SELECT username FROM users WHERE id = $1",
-                    [100],
-                )
-                dict_result: List[Dict[Any, Any]] = query_result.result()
-            # This way transaction begins and commits by itself.
         ```
         """
     async def execute_many(
@@ -381,7 +379,8 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
             await transaction.begin()
             query_result: QueryResult = await transaction.execute_many(
                 "INSERT INTO users (name, age) VALUES ($1, $2)",
@@ -390,21 +389,6 @@ class Transaction:
             dict_result: List[Dict[Any, Any]] = query_result.result()
             # You must call commit manually
             await transaction.commit()
-
-        # Or you can transaction as a async context manager
-
-        async def main() -> None:
-            db_pool = PSQLPool()
-            await psqlpy.startup()
-
-            transaction = await db_pool.transaction()
-            async with transaction:
-                query_result: QueryResult = await transaction.execute(
-                    "SELECT username FROM users WHERE id = $1",
-                    [100],
-                )
-                dict_result: List[Dict[Any, Any]] = query_result.result()
-            # This way transaction begins and commits by itself.
         ```
         """
     async def fetch_row(
@@ -434,30 +418,16 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
             await transaction.begin()
-            query_result: SingleQueryResult = await transaction.execute(
+            query_result: SingleQueryResult = await transaction.fetch_row(
                 "SELECT username FROM users WHERE id = $1",
                 [100],
             )
             dict_result: Dict[Any, Any] = query_result.result()
             # You must call commit manually
             await transaction.commit()
-
-        # Or you can transaction as a async context manager
-
-        async def main() -> None:
-            db_pool = PSQLPool()
-            await psqlpy.startup()
-
-            transaction = await db_pool.transaction()
-            async with transaction:
-                query_result: SingleQueryResult = await transaction.execute(
-                    "SELECT username FROM users WHERE id = $1 LIMIT 1",
-                    [100],
-                )
-                dict_result: Dict[Any, Any] = query_result.result()
-            # This way transaction begins and commits by itself.
         ```
         """
     async def fetch_val(
@@ -485,27 +455,13 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
             await transaction.begin()
             value: Any | None = await transaction.execute(
                 "SELECT username FROM users WHERE id = $1",
                 [100],
             )
-
-        # Or you can transaction as a async context manager
-
-        async def main() -> None:
-            db_pool = PSQLPool()
-            await psqlpy.startup()
-
-            transaction = await db_pool.transaction()
-            async with transaction:
-                query_result: SingleQueryResult = await transaction.execute(
-                    "SELECT username FROM users WHERE id = $1",
-                    [100],
-                )
-                dict_result: Dict[Any, Any] = query_result.result()
-            # This way transaction begins and commits by itself.
         ```
         """
     async def pipeline(
@@ -547,7 +503,8 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
 
             results: list[QueryResult] = await transaction.pipeline(
                 queries=[
@@ -591,7 +548,8 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
 
             await transaction.savepoint("my_savepoint")
             await transaction.execute(...)
@@ -615,7 +573,8 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
             await transaction.execute(...)
             await transaction.rollback()
         ```
@@ -640,7 +599,8 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
 
             await transaction.savepoint("my_savepoint")
             await transaction.execute(...)
@@ -667,13 +627,14 @@ class Transaction:
             db_pool = PSQLPool()
             await db_pool.startup()
 
-            transaction = await db_pool.transaction()
+            connection = await db_pool.connection()
+            transaction = connection.transaction()
 
             await transaction.savepoint("my_savepoint")
             await transaction.release_savepoint
         ```
         """
-    async def cursor(
+    def cursor(
         self: Self,
         querystring: str,
         parameters: list[Any] | None = None,
@@ -707,15 +668,18 @@ class Transaction:
             connection = await db_pool.connection()
             transaction = await connection.transaction()
 
-            cursor = await transaction.cursor(
+            cursor = transaction.cursor(
                 querystring="SELECT * FROM users WHERE username = $1",
                 parameters=["Some_Username"],
                 fetch_number=5,
             )
+            await cursor.start()
 
             async for fetched_result in cursor:
                 dict_result: List[Dict[Any, Any]] = fetched_result.result()
                 ... # do something with this result.
+
+            await cursor.close()
         ```
         """
 
@@ -772,6 +736,7 @@ class Connection:
         ### Parameters:
         - `isolation_level`: configure isolation level of the transaction.
         - `read_variant`: configure read variant of the transaction.
+        - `deferrable`: configure deferrable of the transaction.
         """
 
 class PSQLPool:
@@ -842,7 +807,7 @@ class PSQLPool:
 
         async def main() -> None:
             db_pool = PSQLPool()
-            await psqlpy.startup()
+            await db_pool.startup()
             query_result: QueryResult = await psqlpy.execute(
                 "SELECT username FROM users WHERE id = $1",
                 [100],

python/tests/conftest.py

@@ -122,8 +122,9 @@ async def test_cursor(
     connection = await psql_pool.connection()
     transaction = connection.transaction()
     await transaction.begin()
-    cursor = await transaction.cursor(
+    cursor = transaction.cursor(
         querystring=f"SELECT * FROM {table_name}",
     )
+    await cursor.start()
     yield cursor
     await transaction.commit()

python/tests/test_cursor.py

@@ -1,6 +1,8 @@
+import math
+
 import pytest
 
-from psqlpy import Cursor
+from psqlpy import Cursor, PSQLPool, QueryResult, Transaction
 
 pytestmark = pytest.mark.anyio
 
@@ -147,3 +149,25 @@ async def test_cursor_fetch_backward_all(
 
     must_not_be_empty = await test_cursor.fetch_backward_all()
     assert len(must_not_be_empty.result()) == default_fetch_number - 1
+
+
+async def test_cursor_as_async_manager(
+    psql_pool: PSQLPool,
+    table_name: str,
+    number_database_records: int,
+) -> None:
+    """Test cursor async manager and async iterator."""
+    connection = await psql_pool.connection()
+    transaction: Transaction
+    cursor: Cursor
+    all_results: list[QueryResult] = []
+    expected_num_results = math.ceil(number_database_records / 3)
+    fetch_number = 3
+    async with connection.transaction() as transaction, transaction.cursor(
+        querystring=f"SELECT * FROM {table_name}",
+        fetch_number=fetch_number,
+    ) as cursor:
+        async for result in cursor:
+            all_results.append(result)
+
+    assert len(all_results) == expected_num_results

python/tests/test_transaction.py

@@ -200,7 +200,7 @@ async def test_transaction_cursor(
     """Test that transaction can create cursor."""
     connection = await psql_pool.connection()
     async with connection.transaction() as transaction:
-        cursor = await transaction.cursor(f"SELECT * FROM {table_name}")
+        cursor = transaction.cursor(f"SELECT * FROM {table_name}")
 
         assert isinstance(cursor, Cursor)
 

src/driver/connection.rs

@@ -76,11 +76,8 @@ impl Connection {
             isolation_level,
             read_variant,
             deferrable,
-            Default::default(),
         );
 
-        Transaction {
-            transaction: Arc::new(tokio::sync::RwLock::new(inner_transaction)),
-        }
+        Transaction::new(Arc::new(inner_transaction), Default::default())
     }
 }