From 22e6be661022186ae35d5040cb89db428feea325 Mon Sep 17 00:00:00 2001
From: Xingdong Zuo <zuoxingdong@users.noreply.github.com>
Date: Sun, 19 Sep 2021 11:06:46 +0900
Subject: [PATCH] Update presto.py

---
 pyhive/presto.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/pyhive/presto.py b/pyhive/presto.py
index a38cd891..8a4e798c 100644
--- a/pyhive/presto.py
+++ b/pyhive/presto.py
@@ -64,6 +64,13 @@ def __init__(self, *args, **kwargs):
         self._args = args
         self._kwargs = kwargs
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self):
+        """Call close"""
+        self.close()
+
     def close(self):
         """Presto does not have anything to close"""
         # TODO cancel outstanding queries?