From 3870bf73a69842d88e8c7bb9e63dbfbbb21cae77 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Mon, 26 Oct 2026 14:21:53 -0700 Subject: [PATCH 1/2] create and add/drop primary keys --- .../iceberg/IcebergAlterTableOps.java | 5 + .../provider/iceberg/IcebergMetastore.java | 30 +++- .../iceberg/BeamSqlCliIcebergAlterTest.java | 34 +++++ .../sql/src/main/codegen/config.fmpp | 2 + .../src/main/codegen/includes/parserImpls.ftl | 76 ++++++++-- .../sql/impl/parser/SqlAlterTable.java | 57 +++++++- .../impl/parser/SqlCreateExternalTable.java | 25 +++- .../sql/impl/parser/SqlDdlNodes.java | 44 ++++++ .../beam/sdk/extensions/sql/meta/Table.java | 5 + .../sql/meta/provider/AlterTableOps.java | 9 ++ .../sql/meta/provider/TableProvider.java | 4 + .../sql/meta/store/InMemoryMetaStore.java | 5 + .../sdk/extensions/sql/BeamSqlCliTest.java | 35 +++-- .../sql/impl/parser/BeamDDLTest.java | 136 ++++++++++++++++++ .../sdk/io/iceberg/IcebergCatalogConfig.java | 45 +++++- 15 files changed, 483 insertions(+), 29 deletions(-) diff --git a/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergAlterTableOps.java b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergAlterTableOps.java index 0c8a5519ee40..31476bc2ab89 100644 --- a/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergAlterTableOps.java +++ b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergAlterTableOps.java @@ -47,4 +47,9 @@ public void updatePartitionSpec( List partitionsToAdd, Collection partitionsToDrop) { table.updatePartitionSpec(partitionsToAdd, partitionsToDrop); } + + @Override + public void updatePrimaryKey(List primaryKey) { + table.updateIdentifierFields(primaryKey); + } } diff --git a/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastore.java b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastore.java index 58b0e4d53c95..f0577083e535 100644 --- a/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastore.java +++ b/sdks/java/extensions/sql/iceberg/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/IcebergMetastore.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.core.type.TypeReference; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.stream.Collectors; import org.apache.beam.sdk.extensions.sql.TableUtils; @@ -44,6 +45,7 @@ public class IcebergMetastore extends InMemoryMetaStore { @VisibleForTesting final IcebergCatalogConfig catalogConfig; private final Map cachedTables = new HashMap<>(); private final String database; + private static final String TYPE = "iceberg"; public IcebergMetastore(String db, IcebergCatalogConfig catalogConfig) { this.database = db; @@ -52,12 +54,12 @@ public IcebergMetastore(String db, IcebergCatalogConfig catalogConfig) { @Override public String getTableType() { - return "iceberg"; + return TYPE; } @Override public void createTable(Table table) { - if (!table.getType().equals("iceberg")) { + if (!table.getType().equals(TYPE)) { getProvider(table.getType()).createTable(table); } else { String identifier = getIdentifier(table); @@ -69,7 +71,12 @@ public void createTable(Table table) { .filter(p -> !p.getKey().startsWith("beam.")) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); try { - catalogConfig.createTable(identifier, table.getSchema(), table.getPartitionFields(), props); + catalogConfig.createTable( + identifier, + table.getSchema(), + table.getPartitionFields(), + props, + table.getPrimaryKey()); } catch (TableAlreadyExistsException e) { LOG.info( "Iceberg table '{}' already exists at location '{}'.", table.getName(), identifier); @@ -136,17 +143,19 @@ private String getIdentifier(Table table) { if (tableInfo == null) { return null; } + List identifierFields = tableInfo.getIdentifierFields(); return Table.builder() .type(getTableType()) .name(tableName(identifier)) .schema(tableInfo.getSchema()) + .primaryKey(identifierFields.isEmpty() ? null : identifierFields) .properties(TableUtils.parseProperties(tableInfo.getProperties())) .build(); } @Override public BeamSqlTable buildBeamSqlTable(Table table) { - if (table.getType().equals("iceberg")) { + if (table.getType().equals(TYPE)) { return new IcebergTable(getIdentifier(table), table, catalogConfig); } return getProvider(table.getType()).buildBeamSqlTable(table); @@ -154,18 +163,27 @@ public BeamSqlTable buildBeamSqlTable(Table table) { @Override public boolean supportsPartitioning(Table table) { - if (table.getType().equals("iceberg")) { + if (table.getType().equals(TYPE)) { return true; } return getProvider(table.getType()).supportsPartitioning(table); } + @Override + public boolean supportsPrimaryKey(Table table) { + if (table.getType().equals(TYPE)) { + return true; + } + return getProvider(table.getType()).supportsPrimaryKey(table); + } + @Override public IcebergAlterTableOps alterTable(String name) { IcebergTableInfo table = checkStateNotNull( catalogConfig.loadTable(getIdentifier(name)), "Could not find table '%s'", name); - + // The ALTER changes the table, so the next lookup reloads it. + cachedTables.remove(name); return new IcebergAlterTableOps(table); } diff --git a/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergAlterTest.java b/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergAlterTest.java index b6fd2b14a5e7..91e961644c81 100644 --- a/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergAlterTest.java +++ b/sdks/java/extensions/sql/iceberg/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/iceberg/BeamSqlCliIcebergAlterTest.java @@ -32,7 +32,9 @@ import java.util.UUID; import org.apache.beam.sdk.extensions.sql.BeamSqlCli; import org.apache.beam.sdk.extensions.sql.meta.catalog.InMemoryCatalogManager; +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.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.Table; @@ -172,6 +174,38 @@ public void testAlterTableSchemaFailsHelpfullyWhenAddingRequiredColumns() { cli.execute("ALTER TABLE my_table ADD COLUMNS (col3 BOOLEAN NOT NULL)"); } + @Test + public void testAlterTablePrimaryKey() { + cli.execute(createCatalog("my_catalog")); + cli.execute("CREATE DATABASE my_catalog.my_db"); + cli.execute("USE DATABASE my_catalog.my_db"); + // Iceberg identifier fields must be required, so this also checks that id becomes NOT NULL. + cli.execute( + "CREATE EXTERNAL TABLE my_table(id BIGINT PRIMARY KEY NOT ENFORCED, " + + "region VARCHAR NOT NULL, name VARCHAR) TYPE 'iceberg'"); + IcebergCatalog catalog = (IcebergCatalog) catalogManager.currentCatalog(); + Table table = + catalog.catalogConfig.catalog().loadTable(TableIdentifier.parse("my_db.my_table")); + assertEquals(ImmutableSet.of("id"), table.schema().identifierFieldNames()); + + cli.execute("ALTER TABLE my_table DROP PRIMARY KEY ADD PRIMARY KEY (id, region) NOT ENFORCED"); + table.refresh(); + assertEquals(ImmutableSet.of("id", "region"), table.schema().identifierFieldNames()); + // The session reloads the altered table, key included. + assertEquals( + ImmutableList.of("id", "region"), + catalog.metaStore("my_db").getTable("my_table").getPrimaryKey()); + + cli.execute("ALTER TABLE my_table DROP PRIMARY KEY"); + table.refresh(); + assertEquals(ImmutableSet.of(), table.schema().identifierFieldNames()); + + // DROP left id NOT NULL, so it can be the key again. + cli.execute("ALTER TABLE my_table ADD PRIMARY KEY (id) NOT ENFORCED"); + table.refresh(); + assertEquals(ImmutableSet.of("id"), table.schema().identifierFieldNames()); + } + @Test public void testAlterTablePartitionSpec() { cli.execute(createCatalog("my_catalog")); diff --git a/sdks/java/extensions/sql/src/main/codegen/config.fmpp b/sdks/java/extensions/sql/src/main/codegen/config.fmpp index c3692430610f..d7486de80caf 100644 --- a/sdks/java/extensions/sql/src/main/codegen/config.fmpp +++ b/sdks/java/extensions/sql/src/main/codegen/config.fmpp @@ -59,6 +59,7 @@ data: { "PARTITIONS" "USE" "UNSET" + "ENFORCED" ] # List of keywords from "keywords" section that are not reserved. @@ -409,6 +410,7 @@ data: { "PROPERTIES" "PARTITIONED" "USE" + "ENFORCED" ] # List of non-reserved keywords to add; 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..8b7e728bf7d9 100644 --- a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl +++ b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl @@ -74,7 +74,7 @@ List FieldListParens() : } { - fields = FieldListBody() + fields = FieldListBody(null) { return fields; @@ -87,35 +87,40 @@ List FieldListAngular() : } { - fields = FieldListBody() + fields = FieldListBody(null) { return fields; } } -List FieldListBody() : +List FieldListBody(SqlNodeList primaryKey) : { final List fields = new ArrayList(); Schema.Field field = null; } { - field = Field() { fields.add(field); } + field = Field(primaryKey) { fields.add(field); } ( - field = Field() { fields.add(field); } + field = Field(primaryKey) { fields.add(field); } )* { return fields; } } -Schema.Field Field() : +/** + * A column definition. Columns declared PRIMARY KEY are added to primaryKey, which is null where + * keys aren't allowed: only CREATE EXTERNAL TABLE's own columns can be keys. + */ +Schema.Field Field(SqlNodeList primaryKey) : { final String name; final Schema.FieldType type; final boolean nullable; Schema.Field field = null; SqlNode comment = null; + final Span s; } { name = Identifier() @@ -130,6 +135,11 @@ Schema.Field Field() : | { field = field.withNullable(true); } ) + [ + LOOKAHEAD(1, { primaryKey != null }) + { s = span(); } NotEnforced(s) + { primaryKey.add(new SqlIdentifier(name, s.pos())); } + ] [ comment = StringLiteral() { @@ -145,6 +155,16 @@ Schema.Field Field() : } } +/** The NOT ENFORCED that must follow a PRIMARY KEY. */ +void NotEnforced(Span s) : +{ +} +{ + +| + { throw SqlDdlNodes.enforcedPrimaryKey(s.end(this)); } +} + SqlNodeList PropertyList() : { SqlNodeList list = new SqlNodeList(getPos()); @@ -546,18 +566,26 @@ SqlNodeList PartitionFieldList() : * Note: This example is probably out of sync with the code. * * CREATE EXTERNAL TABLE ( IF NOT EXISTS )? - * ( catalog_name '.' )? ( database_name '.' )? table_name '(' column_def ( ',' column_def )* ')' + * ( catalog_name '.' )? ( database_name '.' )? table_name + * '(' column_def ( ',' column_def )* + * ( ',' PRIMARY KEY '(' column_name ( ',' column_name )* ')' NOT ENFORCED )? ')' * TYPE type_name * ( PARTITIONED BY '(' partition_field ( ',' partition_field )* ')' )? * ( COMMENT comment_string )? * ( LOCATION location_string )? * ( TBLPROPERTIES tbl_properties )? + * + * column_def: + * column_name type ( NULL | NOT NULL )? ( PRIMARY KEY NOT ENFORCED )? ( COMMENT comment_string )? */ SqlCreate SqlCreateExternalTable(Span s, boolean replace) : { final boolean ifNotExists; final SqlIdentifier id; List fieldList = null; + final SqlNodeList columnKey = new SqlNodeList(SqlParserPos.ZERO); + SqlNodeList primaryKey = null; + final Span primaryKeySpan; final SqlNode type; SqlNodeList partitionFields = null; SqlNode comment = null; @@ -572,7 +600,23 @@ SqlCreate SqlCreateExternalTable(Span s, boolean replace) : ifNotExists = IfNotExistsOpt() id = CompoundIdentifier() - fieldList = FieldListParens() + + fieldList = FieldListBody(columnKey) + [ + { primaryKeySpan = span(); } + primaryKey = ParenthesizedSimpleIdentifierList() + NotEnforced(primaryKeySpan) + ] + + { + if (!columnKey.isEmpty()) { + if (columnKey.size() > 1 || primaryKey != null) { + throw new ParseException("A table can have only one PRIMARY KEY. Declare a key on " + + "several columns as PRIMARY KEY (column, ...) NOT ENFORCED after the columns."); + } + primaryKey = columnKey; + } + } ( type = StringLiteral() @@ -591,6 +635,7 @@ SqlCreate SqlCreateExternalTable(Span s, boolean replace) : ifNotExists, id, fieldList, + primaryKey, type, partitionFields, comment, @@ -606,6 +651,8 @@ SqlCreate SqlCreateExternalTable(Span s, boolean replace) : * [ DROP COLUMNS ] * [ ADD PARTITIONS ] * [ DROP PARTITIONS ] + * [ ADD PRIMARY KEY (column_name, column_name, ...) NOT ENFORCED ] + * [ DROP PRIMARY KEY ] * [ SET (key1=val1, key2=val2, ...) ] * [ (RESET | UNSET) (key1, key2, ...) ] */ @@ -616,6 +663,9 @@ SqlCall SqlAlterTable(Span s, String scope) : List columnsToAdd = null; SqlNodeList partitionsToDrop = null; SqlNodeList partitionsToAdd = null; + SqlNodeList primaryKeyToAdd = null; + boolean dropPrimaryKey = false; + final Span primaryKeySpan; SqlNodeList setProps = null; SqlNodeList resetProps = null; } @@ -630,12 +680,18 @@ SqlCall SqlAlterTable(Span s, String scope) : columnsToDrop = ParenthesizedSimpleIdentifierList() | partitionsToDrop = ParenthesizedLiteralOptionCommaList() + | + { dropPrimaryKey = true; } ) ] [ ( columnsToAdd = FieldListParens() | partitionsToAdd = ParenthesizedLiteralOptionCommaList() + | + { primaryKeySpan = span(); } + primaryKeyToAdd = ParenthesizedSimpleIdentifierList() + NotEnforced(primaryKeySpan) ) ] [ ( | ) resetProps = ArgList() ] @@ -651,7 +707,9 @@ SqlCall SqlAlterTable(Span s, String scope) : partitionsToAdd, partitionsToDrop, setProps, - resetProps); + resetProps, + primaryKeyToAdd, + dropPrimaryKey); } } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlAlterTable.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlAlterTable.java index 2edf8be9aaa3..1c804fa1bd14 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlAlterTable.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlAlterTable.java @@ -19,6 +19,7 @@ import static org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.util.Static.RESOURCE; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects.firstNonNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.util.Collection; @@ -30,7 +31,9 @@ import org.apache.beam.sdk.extensions.sql.impl.CatalogSchema; import org.apache.beam.sdk.extensions.sql.impl.TableName; import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils; +import org.apache.beam.sdk.extensions.sql.meta.Table; import org.apache.beam.sdk.extensions.sql.meta.provider.AlterTableOps; +import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider; import org.apache.beam.sdk.schemas.Schema.Field; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalcitePrepare; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalciteSchema; @@ -60,6 +63,8 @@ public class SqlAlterTable extends SqlAlter implements BeamSqlParser.ExecutableS private final @Nullable SqlNodeList partitionsToDrop; private final @Nullable SqlNodeList setProps; private final @Nullable SqlNodeList resetProps; + private final @Nullable SqlNodeList primaryKeyToAdd; + private final boolean dropPrimaryKey; /** * Called by the auto-generated {@link @@ -75,7 +80,9 @@ public SqlAlterTable( @Nullable SqlNodeList partitionsToAdd, @Nullable SqlNodeList partitionsToDrop, @Nullable SqlNodeList setProps, - @Nullable SqlNodeList resetProps) { + @Nullable SqlNodeList resetProps, + @Nullable SqlNodeList primaryKeyToAdd, + boolean dropPrimaryKey) { super(pos, scope); this.name = SqlDdlNodes.getIdentifier(name, pos); this.columnsToAdd = columnsToAdd; @@ -84,6 +91,8 @@ public SqlAlterTable( this.partitionsToDrop = partitionsToDrop; this.setProps = setProps; this.resetProps = resetProps; + this.primaryKeyToAdd = primaryKeyToAdd; + this.dropPrimaryKey = dropPrimaryKey; } @Override @@ -110,7 +119,9 @@ public void execute(CalcitePrepare.Context context) { + schema.getClass())); } - if (beamCalciteSchema.getTable(pair.right) == null) { + TableProvider provider = beamCalciteSchema.getTableProvider(); + Table table = provider.getTable(pair.right); + if (table == null) { // Table does not exist. throw SqlUtil.newContextException( name.getParserPosition(), RESOURCE.tableNotFound(name.toString())); @@ -121,9 +132,13 @@ public void execute(CalcitePrepare.Context context) { List columnsToDropList = SqlDdlNodes.getStringList(columnsToDrop); List partitionsToAddList = SqlDdlNodes.getStringList(partitionsToAdd); List partitionsToDropList = SqlDdlNodes.getStringList(partitionsToDrop); + List primaryKeyToAddList = SqlDdlNodes.getStringList(primaryKeyToAdd); + boolean alterPrimaryKey = dropPrimaryKey || primaryKeyToAdd != null; + if (alterPrimaryKey) { + validatePrimaryKey(provider, table, primaryKeyToAddList); + } - AlterTableOps alterOps = - beamCalciteSchema.getTableProvider().alterTable(SqlDdlNodes.name(name)); + AlterTableOps alterOps = provider.alterTable(SqlDdlNodes.name(name)); if (!setPropsMap.isEmpty() || !resetPropsList.isEmpty()) { validateNonOverlappingProps(setPropsMap, resetPropsList); @@ -133,11 +148,35 @@ public void execute(CalcitePrepare.Context context) { if (!columnsToDropList.isEmpty() || (columnsToAdd != null && !columnsToAdd.isEmpty())) { alterOps.updateSchema(firstNonNull(columnsToAdd, Collections.emptyList()), columnsToDropList); } + if (alterPrimaryKey) { + // DROP PRIMARY KEY ADD PRIMARY KEY (...) replaces the key in a single update. + alterOps.updatePrimaryKey(primaryKeyToAddList); + } if (!partitionsToDropList.isEmpty() || !partitionsToAddList.isEmpty()) { alterOps.updatePartitionSpec(partitionsToAddList, partitionsToDropList); } } + private void validatePrimaryKey(TableProvider provider, Table table, List columns) { + SqlDdlNodes.checkPrimaryKey(provider, table, columns); + if (dropPrimaryKey) { + checkArgument( + table.getPrimaryKey() != null, "Table '%s' has no primary key to drop.", table.getName()); + } else { + checkArgument( + table.getPrimaryKey() == null, + "Table '%s' already has primary key %s. Drop it first with DROP PRIMARY KEY.", + table.getName(), + table.getPrimaryKey()); + } + for (String column : columns) { + checkArgument( + !table.getSchema().getField(column).getType().getNullable(), + "PRIMARY KEY column '%s' must already be NOT NULL.", + column); + } + } + private void validateNonOverlappingProps( Map setPropsMap, Collection resetPropsList) { ImmutableList.Builder overlappingPropsBuilder = ImmutableList.builder(); @@ -195,6 +234,16 @@ public void unparseAlterOperation(SqlWriter writer, int left, int right) { writer.endList(frame); } + if (dropPrimaryKey) { + writer.keyword("DROP PRIMARY KEY"); + } + + if (primaryKeyToAdd != null) { + writer.keyword("ADD PRIMARY KEY"); + writer.list(SqlWriter.FrameTypeEnum.PARENTHESES, SqlWriter.COMMA, primaryKeyToAdd); + writer.keyword("NOT ENFORCED"); + } + if (resetProps != null && !resetProps.isEmpty()) { writer.keyword("RESET"); SqlWriter.Frame frame = writer.startList("(", ")"); 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..defca3992ac1 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 @@ -57,6 +57,7 @@ public class SqlCreateExternalTable extends SqlCreate implements BeamSqlParser.ExecutableStatement { private final SqlIdentifier name; private final List columnList; + private final @Nullable SqlNodeList primaryKey; private final SqlNode type; private final SqlNode comment; private final SqlNode location; @@ -73,6 +74,7 @@ public SqlCreateExternalTable( boolean ifNotExists, SqlIdentifier name, List columnList, + @Nullable SqlNodeList primaryKey, SqlNode type, SqlNodeList partitionFields, SqlNode comment, @@ -81,6 +83,7 @@ public SqlCreateExternalTable( super(OPERATOR, pos, replace, ifNotExists); this.name = checkNotNull(name); this.columnList = columnList; // may be null + this.primaryKey = primaryKey; this.type = checkNotNull(type); this.partitionFields = partitionFields; this.comment = comment; // may be null @@ -107,6 +110,12 @@ public void unparse(SqlWriter writer, int leftPrec, int rightPrec) { if (columnList != null) { SqlWriter.Frame frame = writer.startList("(", ")"); columnList.forEach(column -> unparseColumn(writer, column)); + if (primaryKey != null) { + writer.sep(","); + writer.keyword("PRIMARY KEY"); + writer.list(SqlWriter.FrameTypeEnum.PARENTHESES, SqlWriter.COMMA, primaryKey); + writer.keyword("NOT ENFORCED"); + } writer.endList(frame); } writer.keyword("TYPE"); @@ -180,6 +189,10 @@ public void execute(CalcitePrepare.Context context) { name(name), SqlDdlNodes.getString(type)); } + List primaryKeyColumns = table.getPrimaryKey(); + if (primaryKeyColumns != null) { + SqlDdlNodes.checkPrimaryKey(beamCalciteSchema.getTableProvider(), table, primaryKeyColumns); + } beamCalciteSchema.getTableProvider().createTable(table); } @@ -207,11 +220,21 @@ private void unparseColumn(SqlWriter writer, Field column) { } private Table toTable() { + List primaryKeyColumns = + primaryKey == null ? null : SqlDdlNodes.getStringList(primaryKey); return Table.builder() .type(SqlDdlNodes.getString(type)) .name(name(name)) - .schema(columnList.stream().collect(toSchema())) + .schema( + columnList.stream() + .map( + column -> + primaryKeyColumns != null && primaryKeyColumns.contains(column.getName()) + ? column.withNullable(false) + : column) + .collect(toSchema())) .partitionFields(parsePartitionFields()) + .primaryKey(primaryKeyColumns) .comment(SqlDdlNodes.getString(comment)) .location(SqlDdlNodes.getString(location)) // .path(path) diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java index f8d7e6f73851..74046523e818 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlDdlNodes.java @@ -19,13 +19,20 @@ import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; 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 static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; +import org.apache.beam.sdk.extensions.sql.meta.Table; +import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalcitePrepare; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalciteSchema; +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.SqlDataTypeSpec; 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.SqlLiteral; @@ -141,6 +148,43 @@ static Map getStringMap(@Nullable SqlNodeList nodeList) { return builder.build(); } + /** + * Checks that {@code table} supports primary keys and that each of {@code primaryKey} is a column + * of the table, listed once. + */ + static void checkPrimaryKey(TableProvider provider, Table table, List primaryKey) { + checkArgument( + provider.supportsPrimaryKey(table), + "Invalid use of 'PRIMARY KEY': Table '%s' of type '%s' does not support primary keys.", + table.getName(), + table.getType()); + Set seen = new HashSet<>(); + for (String column : primaryKey) { + checkArgument( + table.getSchema().hasField(column), + "PRIMARY KEY column '%s' is not a column of table '%s'.", + column, + table.getName()); + checkArgument(seen.add(column), "PRIMARY KEY column '%s' is listed more than once.", column); + } + } + + /** + * The parse error for a {@code PRIMARY KEY} declared without {@code NOT ENFORCED}. Beam SQL never + * checks that key values are unique, so {@code NOT ENFORCED} is the only mode it accepts. + */ + public static CalciteContextException enforcedPrimaryKey(SqlParserPos pos) { + String message = + "Beam SQL doesn't enforce primary keys, so PRIMARY KEY must be declared NOT ENFORCED."; + return new CalciteContextException( + message, + new CalciteException(message, null), + pos.getLineNum(), + pos.getColumnNum(), + pos.getEndLineNum(), + pos.getEndColumnNum()); + } + static SqlIdentifier getIdentifier(SqlNode n, SqlParserPos pos) { if (n instanceof SqlIdentifier) { return (SqlIdentifier) n; diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/Table.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/Table.java index 5c03a2b20b25..a211cc7a72bf 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/Table.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/Table.java @@ -38,6 +38,9 @@ public abstract class Table implements Serializable { public abstract @Nullable List getPartitionFields(); + /** The columns of the table's primary key, or null if it has none. */ + public abstract @Nullable List getPrimaryKey(); + public abstract @Nullable String getComment(); public abstract @Pure @Nullable String getLocation(); @@ -61,6 +64,8 @@ public abstract static class Builder { public abstract Builder partitionFields(List fields); + public abstract Builder primaryKey(@Nullable List columns); + public abstract Builder comment(String name); public abstract Builder location(String location); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/AlterTableOps.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/AlterTableOps.java index 3f1ea1721684..82f6d621da1c 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/AlterTableOps.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/AlterTableOps.java @@ -42,4 +42,13 @@ public interface AlterTableOps { * existing partitions. */ void updatePartitionSpec(List partitionsToAdd, Collection partitionsToDrop); + + /** + * Replaces a table's primary key with the given columns, or drops it if the list is empty. Only + * called for tables that {@link TableProvider#supportsPrimaryKey support primary keys}. + */ + default void updatePrimaryKey(List primaryKey) { + throw new UnsupportedOperationException( + getClass().getSimpleName() + " does not support altering primary keys."); + } } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/TableProvider.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/TableProvider.java index 595d9cba52fa..9bc333771d79 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/TableProvider.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/provider/TableProvider.java @@ -81,6 +81,10 @@ default boolean supportsPartitioning(Table table) { return false; } + default boolean supportsPrimaryKey(Table table) { + return false; + } + default AlterTableOps alterTable(String name) { throw new UnsupportedOperationException( String.format("ALTER is not supported for table '%s' of type '%s'.", name, getTableType())); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/store/InMemoryMetaStore.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/store/InMemoryMetaStore.java index 7fd74b6b85af..d384eaf1f4ae 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/store/InMemoryMetaStore.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/store/InMemoryMetaStore.java @@ -130,6 +130,11 @@ public boolean supportsPartitioning(Table table) { return getProvider(table.getType()).supportsPartitioning(table); } + @Override + public boolean supportsPrimaryKey(Table table) { + return getProvider(table.getType()).supportsPrimaryKey(table); + } + /** * Fetches a {@link TableProvider} for this type. This provider can exist in the current {@link * InMemoryMetaStore} or a nested {@link InMemoryMetaStore}. diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliTest.java index 522d9c35bbf6..b9fb55a2aaa2 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/BeamSqlCliTest.java @@ -402,7 +402,8 @@ public void testUnparseAlter_AddColumns() { List columnsToAdd = Arrays.asList(col1, col2); SqlAlterTable alterTable = - new SqlAlterTable(POS, null, tableName, columnsToAdd, null, null, null, null, null); + new SqlAlterTable( + POS, null, tableName, columnsToAdd, null, null, null, null, null, null, false); String expectedSql = "ALTER TABLE `test_table` " @@ -419,7 +420,8 @@ public void testUnparseAlter_DropColumns() { SqlNodeList columnsToDrop = createStringList("col_to_drop_1", "col_to_drop_2"); SqlAlterTable alterTable = - new SqlAlterTable(POS, null, tableName, null, columnsToDrop, null, null, null, null); + new SqlAlterTable( + POS, null, tableName, null, columnsToDrop, null, null, null, null, null, false); String expectedSql = "ALTER TABLE `test_table` DROP COLUMNS (`col_to_drop_1`, `col_to_drop_2`)"; assertEquals(expectedSql, toSql(alterTable)); @@ -438,7 +440,7 @@ public void testUnparseAlter_AddColumnsAndDropColumns() { SqlAlterTable alterTable = new SqlAlterTable( - POS, null, tableName, columnsToAdd, columnsToDrop, null, null, null, null); + POS, null, tableName, columnsToAdd, columnsToDrop, null, null, null, null, null, false); // unparses DROP before ADD String expectedSql = @@ -457,7 +459,8 @@ public void testUnparseAlter_AddAndDropPartitions() { SqlNodeList partsToDrop = createStringList("p3"); SqlAlterTable alterTable = - new SqlAlterTable(POS, null, tableName, null, null, partsToAdd, partsToDrop, null, null); + new SqlAlterTable( + POS, null, tableName, null, null, partsToAdd, partsToDrop, null, null, null, false); // unparses DROP before ADD String expectedSql = @@ -476,7 +479,8 @@ public void testUnparseAlter_TableProperties() { SqlNodeList resetProps = createStringList("prop2"); SqlAlterTable alterTable = - new SqlAlterTable(POS, null, tableName, null, null, null, null, setProps, resetProps); + new SqlAlterTable( + POS, null, tableName, null, null, null, null, setProps, resetProps, null, false); // unparses RESET before SET String expectedSql = "ALTER TABLE `test_table` RESET ('prop2') SET ('prop1' = 'val1')"; @@ -496,24 +500,39 @@ public void testUnparseAlter_AllOperations() { SqlNodeList setProps = new SqlNodeList(POS); setProps.add(createPropertyPair("k", "v")); SqlNodeList resetProps = createStringList("k_reset"); + SqlNodeList primaryKey = SqlNodeList.of(new SqlIdentifier("c1", POS)); SqlAlterTable alterTable = new SqlAlterTable( - POS, null, tableName, addCols, dropCols, addParts, dropParts, setProps, resetProps); + POS, + null, + tableName, + addCols, + dropCols, + addParts, + dropParts, + setProps, + resetProps, + primaryKey, + true); // Expected Order based on source code: // 1. DROP COLUMNS // 2. ADD COLUMNS // 3. DROP PARTITIONS // 4. ADD PARTITIONS - // 5. RESET - // 6. SET + // 5. DROP PRIMARY KEY + // 6. ADD PRIMARY KEY + // 7. RESET + // 8. SET String expectedSql = "ALTER TABLE `full_table` " + "DROP COLUMNS (`c_old`) " + "ADD COLUMNS (`c1` BOOLEAN NOT NULL) " + "DROP PARTITIONS (`p_old`) " + "ADD PARTITIONS (`p_new`) " + + "DROP PRIMARY KEY " + + "ADD PRIMARY KEY (`c1`) NOT ENFORCED " + "RESET ('k_reset') " + "SET ('k' = 'v')"; 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..68fd5ef948f2 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 @@ -18,12 +18,16 @@ package org.apache.beam.sdk.extensions.sql.impl.parser; import static org.apache.beam.sdk.schemas.Schema.toSchema; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; 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; +import java.util.Arrays; import java.util.stream.Stream; import org.apache.beam.sdk.extensions.sql.TableUtils; import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv; @@ -33,12 +37,17 @@ import org.apache.beam.sdk.extensions.sql.meta.Table; import org.apache.beam.sdk.extensions.sql.meta.provider.bigquery.BeamBigQuerySqlDialect; import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestTableProvider; +import org.apache.beam.sdk.extensions.sql.meta.provider.text.TextTableProvider; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.schemas.Schema; +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.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.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.SqlWriter; +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.sql.parser.SqlParserPos; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.pretty.SqlPrettyWriter; import org.junit.Test; @@ -186,6 +195,125 @@ public void testParseCreateExternalTable_minimal() throws Exception { tableProvider.getTables().get("person")); } + @Test + public void testParseCreateExternalTable_primaryKey() throws Exception { + TestTableProvider tableProvider = new PrimaryKeyTableProvider(); + BeamSqlEnv env = BeamSqlEnv.withTableProvider(tableProvider); + String create = + "CREATE EXTERNAL TABLE person (id INT, region VARCHAR, enforced VARCHAR," + + " PRIMARY KEY (region, id) NOT ENFORCED) TYPE test"; + + env.executeDdl(create); + + Table table = tableProvider.getTables().get("person"); + assertEquals(Arrays.asList("region", "id"), table.getPrimaryKey()); + assertEquals( + Stream.of( + Schema.Field.of("id", CalciteUtils.INTEGER), + Schema.Field.of("region", CalciteUtils.VARCHAR), + Schema.Field.nullable("enforced", CalciteUtils.VARCHAR)) + .collect(toSchema()), + table.getSchema()); + assertThat( + SqlParser.create( + create, + SqlParser.config().withLex(Lex.JAVA).withParserFactory(BeamSqlParser.FACTORY)) + .parseStmt() + .toSqlString(new SqlDialect(SqlDialect.EMPTY_CONTEXT)) + .getSql(), + containsString(", PRIMARY KEY (region, id) NOT ENFORCED) TYPE test")); + + env.executeDdl( + "CREATE EXTERNAL TABLE keyed (name VARCHAR," + + " id INT NOT NULL PRIMARY KEY NOT ENFORCED COMMENT 'key') TYPE test"); + Table keyed = tableProvider.getTables().get("keyed"); + assertEquals(Arrays.asList("id"), keyed.getPrimaryKey()); + assertEquals("key", keyed.getSchema().getField("id").getDescription()); + } + + @Test + public void testInvalidPrimaryKey() { + BeamSqlEnv env = BeamSqlEnv.inMemory(new PrimaryKeyTableProvider(), new TextTableProvider()); + env.executeDdl("CREATE EXTERNAL TABLE keyed (id INT PRIMARY KEY NOT ENFORCED) TYPE test"); + env.executeDdl("CREATE EXTERNAL TABLE unkeyed (id INT) TYPE test"); + env.executeDdl("CREATE EXTERNAL TABLE lines (id INT NOT NULL) TYPE text"); + String[][] ddlAndErrors = { + { + "CREATE EXTERNAL TABLE person (id INT PRIMARY KEY NOT ENFORCED) TYPE text", + "Invalid use of 'PRIMARY KEY': Table 'person' of type 'text' does not support primary keys." + }, + // Text tables don't support ALTER either, so this also checks the key is validated first. + { + "ALTER TABLE lines ADD PRIMARY KEY (id) NOT ENFORCED", + "Invalid use of 'PRIMARY KEY': Table 'lines' of type 'text' does not support primary keys." + }, + { + "CREATE EXTERNAL TABLE person (id INT, PRIMARY KEY (id, name) NOT ENFORCED) TYPE test", + "PRIMARY KEY column 'name' is not a column of table 'person'." + }, + { + "CREATE EXTERNAL TABLE person (id INT, PRIMARY KEY (id, id) NOT ENFORCED) TYPE test", + "PRIMARY KEY column 'id' is listed more than once." + }, + { + "ALTER TABLE keyed ADD PRIMARY KEY (id) NOT ENFORCED", + "Table 'keyed' already has primary key [id]. Drop it first with DROP PRIMARY KEY." + }, + {"ALTER TABLE unkeyed DROP PRIMARY KEY", "Table 'unkeyed' has no primary key to drop."}, + { + "ALTER TABLE unkeyed ADD PRIMARY KEY (id) NOT ENFORCED", + "PRIMARY KEY column 'id' must already be NOT NULL." + } + }; + for (String[] ddlAndError : ddlAndErrors) { + IllegalArgumentException e = + assertThrows( + ddlAndError[0], IllegalArgumentException.class, () -> env.executeDdl(ddlAndError[0])); + assertEquals(ddlAndError[1], e.getMessage()); + } + // A table has one key, declared only among CREATE EXTERNAL TABLE's own columns. + for (String ddl : + new String[] { + "CREATE EXTERNAL TABLE person (r ROW(id INT PRIMARY KEY NOT ENFORCED)) TYPE test", + "ALTER TABLE keyed ADD COLUMNS (k INT NOT NULL PRIMARY KEY NOT ENFORCED)" + }) { + assertThrows(ddl, ParseException.class, () -> env.executeDdl(ddl)); + } + for (String columns : + new String[] { + "(id INT PRIMARY KEY NOT ENFORCED, name VARCHAR PRIMARY KEY NOT ENFORCED)", + "(id INT PRIMARY KEY NOT ENFORCED, PRIMARY KEY (id) NOT ENFORCED)" + }) { + ParseException twoKeys = + assertThrows( + columns, + ParseException.class, + () -> env.executeDdl("CREATE EXTERNAL TABLE person " + columns + " TYPE test")); + assertThat( + columns, + twoKeys.getCause().getMessage(), + containsString("A table can have only one PRIMARY KEY.")); + } + // Beam SQL never checks key uniqueness, so a key without NOT ENFORCED fails at PRIMARY. + for (String ddl : + new String[] { + "CREATE EXTERNAL TABLE person (id INT, PRIMARY KEY (id) ENFORCED) TYPE test", + "CREATE EXTERNAL TABLE person (id INT PRIMARY KEY COMMENT 'id') TYPE test", + "ALTER TABLE keyed ADD PRIMARY KEY (id)" + }) { + SqlParseException enforced = + (SqlParseException) + assertThrows(ddl, ParseException.class, () -> env.executeDdl(ddl)).getCause(); + assertThat( + ddl, + enforced.getMessage(), + containsString( + "Beam SQL doesn't enforce primary keys, so PRIMARY KEY must be declared NOT" + + " ENFORCED.")); + assertEquals(ddl, ddl.indexOf("PRIMARY") + 1, enforced.getPos().getColumnNum()); + } + } + @Test public void testParseCreateExternalTable_withDatabase() throws Exception { TestTableProvider rootProvider = new TestTableProvider(); @@ -250,6 +378,14 @@ public void unparseAggregateFunction() { sqlWriter.toSqlString().getSql()); } + /** A {@link TestTableProvider} whose tables support primary keys. */ + private static class PrimaryKeyTableProvider extends TestTableProvider { + @Override + public boolean supportsPrimaryKey(Table table) { + return true; + } + } + private static Table mockTable(String name, String type, String comment, ObjectNode properties) { return mockTable(name, type, comment, properties, "/home/admin/" + name); } diff --git a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java index fd836dc30ddd..bd8d6ae1b5a1 100644 --- a/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java +++ b/sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergCatalogConfig.java @@ -17,6 +17,9 @@ */ package org.apache.beam.sdk.io.iceberg; +import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import com.google.auto.value.AutoValue; import java.io.Serializable; import java.util.Collection; @@ -46,6 +49,7 @@ import org.apache.iceberg.exceptions.AlreadyExistsException; import org.apache.iceberg.exceptions.NoSuchTableException; import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; import org.checkerframework.checker.nullness.qual.Nullable; import org.checkerframework.dataflow.qual.Pure; import org.slf4j.Logger; @@ -163,8 +167,22 @@ public void createTable( Schema tableSchema, @Nullable List partitionFields, @Nullable Map properties) { + createTable(tableIdentifier, tableSchema, partitionFields, properties, null); + } + + /** + * Like {@link #createTable(String, Schema, List, Map)}, also setting the table's identifier + * (primary key) fields, which must be non-nullable columns. + */ + public void createTable( + String tableIdentifier, + Schema tableSchema, + @Nullable List partitionFields, + @Nullable Map properties, + @Nullable List identifierFields) { TableIdentifier icebergIdentifier = IcebergUtils.parseTableIdentifier(tableIdentifier); - org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(tableSchema); + org.apache.iceberg.Schema icebergSchema = + withIdentifierFields(IcebergUtils.beamSchemaToIcebergSchema(tableSchema), identifierFields); PartitionSpec icebergSpec = PartitionUtils.toPartitionSpec(partitionFields, tableSchema); try { LOG.info( @@ -185,6 +203,19 @@ public void createTable( } } + private static org.apache.iceberg.Schema withIdentifierFields( + org.apache.iceberg.Schema schema, @Nullable List identifierFields) { + if (identifierFields == null || identifierFields.isEmpty()) { + return schema; + } + Set identifierFieldIds = + identifierFields.stream() + .map(name -> checkArgumentNotNull(schema.findField(name), "Unknown column '%s'", name)) + .map(Types.NestedField::fieldId) + .collect(Collectors.toSet()); + return new org.apache.iceberg.Schema(schema.columns(), identifierFieldIds); + } + public @Nullable IcebergTableInfo loadTable(String tableIdentifier) { TableIdentifier icebergIdentifier = IcebergUtils.parseTableIdentifier(tableIdentifier); try { @@ -217,6 +248,14 @@ public Map getProperties() { return table.properties(); } + /** The table's identifier (primary key) fields. */ + public List getIdentifierFields() { + org.apache.iceberg.Schema schema = table.schema(); + return schema.identifierFieldIds().stream() + .map(id -> checkStateNotNull(schema.findColumnName(id))) + .collect(Collectors.toList()); + } + public void updateTableProps(Map setProps, List resetProps) { if (setProps.isEmpty() && resetProps.isEmpty()) { return; @@ -263,6 +302,10 @@ public void updateSchema(List columnsToAdd, Collection col update.commit(); } + public void updateIdentifierFields(List identifierFields) { + table.updateSchema().setIdentifierFields(identifierFields).commit(); + } + public void updatePartitionSpec( List partitionsToAdd, Collection partitionsToDrop) { if (partitionsToAdd.isEmpty() && partitionsToDrop.isEmpty()) { From 696a4df55050d0fb70b5e81471cb35f540b73fdb Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 27 Oct 2026 11:34:13 -0700 Subject: [PATCH 2/2] CHANGES --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index 9d3f6fbac241..917fd1ccd5b2 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -85,6 +85,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 `PRIMARY KEY (...) NOT ENFORCED` to `CREATE EXTERNAL TABLE`, and `ALTER TABLE ... ADD/DROP PRIMARY KEY` ([#40826](https://github.com/apache/beam/issues/40826)). ## Breaking Changes