Skip to content

Commit 310775c

Browse files
mharrisb1mharrisb1
committed
feat(routes): add support for custom response models
1 parent f71ee66 commit 310775c

File tree

6 files changed

+679
-6
lines changed

6 files changed

+679
-6
lines changed

‎README.md

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,63 @@ for result in load_response.results:
203203
print("Used pre-aggregations:", result.used_pre_aggregations)
204204
```
205205

206+
### Custom Response Models
207+
208+
You can provide custom response models to extend the default response models:
209+
210+
```python
211+
from pydantic import Field
212+
from cube_http.types.v1.load_response import V1LoadResponse, V1LoadResult
213+
214+
# Extend the default response model with additional fields
215+
class CustomLoadResult(V1LoadResult):
216+
custom_field: str = Field(default=None, alias="customField")
217+
218+
class CustomLoadResponse(V1LoadResponse):
219+
custom_metadata: dict = Field(default_factory=dict, alias="customMetadata")
220+
221+
# Use the custom response model
222+
response = cube.v1.load(
223+
{
224+
"query": {
225+
"measures": ["tasks.count"],
226+
"dimensions": ["tasks.status"],
227+
}
228+
},
229+
response_model=CustomLoadResponse
230+
)
231+
232+
# Access custom fields
233+
custom_metadata = response.custom_metadata
234+
```
235+
236+
This feature is available for all endpoints:
237+
238+
```python
239+
# For SQL endpoint
240+
from cube_http.types.v1.sql_response import V1SqlResponse
241+
242+
class CustomSqlResponse(V1SqlResponse):
243+
extra_data: dict = Field(default_factory=dict)
244+
245+
sql_response = cube.v1.sql(
246+
{"query": {"measures": ["tasks.count"]}},
247+
response_model=CustomSqlResponse
248+
)
249+
250+
# For Meta endpoint
251+
from cube_http.types.v1.meta_response import V1MetaResponse
252+
253+
class CustomMetaResponse(V1MetaResponse):
254+
extended_info: dict = Field(default_factory=dict)
255+
256+
meta_response = cube.v1.meta(
257+
response_model=CustomMetaResponse
258+
)
259+
```
260+
261+
Custom response models are particularly useful when working with custom or extended Cube instances that return additional fields not covered by the default models.
262+
206263
### SQL Query Compilation
207264

208265
You can retrieve the SQL that Cube would execute:

‎examples/custom_response_models.py

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
"""
2+
Cube.dev HTTP Client - Custom Response Models Example
3+
4+
This example demonstrates how to create and use custom response models
5+
with the cube-http-client library to handle extended Cube.js responses.
6+
"""
7+
8+
import asyncio
9+
from typing import Any
10+
11+
from pydantic import Field
12+
13+
import cube_http
14+
from cube_http.types.v1 import V1LoadResponse, V1MetaResponse, V1SqlResponse
15+
16+
# ------ Custom Load Response Models ------
17+
18+
19+
class CustomLoadResponse(V1LoadResponse):
20+
"""Extended load response model with additional custom fields."""
21+
22+
# Add a top-level custom field
23+
execution_details: dict[str, Any] = Field(
24+
default_factory=dict,
25+
alias="executionDetails",
26+
description="Extended execution details from a modified Cube.js server.",
27+
)
28+
29+
30+
# ------ Custom SQL Response Model ------
31+
32+
33+
class CustomSqlResponse(V1SqlResponse):
34+
"""Extended SQL response model with additional custom fields."""
35+
36+
optimizer_info: dict[str, Any] = Field(
37+
default_factory=dict,
38+
alias="optimizerInfo",
39+
description="SQL query optimizer information from an extended Cube.js server.",
40+
)
41+
42+
43+
# ------ Custom Meta Response Model ------
44+
45+
46+
class CustomMetaResponse(V1MetaResponse):
47+
"""Extended metadata response model with additional custom fields."""
48+
49+
schema_version: str | None = Field(
50+
default=None,
51+
alias="schemaVersion",
52+
description="Version information for the data schema.",
53+
)
54+
deployment_info: dict[str, Any] = Field(
55+
default_factory=dict,
56+
alias="deploymentInfo",
57+
description="Additional deployment information from an extended Cube.js server.",
58+
)
59+
60+
61+
def synchronous_example():
62+
"""Demonstrate custom response models with synchronous client."""
63+
print("\n===== Synchronous Custom Response Models Example =====")
64+
65+
# Initialize client
66+
cube = cube_http.Client(
67+
{
68+
"url": "http://localhost:4000/cubejs-api",
69+
"token": "i9f5e76b519a44b060daa33e78c5de170",
70+
}
71+
)
72+
73+
# 1. Load query with custom response model
74+
print("\n----- Load Query with Custom Response Model -----")
75+
try:
76+
response = cube.v1.load(
77+
{
78+
"query": {
79+
"measures": ["tasks.count"],
80+
"dimensions": ["tasks.status"],
81+
}
82+
},
83+
response_model=CustomLoadResponse,
84+
)
85+
86+
print(f"Data rows: {len(response.results[0].data)}")
87+
88+
# Access standard fields
89+
print("Standard fields available:")
90+
print(f" - Query type: {response.query_type}")
91+
print(f" - Results count: {len(response.results)}")
92+
93+
# Access custom fields (will use default values if not provided by server)
94+
print("Custom fields (with default values if not in response):")
95+
print(f" - Execution details: {response.execution_details}")
96+
97+
except Exception as e:
98+
print(f"Error: {e}")
99+
100+
# 2. SQL query with custom response model
101+
print("\n----- SQL Query with Custom Response Model -----")
102+
try:
103+
sql_response = cube.v1.sql(
104+
{
105+
"query": {
106+
"measures": ["tasks.count"],
107+
"dimensions": ["tasks.status"],
108+
}
109+
},
110+
response_model=CustomSqlResponse,
111+
)
112+
113+
# Access standard fields
114+
sql_query, params = sql_response.sql.sql
115+
print(f"SQL query generated: {sql_query[:60]}...")
116+
print(f"Parameters: {params}")
117+
118+
# Access custom fields
119+
print(f"Optimizer info: {sql_response.optimizer_info}")
120+
121+
except Exception as e:
122+
print(f"Error: {e}")
123+
124+
# 3. Meta query with custom response model
125+
print("\n----- Meta Query with Custom Response Model -----")
126+
try:
127+
meta_response = cube.v1.meta(response_model=CustomMetaResponse)
128+
assert meta_response.cubes
129+
130+
# Access standard fields
131+
print(f"Available cubes: {[c.name for c in meta_response.cubes[:3]]}...")
132+
133+
# Access custom fields
134+
print(f"Schema version: {meta_response.schema_version}")
135+
print(f"Deployment info: {meta_response.deployment_info}")
136+
137+
except Exception as e:
138+
print(f"Error: {e}")
139+
140+
141+
async def async_example():
142+
"""Demonstrate custom response models with asynchronous client."""
143+
print("\n===== Asynchronous Custom Response Models Example =====")
144+
145+
# Initialize client
146+
cube = cube_http.AsyncClient(
147+
{
148+
"url": "http://localhost:4000/cubejs-api",
149+
"token": "i9f5e76b519a44b060daa33e78c5de170",
150+
}
151+
)
152+
153+
# 1. Load query with custom response model
154+
print("\n----- Load Query with Custom Response Model -----")
155+
try:
156+
response = await cube.v1.load(
157+
{
158+
"query": {
159+
"measures": ["tasks.count"],
160+
"dimensions": ["tasks.priority"],
161+
}
162+
},
163+
response_model=CustomLoadResponse,
164+
)
165+
166+
print(f"Data rows: {len(response.results[0].data)}")
167+
168+
# Access standard fields
169+
print("Standard fields available:")
170+
print(f" - Query type: {response.query_type}")
171+
print(f" - Results count: {len(response.results)}")
172+
173+
# Access custom fields (will use default values if not provided by server)
174+
print("Custom fields (with default values if not in response):")
175+
print(f" - Execution details: {response.execution_details}")
176+
177+
except Exception as e:
178+
print(f"Error: {e}")
179+
180+
# 2. Run multiple queries in parallel with custom response models
181+
print("\n----- Parallel Queries with Custom Response Models -----")
182+
try:
183+
# Define three queries with different custom response models
184+
load_query = cube.v1.load(
185+
{
186+
"query": {
187+
"measures": ["tasks.count"],
188+
"dimensions": ["tasks.status"],
189+
}
190+
},
191+
response_model=CustomLoadResponse,
192+
)
193+
194+
sql_query = cube.v1.sql(
195+
{
196+
"query": {
197+
"measures": ["tasks.count"],
198+
}
199+
},
200+
response_model=CustomSqlResponse,
201+
)
202+
203+
meta_query = cube.v1.meta(response_model=CustomMetaResponse)
204+
205+
# Run all three queries in parallel
206+
load_response, sql_response, meta_response = await asyncio.gather(
207+
load_query, sql_query, meta_query
208+
)
209+
210+
print("Successfully ran 3 parallel queries with custom response models")
211+
print(f"Load data rows: {len(load_response.results[0].data)}")
212+
print(f"SQL query: {sql_response.sql.sql[0][:40]}...")
213+
print(f"Meta cubes count: {len(meta_response.cubes or [])}")
214+
215+
except Exception as e:
216+
print(f"Error in parallel queries: {e}")
217+
218+
219+
if __name__ == "__main__":
220+
# Run synchronous example
221+
synchronous_example()
222+
223+
# Run asynchronous example
224+
asyncio.run(async_example())

‎src/cube_http/routes/v1/load.py

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,80 @@
1+
from typing import TypeVar, overload
2+
13
from ...exc import V1LoadError
24
from ...types.v1.load_request import V1LoadRequest
35
from ...types.v1.load_response import V1LoadResponse
46
from .._base import AsyncRoute, SyncRoute
57

8+
T = TypeVar("T", bound=V1LoadResponse)
9+
610

711
class SyncLoadRoute(SyncRoute):
8-
def load(self, request: V1LoadRequest) -> V1LoadResponse:
12+
@overload
13+
def load(
14+
self, request: V1LoadRequest, *, response_model: None = None
15+
) -> V1LoadResponse: ...
16+
17+
@overload
18+
def load(self, request: V1LoadRequest, *, response_model: type[T]) -> T: ...
19+
20+
def load(
21+
self, request: V1LoadRequest, *, response_model: type[T] | None = None
22+
) -> T | V1LoadResponse:
23+
"""
24+
Execute a load query.
25+
26+
Args:
27+
request: The load request parameters
28+
response_model: Optional custom response model class to use instead of the default
29+
Must inherit from `V1LoadResponse` model.
30+
31+
Returns:
32+
The response model instance
33+
34+
Raises:
35+
V1LoadError: If the request failed
36+
"""
937
res = self._post("/v1/load", request | {"queryType": "multi"})
1038
if res.status_code == 200:
39+
if response_model is not None:
40+
return response_model.from_response(res)
1141
return V1LoadResponse.from_response(res)
1242
else:
1343
raise V1LoadError.from_response(res)
1444

1545

1646
class AsyncLoadRoute(AsyncRoute):
17-
async def load(self, request: V1LoadRequest) -> V1LoadResponse:
47+
@overload
48+
async def load(
49+
self, request: V1LoadRequest, *, response_model: None = None
50+
) -> V1LoadResponse: ...
51+
52+
@overload
53+
async def load(
54+
self, request: V1LoadRequest, *, response_model: type[T]
55+
) -> T: ...
56+
57+
async def load(
58+
self, request: V1LoadRequest, *, response_model: type[T] | None = None
59+
) -> T | V1LoadResponse:
60+
"""
61+
Execute a load query asynchronously.
62+
63+
Args:
64+
request: The load request parameters
65+
response_model: Optional custom response model class to use instead of the default
66+
Must inherit from `V1LoadResponse` model.
67+
68+
Returns:
69+
The response model instance
70+
71+
Raises:
72+
V1LoadError: If the request failed
73+
"""
1874
res = await self._post("/v1/load", request | {"queryType": "multi"})
1975
if res.status_code == 200:
76+
if response_model is not None:
77+
return response_model.from_response(res)
2078
return V1LoadResponse.from_response(res)
2179
else:
2280
raise V1LoadError.from_response(res)

0 commit comments

Comments
 (0)