Merge pull request #3 from wherobots/cursor-add-close

Add close and exit method for cursor

Author: mpetazzoni

wherobots/db/constants.py

@@ -1,5 +1,5 @@
 from enum import auto
-from strenum import StrEnum
+from strenum import LowercaseStrEnum
 
 from .region import Region
 from .runtime import Runtime
@@ -14,7 +14,7 @@
 MAX_MESSAGE_SIZE: int = 100 * 2**20  # 100MiB
 
 
-class ExecutionState(StrEnum):
+class ExecutionState(LowercaseStrEnum):
     IDLE = auto()
     "Not executing any operation."
 
@@ -40,21 +40,21 @@ def is_terminal_state(self):
         return self in (ExecutionState.COMPLETED, ExecutionState.FAILED)
 
 
-class RequestKind(StrEnum):
+class RequestKind(LowercaseStrEnum):
     EXECUTE_SQL = auto()
     RETRIEVE_RESULTS = auto()
 
 
-class EventKind(StrEnum):
+class EventKind(LowercaseStrEnum):
     STATE_UPDATED = auto()
     EXECUTION_RESULT = auto()
     ERROR = auto()
 
 
-class ResultsFormat(StrEnum):
+class ResultsFormat(LowercaseStrEnum):
     JSON = auto()
     ARROW = auto()
 
 
-class DataCompression(StrEnum):
+class DataCompression(LowercaseStrEnum):
     BROTLI = auto()

wherobots/db/cursor.py

@@ -77,8 +77,18 @@ def fetchmany(self, size: int = None):
     def fetchall(self):
         return self.__get_results()[self.__current_row :]
 
+    def close(self):
+        """Close the cursor."""
+        pass
+
     def __iter__(self):
         return self
 
     def __next__(self):
         raise StopIteration
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()