PoolProxiedConnection support context protocol

The connection object returned by :meth:`_engine.Engine.raw_connection`
now supports the context manager protocol, automatically returning the
connection to the pool when exiting the context.

Fixes: #13116
Change-Id: I51eb1fd61b772368f12a787e5f60db153a839e70
(cherry picked from commit 6fa097ed100dbe5553796c3b45efc50f614c2371)

diff --git a/doc/build/changelog/unreleased_20/13116.rst b/doc/build/changelog/unreleased_20/13116.rst
new file mode 100644 (file)
index 0000000..3e4f98c
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/13116.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: usecase, engine
+    :tickets: 13116
+
+    The connection object returned by :meth:`_engine.Engine.raw_connection`
+    now supports the context manager protocol, automatically returning the
+    connection to the pool when exiting the context.

diff --git a/lib/sqlalchemy/pool/base.py b/lib/sqlalchemy/pool/base.py
index d3527a4fff11fa2cd385bcc30f999421c3da04d1..d2090af6db8db0fed62635ea1303d5fc148f2f0b 100644 (file)
--- a/lib/sqlalchemy/pool/base.py
+++ b/lib/sqlalchemy/pool/base.py
@@ -34,6 +34,7 @@ from .. import log
 from .. import util
 from ..util.typing import Literal
 from ..util.typing import Protocol
+from ..util.typing import Self
 
 if TYPE_CHECKING:
     from ..engine.interfaces import DBAPIConnection
@@ -1124,6 +1125,13 @@ class PoolProxiedConnection(ManagesConnection):
         """
         raise NotImplementedError()
 
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+        return None
+
 
 class _AdhocProxiedConnection(PoolProxiedConnection):
     """provides the :class:`.PoolProxiedConnection` interface for cases where

diff --git a/test/engine/test_pool.py b/test/engine/test_pool.py
index f7016a02e7b55476180f921cd694b246162bd03e..b034e3be78dd0cba6344f02c2bb625b684f8e1b3 100644 (file)
--- a/test/engine/test_pool.py
+++ b/test/engine/test_pool.py
@@ -31,6 +31,7 @@ from sqlalchemy.testing import is_not
 from sqlalchemy.testing import is_not_none
 from sqlalchemy.testing import is_true
 from sqlalchemy.testing import mock
+from sqlalchemy.testing.assertions import expect_raises_message
 from sqlalchemy.testing.engines import testing_engine
 from sqlalchemy.testing.util import gc_collect
 from sqlalchemy.testing.util import lazy_gc
@@ -164,6 +165,25 @@ class PoolTest(PoolTestBase):
         assert not c2.info
         assert "foo2" in c.info
 
+    def test_raw_connection_context(self):
+        p = self._queuepool_fixture(pool_size=1, max_overflow=0)
+        e = create_engine("sqlite://", pool=p, _initialize=False)
+
+        eq_(p.checkedout(), 0)
+        conn = e.raw_connection()
+        eq_(p.checkedout(), 1)
+        with conn as wc:
+            eq_(p.checkedout(), 1)
+            is_(conn, wc)
+
+        eq_(p.checkedout(), 0)
+
+        with expect_raises_message(ValueError, "an error"):
+            with e.raw_connection() as wc:
+                eq_(p.checkedout(), 1)
+                raise ValueError("an error")
+        eq_(p.checkedout(), 0)
+
     def test_rec_info(self):
         p = self._queuepool_fixture(pool_size=1, max_overflow=0)