From 643220972e0164aaeaf018220ff738fb0e31bfa8 Mon Sep 17 00:00:00 2001
From: Inada Naoki <songofacandy@gmail.com>
Date: Fri, 29 Nov 2019 21:15:17 +0900
Subject: [PATCH] Add Context Manager Interface to Connection. (#413)

Fixes #400.
---
 MySQLdb/connections.py            | 6 ++++++
 tests/test_MySQLdb_nonstandard.py | 5 +++++
 2 files changed, 11 insertions(+)

diff --git a/MySQLdb/connections.py b/MySQLdb/connections.py
index 27a0477..020b4a9 100644
--- a/MySQLdb/connections.py
+++ b/MySQLdb/connections.py
@@ -207,6 +207,12 @@ class Connection(_mysql.connection):
                 self.autocommit(autocommit)
         self.messages = []
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def autocommit(self, on):
         on = bool(on)
         if self.get_autocommit() != on:
diff --git a/tests/test_MySQLdb_nonstandard.py b/tests/test_MySQLdb_nonstandard.py
index fa4692e..c5cacbe 100644
--- a/tests/test_MySQLdb_nonstandard.py
+++ b/tests/test_MySQLdb_nonstandard.py
@@ -99,3 +99,8 @@ class CoreAPI(unittest.TestCase):
 
     def test_fileno(self):
         self.assertGreaterEqual(self.conn.fileno(), 0)
+
+    def test_context_manager(self):
+        with connection_factory() as conn:
+            self.assertFalse(conn.closed)
+        self.assertTrue(conn.closed)