diff --git a/_duckdb-stubs/__init__.pyi b/_duckdb-stubs/__init__.pyi index 8770483f..c0f9d714 100644 --- a/_duckdb-stubs/__init__.pyi +++ b/_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: diff --git a/duckdb/__init__.py b/duckdb/__init__.py index d17c530f..52f10523 100644 --- a/duckdb/__init__.py +++ b/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] + diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 638b0a4b..6f7cd8d5 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -204,11 +204,6 @@ struct DuckDBPyConnection : public std::enable_shared_from_this 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(); diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index 61bbe293..f5e83ca1 100644 --- a/src/pyconnection.cpp +++ b/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_(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::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(); diff --git a/tests/fast/api/test_with_propagating_exceptions.py b/tests/fast/api/test_with_propagating_exceptions.py index edf335b6..95ab7a95 100644 --- a/tests/fast/api/test_with_propagating_exceptions.py +++ b/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: diff --git a/tests/fast/test_context_manager.py b/tests/fast/test_context_manager.py index b6a9ebb2..c4bc2892 100644 --- a/tests/fast/test_context_manager.py +++ b/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() +