add enter/exit to cursor and connection classes

test_dbapi/test_dbapi.py

@@ -94,31 +94,29 @@ def _test_errors(self, connection: dbapi.Connection):
         with pytest.raises(dbapi.InterfaceError):
             dbapi.YdbDBApi().connect("localhost:2136", database="/local666")
 
-        cur = connection.cursor()
-
-        with suppress(dbapi.DatabaseError):
-            cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
+        with connection.cursor() as cur:
+            with suppress(dbapi.DatabaseError):
+                cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
 
-        with pytest.raises(dbapi.DataError):
-            cur.execute(dbapi.YdbQuery("SELECT 18446744073709551616"))
+            with pytest.raises(dbapi.DataError):
+                cur.execute(dbapi.YdbQuery("SELECT 18446744073709551616"))
 
-        with pytest.raises(dbapi.DataError):
-            cur.execute(dbapi.YdbQuery("SELECT * FROM 拉屎"))
+            with pytest.raises(dbapi.DataError):
+                cur.execute(dbapi.YdbQuery("SELECT * FROM 拉屎"))
 
-        with pytest.raises(dbapi.DataError):
-            cur.execute(dbapi.YdbQuery("SELECT floor(5 / 2)"))
+            with pytest.raises(dbapi.DataError):
+                cur.execute(dbapi.YdbQuery("SELECT floor(5 / 2)"))
 
-        with pytest.raises(dbapi.ProgrammingError):
-            cur.execute(dbapi.YdbQuery("SELECT * FROM test"))
+            with pytest.raises(dbapi.ProgrammingError):
+                cur.execute(dbapi.YdbQuery("SELECT * FROM test"))
 
-        cur.execute(dbapi.YdbQuery("CREATE TABLE test(id Int64, PRIMARY KEY (id))", is_ddl=True))
+            cur.execute(dbapi.YdbQuery("CREATE TABLE test(id Int64, PRIMARY KEY (id))", is_ddl=True))
 
-        cur.execute(dbapi.YdbQuery("INSERT INTO test(id) VALUES(1)"))
-        with pytest.raises(dbapi.IntegrityError):
             cur.execute(dbapi.YdbQuery("INSERT INTO test(id) VALUES(1)"))
+            with pytest.raises(dbapi.IntegrityError):
+                cur.execute(dbapi.YdbQuery("INSERT INTO test(id) VALUES(1)"))
 
-        cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
-        cur.close()
+            cur.execute(dbapi.YdbQuery("DROP TABLE test", is_ddl=True))
 
 
 class TestSyncConnection(BaseDBApiTestSuit):
@@ -130,6 +128,11 @@ def sync_connection(self) -> dbapi.Connection:
         finally:
             conn.close()
 
+    @pytest.fixture
+    def sync_connection_with_cm(self) -> dbapi.Connection:
+        with dbapi.YdbDBApi().connect(host="localhost", port="2136", database="/local") as conn:
+            yield conn
+
     @pytest.mark.parametrize(
         "isolation_level, read_only",
         [
@@ -141,8 +144,10 @@ def sync_connection(self) -> dbapi.Connection:
             (dbapi.IsolationLevel.SNAPSHOT_READONLY, True),
         ],
     )
-    def test_isolation_level_read_only(self, isolation_level: str, read_only: bool, sync_connection: dbapi.Connection):
-        self._test_isolation_level_read_only(sync_connection, isolation_level, read_only)
+    def test_isolation_level_read_only(
+        self, isolation_level: str, read_only: bool, sync_connection_with_cm: dbapi.Connection
+    ):
+        self._test_isolation_level_read_only(sync_connection_with_cm, isolation_level, read_only)
 
     def test_connection(self, sync_connection: dbapi.Connection):
         self._test_connection(sync_connection)

ydb_sqlalchemy/dbapi/connection.py

@@ -134,6 +134,12 @@ def rollback(self):
             self._maybe_await(self.session_pool.release, self.tx_context.session)
             self.tx_context = None
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
     def close(self):
         self.rollback()
         if not self._shared_session_pool:

ydb_sqlalchemy/dbapi/cursor.py

@@ -281,6 +281,12 @@ def setinputsizes(self, sizes):
     def setoutputsize(self, column=None):
         pass
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
     def close(self):
         self.rows = None
         self._rows_prefetched = None