feat(python): implement runtime connection context manager hooks #3152

_duckdb-stubs/__init__.pyi

@@ -199,7 +199,12 @@ class DependencyException(DatabaseError): ...
 class DuckDBPyConnection:
     def __del__(self) -> None: ...
     def __enter__(self) -> Self: ...
-    def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: ...
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None = ...,
+        exc: BaseException | None = ...,
+        traceback: object | None = ...,
+    ) -> bool | None: ...
     def append(self, table_name: str, df: pandas.DataFrame, *, by_name: bool = False) -> DuckDBPyConnection: ...
     def array_type(self, type: IntoPyType, size: typing.SupportsInt) -> sqltypes.DuckDBPyType: ...
     def arrow(self, rows_per_batch: typing.SupportsInt = 1000000) -> pyarrow.lib.RecordBatchReader:

duckdb/__init__.py

@@ -389,3 +389,24 @@
     "view",
     "write_csv",
 ]
+
+
+def _duckdb_pyconnection_enter(self: DuckDBPyConnection) -> DuckDBPyConnection:
+    """Enter the connection context manager."""
+    return self
+
+
+def _duckdb_pyconnection_exit(
+    self: DuckDBPyConnection,
+    exc_type: object,
+    exc_val: object,
+    exc_tb: object,
+) -> bool | None:
+    """Exit the connection context manager and close the connection."""
+    self.close()
+    return False
+
+
+DuckDBPyConnection.__enter__ = _duckdb_pyconnection_enter  # type: ignore[assignment]
+DuckDBPyConnection.__exit__ = _duckdb_pyconnection_exit  # type: ignore[assignment]
+

src/include/duckdb_python/pyconnection/pyconnection.hpp

@@ -204,11 +204,6 @@ struct DuckDBPyConnection : public std::enable_shared_from_this<DuckDBPyConnecti
 	static void Initialize(nb::handle &m);
 	static void Cleanup();
 
-	std::shared_ptr<DuckDBPyConnection> Enter();
-
-	static void Exit(DuckDBPyConnection &self, const nb::object &exc_type, const nb::object &exc,
-	                 const nb::object &traceback);
-
 	static bool DetectAndGetEnvironment();
 	static bool IsJupyter();
 	static std::string FormattedPythonVersion();

src/pyconnection.cpp

@@ -493,12 +493,6 @@ void DuckDBPyConnection::Initialize(nb::handle &m) {
 	// otherwise weakref.ref/proxy/finalize on a connection raises TypeError.
 	auto connection_module = nb::class_<DuckDBPyConnection>(m, "DuckDBPyConnection", nb::is_weak_referenceable());
 
-	connection_module.def("__enter__", &DuckDBPyConnection::Enter)
-	    .def(
-	        "__exit__",
-	        [](DuckDBPyConnection *self, const nb::object &exc_type, const nb::object &exc,
-	           const nb::object &traceback) { DuckDBPyConnection::Exit(*self, exc_type, exc, traceback); },
-	        nb::arg("exc_type").none(), nb::arg("exc").none(), nb::arg("traceback").none());
 	connection_module.def("__del__", &DuckDBPyConnection::Close);
 
 	InitializeConnectionMethods(connection_module);
@@ -2368,20 +2362,6 @@ bool DuckDBPyConnection::IsInteractive() {
 	return GetModuleState().environment != PythonEnvironmentType::NORMAL;
 }
 
-std::shared_ptr<DuckDBPyConnection> DuckDBPyConnection::Enter() {
-	return shared_from_this();
-}
-
-void DuckDBPyConnection::Exit(DuckDBPyConnection &self, const nb::object &exc_type, const nb::object &exc,
-                              const nb::object &traceback) {
-	self.Close();
-	if (exc_type.ptr() != Py_None) {
-		// Propagate the exception if any occurred
-		PyErr_SetObject(exc_type.ptr(), exc.ptr());
-		throw nb::python_error();
-	}
-}
-
 void DuckDBPyConnection::Cleanup() {
 	GetModuleState().default_connection.Set(nullptr);
 	GetModuleState().import_cache.reset();

tests/fast/api/test_with_propagating_exceptions.py

@@ -10,7 +10,7 @@ def test_with(self):
             pytest.raises(duckdb.CatalogException, match="Table with name invalid does not exist"),
             duckdb.connect() as con,
         ):
-            con.execute("invalid")
+            con.execute("select * from invalid")
 
         # Does not raise an exception
         with duckdb.connect() as con:

tests/fast/test_context_manager.py

@@ -1,7 +1,61 @@
+import pytest
 import duckdb
 
 
 class TestContextManager:
-    def test_context_manager(self, duckdb_cursor):
+    def test_context_manager(self):
         with duckdb.connect(database=":memory:", read_only=False) as con:
             assert con.execute("select 1").fetchall() == [(1,)]
+        with pytest.raises(duckdb.ConnectionException, match="Connection already closed"):
+            con.execute("select 1")
+
+    def test_existing_connection_context_manager(self):
+        con = duckdb.connect()
+        with con as c:
+            assert c is con
+            assert c.execute("select 42").fetchall() == [(42,)]
+        with pytest.raises(duckdb.ConnectionException, match="Connection already closed"):
+            con.execute("select 42")
+
+    def test_cursor_context_manager(self):
+        con = duckdb.connect()
+        with con.cursor() as cur:
+            assert cur.execute("select 'hello'").fetchall() == [("hello",)]
+        with pytest.raises(duckdb.ConnectionException, match="Connection already closed"):
+            cur.execute("select 'hello'")
+        # Parent connection should still be open
+        assert con.execute("select 'world'").fetchall() == [("world",)]
+        con.close()
+
+    def test_nested_context_managers(self):
+        with duckdb.connect() as con:
+            assert con.execute("select 1").fetchall() == [(1,)]
+            with con.cursor() as cur:
+                assert cur.execute("select 2").fetchall() == [(2,)]
+            with pytest.raises(duckdb.ConnectionException, match="Connection already closed"):
+                cur.execute("select 2")
+            # con is still valid
+            assert con.execute("select 3").fetchall() == [(3,)]
+        with pytest.raises(duckdb.ConnectionException, match="Connection already closed"):
+            con.execute("select 3")
+
+    def test_exception_propagation_and_cleanup(self):
+        con_ref = None
+        with pytest.raises(ValueError, match="test error"):
+            with duckdb.connect() as con:
+                con_ref = con
+                con.execute("select 1")
+                raise ValueError("test error")
+
+        assert con_ref is not None
+        with pytest.raises(duckdb.ConnectionException, match="Connection already closed"):
+            con_ref.execute("select 1")
+
+    def test_pure_python_context_manager_binding(self):
+        con = duckdb.connect()
+        try:
+            assert con.__enter__.__module__ == "duckdb"
+            assert con.__exit__.__module__ == "duckdb"
+        finally:
+            con.close()
+