feat(context-manager): support to context manager for transactions

Author: danielebriggi

src/sqlitecloud/dbapi2.py

@@ -330,6 +330,28 @@ def _apply_adapter(self, value: Any) -> SQLiteTypes:
 
         return value
 
+    def __enter__(self):
+        """
+        Context manager to handle transactions.
+
+        In sqlite3 module the control of the autocommit mode is governed by
+        the `isolation_level` of the connection. To follow this behavior, the
+        context manager does't start a new transaction implicitly. Instead,
+        it handles the commit or rollback of transactions that are explicitly opened.
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            self.commit()
+        else:
+            self.rollback()
+            logging.error(
+                f"Rolling back transaction - error '{exc_value}'",
+                exc_info=True,
+                extra={"traceback": traceback},
+            )
+
     def __del__(self) -> None:
         self.close()
 

src/tests/integration/test_sqlite3_parity.py

@@ -1105,3 +1105,57 @@ def test_parse_colnames_and_decltypes_when_both_are_not_specified(
 
         assert result[0] == 12
         assert result[1] == 25
+
+    @pytest.mark.parametrize(
+        "connection",
+        [
+            "sqlitecloud_dbapi2_connection",
+            "sqlite3_connection",
+        ],
+    )
+    def test_transaction_context_manager_on_success(self, connection, request):
+        connection = request.getfixturevalue(connection)
+
+        connection.execute("BEGIN")
+        with connection:
+            cursor = connection.execute(
+                "INSERT INTO albums (Title, ArtistId) VALUES ('Test Album 1', 1)"
+            )
+            id1 = cursor.lastrowid
+            cursor = connection.execute(
+                "INSERT INTO albums (Title, ArtistId) VALUES ('Test Album 2', 1)"
+            )
+            id2 = cursor.lastrowid
+
+        cursor = connection.execute(
+            "SELECT * FROM albums WHERE AlbumId IN (?, ?)", (id1, id2)
+        )
+        result = cursor.fetchall()
+
+        assert len(result) == 2
+
+    @pytest.mark.parametrize(
+        "connection",
+        [
+            "sqlitecloud_dbapi2_connection",
+            "sqlite3_connection",
+        ],
+    )
+    def test_transaction_context_manager_on_failure(self, connection, request):
+        connection = request.getfixturevalue(connection)
+
+        try:
+            connection.execute("BEGIN")
+            with connection:
+                cursor = connection.execute(
+                    "INSERT INTO albums (Title, ArtistId) VALUES ('Test Album 1', 1)"
+                )
+                id1 = cursor.lastrowid
+                connection.execute("INVALID COMMAND")
+        except Exception:
+            assert True
+
+        cursor = connection.execute("SELECT * FROM albums WHERE AlbumId = ?", (id1,))
+        result = cursor.fetchone()
+
+        assert result is None