refactor: removed session._ensure_transaction (#68)

* refactor: removed session._ensure_transaction

* chore: address pr comments

Author: mplemay

iceaxe/__tests__/benchmarks/test_select.py

@@ -107,8 +107,8 @@ async def test_benchmark(
     CONSOLE.print(f"Performance difference: {performance_diff:.2f}%")
 
     # Assert that DBConnection.exec is at most X% slower than raw query
-    assert (
-        performance_diff <= allowed_overhead
-    ), f"DBConnection.exec is {performance_diff:.2f}% slower than raw query, which exceeds the {allowed_overhead}% threshold"
+    assert performance_diff <= allowed_overhead, (
+        f"DBConnection.exec is {performance_diff:.2f}% slower than raw query, which exceeds the {allowed_overhead}% threshold"
+    )
 
     LOGGER.info("Benchmark completed successfully.")

iceaxe/__tests__/schemas/test_db_memory_serializer.py

@@ -85,7 +85,7 @@ class ModelA(TableBase):
             },
         ),
         DryRunComment(
-            text="\n" "NEW TABLE: modela\n",
+            text="\nNEW TABLE: modela\n",
             previous_line=False,
         ),
         DryRunAction(
@@ -294,7 +294,7 @@ class Model2(TableBase):
             },
         ),
         DryRunComment(
-            text="\n" "NEW TABLE: model1\n",
+            text="\nNEW TABLE: model1\n",
             previous_line=False,
         ),
         DryRunAction(
@@ -350,7 +350,7 @@ class Model2(TableBase):
             },
         ),
         DryRunComment(
-            text="\n" "NEW TABLE: model2\n",
+            text="\nNEW TABLE: model2\n",
             previous_line=False,
         ),
         DryRunAction(
@@ -998,7 +998,7 @@ class ModelA(TableBase, GenericSuperclass[OldValues]):
             },
         ),
         DryRunComment(
-            text="\n" "NEW TABLE: modela\n",
+            text="\nNEW TABLE: modela\n",
             previous_line=False,
         ),
         DryRunAction(
@@ -1377,18 +1377,18 @@ class SourceModel(TableBase):
     )
 
     # The foreign key constraint should come after both tables and the target column are created
-    assert (
-        target_table_pos < fk_constraint_pos
-    ), "Foreign key constraint should be created after target table"
-    assert (
-        source_table_pos < fk_constraint_pos
-    ), "Foreign key constraint should be created after source table"
-    assert (
-        target_column_pos < fk_constraint_pos
-    ), "Foreign key constraint should be created after target column"
-    assert (
-        target_pk_pos < fk_constraint_pos
-    ), "Foreign key constraint should be created after target primary key"
+    assert target_table_pos < fk_constraint_pos, (
+        "Foreign key constraint should be created after target table"
+    )
+    assert source_table_pos < fk_constraint_pos, (
+        "Foreign key constraint should be created after source table"
+    )
+    assert target_column_pos < fk_constraint_pos, (
+        "Foreign key constraint should be created after target column"
+    )
+    assert target_pk_pos < fk_constraint_pos, (
+        "Foreign key constraint should be created after target primary key"
+    )
 
     # Verify the actual migration actions
     actor = DatabaseActions()

iceaxe/__tests__/test_session.py

@@ -1128,13 +1128,13 @@ async def test_db_connection_update_batched(db_connection: DBConnection):
 
     # Check group 2 (only emails updated)
     for i, row in enumerate(result[10:20]):
-        assert row["name"] == f"User{i+10}"
-        assert row["email"] == f"updated_user{i+10}@example.com"
+        assert row["name"] == f"User{i + 10}"
+        assert row["email"] == f"updated_user{i + 10}@example.com"
 
     # Check group 3 (both fields updated)
     for i, row in enumerate(result[20:30]):
-        assert row["name"] == f"UpdatedUser{i+20}"
-        assert row["email"] == f"updated_user{i+20}@example.com"
+        assert row["name"] == f"UpdatedUser{i + 20}"
+        assert row["email"] == f"updated_user{i + 20}@example.com"
 
     # Verify all modifications were cleared
     assert all(user.get_modified_attributes() == {} for user in all_users)
@@ -1477,3 +1477,33 @@ async def test_get_dsn(db_connection: DBConnection):
     assert "localhost" in dsn
     assert "5438" in dsn
     assert "iceaxe_test_db" in dsn
+
+
+@pytest.mark.asyncio
+async def test_nested_transactions(db_connection):
+    """
+    Test that nested transactions raise an error by default, but work with ensure=True.
+    """
+    # Start an outer transaction
+    async with db_connection.transaction():
+        # This should work fine
+        assert db_connection.in_transaction is True
+
+        # Nested transaction with ensure=True should work
+        async with db_connection.transaction(ensure=True):
+            assert db_connection.in_transaction is True
+
+        # Nested transaction without ensure should fail
+        with pytest.raises(
+            RuntimeError,
+            match="Cannot start a new transaction while already in a transaction",
+        ):
+            async with db_connection.transaction():
+                pass  # Should not reach here
+
+    # After outer transaction ends, we should be out of transaction
+    assert db_connection.in_transaction is False
+
+    # Now a new transaction should start without error
+    async with db_connection.transaction():
+        assert db_connection.in_transaction is True

iceaxe/logging.py

@@ -79,7 +79,7 @@ def log_time_duration(message: str):
     """
     start = monotonic_ns()
     yield
-    LOGGER.debug(f"{message} : Took {(monotonic_ns() - start)/1e9:.2f}s")
+    LOGGER.debug(f"{message} : Took {(monotonic_ns() - start) / 1e9:.2f}s")
 
 
 # Our global logger should only surface warnings and above by default

iceaxe/session.py

@@ -188,11 +188,14 @@ def get_dsn(self) -> str:
         return "".join(dsn_parts)
 
     @asynccontextmanager
-    async def transaction(self):
+    async def transaction(self, *, ensure: bool = False):
         """
         Context manager for managing database transactions. Ensures that a series of database
         operations are executed atomically.
 
+        :param ensure: If True and already in a transaction, the context manager will yield without creating a new transaction.
+                       If False (default) and already in a transaction, raises a RuntimeError.
+
         ```python {{sticky: True}}
         async with conn.transaction():
             # All operations here are executed in a transaction
@@ -205,6 +208,17 @@ async def transaction(self):
             # If any operation fails, all changes are rolled back
         ```
         """
+        # If ensure is True and we're already in a transaction, just yield
+        if self.in_transaction:
+            if ensure:
+                yield
+                return
+            else:
+                raise RuntimeError(
+                    "Cannot start a new transaction while already in a transaction. Use ensure=True if this is intentional."
+                )
+
+        # Otherwise, start a new transaction
         self.in_transaction = True
         async with self.conn.transaction():
             try:
@@ -325,7 +339,7 @@ async def insert(self, objects: Sequence[TableBase]):
             return
 
         # Reuse a single transaction for all inserts
-        async with self._ensure_transaction():
+        async with self.transaction(ensure=True):
             for model, model_objects in self._aggregate_models_by_table(objects):
                 # For each table, build batched insert queries
                 table_name = QueryIdentifier(model.get_table_name())
@@ -466,7 +480,7 @@ async def upsert(
                 raise ValueError(f"Field {field} is not a column")
 
         results: list[tuple[T, *Ts]] = []
-        async with self._ensure_transaction():
+        async with self.transaction(ensure=True):
             for model, model_objects in self._aggregate_models_by_table(objects):
                 table_name = QueryIdentifier(model.get_table_name())
                 fields = {
@@ -559,7 +573,7 @@ async def update(self, objects: Sequence[TableBase]):
         if not objects:
             return
 
-        async with self._ensure_transaction():
+        async with self.transaction(ensure=True):
             for model, model_objects in self._aggregate_models_by_table(objects):
                 table_name = QueryIdentifier(model.get_table_name())
                 primary_key = self._get_primary_key(model)
@@ -638,7 +652,7 @@ async def delete(self, objects: Sequence[TableBase]):
         :param objects: A sequence of TableBase instances to delete
 
         """
-        async with self._ensure_transaction():
+        async with self.transaction(ensure=True):
             for model, model_objects in self._aggregate_models_by_table(objects):
                 table_name = QueryIdentifier(model.get_table_name())
                 primary_key = self._get_primary_key(model)
@@ -797,19 +811,6 @@ def _get_primary_key(self, obj: Type[TableBase]) -> str | None:
             )
         return self.obj_to_primary_key[table_name]
 
-    @asynccontextmanager
-    async def _ensure_transaction(self):
-        """
-        Context manager that ensures operations are executed within a transaction.
-        If no transaction is active, creates a new one for the duration of the context.
-        If a transaction is already active, uses the existing transaction.
-        """
-        if not self.in_transaction:
-            async with self.transaction():
-                yield
-        else:
-            yield
-
     def _batch_objects_and_values(
         self,
         objects: Sequence[TableBase],