feat: make result store a dataclass

Author: peterfoldes

pyproject.toml

@@ -31,6 +31,7 @@ dev = [
     "pre-commit",
     "pytest>=8.0.2",
     "pyarrow>=14.0.2",
+    "pyarrow-stubs>=17.0",
     "pandas",
     "rich>=13.7.1",
     "mypy>=1.14.1",

uv.lock

@@ -72,10 +72,10 @@ def __init__(
         )
         self.__thread.start()
 
-    def __enter__(self):
+    def __enter__(self) -> "Connection":
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         self.close()
 
     def close(self) -> None:
@@ -175,7 +175,7 @@ def _handle_results(self, execution_id: str, results: Dict[str, Any]) -> Any:
         result_compression = results.get("compression")
         logging.info(
             "Received %d bytes of %s-compressed %s results from %s.",
-            len(result_bytes),
+            len(result_bytes) if result_bytes else 0,
             result_compression,
             result_format,
             execution_id,
@@ -216,15 +216,11 @@ def __execute_sql(self, sql: str, handler: Callable[[Any], None]) -> str:
         }
 
         if self.__store:
-            request["store"] = {}
-            if self.__store.format:
-                request["store"]["format"] = self.__store.format.value
-            if self.__store.single:
-                request["store"]["single"] = str(self.__store.single)
-            if self.__store.generate_presigned_url:
-                request["store"]["generate_presigned_url"] = str(
-                    self.__store.generate_presigned_url
-                )
+            request["store"] = {
+                "format": self.__store.format.value if self.__store.format else None,
+                "single": str(self.__store.single),
+                "generate_presigned_url": str(self.__store.generate_presigned_url),
+            }
 
         self.__queries[execution_id] = Query(
             sql=sql,

wherobots/db/connection.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from enum import auto
 from strenum import LowercaseStrEnum
 from typing import Union
@@ -9,22 +10,13 @@ class StorageFormat(LowercaseStrEnum):
     GEOJSON = auto()
 
 
+@dataclass
 class Store:
-    def __init__(
-        self,
-        format: Union[StorageFormat, None] = None,
-        single: bool = False,
-        generate_presigned_url: bool = False,
-    ):
-        self.format = format
-        self.single = single
-        self.generate_presigned_url = generate_presigned_url
+    format: Union[StorageFormat, None] = None
+    single: bool = False
+    generate_presigned_url: bool = False
+
+    def __post_init__(self) -> None:
         assert (
-            single or not generate_presigned_url
+            self.single or not self.generate_presigned_url
         ), "Presigned URL can only be generated when single part file is requested."
-
-    def __repr__(self):
-        return f"Store(format={self.format}, single={self.single}, generate_presigned_url={self.generate_presigned_url})"
-
-    def __str__(self):
-        return f"Store(format={self.format}, single={self.single}, generate_presigned_url={self.generate_presigned_url})"