From dd47d4b082d7881b4f79c3dd70ca62f1d60506f1 Mon Sep 17 00:00:00 2001 From: Brijesh-Thakkar Date: Sat, 31 Oct 2026 18:06:58 +0530 Subject: [PATCH] Make SqlTransformSchemaTransformProvider config schema portable (#34613) Co-Authored-By: Claude Sonnet 5.5 --- CHANGES.md | 1 + .../sql/expansion-service/build.gradle | 4 +- .../SqlTransformSchemaTransformProvider.java | 108 ++++++----- ...lTransformSchemaTransformProviderTest.java | 170 ++++++++++++++++++ .../python/apache_beam/transforms/sql_test.py | 70 ++++++++ 5 files changed, 296 insertions(+), 57 deletions(-) create mode 100644 sdks/java/extensions/sql/expansion-service/src/test/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProviderTest.java diff --git a/CHANGES.md b/CHANGES.md index cd87a9609674..84f127a28287 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -82,6 +82,7 @@ ## Bugfixes * (Python) RunInference now measures model memory (the `model_byte_size` metric) when `torch` is not installed. It was always 0, because a failed `ModelManager` import also disabled the `resource` module ([#40854](https://github.com/apache/beam/issues/40854)). +* (Java/Python) The SQL `SqlTransformSchemaTransformProvider` configuration schema (`schematransform:org.apache.beam:sql_transform:v1`) now uses only portable types so it can be discovered from Python: `dialect` is a string, `tableproviders` is an array of strings, and the unused `parameters` field was removed. Invalid dialects and unknown table providers are rejected with a clear error ([#34613](https://github.com/apache/beam/issues/34613)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Security Fixes diff --git a/sdks/java/extensions/sql/expansion-service/build.gradle b/sdks/java/extensions/sql/expansion-service/build.gradle index 562c1ac8dc76..568b642ebc74 100644 --- a/sdks/java/extensions/sql/expansion-service/build.gradle +++ b/sdks/java/extensions/sql/expansion-service/build.gradle @@ -43,7 +43,9 @@ dependencies { permitUnusedDeclared project(path: ":sdks:java:expansion-service") // BEAM-11761 implementation project(path: ":sdks:java:extensions:sql") implementation library.java.vendored_guava_32_1_2_jre - + testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") + testImplementation library.java.junit + testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow") } task runExpansionService (type: JavaExec) { diff --git a/sdks/java/extensions/sql/expansion-service/src/main/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProvider.java b/sdks/java/extensions/sql/expansion-service/src/main/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProvider.java index 3ac890a7370d..7819dc151b4f 100644 --- a/sdks/java/extensions/sql/expansion-service/src/main/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProvider.java +++ b/sdks/java/extensions/sql/expansion-service/src/main/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProvider.java @@ -25,14 +25,13 @@ import java.util.List; import java.util.Map; import java.util.ServiceLoader; -import java.util.stream.Collectors; +import java.util.Set; +import java.util.TreeSet; import org.apache.beam.sdk.extensions.sql.SqlTransform; import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils; import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider; import org.apache.beam.sdk.schemas.Schema; -import org.apache.beam.sdk.schemas.logicaltypes.EnumerationType; -import org.apache.beam.sdk.schemas.logicaltypes.OneOfType; import org.apache.beam.sdk.schemas.transforms.SchemaTransform; import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider; import org.apache.beam.sdk.transforms.Create; @@ -51,18 +50,6 @@ public class SqlTransformSchemaTransformProvider implements SchemaTransformProvi private static final Map> QUERY_PLANNERS = ImmutableMap.of("calcite", org.apache.beam.sdk.extensions.sql.impl.CalciteQueryPlanner.class); - private static final EnumerationType QUERY_ENUMERATION = - EnumerationType.create(QUERY_PLANNERS.keySet().stream().collect(Collectors.toList())); - - private static final OneOfType QUERY_PARAMETER = - OneOfType.create( - Schema.Field.nullable("string", Schema.FieldType.STRING), - Schema.Field.nullable("short", Schema.FieldType.INT16), - Schema.Field.nullable("int", Schema.FieldType.INT32), - Schema.Field.nullable("long", Schema.FieldType.INT64), - Schema.Field.nullable("float", Schema.FieldType.FLOAT), - Schema.Field.nullable("double", Schema.FieldType.DOUBLE), - Schema.Field.nullable("datetime", Schema.FieldType.DATETIME)); @Override public String identifier() { @@ -84,38 +71,58 @@ public String description() { + "[the Beam SQL documentation](https://beam.apache.org/documentation/dsls/sql/overview/)."; } + /** + * The configuration schema deliberately uses only primitive field types (strings, booleans and + * arrays of strings) so that it can be decoded by every SDK. Logical types such as {@code + * EnumerationType} and {@code OneOfType} are Java-specific and cannot be decoded from Python. + * + *

{@code dialect} must be one of the supported query planner names and each entry of {@code + * tableproviders} must be the table type of an available {@link TableProvider}; both are + * validated in {@link #from(Row)}. + */ @Override public Schema configurationSchema() { - List providers = new ArrayList<>(); - ServiceLoader.load(TableProvider.class) - .forEach( - (provider) -> { - providers.add(provider.getTableType()); - }); - EnumerationType providerEnum = EnumerationType.create(providers); - return Schema.of( Schema.Field.of("query", Schema.FieldType.STRING).withDescription("SQL query to execute"), Schema.Field.nullable( "ddl", Schema.FieldType.STRING), // TODO: Underlying builder seems more capable? - Schema.Field.nullable("dialect", Schema.FieldType.logicalType(QUERY_ENUMERATION)), + Schema.Field.nullable("dialect", Schema.FieldType.STRING) + .withDescription( + "Query planner to use. Supported values: " + + String.join(", ", QUERY_PLANNERS.keySet())), Schema.Field.nullable("autoload", Schema.FieldType.BOOLEAN), - Schema.Field.nullable( - "tableproviders", Schema.FieldType.array(Schema.FieldType.logicalType(providerEnum))), - Schema.Field.nullable( - "parameters", - Schema.FieldType.logicalType( - OneOfType.create( - Schema.Field.of( - "positional", - Schema.FieldType.array(Schema.FieldType.logicalType(QUERY_PARAMETER))), - Schema.Field.of( - "named", - Schema.FieldType.array(Schema.FieldType.logicalType(QUERY_PARAMETER))))))); + Schema.Field.nullable("tableproviders", Schema.FieldType.array(Schema.FieldType.STRING)) + .withDescription("Names (table types) of the table providers to make available.")); + } + + private static Map availableTableProviders() { + Map tableProviders = new HashMap<>(); + ServiceLoader.load(TableProvider.class) + .forEach((provider) -> tableProviders.put(provider.getTableType(), provider)); + return tableProviders; } @Override public SchemaTransform from(Row configuration) { + String dialect = configuration.getString("dialect"); + if (dialect != null && !QUERY_PLANNERS.containsKey(dialect)) { + throw new IllegalArgumentException( + String.format( + "Unsupported dialect '%s'. Supported dialects: %s.", + dialect, QUERY_PLANNERS.keySet())); + } + Collection requested = configuration.getArray("tableproviders"); + if (requested != null) { + Set available = availableTableProviders().keySet(); + for (String name : requested) { + if (name != null && !available.contains(name)) { + throw new IllegalArgumentException( + String.format( + "Unknown table provider '%s'. Available table providers: %s.", + name, new TreeSet<>(available))); + } + } + } return new SqlSchemaTransform(configuration); } @@ -172,11 +179,9 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { SqlTransform transform = SqlTransform.query(queryString); // Allow setting the query planner class via the dialect name. - EnumerationType.Value dialect = - config.getLogicalTypeValue("dialect", EnumerationType.Value.class); + String dialect = config.getString("dialect"); if (dialect != null) { - Class queryPlannerClass = - QUERY_PLANNERS.get(QUERY_ENUMERATION.toString(dialect)); + Class queryPlannerClass = QUERY_PLANNERS.get(dialect); if (queryPlannerClass != null) { transform = transform.withQueryPlannerClass(queryPlannerClass); } @@ -196,28 +201,19 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) { transform = transform.withAutoLoading(false); // Add any user specified table providers from the set of available tableproviders. - Map tableProviders = new HashMap<>(); - ServiceLoader.load(TableProvider.class) - .forEach( - (provider) -> { - tableProviders.put(provider.getTableType(), provider); - }); - Collection tableproviderList = config.getArray("tableproviders"); + Map tableProviders = availableTableProviders(); + Collection tableproviderList = config.getArray("tableproviders"); if (tableproviderList != null) { - for (Object nameObj : tableproviderList) { - if (nameObj != null) { // This actually could in theory be null... - TableProvider p = tableProviders.get(nameObj); - if (p != null) { // TODO: We ignore tableproviders that don't exist, we could change - // that. - transform = transform.withTableProvider(p.getTableType(), p); - } + for (String name : tableproviderList) { + TableProvider p = name == null ? null : tableProviders.get(name); + if (p != null) { // Names are validated in from(Row). + transform = transform.withTableProvider(p.getTableType(), p); } } } } - // TODO: Process query parameters. This is not necessary for Syndeo GA but would be - // really nice to have. + // TODO: Support query parameters (positional/named) using primitive field types. // TODO: See about reimplementing a correct version of SqlTransform ErrorCapture errors = new ErrorCapture(); diff --git a/sdks/java/extensions/sql/expansion-service/src/test/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProviderTest.java b/sdks/java/extensions/sql/expansion-service/src/test/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProviderTest.java new file mode 100644 index 000000000000..e3e51f072ca7 --- /dev/null +++ b/sdks/java/extensions/sql/expansion-service/src/test/java/org/apache/beam/sdk/extensions/sql/expansion/SqlTransformSchemaTransformProviderTest.java @@ -0,0 +1,170 @@ +/* + * 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. + */ +package org.apache.beam.sdk.extensions.sql.expansion; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.Collections; +import java.util.ServiceLoader; +import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.FieldType; +import org.apache.beam.sdk.schemas.Schema.TypeName; +import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.Row; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link SqlTransformSchemaTransformProvider}. */ +@RunWith(JUnit4.class) +public class SqlTransformSchemaTransformProviderTest { + + @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + @Rule public final transient TemporaryFolder tmp = new TemporaryFolder(); + + private final SqlTransformSchemaTransformProvider provider = + new SqlTransformSchemaTransformProvider(); + + private Row.FieldValueBuilder config(String query) { + return Row.withSchema(provider.configurationSchema()).withFieldValue("query", query); + } + + /** Python (and other SDKs) can only decode schemas made of non-Java-specific types. */ + @Test + public void testConfigurationSchemaUsesOnlyPortableTypes() { + Schema schema = provider.configurationSchema(); + assertEquals( + Arrays.asList("query", "ddl", "dialect", "autoload", "tableproviders"), + schema.getFieldNames()); + assertEquals(FieldType.STRING, schema.getField("query").getType()); + assertEquals(FieldType.STRING.withNullable(true), schema.getField("ddl").getType()); + assertEquals(FieldType.STRING.withNullable(true), schema.getField("dialect").getType()); + assertEquals(FieldType.BOOLEAN.withNullable(true), schema.getField("autoload").getType()); + assertEquals( + FieldType.array(FieldType.STRING).withNullable(true), + schema.getField("tableproviders").getType()); + for (Schema.Field field : schema.getFields()) { + assertTrue( + field.getName() + " must not be a logical type", + !field.getType().getTypeName().equals(TypeName.LOGICAL_TYPE) + && (field.getType().getCollectionElementType() == null + || !field + .getType() + .getCollectionElementType() + .getTypeName() + .equals(TypeName.LOGICAL_TYPE))); + } + } + + @Test + public void testInvalidDialectIsRejected() { + Row row = config("SELECT 1").withFieldValue("dialect", "not_a_dialect").build(); + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> provider.from(row)); + assertTrue(e.getMessage(), e.getMessage().contains("not_a_dialect")); + assertTrue(e.getMessage(), e.getMessage().contains("calcite")); + } + + @Test + public void testUnknownTableProviderIsRejected() { + Row row = + config("SELECT 1") + .withFieldValue("tableproviders", Collections.singletonList("no_such_provider")) + .build(); + IllegalArgumentException e = + assertThrows(IllegalArgumentException.class, () -> provider.from(row)); + assertTrue(e.getMessage(), e.getMessage().contains("no_such_provider")); + } + + @Test + public void testValidConfigurationIsAccepted() { + String tableType = ServiceLoader.load(TableProvider.class).iterator().next().getTableType(); + provider.from(config("SELECT 1").build()); + provider.from( + config("SELECT 1") + .withFieldValue("dialect", "calcite") + .withFieldValue("autoload", false) + .withFieldValue("tableproviders", Collections.singletonList(tableType)) + .build()); + } + + @Test + public void testExpandWithCalciteDialect() { + Schema inputSchema = Schema.builder().addInt64Field("id").addStringField("name").build(); + PCollection input = + pipeline.apply( + Create.of( + Row.withSchema(inputSchema).addValues(1L, "a").build(), + Row.withSchema(inputSchema).addValues(2L, "b").build()) + .withRowSchema(inputSchema)); + + PCollectionRowTuple result = + PCollectionRowTuple.of("PCOLLECTION", input) + .apply( + provider.from( + config("SELECT name FROM PCOLLECTION WHERE id > 1") + .withFieldValue("dialect", "calcite") + .build())); + + Schema outputSchema = Schema.builder().addStringField("name").build(); + PAssert.that(result.get("output")) + .containsInAnyOrder(Row.withSchema(outputSchema).addValues("b").build()); + pipeline.run().waitUntilFinish(); + } + + @Test + public void testExpandWithDdl() throws Exception { + // DDL table types (here 'text') are resolved from the auto-loaded table providers. + File data = tmp.newFile("people.csv"); + Files.write(data.toPath(), Arrays.asList("1,a", "2,b", "3,c"), StandardCharsets.UTF_8); + + String ddl = + "CREATE EXTERNAL TABLE people (id BIGINT, name VARCHAR) TYPE 'text' LOCATION '" + + data.getAbsolutePath() + + "'"; + + PCollectionRowTuple result = + PCollectionRowTuple.empty(pipeline) + .apply( + provider.from( + config("SELECT name FROM people WHERE id > 1") + .withFieldValue("ddl", ddl) + .withFieldValue("autoload", true) + .build())); + + Schema outputSchema = Schema.builder().addStringField("name").build(); + PAssert.that(result.get("output")) + .containsInAnyOrder( + Row.withSchema(outputSchema).addValues("b").build(), + Row.withSchema(outputSchema).addValues("c").build()); + pipeline.run().waitUntilFinish(); + } +} diff --git a/sdks/python/apache_beam/transforms/sql_test.py b/sdks/python/apache_beam/transforms/sql_test.py index 420a4121e35b..fb26e0eb3282 100644 --- a/sdks/python/apache_beam/transforms/sql_test.py +++ b/sdks/python/apache_beam/transforms/sql_test.py @@ -28,6 +28,7 @@ import apache_beam as beam from apache_beam import coders from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.portability.api import schema_pb2 from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to @@ -240,6 +241,75 @@ def test_tuple_field(self): assert_that(out, equal_to([(2, "bar", 200, ("bar", 200))])) +class SqlTransformSchemaTransformConfigTest(unittest.TestCase): + """Regression guard for the SQL SchemaTransform configuration schema. + + This is NOT an end-to-end test: it does not start an expansion service or + talk to Java. It hand-builds the schema proto that + SqlTransformSchemaTransformProvider.configurationSchema() emits in Java for + ``schematransform:org.apache.beam:sql_transform:v1`` and checks that Python + can decode it. The Java side is pinned separately by + SqlTransformSchemaTransformProviderTest, so if the Java schema changes, this + hand-built copy must be updated by hand. + + Previously ``dialect`` and ``tableproviders`` used the Java-only + ``javasdk_enum`` logical type and ``parameters`` used ``javasdk_oneof``, + which made schema discovery fail with "No logical type registered for URN". + """ + @staticmethod + def _field(name, atomic_type, nullable=True): + return schema_pb2.Field( + name=name, + type=schema_pb2.FieldType(atomic_type=atomic_type, nullable=nullable)) + + def _config_schema(self): + string_array = schema_pb2.FieldType( + nullable=True, + array_type=schema_pb2.ArrayType( + element_type=schema_pb2.FieldType(atomic_type=schema_pb2.STRING))) + return schema_pb2.Schema( + fields=[ + self._field('query', schema_pb2.STRING, nullable=False), + self._field('ddl', schema_pb2.STRING), + self._field('dialect', schema_pb2.STRING), + self._field('autoload', schema_pb2.BOOLEAN), + schema_pb2.Field(name='tableproviders', type=string_array), + ]) + + def test_config_schema_decodes(self): + from apache_beam.typehints.schemas import named_tuple_from_schema + config = named_tuple_from_schema(self._config_schema()) + self.assertEqual( + list(config._fields), + ['query', 'ddl', 'dialect', 'autoload', 'tableproviders']) + hints = config.__annotations__ + self.assertEqual(hints['query'], str) + self.assertEqual(hints['dialect'], typing.Optional[str]) + self.assertEqual(hints['autoload'], typing.Optional[bool]) + # Compare structurally; the repr of typing.Sequence varies by Python version. + non_none = [ + arg for arg in typing.get_args(hints['tableproviders']) + if arg is not type(None) + ] + self.assertEqual(len(non_none), 1) + self.assertEqual(typing.get_args(non_none[0]), (str, )) + self.assertIn(type(None), typing.get_args(hints['tableproviders'])) + + def test_java_enum_logical_type_is_not_decodable(self): + # Documents why the Java schema must not use EnumerationType. + from apache_beam.typehints.schemas import named_tuple_from_schema + enum_field = schema_pb2.Field( + name='dialect', + type=schema_pb2.FieldType( + nullable=True, + logical_type=schema_pb2.LogicalType( + urn='beam:logical_type:javasdk_enum:v1', + representation=schema_pb2.FieldType( + atomic_type=schema_pb2.INT32)))) + with self.assertRaisesRegex(ValueError, 'javasdk_enum'): + named_tuple_from_schema(schema_pb2.Schema(fields=[enum_field])) + + if __name__ == "__main__": logging.getLogger().setLevel(logging.INFO) unittest.main()