Continue adding tests

Author: chandr-andr

python/psqlpy/__init__.py

@@ -1,5 +1,6 @@
 from ._internal import (
     Connection,
+    Cursor,
     IsolationLevel,
     PSQLPool,
     QueryResult,
@@ -14,4 +15,5 @@
     "IsolationLevel",
     "ReadVariant",
     "Connection",
+    "Cursor",
 ]

python/psqlpy/_internal/__init__.pyi

@@ -1,3 +1,4 @@
+import types
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
@@ -153,6 +154,13 @@ class Transaction:
     `.transaction()`.
     """
 
+    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 begin(self: Self) -> None:
         """Start the transaction.
 

python/tests/test_transaction.py

@@ -1,6 +1,6 @@
 import pytest
 
-from psqlpy import PSQLPool
+from psqlpy import Cursor, PSQLPool
 from psqlpy.exceptions import DBTransactionError
 
 
@@ -93,3 +93,77 @@ async def test_transaction_savepoint(
     assert not len(result.result())
 
     await transaction.commit()
+
+
+@pytest.mark.anyio
+async def test_transaction_rollback(
+    psql_pool: PSQLPool,
+    table_name: str,
+) -> None:
+    """Test that ROLLBACK works correctly."""
+    connection = await psql_pool.connection()
+    transaction = connection.transaction()
+    await transaction.begin()
+
+    test_name = "test_name"
+    await transaction.execute(
+        f"INSERT INTO {table_name} VALUES ($1, $2)",
+        parameters=[100, test_name],
+    )
+
+    result = await transaction.execute(
+        f"SELECT * FROM {table_name} WHERE name = $1",
+        parameters=[test_name],
+    )
+    assert result.result()
+
+    await transaction.rollback()
+
+    with pytest.raises(expected_exception=DBTransactionError):
+        await transaction.execute(
+            f"SELECT * FROM {table_name} WHERE name = $1",
+            parameters=[test_name],
+        )
+
+    result_from_conn = await psql_pool.execute(
+        f"INSERT INTO {table_name} VALUES ($1, $2)",
+        parameters=[100, test_name],
+    )
+
+    assert not (result_from_conn.result())
+
+
+@pytest.mark.anyio
+async def test_transaction_release_savepoint(
+    psql_pool: PSQLPool,
+) -> None:
+    """Test that it is possible to acquire and release savepoint."""
+    connection = await psql_pool.connection()
+    transaction = connection.transaction()
+    await transaction.begin()
+
+    sp_name_1 = "sp1"
+    sp_name_2 = "sp2"
+
+    await transaction.savepoint(sp_name_1)
+
+    with pytest.raises(expected_exception=DBTransactionError):
+        await transaction.savepoint(sp_name_1)
+
+    await transaction.savepoint(sp_name_2)
+
+    await transaction.release_savepoint(sp_name_1)
+    await transaction.savepoint(sp_name_1)
+
+
+@pytest.mark.anyio
+async def test_transaction_cursor(
+    psql_pool: PSQLPool,
+    table_name: str,
+) -> None:
+    """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}")
+
+        assert isinstance(cursor, Cursor)

src/driver/transaction.rs

@@ -396,11 +396,9 @@ impl RustTransaction {
             ));
         };
 
-        let rollback_savepoint_arc = self.rollback_savepoint.clone();
-        let is_rollback_exists = {
-            let rollback_savepoint_guard = rollback_savepoint_arc.read().await;
-            rollback_savepoint_guard.contains(&rollback_name)
-        };
+        let mut rollback_savepoint_guard = self.rollback_savepoint.write().await;
+        let is_rollback_exists = rollback_savepoint_guard.remove(&rollback_name);
+
         if !is_rollback_exists {
             return Err(RustPSQLDriverError::DataBaseTransactionError(
                 "Don't have rollback with this name".into(),