Skip to content

Commit 5a4b1f8

Browse files
mpetazzonimpetazzoni
committed
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
1 parent 4df6626 commit 5a4b1f8

8 files changed

Lines changed: 152 additions & 151 deletions

File tree

‎wherobots/db/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,10 @@
1010
ProgrammingError,
1111
NotSupportedError,
1212
)
13+
from .models import Store, StoreResult
1314
from .region import Region
1415
from .runtime import Runtime
15-
from .store import Store, StorageFormat, StoreResult
16+
from .types import StorageFormat
1617

1718
__all__ = [
1819
"Connection",

‎wherobots/db/connection.py‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,18 @@
1313
import websockets.protocol
1414
import websockets.sync.client
1515

16-
from wherobots.db.constants import (
17-
DEFAULT_READ_TIMEOUT_SECONDS,
16+
from .constants import DEFAULT_READ_TIMEOUT_SECONDS
17+
from .cursor import Cursor
18+
from .errors import NotSupportedError, OperationalError
19+
from .models import ExecutionResult, Store, StoreResult
20+
from .types import (
1821
RequestKind,
1922
EventKind,
2023
ExecutionState,
2124
ResultsFormat,
2225
DataCompression,
2326
GeometryRepresentation,
2427
)
25-
from wherobots.db.cursor import Cursor
26-
from wherobots.db.errors import NotSupportedError, OperationalError
27-
from wherobots.db.results import ExecutionResult
28-
from wherobots.db.store import Store, StoreResult
2928

3029

3130
@dataclass
@@ -239,7 +238,7 @@ def __execute_sql(
239238

240239
if store:
241240
request["store"] = {
242-
"format": store.format.value if store.format else None,
241+
"format": store.format.value,
243242
"single": str(store.single).lower(),
244243
"generate_presigned_url": str(store.generate_presigned_url).lower(),
245244
}

‎wherobots/db/constants.py‎

Lines changed: 1 addition & 101 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
1-
from enum import auto
21
from packaging.version import Version
3-
from strenum import LowercaseStrEnum, StrEnum
42

53
from .region import Region
64
from .runtime import Runtime
75
from .session_type import SessionType
8-
from .store import StorageFormat
6+
from .types import StorageFormat
97

108

119
DEFAULT_ENDPOINT: str = "api.cloud.wherobots.com" # "api.cloud.wherobots.com"
@@ -22,101 +20,3 @@
2220
PROTOCOL_VERSION: Version = Version("1.0.0")
2321

2422
PARAM_STYLE = "pyformat"
25-
26-
27-
class ExecutionState(LowercaseStrEnum):
28-
IDLE = auto()
29-
"Not executing any operation."
30-
31-
EXECUTION_REQUESTED = auto()
32-
"Execution of a query has been requested by the driver."
33-
34-
RUNNING = auto()
35-
"The SQL session has reported the query is running."
36-
37-
SUCCEEDED = auto()
38-
"The SQL session has reported the query has completed successfully."
39-
40-
CANCELLED = auto()
41-
"The SQL session has reported the query has been cancelled."
42-
43-
FAILED = auto()
44-
"The SQL session has reported the query has failed."
45-
46-
RESULTS_REQUESTED = auto()
47-
"The driver has requested the query results from the SQL session."
48-
49-
COMPLETED = auto()
50-
"The driver has completed processing the query results."
51-
52-
def is_terminal_state(self) -> bool:
53-
return self in (
54-
ExecutionState.COMPLETED,
55-
ExecutionState.CANCELLED,
56-
ExecutionState.FAILED,
57-
)
58-
59-
60-
class RequestKind(LowercaseStrEnum):
61-
EXECUTE_SQL = auto()
62-
RETRIEVE_RESULTS = auto()
63-
CANCEL = auto()
64-
65-
66-
class EventKind(LowercaseStrEnum):
67-
STATE_UPDATED = auto()
68-
EXECUTION_RESULT = auto()
69-
ERROR = auto()
70-
71-
72-
class ResultsFormat(LowercaseStrEnum):
73-
JSON = auto()
74-
ARROW = auto()
75-
76-
77-
class DataCompression(LowercaseStrEnum):
78-
BROTLI = auto()
79-
80-
81-
class GeometryRepresentation(LowercaseStrEnum):
82-
WKT = auto()
83-
WKB = auto()
84-
EWKT = auto()
85-
EWKB = auto()
86-
GEOJSON = auto()
87-
88-
89-
class AppStatus(StrEnum):
90-
PENDING = auto()
91-
PREPARING = auto()
92-
PREPARE_FAILED = auto()
93-
REQUESTED = auto()
94-
DEPLOYING = auto()
95-
DEPLOY_FAILED = auto()
96-
DEPLOYED = auto()
97-
INITIALIZING = auto()
98-
INIT_FAILED = auto()
99-
READY = auto()
100-
DESTROY_REQUESTED = auto()
101-
DESTROYING = auto()
102-
DESTROY_FAILED = auto()
103-
DESTROYED = auto()
104-
105-
def is_starting(self) -> bool:
106-
return self in (
107-
AppStatus.PENDING,
108-
AppStatus.PREPARING,
109-
AppStatus.REQUESTED,
110-
AppStatus.DEPLOYING,
111-
AppStatus.DEPLOYED,
112-
AppStatus.INITIALIZING,
113-
)
114-
115-
def is_terminal_state(self) -> bool:
116-
return self in (
117-
AppStatus.PREPARE_FAILED,
118-
AppStatus.DEPLOY_FAILED,
119-
AppStatus.INIT_FAILED,
120-
AppStatus.DESTROY_FAILED,
121-
AppStatus.DESTROYED,
122-
)

‎wherobots/db/cursor.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,7 @@
22
from typing import Any, List, Tuple, Dict
33

44
from .errors import ProgrammingError
5-
from .results import ExecutionResult
6-
from .store import Store, StoreResult
5+
from .models import ExecutionResult, Store, StoreResult
76

87
_TYPE_MAP = {
98
"object": "STRING",

‎wherobots/db/driver.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,20 @@
2727
MAX_MESSAGE_SIZE,
2828
PARAM_STYLE,
2929
PROTOCOL_VERSION,
30-
AppStatus,
31-
DataCompression,
32-
GeometryRepresentation,
33-
ResultsFormat,
34-
SessionType,
3530
)
3631
from .errors import (
3732
InterfaceError,
3833
OperationalError,
3934
)
4035
from .region import Region
4136
from .runtime import Runtime
37+
from .session_type import SessionType
38+
from .types import (
39+
AppStatus,
40+
DataCompression,
41+
GeometryRepresentation,
42+
ResultsFormat,
43+
)
4244

4345
apilevel = "2.0"
4446
threadsafety = 1

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

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,9 @@
11
from dataclasses import dataclass
2-
from enum import auto
3-
from strenum import LowercaseStrEnum
42

3+
import pandas
54

6-
class StorageFormat(LowercaseStrEnum):
7-
"""Storage formats for storing query results to cloud storage."""
8-
9-
PARQUET = auto()
10-
CSV = auto()
11-
GEOJSON = auto()
5+
from .constants import DEFAULT_STORAGE_FORMAT
6+
from .types import StorageFormat
127

138

149
@dataclass(frozen=True)
@@ -38,7 +33,7 @@ class Store:
3833
Requires single=True.
3934
"""
4035

41-
format: StorageFormat | None = None
36+
format: StorageFormat
4237
single: bool = False
4338
generate_presigned_url: bool = False
4439

@@ -54,9 +49,32 @@ def for_download(cls, format: StorageFormat | None = None) -> "Store":
5449
single file mode and presigned URL generation enabled.
5550
5651
Args:
57-
format: The storage format. Defaults to parquet if not specified.
52+
format: The storage format.
5853
5954
Returns:
6055
A Store configured for single-file download with presigned URL.
6156
"""
62-
return cls(format=format, single=True, generate_presigned_url=True)
57+
return cls(
58+
format=format or DEFAULT_STORAGE_FORMAT,
59+
single=True,
60+
generate_presigned_url=True,
61+
)
62+
63+
64+
@dataclass
65+
class ExecutionResult:
66+
"""Result of a query execution.
67+
68+
This class encapsulates all possible outcomes of a query execution:
69+
a DataFrame result, an error, or a store result (when results are
70+
written to cloud storage).
71+
72+
Attributes:
73+
results: The query results as a pandas DataFrame, or None if an error occurred.
74+
error: The error that occurred during execution, or None if successful.
75+
store_result: The store result if results were written to cloud storage.
76+
"""
77+
78+
results: pandas.DataFrame | None = None
79+
error: Exception | None = None
80+
store_result: StoreResult | None = None

‎wherobots/db/results.py‎

Lines changed: 0 additions & 24 deletions
This file was deleted.

‎wherobots/db/types.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
from enum import auto
2+
from strenum import LowercaseStrEnum, StrEnum
3+
4+
5+
class ExecutionState(LowercaseStrEnum):
6+
IDLE = auto()
7+
"Not executing any operation."
8+
9+
EXECUTION_REQUESTED = auto()
10+
"Execution of a query has been requested by the driver."
11+
12+
RUNNING = auto()
13+
"The SQL session has reported the query is running."
14+
15+
SUCCEEDED = auto()
16+
"The SQL session has reported the query has completed successfully."
17+
18+
CANCELLED = auto()
19+
"The SQL session has reported the query has been cancelled."
20+
21+
FAILED = auto()
22+
"The SQL session has reported the query has failed."
23+
24+
RESULTS_REQUESTED = auto()
25+
"The driver has requested the query results from the SQL session."
26+
27+
COMPLETED = auto()
28+
"The driver has completed processing the query results."
29+
30+
def is_terminal_state(self) -> bool:
31+
return self in (
32+
ExecutionState.COMPLETED,
33+
ExecutionState.CANCELLED,
34+
ExecutionState.FAILED,
35+
)
36+
37+
38+
class RequestKind(LowercaseStrEnum):
39+
EXECUTE_SQL = auto()
40+
RETRIEVE_RESULTS = auto()
41+
CANCEL = auto()
42+
43+
44+
class EventKind(LowercaseStrEnum):
45+
STATE_UPDATED = auto()
46+
EXECUTION_RESULT = auto()
47+
ERROR = auto()
48+
49+
50+
class ResultsFormat(LowercaseStrEnum):
51+
JSON = auto()
52+
ARROW = auto()
53+
54+
55+
class DataCompression(LowercaseStrEnum):
56+
BROTLI = auto()
57+
58+
59+
class GeometryRepresentation(LowercaseStrEnum):
60+
WKT = auto()
61+
WKB = auto()
62+
EWKT = auto()
63+
EWKB = auto()
64+
GEOJSON = auto()
65+
66+
67+
class StorageFormat(LowercaseStrEnum):
68+
PARQUET = auto()
69+
CSV = auto()
70+
GEOJSON = auto()
71+
72+
73+
class AppStatus(StrEnum):
74+
PENDING = auto()
75+
PREPARING = auto()
76+
PREPARE_FAILED = auto()
77+
REQUESTED = auto()
78+
DEPLOYING = auto()
79+
DEPLOY_FAILED = auto()
80+
DEPLOYED = auto()
81+
INITIALIZING = auto()
82+
INIT_FAILED = auto()
83+
READY = auto()
84+
DESTROY_REQUESTED = auto()
85+
DESTROYING = auto()
86+
DESTROY_FAILED = auto()
87+
DESTROYED = auto()
88+
89+
def is_starting(self) -> bool:
90+
return self in (
91+
AppStatus.PENDING,
92+
AppStatus.PREPARING,
93+
AppStatus.REQUESTED,
94+
AppStatus.DEPLOYING,
95+
AppStatus.DEPLOYED,
96+
AppStatus.INITIALIZING,
97+
)
98+
99+
def is_terminal_state(self) -> bool:
100+
return self in (
101+
AppStatus.PREPARE_FAILED,
102+
AppStatus.DEPLOY_FAILED,
103+
AppStatus.INIT_FAILED,
104+
AppStatus.DESTROY_FAILED,
105+
AppStatus.DESTROYED,
106+
)

0 commit comments

Comments
 (0)