refactor: reorganize types and models into dedicated modules

- Move enums (ExecutionState, RequestKind, EventKind, etc.) from constants.py to types.py
- Move dataclasses (Store, StoreResult, ExecutionResult) from store.py/results.py to models.py
- Delete store.py and results.py
- Update all imports across the codebase

Author: mpetazzoni

wherobots/db/__init__.py

@@ -10,9 +10,10 @@
     ProgrammingError,
     NotSupportedError,
 )
+from .models import Store, StoreResult
 from .region import Region
 from .runtime import Runtime
-from .store import Store, StorageFormat, StoreResult
+from .types import StorageFormat
 
 __all__ = [
     "Connection",

wherobots/db/connection.py

@@ -13,19 +13,18 @@
 import websockets.protocol
 import websockets.sync.client
 
-from wherobots.db.constants import (
-    DEFAULT_READ_TIMEOUT_SECONDS,
+from .constants import DEFAULT_READ_TIMEOUT_SECONDS
+from .cursor import Cursor
+from .errors import NotSupportedError, OperationalError
+from .models import ExecutionResult, Store, StoreResult
+from .types import (
     RequestKind,
     EventKind,
     ExecutionState,
     ResultsFormat,
     DataCompression,
     GeometryRepresentation,
 )
-from wherobots.db.cursor import Cursor
-from wherobots.db.errors import NotSupportedError, OperationalError
-from wherobots.db.results import ExecutionResult
-from wherobots.db.store import Store, StoreResult
 
 
 @dataclass
@@ -239,7 +238,7 @@ def __execute_sql(
 
         if store:
             request["store"] = {
-                "format": store.format.value if store.format else None,
+                "format": store.format.value,
                 "single": str(store.single).lower(),
                 "generate_presigned_url": str(store.generate_presigned_url).lower(),
             }

wherobots/db/constants.py

@@ -1,11 +1,9 @@
-from enum import auto
 from packaging.version import Version
-from strenum import LowercaseStrEnum, StrEnum
 
 from .region import Region
 from .runtime import Runtime
 from .session_type import SessionType
-from .store import StorageFormat
+from .types import StorageFormat
 
 
 DEFAULT_ENDPOINT: str = "api.cloud.wherobots.com"  # "api.cloud.wherobots.com"
@@ -22,101 +20,3 @@
 PROTOCOL_VERSION: Version = Version("1.0.0")
 
 PARAM_STYLE = "pyformat"
-
-
-class ExecutionState(LowercaseStrEnum):
-    IDLE = auto()
-    "Not executing any operation."
-
-    EXECUTION_REQUESTED = auto()
-    "Execution of a query has been requested by the driver."
-
-    RUNNING = auto()
-    "The SQL session has reported the query is running."
-
-    SUCCEEDED = auto()
-    "The SQL session has reported the query has completed successfully."
-
-    CANCELLED = auto()
-    "The SQL session has reported the query has been cancelled."
-
-    FAILED = auto()
-    "The SQL session has reported the query has failed."
-
-    RESULTS_REQUESTED = auto()
-    "The driver has requested the query results from the SQL session."
-
-    COMPLETED = auto()
-    "The driver has completed processing the query results."
-
-    def is_terminal_state(self) -> bool:
-        return self in (
-            ExecutionState.COMPLETED,
-            ExecutionState.CANCELLED,
-            ExecutionState.FAILED,
-        )
-
-
-class RequestKind(LowercaseStrEnum):
-    EXECUTE_SQL = auto()
-    RETRIEVE_RESULTS = auto()
-    CANCEL = auto()
-
-
-class EventKind(LowercaseStrEnum):
-    STATE_UPDATED = auto()
-    EXECUTION_RESULT = auto()
-    ERROR = auto()
-
-
-class ResultsFormat(LowercaseStrEnum):
-    JSON = auto()
-    ARROW = auto()
-
-
-class DataCompression(LowercaseStrEnum):
-    BROTLI = auto()
-
-
-class GeometryRepresentation(LowercaseStrEnum):
-    WKT = auto()
-    WKB = auto()
-    EWKT = auto()
-    EWKB = auto()
-    GEOJSON = auto()
-
-
-class AppStatus(StrEnum):
-    PENDING = auto()
-    PREPARING = auto()
-    PREPARE_FAILED = auto()
-    REQUESTED = auto()
-    DEPLOYING = auto()
-    DEPLOY_FAILED = auto()
-    DEPLOYED = auto()
-    INITIALIZING = auto()
-    INIT_FAILED = auto()
-    READY = auto()
-    DESTROY_REQUESTED = auto()
-    DESTROYING = auto()
-    DESTROY_FAILED = auto()
-    DESTROYED = auto()
-
-    def is_starting(self) -> bool:
-        return self in (
-            AppStatus.PENDING,
-            AppStatus.PREPARING,
-            AppStatus.REQUESTED,
-            AppStatus.DEPLOYING,
-            AppStatus.DEPLOYED,
-            AppStatus.INITIALIZING,
-        )
-
-    def is_terminal_state(self) -> bool:
-        return self in (
-            AppStatus.PREPARE_FAILED,
-            AppStatus.DEPLOY_FAILED,
-            AppStatus.INIT_FAILED,
-            AppStatus.DESTROY_FAILED,
-            AppStatus.DESTROYED,
-        )

wherobots/db/cursor.py

@@ -2,8 +2,7 @@
 from typing import Any, List, Tuple, Dict
 
 from .errors import ProgrammingError
-from .results import ExecutionResult
-from .store import Store, StoreResult
+from .models import ExecutionResult, Store, StoreResult
 
 _TYPE_MAP = {
     "object": "STRING",

wherobots/db/driver.py

@@ -27,18 +27,20 @@
     MAX_MESSAGE_SIZE,
     PARAM_STYLE,
     PROTOCOL_VERSION,
-    AppStatus,
-    DataCompression,
-    GeometryRepresentation,
-    ResultsFormat,
-    SessionType,
 )
 from .errors import (
     InterfaceError,
     OperationalError,
 )
 from .region import Region
 from .runtime import Runtime
+from .session_type import SessionType
+from .types import (
+    AppStatus,
+    DataCompression,
+    GeometryRepresentation,
+    ResultsFormat,
+)
 
 apilevel = "2.0"
 threadsafety = 1

wherobots/db/store.py wherobots/db/models.pywherobots/db/store.py renamed to wherobots/db/models.py

@@ -1,14 +1,9 @@
 from dataclasses import dataclass
-from enum import auto
-from strenum import LowercaseStrEnum
 
+import pandas
 
-class StorageFormat(LowercaseStrEnum):
-    """Storage formats for storing query results to cloud storage."""
-
-    PARQUET = auto()
-    CSV = auto()
-    GEOJSON = auto()
+from .constants import DEFAULT_STORAGE_FORMAT
+from .types import StorageFormat
 
 
 @dataclass(frozen=True)
@@ -38,7 +33,7 @@ class Store:
             Requires single=True.
     """
 
-    format: StorageFormat | None = None
+    format: StorageFormat
     single: bool = False
     generate_presigned_url: bool = False
 
@@ -54,9 +49,32 @@ def for_download(cls, format: StorageFormat | None = None) -> "Store":
         single file mode and presigned URL generation enabled.
 
         Args:
-            format: The storage format. Defaults to parquet if not specified.
+            format: The storage format.
 
         Returns:
             A Store configured for single-file download with presigned URL.
         """
-        return cls(format=format, single=True, generate_presigned_url=True)
+        return cls(
+            format=format or DEFAULT_STORAGE_FORMAT,
+            single=True,
+            generate_presigned_url=True,
+        )
+
+
+@dataclass
+class ExecutionResult:
+    """Result of a query execution.
+
+    This class encapsulates all possible outcomes of a query execution:
+    a DataFrame result, an error, or a store result (when results are
+    written to cloud storage).
+
+    Attributes:
+        results: The query results as a pandas DataFrame, or None if an error occurred.
+        error: The error that occurred during execution, or None if successful.
+        store_result: The store result if results were written to cloud storage.
+    """
+
+    results: pandas.DataFrame | None = None
+    error: Exception | None = None
+    store_result: StoreResult | None = None

wherobots/db/results.py

@@ -0,0 +1,106 @@
+from enum import auto
+from strenum import LowercaseStrEnum, StrEnum
+
+
+class ExecutionState(LowercaseStrEnum):
+    IDLE = auto()
+    "Not executing any operation."
+
+    EXECUTION_REQUESTED = auto()
+    "Execution of a query has been requested by the driver."
+
+    RUNNING = auto()
+    "The SQL session has reported the query is running."
+
+    SUCCEEDED = auto()
+    "The SQL session has reported the query has completed successfully."
+
+    CANCELLED = auto()
+    "The SQL session has reported the query has been cancelled."
+
+    FAILED = auto()
+    "The SQL session has reported the query has failed."
+
+    RESULTS_REQUESTED = auto()
+    "The driver has requested the query results from the SQL session."
+
+    COMPLETED = auto()
+    "The driver has completed processing the query results."
+
+    def is_terminal_state(self) -> bool:
+        return self in (
+            ExecutionState.COMPLETED,
+            ExecutionState.CANCELLED,
+            ExecutionState.FAILED,
+        )
+
+
+class RequestKind(LowercaseStrEnum):
+    EXECUTE_SQL = auto()
+    RETRIEVE_RESULTS = auto()
+    CANCEL = auto()
+
+
+class EventKind(LowercaseStrEnum):
+    STATE_UPDATED = auto()
+    EXECUTION_RESULT = auto()
+    ERROR = auto()
+
+
+class ResultsFormat(LowercaseStrEnum):
+    JSON = auto()
+    ARROW = auto()
+
+
+class DataCompression(LowercaseStrEnum):
+    BROTLI = auto()
+
+
+class GeometryRepresentation(LowercaseStrEnum):
+    WKT = auto()
+    WKB = auto()
+    EWKT = auto()
+    EWKB = auto()
+    GEOJSON = auto()
+
+
+class StorageFormat(LowercaseStrEnum):
+    PARQUET = auto()
+    CSV = auto()
+    GEOJSON = auto()
+
+
+class AppStatus(StrEnum):
+    PENDING = auto()
+    PREPARING = auto()
+    PREPARE_FAILED = auto()
+    REQUESTED = auto()
+    DEPLOYING = auto()
+    DEPLOY_FAILED = auto()
+    DEPLOYED = auto()
+    INITIALIZING = auto()
+    INIT_FAILED = auto()
+    READY = auto()
+    DESTROY_REQUESTED = auto()
+    DESTROYING = auto()
+    DESTROY_FAILED = auto()
+    DESTROYED = auto()
+
+    def is_starting(self) -> bool:
+        return self in (
+            AppStatus.PENDING,
+            AppStatus.PREPARING,
+            AppStatus.REQUESTED,
+            AppStatus.DEPLOYING,
+            AppStatus.DEPLOYED,
+            AppStatus.INITIALIZING,
+        )
+
+    def is_terminal_state(self) -> bool:
+        return self in (
+            AppStatus.PREPARE_FAILED,
+            AppStatus.DEPLOY_FAILED,
+            AppStatus.INIT_FAILED,
+            AppStatus.DESTROY_FAILED,
+            AppStatus.DESTROYED,
+        )