From e6a25206ee35937aa21423d6a12a62f3118cbc21 Mon Sep 17 00:00:00 2001
From: Marcin Szczepanski <szczepanski.m.3@pg.com>
Date: Wed, 20 May 2020 11:00:40 +0200
Subject: [PATCH 1/2] Support with statement by adding __enter__ and __exit__
 to Cursor and Connection classes

---
 jaydebeapi/__init__.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/jaydebeapi/__init__.py b/jaydebeapi/__init__.py
index 0abd4c1..3b105e2 100644
--- a/jaydebeapi/__init__.py
+++ b/jaydebeapi/__init__.py
@@ -156,7 +156,7 @@ def _handle_sql_exception_jpype():
     else:
         exc_type = InterfaceError
     reraise(exc_type, exc_info[1], exc_info[2])
-    
+
 def _jdbc_connect_jpype(jclassname, url, driver_args, jars, libs):
     import jpype
     if not jpype.isJVMStarted():
@@ -441,6 +441,12 @@ class Connection(object):
     def cursor(self):
         return Cursor(self, self._converters)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
 # DB-API 2.0 Cursor Object
 class Cursor(object):
 
@@ -594,6 +600,12 @@ class Cursor(object):
     def setoutputsize(self, size, column=None):
         pass
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
 def _unknownSqlTypeConverter(rs, col):
     return rs.getObject(col)
 

From 34be9556a90cca3b7e44b2e93277aa08fe361265 Mon Sep 17 00:00:00 2001
From: Marcin Szczepanski <szczepanski.m.3@pg.com>
Date: Wed, 20 May 2020 15:32:33 +0200
Subject: [PATCH 2/2] Add integration tests for with statement

---
 test/test_integration.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/test/test_integration.py b/test/test_integration.py
index 8818975..a99c69d 100644
--- a/test/test_integration.py
+++ b/test/test_integration.py
@@ -207,6 +207,17 @@ class IntegrationTestBase(object):
         cursor.execute("select * from ACCOUNT")
         self.assertEqual(cursor.rowcount, -1)
 
+    def test_connection_with_statement(self):
+        with self.connect() as conn:
+            self.assertEqual(conn._closed, False)
+        self.assertEqual(conn._closed, True)
+
+    def test_cursor_with_statement(self):
+        with self.conn.cursor() as cursor:
+            cursor.execute("select 1 from ACCOUNT")
+            self.assertIsNotNone(cursor._connection)
+        self.assertIsNone(cursor._connection)
+
 class SqliteTestBase(IntegrationTestBase):
 
     def setUpSql(self):