feat: add support for sql query and export and catalog api endpoints in experimental - #910
feat: add support for sql query and export and catalog api endpoints in experimental#910smithednominal wants to merge 6 commits into
Conversation
…to python /experimental/
…eys (presigned_url => presignedUrl for example)
…omatically fixable -- this conditionally imports pyarrow since its an optional dependency.
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
|
Claude finished @smithednominal's task in 4m 10s —— View job ReviewSolid, well-scoped addition. The free-function-taking- Stability-tier leak. Experimental service in core. Docs. Verification gap: this was a static read, so I did not confirm that conjure's |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| catalog = raw["sqlCatalog"] | ||
| if not isinstance(catalog, dict): | ||
| raise ValueError(f"Expected sql_catalog to be a dict, got {type(catalog)}") | ||
|
|
||
| tables = [ | ||
| SqlCatalogTable( | ||
| name=t["name"], | ||
| columns=[SqlCatalogColumn(**c) for c in t.get("columns", [])], | ||
| ) | ||
| for t in catalog.get("tables", []) | ||
| ] | ||
| functions = [SqlCatalogFunction(name=f["name"]) for f in catalog.get("functions", [])] |
There was a problem hiding this comment.
The module docstring says this JSON shape is inferred from live testing, not codegen — but the parser is written as if the shape is guaranteed. SqlCatalogColumn(**c) raises TypeError: unexpected keyword argument the moment the backend adds a field to a column, and raw["sqlCatalog"] / t["name"] raise bare KeyErrors with no context. For a shim whose whole premise is that the wire format may shift, parse defensively and pull fields explicitly:
| catalog = raw["sqlCatalog"] | |
| if not isinstance(catalog, dict): | |
| raise ValueError(f"Expected sql_catalog to be a dict, got {type(catalog)}") | |
| tables = [ | |
| SqlCatalogTable( | |
| name=t["name"], | |
| columns=[SqlCatalogColumn(**c) for c in t.get("columns", [])], | |
| ) | |
| for t in catalog.get("tables", []) | |
| ] | |
| functions = [SqlCatalogFunction(name=f["name"]) for f in catalog.get("functions", [])] | |
| catalog = raw.get("sqlCatalog") | |
| if not isinstance(catalog, dict): | |
| raise ValueError(f"Expected 'sqlCatalog' to be a dict, got {type(catalog)}") | |
| tables = [ | |
| SqlCatalogTable( | |
| name=t["name"], | |
| columns=[ | |
| SqlCatalogColumn(name=c["name"], type=c["type"], nullable=c["nullable"]) | |
| for c in t.get("columns", []) | |
| ], | |
| ) | |
| for t in catalog.get("tables", []) | |
| ] | |
| functions = [SqlCatalogFunction(name=f["name"]) for f in catalog.get("functions", [])] |
(this also fixes the message naming the key sql_catalog when the lookup is sqlCatalog.)
| conjure_python_client.ConjureHTTPError: On invalid query, missing datasets, execution failure, | ||
| resource exhaustion, or timeout. | ||
| """ | ||
| import pyarrow as pa |
There was a problem hiding this comment.
pyarrow is behind the new sql extra, so a user without it gets a bare ModuleNotFoundError: No module named 'pyarrow' with no hint at the fix. The repo convention for optional extras is to translate it (nominal/experimental/video/__init__.py:20, nominal/cli/mis.py:28, nominal/core/datasource.py:481):
| import pyarrow as pa | |
| try: | |
| import pyarrow as pa | |
| except ImportError as ex: | |
| raise ImportError("nominal[sql] is required for SQL queries. Install it with: pip install 'nominal[sql]'") from ex |
The pandas/polars wrappers inherit this path, so fixing it here covers all three entry points.
| resolved_workspace_rid, | ||
| query, | ||
| ) | ||
| return cast(str, response["presignedUrl"]) |
There was a problem hiding this comment.
The module docstring (line 11) and README.md:299 both tell users "use export_sql() if you need a query ID" — but this discards query_id and returns only the URL, so the documented escape hatch doesn't exist. Either return both (a small frozen dataclass, consistent with the SqlCatalog* types above) or drop the claim from both docs.
| def query(self, auth_header: str, workspace_rid: str, query: str, max_rows: int | None = None) -> bytes: | ||
| _headers = { | ||
| "Accept": "application/octet-stream", | ||
| "Content-Type": "application/json", | ||
| "Authorization": auth_header, | ||
| } | ||
| _json: dict[str, object] = {"query": query, "workspace_rid": workspace_rid} | ||
| if max_rows is not None: | ||
| _json["max_rows"] = max_rows | ||
| _response = self._request("POST", self._uri + "/sql/v1/query", params={}, headers=_headers, json=_json) | ||
| return _response.content | ||
|
|
||
| def export(self, auth_header: str, workspace_rid: str, query: str) -> dict[str, object]: | ||
| _headers = { | ||
| "Accept": "application/json", | ||
| "Content-Type": "application/json", | ||
| "Authorization": auth_header, | ||
| } | ||
| _json = {"query": query, "workspace_rid": workspace_rid} | ||
| _response = self._request("POST", self._uri + "/sql/v1/query/export", params={}, headers=_headers, json=_json) | ||
| return _response.json() # type: ignore[no-any-return] | ||
|
|
||
| def get_sql_catalog(self, auth_header: str) -> dict[str, object]: | ||
| _headers = {"Accept": "application/json", "Authorization": auth_header} | ||
| _response = self._request("GET", self._uri + "/sql/v1/catalog", params={}, headers=_headers) | ||
| return _response.json() # type: ignore[no-any-return] |
There was a problem hiding this comment.
Nothing in tests/experimental/test_sql.py exercises this class — every test mocks clients.sql wholesale. That leaves the one part of the PR the docstring explicitly calls unverified (URL paths, Accept headers, and the workspace_rid/max_rows snake_case JSON keys) with zero coverage, while the thin _sql.py plumbing on top of it is tested five times over.
A test subclassing/MagicMock-ing _request and asserting the method, URL suffix, headers, and request body for each of the three calls would pin the wire contract, so a backend casing change fails a test instead of failing at runtime in a user's notebook.
| def test_query_sql_calls_with_default_workspace_when_workspace_rid_none() -> None: | ||
| """Test that query_sql explicitly passes None for workspace_rid in service call when max_rows is not set.""" | ||
| clients = MagicMock() | ||
| clients.auth_header = "Bearer token" | ||
| clients.resolve_default_workspace_rid.return_value = "ri.workspace.default" | ||
| client = NominalClient(_clients=clients) | ||
|
|
||
| table = pa.table({"x": [1]}) | ||
| clients.sql.query.return_value = _arrow_ipc_bytes(table) | ||
|
|
||
| query_sql(client, "SELECT 1") | ||
|
|
||
| clients.sql.query.assert_called_once_with("Bearer token", "ri.workspace.default", "SELECT 1", None) |
There was a problem hiding this comment.
This is a strict subset of test_query_sql_round_trips_arrow_payload_and_resolves_default_workspace — same call, same assert_called_once_with(..., None) assertion. The docstring claims it covers "when max_rows is not set", but so does the first test. Delete it.
| def test_query_sql_propagates_conjure_http_error() -> None: | ||
| """Test that ConjureHTTPError from the service propagates unchanged through query_sql.""" | ||
| clients = MagicMock() | ||
| clients.auth_header = "Bearer token" | ||
| clients.resolve_default_workspace_rid.return_value = "ri.workspace.default" | ||
| client = NominalClient(_clients=clients) | ||
|
|
||
| # Construct a real ConjureHTTPError with a minimal fake requests.Response | ||
| fake_response = MagicMock() | ||
| fake_response.status_code = 400 | ||
| fake_response.json.return_value = { | ||
| "errorCode": "INVALID_ARGUMENT", | ||
| "errorName": "Invalid query syntax", | ||
| "errorInstanceId": "e-123", | ||
| "parameters": {}, | ||
| } | ||
| fake_response.headers.get.return_value = "trace-id-123" | ||
| fake_response.request = MagicMock() | ||
|
|
||
| http_error = HTTPError(response=fake_response) | ||
| conjure_error = ConjureHTTPError(http_error) | ||
| clients.sql.query.side_effect = conjure_error | ||
|
|
||
| with pytest.raises(ConjureHTTPError) as exc_info: | ||
| query_sql(client, "SELECT * FROM nonexistent_table") | ||
|
|
||
| assert exc_info.value is conjure_error | ||
|
|
||
|
|
||
| def test_export_sql_propagates_conjure_http_error() -> None: | ||
| """Test that ConjureHTTPError from export_sql propagates unchanged.""" | ||
| clients = MagicMock() | ||
| clients.auth_header = "Bearer token" | ||
| clients.resolve_default_workspace_rid.return_value = "ri.workspace.default" | ||
| client = NominalClient(_clients=clients) | ||
|
|
||
| fake_response = MagicMock() | ||
| fake_response.status_code = 412 | ||
| fake_response.json.return_value = { | ||
| "errorCode": "FAILED_PRECONDITION", | ||
| "errorName": "SQL export is not configured for this deployment", | ||
| "errorInstanceId": "e-124", | ||
| "parameters": {}, | ||
| } | ||
| fake_response.headers.get.return_value = "trace-id-124" | ||
| fake_response.request = MagicMock() | ||
|
|
||
| http_error = HTTPError(response=fake_response) | ||
| conjure_error = ConjureHTTPError(http_error) | ||
| clients.sql.export.side_effect = conjure_error | ||
|
|
||
| with pytest.raises(ConjureHTTPError) as exc_info: | ||
| export_sql(client, "SELECT * FROM datasets") | ||
|
|
||
| assert exc_info.value is conjure_error |
There was a problem hiding this comment.
Neither query_sql nor export_sql has a try/except, so these 55 lines only assert that MagicMock.side_effect propagates — they test the mock library, not this module, and the elaborate ConjureHTTPError construction implies a translation layer that doesn't exist. They'd also keep passing if error handling were later added and got it wrong (the assertion is is conjure_error, but any raised ConjureHTTPError satisfies the pytest.raises).
Higher-value coverage for the same budget: the _from_json ValueError branch (_sql.py:64) and a non-Arrow/empty query payload, which currently surfaces as an opaque ArrowInvalid rather than the ConjureHTTPError the docstring promises.
1 participant
feat(whs): add support to python client /experimental/ repo section to support beta sql endpoints (query, export, catalog).
Intended flow with examples is in nominal/experimental/sql/README.md including step-by-step snippets and an overall superset example. Documentation on the feature is in progress in this pr: https://github.com/nominal-io/fern-docs/pull/1615