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

tests/smoke.py

@@ -16,6 +16,8 @@
 from wherobots.db.region import Region
 from wherobots.db.runtime import Runtime
 from wherobots.db.session_type import SessionType
+from wherobots.db.models import Store, StoreResult
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -24,6 +26,8 @@
     parser.add_argument("--region", help="Region to connect to (ie. aws-us-west-2)")
     parser.add_argument("--runtime", help="Runtime type (ie. tiny)")
     parser.add_argument("--version", help="Runtime version (ie. latest)")
+    parser.add_argument("--force-new", action="store_true")
+    parser.add_argument("--store", action="store_true")
     parser.add_argument(
         "--session-type",
         help="Type of session to create",
@@ -64,6 +68,7 @@
     api_key = None
     token = None
     headers = None
+    store = None
 
     if args.api_key_file:
         with open(args.api_key_file) as f:
@@ -75,6 +80,10 @@
             token = f.read().strip()
         headers = {"Authorization": f"Bearer {token}"}
 
+    if args.store:
+        store = Store.for_download()
+        logging.info("Will requests for results to be stored in cloud storage.")
+
     if args.ws_url:
         conn_func = functools.partial(connect_direct, uri=args.ws_url, headers=headers)
     else:
@@ -88,23 +97,40 @@
             runtime=Runtime(args.runtime) if args.runtime else Runtime.MICRO,
             region=Region(args.region) if args.region else Region.AWS_US_WEST_2,
             version=args.version,
+            force_new=args.force_new,
             session_type=SessionType(args.session_type),
         )
 
-    def render(results: pandas.DataFrame) -> None:
+    def render_df(df: pandas.DataFrame) -> Table:
         table = Table()
         table.add_column("#")
-        for column in results.columns:
+        for column in df.columns:
             table.add_column(column, max_width=args.wide, no_wrap=True)
-        for row in results.itertuples(name=None):
+        for row in df.itertuples(name=None):
             r = [str(x) for x in row]
             table.add_row(*r)
-        Console().print(table)
+        return Table
+
+    def render_stored(sr: StoreResult) -> Table:
+        table = Table()
+        table.add_column("URI")
+        table.add_column("Size")
+        table.add_row(sr.result_uri, str(sr.size))
+        return Table
+
+    def render(results: pandas.DataFrame | StoreResult) -> None:
+        if isinstance(results, StoreResult):
+            Console().print(render_stored(results))
+        else:
+            Console().print(render_df(results))
 
-    def execute(conn: Connection, sql: str) -> pandas.DataFrame:
+    def execute(conn: Connection, sql: str) -> pandas.DataFrame | StoreResult:
         with conn.cursor() as cursor:
-            cursor.execute(sql)
-            return cursor.fetchall()
+            cursor.execute(sql, store=store)
+            if args.store:
+                return cursor.get_store_result()
+            else:
+                return cursor.fetchall()
 
     try:
         with conn_func() as conn:

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