Skip to content

Commit 92562d0

Browse files
author
Jaeho Yoo
committed
Add support for using Cursor as a context manager
1 parent 1e95bbf commit 92562d0

File tree

2 files changed

+8
-2
lines changed

2 files changed

+8
-2
lines changed

tests/unit/test_dbapi.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -312,5 +312,5 @@ def test_hostname_parsing():
312312

313313
def test_description_is_none_when_cursor_is_not_executed():
314314
connection = Connection("sample_trino_cluster:443")
315-
cursor = connection.cursor()
316-
assert hasattr(cursor, 'description')
315+
with connection.cursor() as cursor:
316+
assert hasattr(cursor, 'description')

trino/dbapi.py

+6
Original file line numberDiff line numberDiff line change
@@ -382,6 +382,12 @@ def __init__(
382382
def __iter__(self):
383383
return self._iterator
384384

385+
def __enter__(self):
386+
return self
387+
388+
def __exit__(self, exc_type, exc_value, traceback):
389+
self.close()
390+
385391
@property
386392
def connection(self):
387393
return self._connection

0 commit comments

Comments
 (0)