From 5eccba7ae9b5f8434ae6a3910548f885cfb16308 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 27 Oct 2026 11:27:09 -0700 Subject: [PATCH 1/2] WITH clause --- sdks/java/extensions/sql/build.gradle | 28 ++ .../meta/provider/iceberg/IcebergTable.java | 83 +++- .../meta/provider/iceberg/ManagedOptions.java | 143 +++++++ .../iceberg/BeamSqlCliIcebergTest.java | 61 +++ .../iceberg/IcebergMetastoreTest.java | 15 +- .../sql/src/main/codegen/config.fmpp | 2 + .../src/main/codegen/includes/parserImpls.ftl | 56 +++ .../sql/TableNameExtractionUtils.java | 8 +- .../sql/impl/BeamCalciteSchema.java | 26 +- .../extensions/sql/impl/CatalogSchema.java | 3 +- .../impl/parser/SqlCreateExternalTable.java | 4 + .../sql/impl/parser/TableClause.java | 372 ++++++++++++++++++ .../sdk/extensions/sql/meta/BeamSqlTable.java | 25 ++ .../sql/impl/TableClauseResolutionTest.java | 324 +++++++++++++++ .../sql/impl/parser/BeamDDLTest.java | 8 + .../sql/impl/parser/TableClauseTest.java | 310 +++++++++++++++ 16 files changed, 1462 insertions(+), 6 deletions(-) create mode 100644 sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/ManagedOptions.java create mode 100644 sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClause.java create mode 100644 sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/TableClauseResolutionTest.java create mode 100644 sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClauseTest.java diff --git a/sdks/java/extensions/sql/build.gradle b/sdks/java/extensions/sql/build.gradle index 49633a0a47a1..79ae40d152c6 100644 --- a/sdks/java/extensions/sql/build.gradle +++ b/sdks/java/extensions/sql/build.gradle @@ -150,8 +150,24 @@ task copyFmppTemplatesFromSrc(type: Copy) { from "src/main/codegen" into "${project.buildDir}/templates-fmpp/codegen" } +// Calcite's Parser.jj has no extension points for the WITH clause of a table, so each anchor +// below gets Beam's grammar (parserImpls.ftl) appended on the same line. +def calciteParserExtensions = [ + // TableRef3, after FOR SYSTEM_TIME AS OF: WITH of a table in FROM. + '[ tableRef = Snapshot(tableRef) ]': 'tableRef = TableClauseOpt(tableRef, TableClause.Role.READ)', + // SqlInsert, before the target column list: WITH of the INSERT target. + '[ LOOKAHEAD(5) tableRef = ExtendTable(tableRef) ]': 'tableRef = TableClauseOpt(tableRef, TableClause.Role.WRITE)', + // End of TableRef3: let a user alias replace the implicit alias of a table with WITH. + '[ tableRef = Tablesample(tableRef) ]': '{ tableRef = TableClause.collapseImplicitAlias(tableRef); }', + // Every WITH query: reject WITH options on its common table expressions. + 'JAVACODE SqlNode addWith(SqlNodeList withList, SqlNode e) {': 'TableClause.checkCteReferences(withList, e);', + // OrderByLimitOpt, which attaches an ORDER BY to a WITH query after addWith: check it too. + 'if (orderBy != null || offsetFetch[0] != null || offsetFetch[1] != null) {': 'TableClause.checkCteReferencesInOrderBy(e, orderBy);', +] + task copyFmppTemplatesFromCalciteCore(type: Copy) { dependsOn configurations.fmppTemplates + inputs.property("calciteParserExtensions", calciteParserExtensions) File calciteCoreJar = files(configurations.fmppTemplates.files).filter { it.name.startsWith("beam-vendor-calcite") }.singleFile @@ -170,6 +186,18 @@ task copyFmppTemplatesFromCalciteCore(type: Copy) { line -> line.replace('import com.google.common.', 'import org.apache.beam.vendor.calcite.v1_40_0.com.google.common.') } + doLast { + def parserJj = file("${project.buildDir}/templates-fmpp/codegen/templates/Parser.jj") + def grammar = parserJj.text + calciteParserExtensions.each { anchor, extension -> + def count = grammar.count(anchor) + if (count != 1) { + throw new GradleException("Expected '${anchor}' once in Calcite's Parser.jj, found it ${count} times") + } + grammar = grammar.replace(anchor, "${anchor} ${extension}") + } + parserJj.text = grammar + } } // Generate the FMPP sources from the FMPP templates. diff --git a/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergTable.java b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergTable.java index 4c198e8e8f6c..13f59b05602a 100644 --- a/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergTable.java +++ b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergTable.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.extensions.sql.meta.provider.iceberg; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -25,6 +26,7 @@ import java.util.List; import java.util.Map; import java.util.function.IntFunction; +import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable; import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTableFilter; import org.apache.beam.sdk.extensions.sql.meta.DefaultTableFilter; import org.apache.beam.sdk.extensions.sql.meta.ProjectSupport; @@ -32,6 +34,8 @@ import org.apache.beam.sdk.extensions.sql.meta.Table; import org.apache.beam.sdk.extensions.sql.meta.provider.bigquery.BeamSqlUnparseContext; import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig; +import org.apache.beam.sdk.io.iceberg.IcebergWriteSchemaTransformProvider; +import org.apache.beam.sdk.io.iceberg.cdc.IcebergCdcMetadataColumns; import org.apache.beam.sdk.managed.Managed; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.values.PBegin; @@ -63,11 +67,37 @@ final class IcebergTable extends SchemaBaseBeamTable { @VisibleForTesting static final String TRIGGERING_FREQUENCY_FIELD = "triggering_frequency_seconds"; + private static final Schema WRITE_CONFIG_SCHEMA = + new IcebergWriteSchemaTransformProvider().configurationSchema(); + private static final Map WRITE_OWNED_OPTIONS = + ImmutableMap.builder() + .putAll( + ManagedOptions.owned( + "the SQL table sets it", + "table", + "catalog_name", + "catalog_properties", + "config_properties")) + .putAll( + ManagedOptions.owned( + "the INSERT writes all of the table's columns", "keep", "drop", "only")) + .putAll( + ManagedOptions.owned( + "it applies only when the write creates the table", + "partition_fields", + "sort_fields", + "table_properties")) + .putAll(ManagedOptions.owned("an INSERT has no error output", "error_handling")) + .build(); + @VisibleForTesting final String tableIdentifier; @VisibleForTesting final IcebergCatalogConfig catalogConfig; @VisibleForTesting @Nullable Integer triggeringFrequency; @VisibleForTesting final @Nullable List partitionFields; + /** The options of an INSERT's {@code WITH} clause, or empty for the table itself. */ + private final Map writeOptions; + IcebergTable(String tableIdentifier, Table table, IcebergCatalogConfig catalogConfig) { super(table.getSchema()); this.schema = table.getSchema(); @@ -90,19 +120,70 @@ final class IcebergTable extends SchemaBaseBeamTable { } this.partitionFields = table.getPartitionFields(); + this.writeOptions = ImmutableMap.of(); + } + + private IcebergTable(IcebergTable table, Schema schema, Map writeOptions) { + super(schema); + this.tableIdentifier = table.tableIdentifier; + this.catalogConfig = table.catalogConfig; + this.triggeringFrequency = table.triggeringFrequency; + this.partitionFields = table.partitionFields; + this.writeOptions = writeOptions; + } + + /** + * Returns this table as the target of an INSERT with {@code WITH} options, which are options of + * the Managed Iceberg write. With {@code 'mode' = 'merge-on-read'}, the INSERT also supplies the + * change type column, if {@code change_type_column} names one, and the sequence number column. + */ + @Override + public BeamSqlTable forWrite(Map options) { + Map checked = + ManagedOptions.check(options, WRITE_CONFIG_SCHEMA, WRITE_OWNED_OPTIONS, "an Iceberg write"); + Schema schema = getSchema(); + if ("merge-on-read".equalsIgnoreCase((String) checked.get("mode"))) { + @Nullable String changeTypeColumn = (String) checked.get("change_type_column"); + if (changeTypeColumn != null) { + schema = withControlColumn(schema, changeTypeColumn, Schema.FieldType.STRING); + } + Object sequenceNumberColumn = + checked.getOrDefault( + "sequence_number_column", IcebergCdcMetadataColumns.COMMIT_SNAPSHOT_SEQUENCE_NUMBER); + schema = withControlColumn(schema, (String) sequenceNumberColumn, Schema.FieldType.INT64); + } + return new IcebergTable(this, schema, checked); + } + + private static Schema withControlColumn(Schema schema, String name, Schema.FieldType type) { + checkArgument( + !schema.hasField(name), + "Merge-on-read control column '%s' collides with another column", + name); + return Schema.builder().addFields(schema.getFields()).addField(name, type).build(); } @Override public POutput buildIOWriter(PCollection input) { + return input.apply(Managed.write(Managed.ICEBERG).withConfig(writeConfig())); + } + + /** + * The Managed write configuration: the table's, its stored triggering frequency, the INSERT's + * {@code WITH} options, which override it, and its partition fields. + */ + @VisibleForTesting + Map writeConfig() { ImmutableMap.Builder configBuilder = ImmutableMap.builder(); configBuilder.putAll(getBaseConfig()); if (triggeringFrequency != null) { configBuilder.put(TRIGGERING_FREQUENCY_FIELD, triggeringFrequency); } + configBuilder.putAll(writeOptions); if (partitionFields != null) { configBuilder.put("partition_fields", partitionFields); } - return input.apply(Managed.write(Managed.ICEBERG).withConfig(configBuilder.build())); + return configBuilder.buildKeepingLast(); } @Override diff --git a/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/ManagedOptions.java b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/ManagedOptions.java new file mode 100644 index 000000000000..25d995c73f88 --- /dev/null +++ b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/ManagedOptions.java @@ -0,0 +1,143 @@ +/* + * 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.meta.provider.iceberg; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.Schema.FieldType; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** Checks the options of a {@code WITH} clause against a Managed configuration schema. */ +final class ManagedOptions { + private ManagedOptions() {} + + /** Maps each key to {@code reason}, why {@code WITH} cannot set it. */ + static Map owned(String reason, String... keys) { + ImmutableMap.Builder owned = ImmutableMap.builder(); + for (String key : keys) { + owned.put(key, reason); + } + return owned.build(); + } + + /** + * Returns the options as the Java values {@code configSchema} expects. + * + * @param owned the keys that {@code WITH} cannot set, each mapped to the reason + * @param use what the options configure, for error messages + * @throws IllegalArgumentException for an owned or unknown key, or a value of the wrong type + */ + static Map check( + Map options, Schema configSchema, Map owned, String use) { + ImmutableMap.Builder checked = ImmutableMap.builder(); + for (Map.Entry option : options.entrySet()) { + String key = option.getKey(); + @Nullable String reason = owned.get(key); + checkArgument(reason == null, "Option '%s' cannot be set in WITH: %s", key, reason); + if (!configSchema.hasField(key)) { + List accepted = + configSchema.getFieldNames().stream() + .filter(name -> !owned.containsKey(name)) + .sorted() + .collect(Collectors.toList()); + throw new IllegalArgumentException( + String.format("Unknown option '%s' for %s. Accepted options: %s", key, use, accepted)); + } + FieldType type = configSchema.getField(key).getType(); + @Nullable Object value = convert(option.getValue(), type); + if (value == null) { + throw new IllegalArgumentException( + String.format("Option '%s' must be of type %s", key, sqlType(type))); + } + checked.put(key, value); + } + return checked.build(); + } + + /** Returns {@code value} as the Java type of {@code type}, or null if it is not of that type. */ + private static @Nullable Object convert(@Nullable Object value, FieldType type) { + switch (type.getTypeName()) { + case STRING: + return value instanceof String ? value : null; + case BOOLEAN: + return value instanceof Boolean ? value : null; + case INT32: + return value instanceof Long && (Long) value == ((Long) value).intValue() + ? ((Long) value).intValue() + : null; + case ARRAY: + if (!(value instanceof List)) { + return null; + } + FieldType elementType = checkStateNotNull(type.getCollectionElementType()); + ImmutableList.Builder list = ImmutableList.builder(); + for (Object element : (List) value) { + @Nullable Object converted = convert(element, elementType); + if (converted == null) { + return null; + } + list.add(converted); + } + return list.build(); + case MAP: + if (!(value instanceof Map)) { + return null; + } + FieldType keyType = checkStateNotNull(type.getMapKeyType()); + FieldType valueType = checkStateNotNull(type.getMapValueType()); + ImmutableMap.Builder map = ImmutableMap.builder(); + for (Map.Entry entry : ((Map) value).entrySet()) { + @Nullable Object convertedKey = convert(entry.getKey(), keyType); + @Nullable Object convertedValue = convert(entry.getValue(), valueType); + if (convertedKey == null || convertedValue == null) { + return null; + } + map.put(convertedKey, convertedValue); + } + return map.build(); + default: + return null; + } + } + + private static String sqlType(FieldType type) { + switch (type.getTypeName()) { + case STRING: + return "VARCHAR"; + case INT32: + return "INTEGER"; + case ARRAY: + return "ARRAY<" + sqlType(checkStateNotNull(type.getCollectionElementType())) + ">"; + case MAP: + return "MAP<" + + sqlType(checkStateNotNull(type.getMapKeyType())) + + ", " + + sqlType(checkStateNotNull(type.getMapValueType())) + + ">"; + default: + return type.getTypeName().name(); + } + } +} diff --git a/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergTest.java b/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergTest.java index 567c3bbc7fae..dabf017823cd 100644 --- a/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergTest.java +++ b/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergTest.java @@ -18,9 +18,12 @@ package org.apache.beam.sdk.extensions.sql.meta.provider.iceberg; import static java.lang.String.format; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import java.io.File; @@ -38,6 +41,7 @@ import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.Row; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.runtime.CalciteContextException; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.DateTime; import org.junit.Before; @@ -233,4 +237,61 @@ public void testCrossCatalogTableWriteAndRead() throws IOException { assertEquals("catalog_1", catalogManager.currentCatalog().name()); assertEquals("my_namespace", catalogManager.currentCatalog().currentDatabase()); } + + @Test + public void testInsertWithOptions() { + sqlEnv.executeDdl(createCatalog("my_catalog")); + sqlEnv.executeDdl("USE CATALOG my_catalog"); + sqlEnv.executeDdl( + "CREATE EXTERNAL TABLE orders (id BIGINT NOT NULL, name VARCHAR) TYPE 'iceberg'"); + // Merge-on-read applies each row as a change, given by the extra change type and sequence + // columns. A colon in an option must not reach the Hadoop catalog as a database name. + String applyChanges = + "INSERT INTO orders WITH ('mode' = 'merge-on-read', 'equality_columns' = ARRAY['id']," + + " 'num_shards' = 2, 'snapshot_properties' = MAP['source', 'sql:test']," + + " 'change_type_column' = 'op', 'sequence_number_column' = 'seq') VALUES "; + for (String changes : + new String[] { + "(1, 'a', 'INSERT', 1), (2, 'b', 'INSERT', 1)", + "(1, 'a', 'UPDATE_BEFORE', 2), (1, 'aa', 'UPDATE_AFTER', 2), (2, 'b', 'DELETE', 2)" + }) { + Pipeline p = Pipeline.create(); + BeamSqlRelUtils.toPCollection(p, sqlEnv.parseQuery(applyChanges + changes)); + p.run().waitUntilFinish(); + } + + Pipeline p = Pipeline.create(); + PCollection output = + BeamSqlRelUtils.toPCollection(p, sqlEnv.parseQuery("SELECT * FROM orders")); + PAssert.that(output) + .containsInAnyOrder(Row.withSchema(output.getSchema()).addValues(1L, "aa").build()); + p.run().waitUntilFinish(); + + String[][] invalidStatements = { + {"SELECT * FROM orders WITH ('mode' = 'append')", "does not support WITH options for reads"}, + { + "INSERT INTO orders WITH ('partition_fields' = ARRAY['id']) SELECT * FROM orders", + "Option 'partition_fields' cannot be set in WITH: it applies only when the write creates" + + " the table" + }, + { + "INSERT INTO orders WITH ('streaming' = TRUE) SELECT * FROM orders", + "Unknown option 'streaming' for an Iceberg write. Accepted options: [" + }, + { + "INSERT INTO orders WITH ('num_shards' = 'two') SELECT * FROM orders", + "Option 'num_shards' must be of type INTEGER" + }, + { + "INSERT INTO orders WITH ('mode' = 'merge-on-read', 'sequence_number_column' = 'name')" + + " SELECT * FROM orders", + "Merge-on-read control column 'name' collides with another column" + }, + }; + for (String[] invalid : invalidStatements) { + Exception e = + assertThrows(invalid[0], RuntimeException.class, () -> sqlEnv.parseQuery(invalid[0])); + assertThat(invalid[0], Throwables.getStackTraceAsString(e), containsString(invalid[1])); + } + } } diff --git a/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastoreTest.java b/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastoreTest.java index 5e76b536134d..703e34f3245b 100644 --- a/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastoreTest.java +++ b/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastoreTest.java @@ -25,6 +25,7 @@ import java.io.IOException; import java.util.Map; import java.util.UUID; +import org.apache.beam.sdk.extensions.sql.TableUtils; import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable; import org.apache.beam.sdk.extensions.sql.meta.Table; import org.apache.beam.sdk.schemas.Schema; @@ -61,7 +62,14 @@ public void testGetTableType() { @Test public void testBuildBeamSqlTable() { - Table table = Table.builder().name("my_table").schema(Schema.of()).type("iceberg").build(); + Table table = + Table.builder() + .name("my_table") + .schema(Schema.of()) + .type("iceberg") + .properties( + TableUtils.parseProperties("{\"beam.write.triggering_frequency_seconds\": 10}")) + .build(); BeamSqlTable sqlTable = metastore().buildBeamSqlTable(table); assertNotNull(sqlTable); @@ -70,6 +78,11 @@ public void testBuildBeamSqlTable() { IcebergTable icebergTable = (IcebergTable) sqlTable; assertEquals(catalog.currentDatabase() + ".my_table", icebergTable.tableIdentifier); assertEquals(catalog.catalogConfig, icebergTable.catalogConfig); + assertEquals(10, icebergTable.writeConfig().get("triggering_frequency_seconds")); + // An INSERT's WITH options override the table's + IcebergTable insertTarget = + (IcebergTable) icebergTable.forWrite(ImmutableMap.of("triggering_frequency_seconds", 30L)); + assertEquals(30, insertTarget.writeConfig().get("triggering_frequency_seconds")); } @Test diff --git a/sdks/java/extensions/sql/src/main/codegen/config.fmpp b/sdks/java/extensions/sql/src/main/codegen/config.fmpp index c3692430610f..1b1da14f5989 100644 --- a/sdks/java/extensions/sql/src/main/codegen/config.fmpp +++ b/sdks/java/extensions/sql/src/main/codegen/config.fmpp @@ -38,8 +38,10 @@ data: { "org.apache.beam.sdk.extensions.sql.impl.parser.SqlSetOptionBeam" "org.apache.beam.sdk.extensions.sql.impl.parser.SqlAlterCatalog" "org.apache.beam.sdk.extensions.sql.impl.parser.SqlAlterTable" + "org.apache.beam.sdk.extensions.sql.impl.parser.TableClause" "org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils" "org.apache.beam.sdk.schemas.Schema" + "java.util.LinkedHashMap" ] # List of keywords. diff --git a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl index cb8eec438728..fbe29c77b7e8 100644 --- a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl +++ b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl @@ -899,4 +899,60 @@ SqlSetOptionBeam SqlSetOptionBeam(Span s, String scope) : ) } +/** + * Parses the WITH clause of a table in FROM or of an INSERT target and folds it into the table's + * name; see TableClause. (see also build.gradle) + */ +SqlNode TableClauseOpt(SqlNode tableRef, TableClause.Role role) : +{ + final Map withOptions = new LinkedHashMap(); +} +{ + [ AddTableOptions(withOptions) ] + { + return TableClause.fold(tableRef, role, withOptions); + } +} + +/** Parses a name component written by TableClause.encode(). */ +TableClause TableClauseEof() : +{ + TableClause.Role role = TableClause.Role.READ; + final String name; + final Map withOptions = new LinkedHashMap(); +} +{ + [ { role = TableClause.Role.WRITE; } ] + name = Identifier() AddTableOptions(withOptions) { + return TableClause.create(name, role, withOptions); + } +} + +/** + * WITH '(' 'key' = value ( ',' 'key' = value )* ')' + */ +void AddTableOptions(Map withOptions) : +{ +} +{ + + AddTableOption(withOptions) + ( AddTableOption(withOptions) )* + +} + +void AddTableOption(Map withOptions) : +{ + final String key; + final SqlParserPos keyPos; + final SqlNode value; +} +{ + key = SimpleStringLiteral() { keyPos = getPos(); } + + value = Expression(ExprContext.ACCEPT_NON_QUERY) { + TableClause.addOption(withOptions, key, keyPos, value); + } +} + // End parserImpls.ftl diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/TableNameExtractionUtils.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/TableNameExtractionUtils.java index efd3d8dfc353..96688a27565a 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/TableNameExtractionUtils.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/TableNameExtractionUtils.java @@ -23,6 +23,7 @@ import java.util.Collections; import java.util.List; import org.apache.beam.sdk.extensions.sql.impl.TableName; +import org.apache.beam.sdk.extensions.sql.impl.parser.TableClause; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlAsOperator; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlCall; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlIdentifier; @@ -44,6 +45,8 @@ * ... FROM table1 JOIN (LEFT, INNER, OUTER etc) table2 JOIN table3 ... * ... FROM table1 UNION (INTERSECT etc) SELECT ... * + * + *

A table with a {@code WITH} clause is extracted by its name alone. */ public class TableNameExtractionUtils { @@ -94,6 +97,9 @@ private static List extractFromJoin(SqlJoin join) { } private static List extractFromIdentifier(SqlIdentifier identifier) { - return ImmutableList.of(TableName.create(identifier.names)); + int last = identifier.names.size() - 1; + @Nullable TableClause clause = TableClause.decode(identifier.names.get(last)); + SqlIdentifier name = clause == null ? identifier : identifier.setName(last, clause.baseName()); + return ImmutableList.of(TableName.create(name.names)); } } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamCalciteSchema.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamCalciteSchema.java index f7783e7c3eca..1205a47643e3 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamCalciteSchema.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/BeamCalciteSchema.java @@ -24,6 +24,8 @@ import java.util.HashMap; import java.util.Map; import java.util.Set; +import org.apache.beam.sdk.extensions.sql.impl.parser.TableClause; +import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable; import org.apache.beam.sdk.extensions.sql.meta.Table; import org.apache.beam.sdk.extensions.sql.meta.catalog.Catalog; import org.apache.beam.sdk.extensions.sql.meta.catalog.CatalogManager; @@ -116,19 +118,36 @@ public Set getTypeNames() { return Collections.emptySet(); } + /** + * Returns the named table. For a name that encodes a {@link TableClause}, returns its base table + * as derived by {@link BeamSqlTable#forRead} or {@link BeamSqlTable#forWrite}. + */ @Override public org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.schema.@Nullable Table getTable( String name) { - Table table = tableProvider.getTable(name); + @Nullable TableClause clause = TableClause.decode(name); + Table table = tableProvider.getTable(clause == null ? name : clause.baseName()); if (table == null) { return null; } + BeamSqlTable beamTable = tableProvider.buildBeamSqlTable(table); return new BeamCalciteTable( - tableProvider.buildBeamSqlTable(table), + clause == null ? beamTable : derive(beamTable, clause), connection.getPipelineOptionsMap(), connection.getPipelineOptions()); } + private static BeamSqlTable derive(BeamSqlTable table, TableClause clause) { + try { + return clause.role() == TableClause.Role.READ + ? table.forRead(clause.options()) + : table.forWrite(clause.options()); + } catch (RuntimeException e) { + throw new IllegalArgumentException( + String.format("Cannot resolve %s: %s", clause.encode(), e.getMessage()), e); + } + } + public Collection getTables() { return tableProvider.getTables().values(); } @@ -156,6 +175,9 @@ public Set getSubSchemaNames() { */ @Override public @Nullable Schema getSubSchema(String name) { + if (TableClause.decode(name) != null) { + return null; + } BeamCalciteSchema subSchema = subSchemas.get(name); if (subSchema != null) { diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CatalogSchema.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CatalogSchema.java index b8a2b0031135..3aac20cb2fa5 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CatalogSchema.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CatalogSchema.java @@ -28,6 +28,7 @@ import java.util.Map; import java.util.Set; import org.apache.beam.sdk.extensions.sql.impl.parser.SqlDdlNodes; +import org.apache.beam.sdk.extensions.sql.impl.parser.TableClause; import org.apache.beam.sdk.extensions.sql.meta.catalog.Catalog; import org.apache.beam.sdk.extensions.sql.meta.catalog.CatalogManager; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.linq4j.tree.Expression; @@ -178,7 +179,7 @@ public Set getTableNames() { @Override public @Nullable BeamCalciteSchema getSubSchema(@Nullable String name) { - if (name == null) { + if (name == null || TableClause.decode(name) != null) { return null; } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCreateExternalTable.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCreateExternalTable.java index 499a12ee9e04..74f639fa2a69 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCreateExternalTable.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCreateExternalTable.java @@ -140,6 +140,10 @@ public void unparse(SqlWriter writer, int leftPrec, int rightPrec) { @Override public void execute(CalcitePrepare.Context context) { + checkArgument( + TableClause.decode(name(name)) == null, + "Table name '%s' is reserved for WITH options.", + name(name)); final Pair pair = SqlDdlNodes.schema(context, true, name); if (pair.left.plus().getTable(pair.right) != null) { // Table exists. diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClause.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClause.java new file mode 100644 index 000000000000..0c64b9360ec6 --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClause.java @@ -0,0 +1,372 @@ +/* + * 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.impl.parser; + +import static java.util.stream.Collectors.joining; + +import com.google.auto.value.AutoValue; +import java.io.StringReader; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.extensions.sql.impl.parser.impl.BeamSqlParserImpl; +import org.apache.beam.sdk.extensions.sql.impl.parser.impl.TokenMgrError; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.avatica.util.Casing; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.runtime.CalciteContextException; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.runtime.CalciteException; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlCharStringLiteral; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlIdentifier; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlKind; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlLiteral; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlNode; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlNodeList; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlNumericLiteral; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlWith; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlWithItem; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlAbstractParserImpl; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.util.SqlShuttle; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.validate.SqlConformanceEnum; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.util.Util; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * The {@code WITH} clause of one table occurrence in a statement. + * + *

Calcite resolves a table occurrence by its name alone, so the parser folds the clause into the + * last component of the table's name as canonical SQL text. {@link #decode} reads it back and + * breaks it apart. {@code FROM cat.db.orders WITH ('k' = 1)} becomes the name {@code cat}, {@code + * db}, {@code `orders` WITH ('k' = 1)} aliased as {@code orders}, and the INSERT target {@code INTO + * orders WITH ('k' = 1)} becomes the name {@code INTO `orders` WITH ('k' = 1)}. + * + *

This is not a {@code SqlCall} like the DDL nodes in this package: Beam executes those itself, + * while a table occurrence, including an INSERT target, goes through Calcite's validator and + * SqlToRelConverter, which accept only their own node kinds and look tables up by name. + */ +@AutoValue +public abstract class TableClause { + /** Whether a table occurrence is read (in FROM) or written (as an INSERT target). */ + public enum Role { + READ, + WRITE + } + + /** The last component of the original table name. */ + public abstract String baseName(); + + public abstract Role role(); + + /** + * The {@code WITH} options in clause order. Values are String, Long, Double, Boolean, List or + * Map. + */ + public abstract Map options(); + + public static TableClause create(String baseName, Role role, Map options) { + return new AutoValue_TableClause(baseName, role, ImmutableMap.copyOf(options)); + } + + /** + * Returns the canonical SQL text of this table reference (e.g. "`orders` WITH ('k' = TRUE)" for a + * read). + */ + public String encode() { + String table = "`" + baseName().replace("`", "``") + "`"; + return (role() == Role.WRITE ? "INTO " + table : table) + + options().entrySet().stream() + .map(option -> quote(option.getKey()) + " = " + toSql(option.getValue())) + .collect(joining(", ", " WITH (", ")")); + } + + /** + * Returns the table reference encoded in a name component, or null if the component is a plain + * name. + */ + public static @Nullable TableClause decode(String name) { + if (!name.startsWith("`") && !name.startsWith("INTO `")) { + return null; + } + BeamSqlParserImpl parser = new BeamSqlParserImpl(new StringReader(name)); + parser.switchTo(SqlAbstractParserImpl.LexicalState.BTID); + parser.setQuotedCasing(Casing.UNCHANGED); + parser.setUnquotedCasing(Casing.UNCHANGED); + parser.setIdentifierMaxLength(Integer.MAX_VALUE); + parser.setConformance(SqlConformanceEnum.DEFAULT); + try { + TableClause clause = parser.TableClauseEof(); + return clause.encode().equals(name) ? clause : null; + } catch (Exception | TokenMgrError e) { + return null; + } + } + + /** + * Folds the {@code WITH} clause of a table in FROM or of an INSERT target into its name. A read + * is aliased as the base name so references like {@code orders.id} keep working. Called by the + * parser. + */ + public static SqlNode fold(SqlNode tableRef, Role role, Map options) { + checkNotEncoded(tableRef); + if (options.isEmpty()) { + return tableRef; + } + SqlIdentifier name = plainTableName(tableRef); + int last = name.names.size() - 1; + String baseName = name.names.get(last); + SqlIdentifier folded = name.setName(last, encode(create(baseName, role, options), name)); + return role == Role.WRITE + ? folded + : SqlStdOperatorTable.AS.createCall( + name.getParserPosition(), + folded, + new SqlIdentifier(baseName, name.getComponentParserPosition(last))); + } + + /** + * Replaces the implicit alias of a folded read with the alias the statement gives it, so that + * {@code AS(AS(folded, base), alias, columns...)} becomes {@code AS(folded, alias, columns...)}, + * also under TABLESAMPLE. Called by the parser at the end of a table reference. + */ + public static SqlNode collapseImplicitAlias(SqlNode tableRef) { + if (tableRef.getKind() == SqlKind.TABLESAMPLE) { + SqlCall sample = (SqlCall) tableRef; + sample.setOperand(0, collapseImplicitAlias(sample.operand(0))); + return sample; + } + if (tableRef.getKind() != SqlKind.AS || !isFoldedRead(((SqlCall) tableRef).operand(0))) { + return tableRef; + } + SqlCall alias = (SqlCall) tableRef; + List operands = new ArrayList<>(alias.getOperandList()); + operands.set(0, ((SqlCall) alias.operand(0)).operand(0)); + return SqlStdOperatorTable.AS.createCall(alias.getParserPosition(), operands); + } + + /** + * Rejects {@code WITH} options on references to the common table expressions of a WITH list, + * which Calcite would resolve to catalog tables instead. Called by the parser for every query + * with a WITH list. + */ + public static void checkCteReferences(@Nullable SqlNodeList withList, SqlNode body) { + if (withList == null) { + return; + } + Set inScope = new HashSet<>(); + for (SqlNode node : withList) { + SqlWithItem item = (SqlWithItem) node; + if (item.recursive.booleanValue()) { + inScope.add(item.name.getSimple()); + } + checkNoClauseOn(inScope, item.query); + inScope.add(item.name.getSimple()); + } + checkNoClauseOn(inScope, body); + } + + /** + * Also rejects them in the ORDER BY of a WITH query, which the parser attaches after {@link + * #checkCteReferences} and Calcite validates in the WITH's scope. Called by the parser. + */ + public static void checkCteReferencesInOrderBy(SqlNode query, @Nullable SqlNodeList orderBy) { + if (query instanceof SqlWith && orderBy != null) { + checkCteReferences(((SqlWith) query).withList, orderBy); + } + } + + /** + * Adds a {@code WITH} option, converting its value to a plain Java value. Called by the parser. + */ + public static void addOption( + Map options, String key, SqlParserPos keyPos, SqlNode value) { + if (options.containsKey(key)) { + throw error(keyPos, String.format("Duplicate option '%s'", key)); + } + options.put(key, toJava(key, value)); + } + + /** Rejects a table name that is written in the encoded form of a {@link TableClause}. */ + private static void checkNotEncoded(SqlNode tableRef) { + SqlNode name = tableRef; + while (name instanceof SqlCall) { // table hints, EXTEND, OVER, FOR SYSTEM_TIME AS OF + name = ((SqlCall) name).operand(0); + } + if (name instanceof SqlIdentifier) { + String last = Util.last(((SqlIdentifier) name).names); + if (decode(last) != null) { + throw error( + name.getParserPosition(), + String.format("Table name '%s' is reserved for WITH options", last)); + } + } + } + + private static String encode(TableClause clause, SqlIdentifier name) { + String encoded = clause.encode(); + if (!clause.equals(decode(encoded))) { + throw error(name.getParserPosition(), "Unsupported WITH options: " + encoded); + } + return encoded; + } + + private static void checkNoClauseOn(Set ctes, SqlNode query) { + query.accept( + new SqlShuttle() { + @Override + public SqlNode visit(SqlIdentifier id) { + TableClause clause = id.isSimple() ? decode(id.getSimple()) : null; + if (clause != null && ctes.contains(clause.baseName())) { + throw error( + id.getParserPosition(), + String.format( + "WITH options cannot be applied to common table expression '%s'", + clause.baseName())); + } + return id; + } + }); + } + + private static SqlIdentifier plainTableName(SqlNode tableRef) { + if (tableRef instanceof SqlIdentifier) { + return (SqlIdentifier) tableRef; + } + throw error( + tableRef.getParserPosition(), + "WITH cannot be combined with table hints, EXTEND, FOR SYSTEM_TIME AS OF or OVER"); + } + + private static boolean isFoldedRead(SqlNode node) { + if (node.getKind() != SqlKind.AS) { + return false; + } + SqlNode table = ((SqlCall) node).operand(0); + if (!(table instanceof SqlIdentifier)) { + return false; + } + TableClause clause = decode(Util.last(((SqlIdentifier) table).names)); + return clause != null && clause.role() == Role.READ; + } + + private static Object toJava(String key, SqlNode node) { + SqlNode value = node.getKind() == SqlKind.LITERAL_CHAIN ? SqlLiteral.unchain(node) : node; + if (value instanceof SqlCharStringLiteral) { + return ((SqlCharStringLiteral) value).getValueAs(String.class); + } + if (value instanceof SqlNumericLiteral) { + return number(key, (SqlNumericLiteral) value); + } + if (value instanceof SqlLiteral) { + Object literal = ((SqlLiteral) value).getValue(); + if (literal instanceof Boolean) { + return literal; + } + } + if (value.getKind() == SqlKind.ARRAY_VALUE_CONSTRUCTOR) { + SqlCall elements = (SqlCall) value; + ImmutableList.Builder list = ImmutableList.builder(); + for (int i = 0; i < elements.operandCount(); i++) { + list.add(toJava(key, elements.operand(i))); + } + return list.build(); + } + if (value.getKind() == SqlKind.MAP_VALUE_CONSTRUCTOR + && ((SqlCall) value).operandCount() % 2 == 0) { + SqlCall entries = (SqlCall) value; + Map map = new LinkedHashMap<>(); + for (int i = 0; i < entries.operandCount(); i += 2) { + Object mapKey = toJava(key, entries.operand(i)); + if (map.put(mapKey, toJava(key, entries.operand(i + 1))) != null) { + throw error( + entries.operand(i).getParserPosition(), + String.format("Duplicate MAP key %s in option '%s'", toSql(mapKey), key)); + } + } + return ImmutableMap.copyOf(map); + } + throw error( + node.getParserPosition(), + String.format( + "Option '%s' must be a string, numeric or boolean literal," + + " or an ARRAY[...] or MAP[...] of them", + key)); + } + + private static Object number(String key, SqlNumericLiteral literal) { + BigDecimal value = literal.getValueAs(BigDecimal.class); + if (literal.isExact() && literal.isInteger()) { + try { + return value.longValueExact(); + } catch (ArithmeticException e) { + throw outOfRange(key, literal); + } + } + double number = value.doubleValue(); + if (Double.isInfinite(number) || (number == 0 && value.signum() != 0)) { + throw outOfRange(key, literal); + } + return number; + } + + private static CalciteContextException outOfRange(String key, SqlNumericLiteral literal) { + return error( + literal.getParserPosition(), + String.format("Option '%s' value %s is out of range", key, literal)); + } + + private static String toSql(@Nullable Object value) { + if (value instanceof String) { + return quote((String) value); + } + if (value instanceof List) { + return ((List) value) + .stream().map(TableClause::toSql).collect(joining(", ", "ARRAY[", "]")); + } + if (value instanceof Map) { + return ((Map) value) + .entrySet().stream() + .map(entry -> toSql(entry.getKey()) + ", " + toSql(entry.getValue())) + .collect(joining(", ", "MAP[", "]")); + } + if (value instanceof Boolean) { + return (Boolean) value ? "TRUE" : "FALSE"; + } + return String.valueOf(value); + } + + private static String quote(String value) { + return "'" + value.replace("'", "''") + "'"; + } + + private static CalciteContextException error(SqlParserPos pos, String message) { + return new CalciteContextException( + message, + new CalciteException(message, null), + pos.getLineNum(), + pos.getColumnNum(), + pos.getEndLineNum(), + pos.getEndColumnNum()); + } +} diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/BeamSqlTable.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/BeamSqlTable.java index e658db8c3d82..4dfa0888c81f 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/BeamSqlTable.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/BeamSqlTable.java @@ -18,6 +18,7 @@ package org.apache.beam.sdk.extensions.sql.meta; import java.util.List; +import java.util.Map; import org.apache.beam.sdk.extensions.sql.impl.BeamTableStatistics; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.schemas.Schema; @@ -55,4 +56,28 @@ public interface BeamSqlTable { * estimate the row count or rate it will return BeamTableStatistics.BOUNDED_UNKNOWN. */ BeamTableStatistics getTableStatistics(PipelineOptions options); + + /** + * Returns a new table for one reference in FROM with a {@code WITH} clause. The new table may + * differ in schema, boundedness and push-down support. + * + *

Must not mutate this table or its {@link Table}, and must be cheap with no catalog calls, + * because planning calls it several times per reference. + * + * @throws IllegalArgumentException for an unknown option key, naming it and the accepted keys + * @throws UnsupportedOperationException if this table does not support options + */ + default BeamSqlTable forRead(Map options) { + throw new UnsupportedOperationException( + String.format("%s does not support WITH options for reads", getClass().getSimpleName())); + } + + /** + * Returns a new table for an INSERT target with a {@code WITH} clause. The contract of {@link + * #forRead} applies. + */ + default BeamSqlTable forWrite(Map options) { + throw new UnsupportedOperationException( + String.format("%s does not support WITH options for writes", getClass().getSimpleName())); + } } diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/TableClauseResolutionTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/TableClauseResolutionTest.java new file mode 100644 index 000000000000..7f72b7e610bf --- /dev/null +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/TableClauseResolutionTest.java @@ -0,0 +1,324 @@ +/* + * 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.impl; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.stream.Collectors; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils; +import org.apache.beam.sdk.extensions.sql.meta.BeamSqlTable; +import org.apache.beam.sdk.extensions.sql.meta.SchemaBaseBeamTable; +import org.apache.beam.sdk.extensions.sql.meta.Table; +import org.apache.beam.sdk.extensions.sql.meta.catalog.InMemoryCatalogManager; +import org.apache.beam.sdk.extensions.sql.meta.provider.InMemoryMetaTableProvider; +import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PDone; +import org.apache.beam.sdk.values.POutput; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; + +/** + * Resolves table references with {@code WITH} options through {@link BeamSqlEnv} and through a JDBC + * connection. + */ +@RunWith(Parameterized.class) +public class TableClauseResolutionTest { + private static final List> QUERY_ROWS = new CopyOnWriteArrayList<>(); + private static final List>> WRITES = new CopyOnWriteArrayList<>(); + + @Parameters(name = "{0}") + public static List sessions() { + return Arrays.asList("BeamSqlEnv", "JDBC"); + } + + @Parameter public String sessionType; + + private Session session; + + @Before + public void setUp() throws Exception { + QUERY_ROWS.clear(); + WRITES.clear(); + InMemoryCatalogManager catalogManager = new InMemoryCatalogManager(); + catalogManager.registerTableProvider(new ClauseTableProvider()); + catalogManager.registerTableProvider(new TestTableProvider()); + session = + sessionType.equals("JDBC") + ? new JdbcSession(JdbcDriver.connect(catalogManager, PipelineOptionsFactory.create())) + : new EnvSession( + BeamSqlEnv.builder(catalogManager) + .setPipelineOptions(PipelineOptionsFactory.create()) + .build()); + session.execute("CREATE EXTERNAL TABLE orders (id BIGINT NOT NULL, name VARCHAR) TYPE clause"); + session.execute("CREATE EXTERNAL TABLE plain (id BIGINT, name VARCHAR) TYPE test"); + } + + @Test + public void eachReferenceReadsWithItsOwnOptions() throws Exception { + assertEquals( + Collections.singletonList( + Arrays.asList( + read(ImmutableMap.of("tag", "a")), read(ImmutableMap.of("tag", "b")), "one")), + session.query( + "SELECT x.clause, y.clause, z.name" + + " FROM orders WITH ('tag' = 'a') AS x" + + " JOIN orders WITH ('tag' = 'b') y ON x.id = y.id" + + " JOIN orders AS z ON y.id = z.id")); + + session.execute("CREATE CATALOG cat TYPE 'local'"); + session.execute("USE CATALOG cat"); + session.execute("CREATE DATABASE db"); + session.execute("USE DATABASE db"); + session.execute("CREATE EXTERNAL TABLE orders (id BIGINT NOT NULL, name VARCHAR) TYPE clause"); + assertEquals( + Collections.singletonList( + Arrays.asList( + read(ImmutableMap.of("tag", "cat.db")), + read(ImmutableMap.of("tag", "db")), + read(ImmutableMap.of("tag", "none")))), + session.query( + "SELECT a.clause, b.clause, orders.clause" + + " FROM cat.db.orders WITH ('tag' = 'cat.db') a" + + " JOIN db.orders WITH ('tag' = 'db') b ON a.id = b.id" + + " JOIN orders WITH ('tag' = 'none') ON a.id = orders.id")); + } + + @Test + public void insertTargetWritesThroughForWrite() throws Exception { + session.execute("INSERT INTO orders WITH ('tag' = 'w') SELECT id, name, 'note' FROM orders"); + + assertEquals( + Collections.singletonList( + KV.of(write(ImmutableMap.of("tag", "w")), Arrays.asList(1L, "one", "note"))), + WRITES); + assertError( + "INSERT INTO orders SELECT id, name, 'note' FROM orders", + "Number of INSERT target columns (2) does not equal number of source items (3)"); + } + + @Test + public void unsupportedOptionsFailWithTheClauseText() { + assertError( + "SELECT * FROM plain WITH ('tag' = 'a')", + "Cannot resolve `plain` WITH ('tag' = 'a'):" + + " InMemoryTable does not support WITH options for reads"); + assertError( + "INSERT INTO plain WITH ('tag' = 'a') SELECT * FROM plain", + "Cannot resolve INTO `plain` WITH ('tag' = 'a'):" + + " InMemoryTable does not support WITH options for writes"); + } + + private void assertError(String sql, String message) { + Exception e = assertThrows(sql, Exception.class, () -> session.execute(sql)); + assertThat( + sql, + Throwables.getCausalChain(e).stream() + .map(Throwable::getMessage) + .collect(Collectors.joining("\n")), + containsString(message)); + } + + private static String read(Map options) { + return "read " + options; + } + + private static String write(Map options) { + return "write " + options; + } + + /** Runs SQL through one of the two execution paths. */ + private interface Session { + void execute(String sql) throws Exception; + + List> query(String sql) throws Exception; + } + + private static class EnvSession implements Session { + private final BeamSqlEnv env; + + EnvSession(BeamSqlEnv env) { + this.env = env; + } + + @Override + public void execute(String sql) throws ParseException { + if (env.isDdl(sql)) { + env.executeDdl(sql); + } else { + run(sql); + } + } + + @Override + public List> query(String sql) throws ParseException { + QUERY_ROWS.clear(); + run(sql); + return new ArrayList<>(QUERY_ROWS); + } + + /** Runs a query or an INSERT, collecting its output rows. */ + private void run(String sql) throws ParseException { + Pipeline pipeline = Pipeline.create(); + BeamSqlRelUtils.toPCollection(pipeline, env.parseQuery(sql)) + .apply(ParDo.of(new CollectRows())); + pipeline.run().waitUntilFinish(); + } + } + + private static class JdbcSession implements Session { + private final Connection connection; + + JdbcSession(Connection connection) { + this.connection = connection; + } + + @Override + public void execute(String sql) throws Exception { + try (Statement statement = connection.createStatement()) { + statement.execute(sql); + } + } + + @Override + public List> query(String sql) throws Exception { + try (Statement statement = connection.createStatement(); + ResultSet result = statement.executeQuery(sql)) { + List> rows = new ArrayList<>(); + while (result.next()) { + List row = new ArrayList<>(); + for (int i = 1; i <= result.getMetaData().getColumnCount(); i++) { + row.add(result.getObject(i)); + } + rows.add(row); + } + return rows; + } + } + } + + private static class CollectRows extends DoFn { + @ProcessElement + public void processElement(@Element Row row) { + QUERY_ROWS.add(row.getValues()); + } + } + + /** Provides tables of type {@code clause}, which support WITH options. */ + private static class ClauseTableProvider extends InMemoryMetaTableProvider { + @Override + public String getTableType() { + return "clause"; + } + + @Override + public BeamSqlTable buildBeamSqlTable(Table table) { + return new ClauseTable(table.getSchema(), null); + } + } + + /** + * A table with the single row {@code (1, 'one')}. WITH options add a {@code clause} column, which + * reads fill with a description of them and the hook that took them, and writes record next to + * each row. + */ + private static class ClauseTable extends SchemaBaseBeamTable { + private final @Nullable String clause; + + ClauseTable(Schema schema, @Nullable String clause) { + super(schema); + this.clause = clause; + } + + @Override + public BeamSqlTable forRead(Map options) { + return withClause(read(options)); + } + + @Override + public BeamSqlTable forWrite(Map options) { + return withClause(write(options)); + } + + private ClauseTable withClause(String clause) { + return new ClauseTable( + Schema.builder().addFields(schema.getFields()).addStringField("clause").build(), clause); + } + + @Override + public PCollection buildIOReader(PBegin begin) { + Row.Builder row = Row.withSchema(schema).addValues(1L, "one"); + if (clause != null) { + row.addValue(clause); + } + return begin.apply(Create.of(row.build()).withRowSchema(schema)); + } + + @Override + public POutput buildIOWriter(PCollection input) { + input.apply(ParDo.of(new RecordWrites(clause))); + return PDone.in(input.getPipeline()); + } + + @Override + public PCollection.IsBounded isBounded() { + return PCollection.IsBounded.BOUNDED; + } + } + + private static class RecordWrites extends DoFn { + private final @Nullable String clause; + + RecordWrites(@Nullable String clause) { + this.clause = clause; + } + + @ProcessElement + public void processElement(@Element Row row) { + WRITES.add(KV.of(clause, row.getValues())); + } + } +} diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/BeamDDLTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/BeamDDLTest.java index e465ce44d056..af23f5e79b92 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/BeamDDLTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/BeamDDLTest.java @@ -21,6 +21,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -184,6 +185,13 @@ public void testParseCreateExternalTable_minimal() throws Exception { .properties(TableUtils.emptyProperties()) .build(), tableProvider.getTables().get("person")); + IllegalArgumentException reservedName = + assertThrows( + IllegalArgumentException.class, + () -> + env.executeDdl("CREATE EXTERNAL TABLE ```t`` WITH ('k' = 1)` (id INT) TYPE test")); + assertEquals( + "Table name '`t` WITH ('k' = 1)' is reserved for WITH options.", reservedName.getMessage()); } @Test diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClauseTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClauseTest.java new file mode 100644 index 000000000000..131d67f4dd1a --- /dev/null +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/TableClauseTest.java @@ -0,0 +1,310 @@ +/* + * 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.impl.parser; + +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; + +import java.util.Arrays; +import java.util.Collections; +import java.util.Map; +import org.apache.beam.sdk.extensions.sql.TableNameExtractionUtils; +import org.apache.beam.sdk.extensions.sql.impl.TableName; +import org.apache.beam.sdk.extensions.sql.impl.parser.TableClause.Role; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.config.Lex; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlDialect; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlIdentifier; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlSelect; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParseException; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParser; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.util.Util; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests {@code WITH} clauses of table references, which the parser folds into table names as the + * canonical text of a {@link TableClause}. Statements are rendered without identifier quoting, so a + * folded name component shows up as text such as {@code cat.db.`orders` WITH ('k' = 1)}. + */ +@RunWith(JUnit4.class) +public class TableClauseTest { + private static final SqlParser.Config CONFIG = + SqlParser.config().withLex(Lex.JAVA).withParserFactory(BeamSqlParser.FACTORY); + private static final SqlDialect UNQUOTED = new SqlDialect(SqlDialect.EMPTY_CONTEXT); + private static final String COMBINED = + "WITH cannot be combined with table hints, EXTEND, FOR SYSTEM_TIME AS OF or OVER"; + private static final String ON_CTE = "WITH options cannot be applied to common table expression"; + private static final String NOT_A_LITERAL = + "Option 'k' must be a string, numeric or boolean literal," + + " or an ARRAY[...] or MAP[...] of them"; + + /** Statements that use the WITH keyword otherwise, and how they render. */ + private static final String[][] UNCHANGED_STATEMENTS = { + {"WITH t AS (SELECT 1 AS a) SELECT * FROM t", "WITH t AS (SELECT 1 AS a) SELECT * FROM t"}, + { + "INSERT INTO t WITH x AS (SELECT 1 AS a) SELECT * FROM x", + "INSERT INTO t WITH x AS (SELECT 1 AS a) SELECT * FROM x" + }, + { + "INSERT INTO t (a) WITH x AS (SELECT 1 AS a) SELECT * FROM x", + "INSERT INTO t (a) WITH x AS (SELECT 1 AS a) SELECT * FROM x" + }, + { + "SELECT * FROM UNNEST(ARRAY[1, 2]) WITH ORDINALITY AS t (x, o)", + "SELECT * FROM UNNEST(ARRAY[1, 2]) WITH ORDINALITY AS t (x, o)" + }, + {"SELECT * FROM ```t```", "SELECT * FROM `t`"}, + }; + + /** Statements with WITH options and how they render once folded. */ + private static final String[][] FOLDED_STATEMENTS = { + { + "SELECT x.a FROM t WITH ('k' = 1) AS x(a, b)", + "SELECT x.a FROM `t` WITH ('k' = 1) AS x (a, b)" + }, + { + "SELECT * FROM t WITH ('k' = 1) TABLESAMPLE BERNOULLI(10)", + "SELECT * FROM `t` WITH ('k' = 1) AS t TABLESAMPLE BERNOULLI(10.0)" + }, + { + "SELECT * FROM t WITH ('k' = 1) MATCH_RECOGNIZE (PATTERN (A) DEFINE A AS true)", + "SELECT * FROM `t` WITH ('k' = 1) AS t MATCH_RECOGNIZE( PATTERN (A) DEFINE A AS TRUE)" + }, + { + "SELECT * FROM a WITH ('k' = 1) x JOIN cat.db.b WITH ('k' = 2) ON x.id = b.id", + "SELECT * FROM `a` WITH ('k' = 1) AS x" + + " INNER JOIN cat.db.`b` WITH ('k' = 2) AS b ON x.id = b.id" + }, + { + "WITH x AS (SELECT * FROM a WITH ('k' = 1)) SELECT * FROM x, b WITH ('k' = 2)", + "WITH x AS (SELECT * FROM `a` WITH ('k' = 1) AS a)" + + " SELECT * FROM x, `b` WITH ('k' = 2) AS b" + }, + { + "WITH t AS (SELECT * FROM t WITH ('k' = 1) AS c) SELECT * FROM t", + "WITH t AS (SELECT * FROM `t` WITH ('k' = 1) AS c) SELECT * FROM t" + }, + { + "WITH t AS (SELECT 1 AS a) SELECT * FROM db.t WITH ('k' = 1) AS t", + "WITH t AS (SELECT 1 AS a) SELECT * FROM db.`t` WITH ('k' = 1) AS t" + }, + { + "SELECT * FROM (WITH t AS (SELECT 1 AS a) SELECT * FROM t), t WITH ('k' = 1) AS c", + "SELECT * FROM (WITH t AS (SELECT 1 AS a) SELECT * FROM t), `t` WITH ('k' = 1) AS c" + }, + { + "INSERT INTO cat.db.mirror WITH ('mode' = 'merge-on-read')" + + " SELECT * FROM cat.db.orders WITH ('k' = ARRAY['a'])", + "INSERT INTO cat.db.INTO `mirror` WITH ('mode' = 'merge-on-read')" + + " SELECT * FROM cat.db.`orders` WITH ('k' = ARRAY['a']) AS orders" + }, + { + "INSERT INTO t WITH ('k' = 1) (a, b) SELECT 1, 2", + "INSERT INTO INTO `t` WITH ('k' = 1) (a, b) SELECT 1, 2" + }, + { + "INSERT INTO t WITH ('k' = 1) WITH x AS (SELECT 1 AS a) SELECT * FROM x", + "INSERT INTO INTO `t` WITH ('k' = 1) WITH x AS (SELECT 1 AS a) SELECT * FROM x" + }, + }; + + /** + * Malformed or misplaced WITH clauses, and part of the parse error each one raises. A ^ marks the + * column the error is reported at. + */ + private static final String[][] REJECTED_STATEMENTS = { + {"SELECT * FROM t WITH ('k' = 1, ^'k' = 2)", "Duplicate option 'k'"}, + { + "SELECT * FROM t WITH ('k' = 9223372036854775808)", + "Option 'k' value 9223372036854775808 is out of range" + }, + {"SELECT * FROM t WITH ('k' = 1e400)", "is out of range"}, + {"SELECT * FROM t WITH ('k' = -1e-400)", "is out of range"}, + {"SELECT * FROM t WITH ('m' = MAP['a', 1, 'a', 2])", "Duplicate MAP key 'a'"}, + {"SELECT * FROM t WITH (k = 1)", "Encountered \"k\""}, + {"SELECT * FROM t /*+ OPTIONS('a'='b') */ WITH ('k' = 1)", COMBINED}, + {"INSERT INTO t EXTEND (c INT) WITH ('k' = 1) SELECT 1, 2", COMBINED}, + {"SELECT * FROM t AS c ^WITH ('k' = 1)", "Encountered \"WITH\""}, + {"SELECT * FROM t WITH ('k' = 1) WITH ('j' = 2)", "'WITH'"}, + {"SELECT * FROM (SELECT * FROM t) WITH ('k' = 1)", "'WITH'"}, + {"INSERT INTO t (a) WITH ('k' = 1) SELECT 1", "'WITH'"}, + { + "SELECT * FROM ^```t`` WITH ('k' = 1)`", + "Table name '`t` WITH ('k' = 1)' is reserved for WITH options" + }, + { + "SELECT * FROM db.```t`` WITH ('k' = 1)` /*+ OPTIONS('a'='b') */ EXTEND (c INT)", + "is reserved for WITH options" + }, + {"INSERT INTO ```t`` WITH ('k' = 1)` SELECT 1", "is reserved for WITH options"}, + {"WITH t AS (SELECT 1 AS a) SELECT * FROM ^t WITH ('k' = 1)", ON_CTE + " 't'"}, + { + "WITH s AS (SELECT 1 AS a), t AS (SELECT * FROM s WITH ('k' = 1)) SELECT * FROM t", + ON_CTE + " 's'" + }, + {"WITH RECURSIVE t AS (SELECT * FROM t WITH ('k' = 1)) SELECT * FROM t", ON_CTE + " 't'"}, + { + "WITH t AS (SELECT 1 AS a) SELECT a FROM t ORDER BY (SELECT a FROM ^t WITH ('k' = 1))", + ON_CTE + " 't'" + }, + }; + + @Test + public void statementsWithoutOptionsParseAsBefore() throws Exception { + for (String[] statement : UNCHANGED_STATEMENTS) { + assertEquals(statement[0], statement[1], render(statement[0])); + } + } + + @Test + public void optionsFoldIntoTheLastNameComponent() throws Exception { + for (String[] statement : FOLDED_STATEMENTS) { + assertEquals(statement[0], statement[1], render(statement[0])); + } + } + + @Test + public void malformedOrMisplacedOptionsAreRejected() { + for (String[] statement : REJECTED_STATEMENTS) { + assertParseError(statement[0], statement[1]); + } + for (String value : + new String[] { + "NULL", "1 + 2", "x", "TIMESTAMP '2020-01-01 00:00:00'", "ARRAY[1, NULL]", "MAP['a']" + }) { + assertParseError("SELECT * FROM t WITH ('k' = " + value + ")", NOT_A_LITERAL); + } + } + + @Test + public void optionValuesBecomeCanonicalJavaValuesInClauseOrder() throws Exception { + SqlCall from = + (SqlCall) + ((SqlSelect) + SqlParser.create( + "SELECT * FROM t WITH ('s' = 'it''s', 'n' = -7, 'd' = 1.50," + + " 'e' = 1e3, 'b' = true, 'c' = 'con' 'cat', 'a' = ARRAY[1, 'x']," + + " 'm' = MAP['k', ARRAY[2.5]], 'h' = 9223372036854775807)", + CONFIG) + .parseQuery()) + .getFrom(); + SqlIdentifier name = from.operand(0); + TableClause clause = checkStateNotNull(TableClause.decode(Util.last(name.names))); + + assertEquals( + ImmutableMap.builder() + .put("s", "it's") + .put("n", -7L) + .put("d", 1.5) + .put("e", 1000.0) + .put("b", true) + .put("c", "concat") + .put("a", ImmutableList.of(1L, "x")) + .put("m", ImmutableMap.of("k", ImmutableList.of(2.5))) + .put("h", Long.MAX_VALUE) + .build(), + clause.options()); + assertEquals( + "`t` WITH ('s' = 'it''s', 'n' = -7, 'd' = 1.5, 'e' = 1000.0, 'b' = TRUE," + + " 'c' = 'concat', 'a' = ARRAY[1, 'x'], 'm' = MAP['k', ARRAY[2.5]]," + + " 'h' = 9223372036854775807)", + clause.encode()); + } + + @Test + public void decodeReadsBackExactlyWhatEncodeWrites() { + Map options = + ImmutableMap.builder() + .put("string", "it's `quoted`\n WITH ('x' = 1) 日本") + .put("empty", "") + .put("min", Long.MIN_VALUE) + .put("max", Long.MAX_VALUE) + .put("double", -0.001) + .put("large", 1.0E20) + .put("small", 1.0E-5) + .put("zero", 0.0) + .put("boolean", false) + .put("list", ImmutableList.of(1L, "x", true, ImmutableList.of())) + .put("map", ImmutableMap.of("k", ImmutableList.of(2.5), 3L, ImmutableMap.of())) + .build(); + TableClause[] clauses = { + TableClause.create("orders", Role.READ, options), + TableClause.create("my `odd` 'name' 日本", Role.READ, ImmutableMap.of("k", 1L)), + TableClause.create("orders", Role.WRITE, options), + }; + for (TableClause clause : clauses) { + assertEquals(clause.encode(), clause, TableClause.decode(clause.encode())); + } + + String[] plainNames = { + "orders", + "`orders`", + "`orders", + "`orders` WITH ('a' = 1", + "`orders` WITH ('a' = 1)", + "`orders` WITH ('a' = NULL)", + "`orders` WITH ('a' = 1e3)", + "INTO db.`orders` WITH ('a' = 1)", + }; + for (String name : plainNames) { + assertNull(name, TableClause.decode(name)); + } + } + + @Test + public void tableNamesAreExtractedWithoutOptions() throws Exception { + assertEquals( + Arrays.asList( + TableName.create(Arrays.asList("cat", "db", "a")), + TableName.create(Collections.singletonList("b")), + TableName.create(Collections.singletonList("c"))), + TableNameExtractionUtils.extractTableNamesFromNode( + SqlParser.create( + "SELECT * FROM cat.db.a WITH ('k' = 1) x" + + " JOIN b WITH ('k' = 1) ON x.id = b.id, c", + CONFIG) + .parseQuery())); + } + + private static String render(String sql) throws SqlParseException { + return SqlParser.create(sql, CONFIG) + .parseStmt() + .toSqlString(UNQUOTED) + .getSql() + .replace('\n', ' '); + } + + /** Asserts that parsing fails with the message, at the column of the ^ in the SQL if any. */ + private static void assertParseError(String sqlWithCaret, String message) { + String sql = sqlWithCaret.replace("^", ""); + SqlParseException e = + assertThrows(sql, SqlParseException.class, () -> SqlParser.create(sql, CONFIG).parseStmt()); + assertThat(sql, e.getMessage(), containsString(message)); + if (sqlWithCaret.contains("^")) { + assertEquals(sql, sqlWithCaret.indexOf('^') + 1, e.getPos().getColumnNum()); + } + } +} From 0fef1730ce1cf020ba6cc5a3baaab82b564c8ff7 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 27 Oct 2026 11:33:09 -0700 Subject: [PATCH 2/2] CHANGES --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index 04ac961fcd93..72e48bc73b56 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -86,6 +86,7 @@ * (Python, Java, Go) Added `--dataflow_container_base_repository` (Python, Go) and `--dataflowContainerBaseRepository` (Java) to set the repository (e.g. `docker.example.com/my-prefix/cloud-dataflow/v1beta3`) from which the Dataflow runner pulls its default SDK container images, including the default images of cross-language Apache Beam environments. Defaults to `gcr.io/cloud-dataflow/v1beta3` ([#40569](https://github.com/apache/beam/issues/40569)). * (Go) GCS and Dataflow clients now support non-default universes such as Trusted Partner Cloud (TPC), using the universe domain from `GOOGLE_CLOUD_UNIVERSE_DOMAIN` or the default credentials. Added a `--gcs_endpoint` pipeline option to override the GCS endpoint ([#40569](https://github.com/apache/beam/issues/40569)). * (Java) `GcsUtil` now uses the java-storage client (GcsUtil V2) for GCS operations by default instead of gcsio (GcsUtil V1). To go back to V1, pass `--experiments=disable_gcsutil_v2`. Pipelines with the `use_grpc_for_gcs` experiment keep using V1. The `use_gcsutil_v2` experiment is no longer needed and is now ignored ([#40712](https://github.com/apache/beam/issues/40712)). +* [Beam SQL] Added a `WITH (...)` clause to set per-statement connector options in `FROM` and `INSERT INTO` ([#40848](https://github.com/apache/beam/issues/40848)). ## Breaking Changes