From dfbbe6ba231c87e51b0dc9eeb200189d773a678c Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Thu, 15 Oct 2026 15:29:39 +0200 Subject: [PATCH 1/3] feat(client): support setting session compute size Co-authored-by: Cursor --- README.md | 11 ++++++ src/altertable_flightsql/__init__.py | 3 +- src/altertable_flightsql/client.py | 53 +++++++++++++++++++++++++++- tests/test_client.py | 49 ++++++++++++++++++++++++- 4 files changed, 113 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 2b59938..1dd5ef9 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,17 @@ client = Client( ) ``` +### Compute size + +```python +from altertable_flightsql import Client, ComputeSize + +client = Client(username="admin", password="secret", compute_size=ComputeSize.L) +``` + +The size is fixed when the session runs its first query, so set it in the +constructor or call `client.set_compute_size(...)` before querying. + ### Queries and Updates ```python diff --git a/src/altertable_flightsql/__init__.py b/src/altertable_flightsql/__init__.py index 7b03997..a31d74b 100644 --- a/src/altertable_flightsql/__init__.py +++ b/src/altertable_flightsql/__init__.py @@ -7,10 +7,11 @@ __version__ = "0.3.3" -from altertable_flightsql.client import Client, PreparedStatement +from altertable_flightsql.client import Client, ComputeSize, PreparedStatement __all__ = [ "__version__", "Client", + "ComputeSize", "PreparedStatement", ] diff --git a/src/altertable_flightsql/client.py b/src/altertable_flightsql/client.py index ef097ac..4447daa 100644 --- a/src/altertable_flightsql/client.py +++ b/src/altertable_flightsql/client.py @@ -46,6 +46,25 @@ class IngestTableMode(Enum): """Create the table if it does not exist, recreate it if it does""" +class ComputeSize(str, Enum): + """Compute size of the worker backing a session.""" + + XS = "XS" + S = "S" + M = "M" + L = "L" + XL = "XL" + XXL = "2XL" + XXXL = "3XL" + XXXXL = "4XL" + + +def _compute_size_value(compute_size: Union[ComputeSize, str]) -> str: + if isinstance(compute_size, ComputeSize): + return compute_size.value + return compute_size + + @dataclass(frozen=True) class IngestIncrementalOptions: """Options for incremental ingestion.""" @@ -113,6 +132,7 @@ def __init__( port: int = 443, tls: bool = True, auto_commit: bool = False, + compute_size: Optional[Union[ComputeSize, str]] = None, ): """ Initialize an Altertable client. @@ -130,6 +150,8 @@ def __init__( port: Server port (default: 443). tls: Whether to use TLS/SSL (default: True). auto_commit: Whether to auto-commit transactions (default: False). + compute_size: Compute size for the session (e.g. ``ComputeSize.L`` + or ``"L"``). When ``None``, the server default is used. """ # Build location URI @@ -162,13 +184,27 @@ def __init__( if schema: options["schema"] = flight_pb2.SessionOptionValue(string_value=schema) + if compute_size: + options["compute_size"] = flight_pb2.SessionOptionValue( + string_value=_compute_size_value(compute_size) + ) + if options: self._set_options(options) def _set_options(self, options: Mapping[str, flight_pb2.SessionOptionValue]): cmd = flight_pb2.SetSessionOptionsRequest(session_options=options) action = flight.Action("SetSessionOptions", cmd.SerializeToString()) - list(self._client.do_action(action)) + results = list(self._client.do_action(action)) + if not results: + return + result = flight_pb2.SetSessionOptionsResult.FromString(bytes(results[0].body)) + if result.errors: + details = ", ".join( + f"{name}: {flight_pb2.SetSessionOptionsResult.ErrorValue.Name(err.value)}" + for name, err in result.errors.items() + ) + raise ValueError(f"Failed to set session options ({details})") def _execute_query_command(self, cmd) -> flight.FlightStreamReader: """Execute a Flight SQL query command and return the result stream.""" @@ -190,6 +226,21 @@ def set_catalog(self, catalog: str): def set_schema(self, schema: str): self._set_options({"schema": flight_pb2.SessionOptionValue(string_value=schema)}) + def set_compute_size(self, compute_size: Union[ComputeSize, str]): + """ + Set the compute size of the session. + + The size is fixed when the session first runs a query, so this only + takes effect if called before the first query. + """ + self._set_options( + { + "compute_size": flight_pb2.SessionOptionValue( + string_value=_compute_size_value(compute_size) + ) + } + ) + def query( self, query: str, diff --git a/tests/test_client.py b/tests/test_client.py index edafdd9..328f9e0 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -5,7 +5,7 @@ import pytest from google.protobuf import any_pb2 -from altertable_flightsql.client import Client, PreparedStatement +from altertable_flightsql.client import Client, ComputeSize, PreparedStatement from altertable_flightsql.generated import arrow_flight_pb2 as flight_pb2 @@ -17,10 +17,16 @@ def __init__(self): self.closed = False close_result = flight_pb2.CloseSessionResult(status=flight_pb2.CloseSessionResult.CLOSED) self.action_results = [SimpleNamespace(body=close_result.SerializeToString())] + self.set_options_result = flight_pb2.SetSessionOptionsResult() def do_action(self, action, options=None): if self.closed: raise RuntimeError("FlightClient is closed") + if action.type == "SetSessionOptions": + self.actions.append(action) + self.options.append(options) + self.events.append(("action", action.type)) + return [SimpleNamespace(body=self.set_options_result.SerializeToString())] self.actions.append(action) self.options.append(options) self.events.append(("action", action.type)) @@ -107,6 +113,47 @@ def test_set_options_serializes_flight_session_request_without_any(): assert _action_body_bytes(action) != wrapped_request.SerializeToString() +@pytest.mark.parametrize( + "value,expected", + [(ComputeSize.XXL, "2XL"), ("m", "m"), (ComputeSize.XS, "XS")], +) +def test_set_compute_size_sends_string_option(value, expected): + flight_client = FakeFlightClient() + client = _client_backed_by(flight_client) + + client.set_compute_size(value) + + request = flight_pb2.SetSessionOptionsRequest.FromString( + _action_body_bytes(flight_client.actions[0]) + ) + assert request.session_options["compute_size"].string_value == expected + + +def test_set_options_raises_on_server_error(): + flight_client = FakeFlightClient() + flight_client.set_options_result.errors["compute_size"].value = ( + flight_pb2.SetSessionOptionsResult.INVALID_VALUE + ) + client = _client_backed_by(flight_client) + + with pytest.raises(ValueError, match="compute_size: INVALID_VALUE"): + client.set_compute_size("AUTO") + + +def test_constructor_sends_compute_size(monkeypatch): + flight_client = FakeFlightClient() + flight_client.authenticate_basic_token = lambda *a: None + monkeypatch.setattr(flight, "FlightClient", lambda *a, **k: flight_client) + + Client("u", "p", catalog="c", compute_size=ComputeSize.L) + + request = flight_pb2.SetSessionOptionsRequest.FromString( + _action_body_bytes(flight_client.actions[0]) + ) + assert request.session_options["compute_size"].string_value == "L" + assert request.session_options["catalog"].string_value == "c" + + def test_close_closes_server_session_before_transport(): flight_client = FakeFlightClient() client = _client_backed_by(flight_client) From b3caca85e9ac801f89156465c0df44b5a5cef3e2 Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Thu, 15 Oct 2026 15:31:14 +0200 Subject: [PATCH 2/3] fix(examples): drop PRIMARY KEY unsupported by DuckLake Co-authored-by: Cursor --- examples/client_usage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/client_usage.py b/examples/client_usage.py index 9e29046..e599fb6 100644 --- a/examples/client_usage.py +++ b/examples/client_usage.py @@ -67,7 +67,7 @@ def example_updates(): ) as client: # Create table client.execute( - "CREATE TABLE IF NOT EXISTS users (id INT PRIMARY KEY, name VARCHAR, age INT)" + "CREATE TABLE IF NOT EXISTS users (id INT, name VARCHAR, age INT)" ) # Insert From 320410efa2bceb6a6d8e3baaab03027e93e66063 Mon Sep 17 00:00:00 2001 From: Yannick Utard Date: Thu, 15 Oct 2026 15:33:05 +0200 Subject: [PATCH 3/3] style: format example Co-authored-by: Cursor --- examples/client_usage.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/client_usage.py b/examples/client_usage.py index e599fb6..84a69a0 100644 --- a/examples/client_usage.py +++ b/examples/client_usage.py @@ -66,9 +66,7 @@ def example_updates(): **CONNECTION_SETTINGS, ) as client: # Create table - client.execute( - "CREATE TABLE IF NOT EXISTS users (id INT, name VARCHAR, age INT)" - ) + client.execute("CREATE TABLE IF NOT EXISTS users (id INT, name VARCHAR, age INT)") # Insert rows = client.execute("INSERT INTO users (id, name, age) VALUES (1, 'Alice', 25)")