fix cursor context and add corresponding test

mara_db/dbs.py

@@ -1,5 +1,6 @@
 """Abstract definition of database connections"""
 
+import contextlib
 import functools
 import pathlib
 
@@ -37,6 +38,7 @@ def connect(self) -> object:
         """
         raise NotImplementedError(f'Please implement connect for type "{self.__class__.__name__}"')
 
+    @contextlib.contextmanager
     def cursor_context(self) -> object:
         """
         A single iteration with a cursor context. When the iteration is
@@ -50,11 +52,10 @@ def cursor_context(self) -> object:
         try:
             cursor = connection.cursor()
             yield cursor
+            connection.commit()
         except Exception:
             connection.rollback()
             raise
-        else:
-            connection.commit()
         finally:
             cursor.close()
             connection.close()

tests/db_test_helper.py

@@ -26,7 +26,7 @@ def db_replace_placeholders(db: dbs.DB, docker_ip: str, docker_port: int) -> dbs
 Basic tests which can be used for different DB engines.
 """
 
-def _test_sqlalchemy(db):
+def _test_sqlalchemy(db: dbs.DB):
     """
     A simple test to check if the SQLAlchemy connection works
     """
@@ -41,7 +41,7 @@ def _test_sqlalchemy(db):
         # appropriately formatted for the backend
         assert conn.scalar(select(1)) == 1
 
-def _test_connect(db):
+def _test_connect(db: dbs.DB):
     connection = db.connect()
     cursor = connection.cursor()
     try:
@@ -55,3 +55,9 @@ def _test_connect(db):
     finally:
         cursor.close()
         connection.close()
+
+def _test_cursor_context(db: dbs.DB):
+    with db.cursor_context() as cursor:
+        cursor.execute('SELECT 1')
+        row = cursor.fetchone()
+        assert row[0] == 1

tests/mssql/test_mssql.py

@@ -64,6 +64,14 @@ def test_mssql_connect(mssql_db):
     _test_connect(mssql_db)
 
 
+def test_mssql_cursor_context(mssql_db):
+    """
+    A simple test to check if the cursor context of the db works.
+    """
+    from ..db_test_helper import _test_cursor_context
+    _test_cursor_context(mssql_db)
+
+
 
 """
 #################################################################################################################################

tests/postgres/test_postgres.py

@@ -158,3 +158,11 @@ def test_postgres_connect(postgres_db):
     """
     from ..db_test_helper import _test_connect
     _test_connect(postgres_db)
+
+
+def test_postgres_cursor_context(postgres_db):
+    """
+    A simple test to check if the cursor context of the db works.
+    """
+    from ..db_test_helper import _test_cursor_context
+    _test_cursor_context(postgres_db)

tests/test_databricks.py

@@ -49,3 +49,11 @@ def test_databricks_connect():
     """
     from .db_test_helper import _test_connect
     _test_connect(DATABRICKS_DB)
+
+
+def test_databricks_cursor_context():
+    """
+    A simple test to check if the cursor context of the db works.
+    """
+    from .db_test_helper import _test_cursor_context
+    _test_cursor_context(DATABRICKS_DB)