From 92562d06b9c33a7a5d1f23f0e2e9c80ef9cb1434 Mon Sep 17 00:00:00 2001
From: Jaeho Yoo <>
Date: Mon, 24 Feb 2025 15:57:30 +0900
Subject: [PATCH] Add support for using Cursor as a context manager

---
 tests/unit/test_dbapi.py | 4 ++--
 trino/dbapi.py           | 6 ++++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py
index ccd44d10..41d96941 100644
--- a/tests/unit/test_dbapi.py
+++ b/tests/unit/test_dbapi.py
@@ -312,5 +312,5 @@ def test_hostname_parsing():
 
 def test_description_is_none_when_cursor_is_not_executed():
     connection = Connection("sample_trino_cluster:443")
-    cursor = connection.cursor()
-    assert hasattr(cursor, 'description')
+    with connection.cursor() as cursor:
+        assert hasattr(cursor, 'description')
diff --git a/trino/dbapi.py b/trino/dbapi.py
index dee7cdb7..f989cced 100644
--- a/trino/dbapi.py
+++ b/trino/dbapi.py
@@ -382,6 +382,12 @@ def __init__(
     def __iter__(self):
         return self._iterator
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     @property
     def connection(self):
         return self._connection