feat(client): support setting session compute size

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

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 PRIMARY KEY, 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)")

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",
 ]

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,

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)