This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-7927-c16e152f2f2b7b239a8660b77b9b1e539c8c2236 in repository https://gitbox.apache.org/repos/asf/texera.git
commit e321f0c717cf5873cebd3ce72885f41976b0282d Author: Eugene Gu <[email protected]> AuthorDate: Mon Sep 21 01:00:25 2026 +0000 feat(amber): register EvaluatedValue in ControlReturn's sealed oneof (#7927) ### What changes were proposed in this PR? Register `EvaluatedValue` as a member of `ControlReturn`'s sealed oneof (`EvaluatedValue evaluatedValue = 53;` in the worker-response range of `controlreturns.proto`), add two Python test files that pin the fix and the underlying invariant, and add a Jackson mix-in in `JSONUtils.scala` that keeps ScalaPB sealed-oneof helper methods out of websocket JSON. The Jackson behavior is also covered directly in `common/workflow-core` by two `JSONUtilsSpec` tests. **Why this is a bug.** The two proto files disagree about `EvaluatePythonExpression`'s reply type: `workerservice.proto` declares the worker's reply as `EvaluatedValue`, but every worker reply must travel inside `ControlReturn`'s sealed oneof. That oneof registers only the coordinator-side wrapper `EvaluatePythonExpressionResponse` (the reply type of the coordinator's RPC, which contains `repeated EvaluatedValue`). `EvaluatedValue` itself is defined in the same file but never joined the oneof, so the declared worker reply has no wire slot. **The failure is silent.** In the Python engine, `set_one_of` assigns the snake_case field name derived from the type name. Assigning a name that is not a oneof field raises no error on a betterproto dataclass, and serialization ignores it, so the worker's reply is packed into an empty `ControlReturn`: `bytes(set_one_of(ControlReturn, EvaluatedValue(...)))` is `b''` and `get_one_of` returns `None`, with no exception or log. The Python worker's handler produces the correct `EvaluatedValue`, but it is lost at the packing step, so the coordinator's `Future.collect` over worker replies can never receive a real value. **Why fix it in the proto.** The bug lives in the contract, not in either engine's code. Both engines' packing and receiving logic assumes that every declared reply type is registered. Registering the type restores that assumption, and both engines regenerate their bindings from the shared proto (generated bindings are not checked in), so no handler code changes on either side. The alternative—changing `workerservice.proto` to reply with the wrapper type—would require handler changes in both engines for no additional benefit. **Follow-up: keep sealed-oneof helper methods out of websocket JSON.** Joining the sealed oneof makes the generated Scala `EvaluatedValue` extend the `ControlReturn` trait, which carries ScalaPB's `isEmpty` and `isDefined` helper methods. Jackson's getter scan then serializes them as `"empty"` and `"defined"` fields in websocket JSON (`EvaluatedValue` is embedded in the `PythonExpressionEvaluateResponse` websocket event). Deserialization rejects those fields because the constructor knows only `value` and `attributes`. This is the round-trip failure that `TexeraWebSocketEventSpec` caught on the first CI run of this PR. The fix adds a `GeneratedSealedOneofMixin` with `@JsonIgnore` on both methods and registers it on `JSONUtils.objectMapper` against the `scalapb.GeneratedSealedOneof` interface, covering every current and future sealed-oneof member. Alternatives were rejected: disabling `FAIL_ON_UNKNOWN_PROPERTIES` globally would hide real deserialization bugs and leave unwanted fields on the wire; loosening the spec would legitimize those fields; and a per-class mix-in would leave the next sealed-oneof member exposed to the same failure. **Not in scope.** The Scala worker's `evaluatePythonExpression` remains a `???` stub (`DataProcessorRPCHandlerInitializer.scala`), and the coordinator fans the request out to all workers of the operator. Evaluating against an operator with Scala workers therefore still fails on the stub; that is pre-existing behavior independent of this change. ### Any related issues, documentation, discussions? Fixes #7924 ### How was this PR tested? Two Python test files were added under `amber/src/test/python`: - `core/util/proto/test_set_one_of.py` contains regression tests showing that `EvaluatedValue` survives `set_one_of`/`get_one_of`, and that a full wire round-trip produces non-empty bytes that parse back to the original value. With the proto change removed and bindings regenerated from the original proto, both tests fail on the empty-bytes/`None` symptoms; with the change restored, they pass. - `core/architecture/rpc/test_reply_types_registered.py` contains invariant tests for the whole bug class. Every reply type declared by `WorkerServiceStub` (21 RPCs) and `CoordinatorServiceStub` (18 RPCs) must be a registered `ControlReturn` oneof member, and every registered member must survive a real `set_one_of`/`get_one_of` round-trip. Non-empty guards prevent reflection from passing silently if the generated-code layout changes. On the unfixed proto, the invariant test identifies `WorkerServiceStub.evaluate_python_expression -> EvaluatedValue` as the missing registration. For the Jackson mix-in, `TexeraWebSocketEventSpec` goes from 9/10 (failing during round-trip deserialization on the unrecognized `"empty"` field, matching this PR's first CI run) to 10/10. Removing only the mix-in reproduces the original failure. Two direct tests were also added to `common/workflow-core/src/test/scala/org/apache/texera/amber/util/JSONUtilsSpec.scala`: - A representative non-empty generated sealed-oneof member, `OpExecWithCode`, serializes without `empty` or `defined` and round-trips to the original Scala value. - The empty generated sealed-oneof value, `OpExecInitInfo.Empty`, serializes to an empty JSON object without helper properties. The direct `workflow-core` tests and formatting checks pass with: ```bash env JAVA_HOME=/Library/Java/JavaVirtualMachines/jdk-17.jdk/Contents/Home sbt -Dsbt.log.noformat=true "WorkflowCore/testOnly org.apache.texera.amber.util.JSONUtilsSpec" env JAVA_HOME=/Library/Java/JavaVirtualMachines/jdk-17.jdk/Contents/Home sbt -Dsbt.log.noformat=true "WorkflowCore/Test/scalafmtCheck" "WorkflowCore/Compile/scalafmtCheck" ``` The first command passes all 23 `JSONUtilsSpec` tests. `git diff --check` also passes. The Python suite (`pytest -m "not integration"`) passes with the two new files included; `ruff check` and `ruff format --check` pass. ScalaPB code generation and a full `sbt compile` on JDK 17 succeed, and the generated Scala `ControlReturn` gains the `SealedValue.EvaluatedValue` case, preserving the coordinator-side `Future[EvaluatedValue]` typing. The websocket serde specs pass after the mix-in, and `scalafix --check` passes. Manual reproduction before and after the proto change: ```python bytes(set_one_of(ControlReturn, EvaluatedValue(value=TypedValue(expression="1+1", value_str="2")))) ``` Before the change, this returns `b''`. After the change, it returns `b'\xaa\x03\n\n\x08\n\x031+1\x1a\x012'` (field 53), and `get_one_of` returns the full value. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Fable 5) and OpenAI Codex --------- Co-authored-by: Xuan Gu <[email protected]> --- .../engine/architecture/rpc/controlreturns.proto | 1 + .../rpc/test_async_rpc_handler_initializer.py | 19 +--- .../rpc/test_reply_types_registered.py | 105 +++++++++++++++++++++ .../test/python/core/util/proto/test_set_one_of.py | 50 ++++++++++ .../org/apache/texera/amber/util/JSONUtils.scala | 16 ++++ .../apache/texera/amber/util/JSONUtilsSpec.scala | 19 ++++ 6 files changed, 195 insertions(+), 15 deletions(-) diff --git a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto index c367bd9dd0..5ce90ebdb6 100644 --- a/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto +++ b/amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlreturns.proto @@ -43,6 +43,7 @@ message ControlReturn { WorkerStateResponse workerStateResponse = 50; WorkerMetricsResponse workerMetricsResponse = 51; FinalizeCheckpointResponse finalizeCheckpointResponse = 52; + EvaluatedValue evaluatedValue = 53; // common responses ControlError controlError = 101; diff --git a/amber/src/test/python/core/architecture/rpc/test_async_rpc_handler_initializer.py b/amber/src/test/python/core/architecture/rpc/test_async_rpc_handler_initializer.py index 89b888ae95..08ed56517e 100644 --- a/amber/src/test/python/core/architecture/rpc/test_async_rpc_handler_initializer.py +++ b/amber/src/test/python/core/architecture/rpc/test_async_rpc_handler_initializer.py @@ -70,15 +70,6 @@ UNIMPLEMENTED_RPCS = frozenset( } ) -# Worker replies with no slot in ControlReturn's oneof. The worker service -# declares EvaluatedValue as EvaluatePythonExpression's reply, but the oneof -# registers only the coordinator-side wrapper, so set_one_of() packs an EMPTY -# ControlReturn: the answer is silently dropped (verified by hand). -# test_async_rpc_server.py pins that swallowing mechanism with a synthetic -# type; this names the real RPC that hits it. The fix is a .proto change in -# its own PR -- this spec only pins the current, broken shape. -REPLIES_MISSING_FROM_CONTROL_RETURN = frozenset({"evaluate_python_expression"}) - # Scanned on disk to catch handler classes the MRO cannot see. HANDLER_PACKAGE = "core.architecture.handlers.control" @@ -218,14 +209,12 @@ class TestTransportOneofs: ) assert unroutable == [] - def test_reply_types_missing_from_control_return_are_exactly_the_known_hole(self): + def test_every_reply_type_is_a_control_return_oneof_member(self): members = _oneof_member_types(ControlReturn) - missing = { + unroutable = sorted( name for name, rpc in RPCS.items() if rpc.handler.reply_type not in members - } - # Both directions: fixing the proto without updating the constant - # fails, and so does padding the constant with a routable RPC. - assert missing == set(REPLIES_MISSING_FROM_CONTROL_RETURN) + ) + assert unroutable == [] class TestHandlerCoverage: diff --git a/amber/src/test/python/core/architecture/rpc/test_reply_types_registered.py b/amber/src/test/python/core/architecture/rpc/test_reply_types_registered.py new file mode 100644 index 0000000000..f21010016c --- /dev/null +++ b/amber/src/test/python/core/architecture/rpc/test_reply_types_registered.py @@ -0,0 +1,105 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import dataclasses +import inspect +import typing + +import proto.org.apache.texera.amber.engine.architecture.rpc as rpc +from core.util import get_one_of, set_one_of +from proto.org.apache.texera.amber.engine.architecture.rpc import ( + ControlReturn, + CoordinatorServiceStub, + WorkerServiceStub, +) + + +def _registered_member_types() -> dict: + """The message types registered in ControlReturn's sealed oneof, by field.""" + hints = typing.get_type_hints(ControlReturn) + return { + field.name: hints[field.name] + for field in dataclasses.fields(ControlReturn) + if getattr(field.metadata["betterproto"], "group", None) == "sealed_value" + } + + +def _declared_reply_types(stub_class) -> dict: + """The declared reply type of every RPC method on a generated service stub.""" + reply_types = {} + for name, method in vars(stub_class).items(): + if not inspect.iscoroutinefunction(method): + continue + annotation = method.__annotations__["return"] + reply_types[name] = ( + getattr(rpc, annotation) if isinstance(annotation, str) else annotation + ) + return reply_types + + +class TestControlReturnRegistry: + def test_every_declared_rpc_reply_type_is_a_registered_oneof_member(self): + # Every reply travels inside ControlReturn's sealed oneof, and + # set_one_of silently packs an unregistered type into an empty + # ControlReturn (see test_async_rpc_server.py, which pins that + # mechanism). A reply type declared by an RPC but missing from the + # oneof is therefore dropped on the wire without any error, so every + # declared reply type must be registered in controlreturns.proto. + registered = set(_registered_member_types().values()) + assert registered, ( + "Reflection found no sealed_value members on ControlReturn; the " + "generated-code layout changed and this test no longer checks " + "anything." + ) + for stub_class in (WorkerServiceStub, CoordinatorServiceStub): + declared = _declared_reply_types(stub_class) + assert declared, ( + f"Reflection found no RPC methods on {stub_class.__name__}; " + f"the generated-stub layout changed and this test no longer " + f"checks anything." + ) + unregistered = { + f"{stub_class.__name__}.{method} -> {reply.__name__}" + for method, reply in declared.items() + if reply not in registered + } + assert not unregistered, ( + f"RPC reply types not registered in ControlReturn's sealed " + f"oneof: {sorted(unregistered)}. Register each type in " + f"controlreturns.proto, otherwise its replies are silently " + f"dropped." + ) + + def test_every_registered_member_survives_set_one_of(self): + # set_one_of derives the oneof field name from the type name, so a + # member whose field name does not follow that convention is packed + # into nothing. Round-tripping every member through set_one_of keeps + # the field names and the conversion logic from drifting apart. + members = _registered_member_types() + assert members, ( + "Reflection found no sealed_value members on ControlReturn; the " + "generated-code layout changed and this test no longer checks " + "anything." + ) + for field_name, member_type in members.items(): + member = member_type() + packed = set_one_of(ControlReturn, member) + assert get_one_of(packed) is member, ( + f"set_one_of failed to pack {member_type.__name__} into " + f"ControlReturn.{field_name}; the oneof field name does not " + f"match the name set_one_of derives from the type." + ) diff --git a/amber/src/test/python/core/util/proto/test_set_one_of.py b/amber/src/test/python/core/util/proto/test_set_one_of.py new file mode 100644 index 0000000000..63a29e688b --- /dev/null +++ b/amber/src/test/python/core/util/proto/test_set_one_of.py @@ -0,0 +1,50 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from core.util import get_one_of, set_one_of +from proto.org.apache.texera.amber.engine.architecture.rpc import ( + ControlReturn, + EvaluatedValue, + TypedValue, +) + + +class TestSetOneOfControlReturn: + def test_evaluated_value_survives_pack_and_unpack(self): + # EvaluatePythonExpression is the one worker RPC whose declared reply + # type is EvaluatedValue; it must be a registered ControlReturn oneof + # member, otherwise set_one_of silently packs an empty ControlReturn + # and the worker's reply is dropped on the wire. + evaluated = EvaluatedValue( + value=TypedValue(expression="1+1", value_str="2"), + attributes=[], + ) + + packed = set_one_of(ControlReturn, evaluated) + + assert get_one_of(packed) == evaluated + + def test_evaluated_value_survives_wire_roundtrip(self): + evaluated = EvaluatedValue( + value=TypedValue(expression="1+1", value_str="2"), + attributes=[], + ) + + wire_bytes = bytes(set_one_of(ControlReturn, evaluated)) + + assert wire_bytes != b"" + assert get_one_of(ControlReturn().parse(wire_bytes)) == evaluated diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/JSONUtils.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/JSONUtils.scala index 7f862b4fbc..8b8d4b3c32 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/JSONUtils.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/JSONUtils.scala @@ -19,6 +19,7 @@ package org.apache.texera.amber.util +import com.fasterxml.jackson.annotation.JsonIgnore import com.fasterxml.jackson.annotation.JsonInclude.Include import com.fasterxml.jackson.databind.module.SimpleModule import com.fasterxml.jackson.databind.{JsonNode, ObjectMapper} @@ -33,6 +34,18 @@ import scala.jdk.CollectionConverters.IteratorHasAsScala object JSONUtils { + // scalapb generates every sealed-oneof trait (a oneof named `sealed_value`, e.g. + // ControlReturn) with concrete `isEmpty`/`isDefined` helper methods that all member + // messages inherit. Jackson treats any public `isXxx` method as a getter, so without + // this mix-in it writes `"empty"`/`"defined"` into the JSON, and deserialization then + // fails because the member's constructor has no such parameters. The mix-in is keyed + // on the `scalapb.GeneratedSealedOneof` interface so every current and future + // sealed-oneof member is exempted at once. + private trait GeneratedSealedOneofMixin { + @JsonIgnore def isEmpty: Boolean + @JsonIgnore def isDefined: Boolean + } + /** * A singleton object for configuring the Jackson `ObjectMapper` to handle JSON serialization and deserialization * in Scala. This custom `ObjectMapper` is tailored for Scala, ensuring compatibility with Scala types @@ -43,6 +56,8 @@ object JSONUtils { * which is common in case classes. * - Registers the `SimpleModule` with pairs of serializer & deserializer to ensure proper handling of serializing * and deserializing the PhysicalPlan + * - Adds the `GeneratedSealedOneofMixin` mix-in for scalapb sealed-oneof types, so their generated + * `isEmpty`/`isDefined` helper methods are not serialized as JSON properties. * - Sets the serialization inclusion rules to exclude `null` and `absent` values: * - `Include.NON_NULL`: Excludes fields with `null` values from the serialized JSON. * - `Include.NON_ABSENT`: Excludes fields with `Option.empty` (or equivalent absent values) from serialization. @@ -61,6 +76,7 @@ object JSONUtils { .addKeySerializer(classOf[PortIdentity], new PortIdentityKeySerializer()) .addKeyDeserializer(classOf[PortIdentity], new PortIdentityKeyDeserializer()) ) + .addMixIn(classOf[scalapb.GeneratedSealedOneof], classOf[GeneratedSealedOneofMixin]) .setSerializationInclusion(Include.NON_NULL) .setSerializationInclusion(Include.NON_ABSENT) .setDateFormat(new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'")) diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/JSONUtilsSpec.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/util/JSONUtilsSpec.scala index 2212393832..59db8f73b9 100644 --- a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/JSONUtilsSpec.scala +++ b/common/workflow-core/src/test/scala/org/apache/texera/amber/util/JSONUtilsSpec.scala @@ -21,6 +21,7 @@ package org.apache.texera.amber.util import com.fasterxml.jackson.databind.JsonNode import com.fasterxml.jackson.databind.node.{JsonNodeFactory, MissingNode} +import org.apache.texera.amber.core.executor.{OpExecInitInfo, OpExecWithCode} import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers @@ -219,4 +220,22 @@ class JSONUtilsSpec extends AnyFlatSpec with Matchers { n.isArray shouldBe true n.size() shouldBe 0 } + + it should "round-trip a generated sealed-oneof member without helper properties" in { + val original = OpExecWithCode("print('hello')", "python") + val json = JSONUtils.objectMapper.writeValueAsString(original) + val root = parse(json) + + root.get("code").asText() shouldBe "print('hello')" + root.get("language").asText() shouldBe "python" + root.has("empty") shouldBe false + root.has("defined") shouldBe false + JSONUtils.objectMapper.readValue(json, classOf[OpExecWithCode]) shouldBe original + } + + it should "exclude helper properties from an empty generated sealed-oneof value" in { + val root = parse(JSONUtils.objectMapper.writeValueAsString(OpExecInitInfo.Empty)) + + root.isEmpty shouldBe true + } }
