chore: Use context manager

Author: paxcodes

clickhouse_backend/driver/connection.py

@@ -1,4 +1,5 @@
 import re
+from contextlib import contextmanager
 
 from clickhouse_driver import connection
 from clickhouse_driver.dbapi import connection as dbapi_connection
@@ -74,7 +75,6 @@ class Cursor(cursor.Cursor):
     # Whether to return data in columnar format. For backwards-compatibility,
     # let's default to None.
     columnar = None
-    _use_numpy = None
 
     def close(self):
         """Push client back to connection pool"""
@@ -89,13 +89,10 @@ def closed(self):
 
     @property
     def use_numpy(self):
-        if self._use_numpy is None:
-            return self._client.client_settings["use_numpy"]
-        return self._use_numpy
+        return self._client.client_settings["use_numpy"]
 
     @use_numpy.setter
     def use_numpy(self, value):
-        self._use_numpy = value
         self._client.client_settings["use_numpy"] = value
         if value:
             try:
@@ -113,6 +110,19 @@ def use_numpy(self, value):
             self._client.query_result_cls = QueryResult
             self._client.iter_query_result_cls = IterQueryResult
             self._client.progress_query_result_cls = ProgressQueryResult
+    
+    @contextmanager
+    def set_query_args(self, columnar: bool, use_numpy: bool):
+        original_use_numpy = self.use_numpy
+        self.use_numpy = use_numpy
+        original_columnar = self.columnar
+        self.columnar = columnar
+        
+        yield self
+        
+        self.use_numpy = original_use_numpy
+        self.columnar = original_columnar
+    
 
     def __del__(self):
         # If someone forgets calling close method,

tests/backends/tests.py

@@ -635,46 +635,43 @@ def test_use_numpy_query(self):
         import numpy as np
 
         with connections["s2r1"].cursor() as cursorWrapper:
-            cursorWrapper.cursor.columnar = True
-            cursorWrapper.cursor.use_numpy = True
-            cursorWrapper.execute(sql)
-            np.testing.assert_equal(
-                cursorWrapper.fetchall(),
+            with cursorWrapper.cursor.set_query_args(columnar=True, use_numpy=True) as cursor:
+                cursor.execute(sql)
+                np.testing.assert_equal(
+                    cursor.fetchall(),
                 [
                     np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
                     np.array([0, 1, 2], dtype=np.uint64),
                     np.array([0, 2.5, 5.0], dtype=np.float64)
                 ],
             )
 
-            cursorWrapper.execute(sql)
-            np.testing.assert_equal(
-                cursorWrapper.fetchmany(2),
-                [
-                    np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
-                    np.array([0, 1, 2], dtype=np.uint64),
+                cursor.execute(sql)
+                np.testing.assert_equal(
+                    cursor.fetchmany(2),
+                    [
+                        np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
+                        np.array([0, 1, 2], dtype=np.uint64),
                 ],
             )
 
-            actual_results = [
-                r
-                for results in iter(lambda: cursorWrapper.fetchmany(2), [])
-                for r in results
-            ]
-            np.testing.assert_equal(
-                actual_results,
-                [
-                    np.array([0, 2.5, 5], dtype=np.float64),
-                ],
-            )
+                actual_results = [
+                    r
+                    for results in iter(lambda: cursor.fetchmany(2), [])
+                    for r in results
+                ]
+                np.testing.assert_equal(
+                    actual_results,
+                    [
+                        np.array([0, 2.5, 5], dtype=np.float64),
+                    ],
+                )
 
-            cursorWrapper.execute(sql)
-            np.testing.assert_equal(
-                cursorWrapper.fetchone(),
-                np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
-            )
-            cursorWrapper.cursor.columnar = False
-            cursorWrapper.cursor.use_numpy = False
+                cursor.execute(sql)
+                np.testing.assert_equal(
+                    cursor.fetchone(),
+                    np.array([np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05'), np.datetime64('2022-01-01T01:00:05')], dtype='datetime64[s]'),
+                )
 
 # These tests aren't conditional because it would require differentiating
 # between MySQL+InnoDB and MySQL+MYISAM (something we currently can't do).