Skip to content

feat(client): support setting session compute size - #20

Merged
utay merged 3 commits into
mainfrom
feat/session-compute-size
Oct 15, 2026
Merged

utay merged 3 commits into
mainfrom
feat/session-compute-size

Conversation

@utay

@utay utay commented Oct 15, 2026

Copy link
Copy Markdown
Member

Summary

  • Add a `ComputeSize` enum and a `compute_size` argument on `Client`, plus `Client.set_compute_size()`, backed by the server's existing `compute_size` session option.
  • `_set_options` now parses the `SetSessionOptionsResult` and raises `ValueError` on rejected options (e.g. `AUTO`), instead of silently ignoring errors. This also applies to `catalog` and `schema`.

Notes

The server fixes the compute size when the session first runs a query, so it must be set before the first query. Per-query sizing would need a backend change and is out of scope.

Test plan

  • Unit tests for enum/string serialization, server error handling, and constructor option

Made with Cursor

Co-authored-by: Cursor <cursoragent@cursor.com>
@utay utay self-assigned this Oct 15, 2026
utay and others added 2 commits October 15, 2026 15:31
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
@utay
utay merged commit ace2ff0 into main Oct 15, 2026
6 checks passed
@utay
utay deleted the feat/session-compute-size branch October 15, 2026 13:43
Comment on lines +62 to +67
def _compute_size_value(compute_size: Union[ComputeSize, str]) -> str:
if isinstance(compute_size, ComputeSize):
return compute_size.value
return compute_size


Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This helper is not needed: ComputeSize mixes in str, so each member already is its string value, and protobuf accepts it as is (ComputeSize.XXL is sent as "2XL").

We can drop _compute_size_value and pass compute_size directly:

options["compute_size"] = flight_pb2.SessionOptionValue(string_value=compute_size)
self._set_options(
    {"compute_size": flight_pb2.SessionOptionValue(string_value=compute_size)}
)

utay pushed a commit that referenced this pull request Oct 15, 2026
🤖 I have created a release *beep* *boop*
---


##
[0.4.0](v0.3.3...v0.4.0)
(2026-10-15)


### Features

* **client:** support setting session compute size
([#20](#20))
([ace2ff0](ace2ff0))

---
This PR was generated with [Release
Please](https://github.com/googleapis/release-please). See
[documentation](https://github.com/googleapis/release-please#release-please).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants