Beam SQL Primary Keys

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
 

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<String> partitionsToAdd, Collection<String> partitionsToDrop) {
     table.updatePartitionSpec(partitionsToAdd, partitionsToDrop);
   }
+
+  @Override
+  public void updatePrimaryKey(List<String> primaryKey) {
+    table.updateIdentifierFields(primaryKey);
+  }
 }

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<String, Table> 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,36 +143,47 @@ private String getIdentifier(Table table) {
     if (tableInfo == null) {
       return null;
     }
+    List<String> 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);
   }
 
   @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);
   }
 

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"));

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;

sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl

@@ -74,7 +74,7 @@ List<Schema.Field> FieldListParens() :
 }
 {
     <LPAREN>
-        fields = FieldListBody()
+        fields = FieldListBody(null)
     <RPAREN>
     {
         return fields;
@@ -87,35 +87,40 @@ List<Schema.Field> FieldListAngular() :
 }
 {
     <LT>
-        fields = FieldListBody()
+        fields = FieldListBody(null)
     <GT>
     {
         return fields;
     }
 }
 
-List<Schema.Field> FieldListBody() :
+List<Schema.Field> FieldListBody(SqlNodeList primaryKey) :
 {
     final List<Schema.Field> fields = new ArrayList<Schema.Field>();
     Schema.Field field = null;
 }
 {
-    field = Field() { fields.add(field); }
+    field = Field(primaryKey) { fields.add(field); }
     (
-        <COMMA> field = Field() { fields.add(field); }
+        <COMMA> 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 })
+        <PRIMARY> { s = span(); } <KEY> NotEnforced(s)
+        { primaryKey.add(new SqlIdentifier(name, s.pos())); }
+    ]
     [
         <COMMENT> comment = StringLiteral()
         {
@@ -145,6 +155,16 @@ Schema.Field Field() :
     }
 }
 
+/** The NOT ENFORCED that must follow a PRIMARY KEY. */
+void NotEnforced(Span s) :
+{
+}
+{
+    <NOT> <ENFORCED>
+|
+    { 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<Schema.Field> 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()
+    <LPAREN>
+    fieldList = FieldListBody(columnKey)
+    [
+        <COMMA> <PRIMARY> { primaryKeySpan = span(); } <KEY>
+        primaryKey = ParenthesizedSimpleIdentifierList()
+        NotEnforced(primaryKeySpan)
+    ]
+    <RPAREN>
+    {
+        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>
     (
         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 <column_name, column_name> ]
  *   [ ADD PARTITIONS <partition_field, partition_field, ...> ]
  *   [ DROP PARTITIONS <partition_field, partition_field, ...> ]
+ *   [ 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<Schema.Field> 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) :
       <COLUMNS> columnsToDrop = ParenthesizedSimpleIdentifierList()
         |
       <PARTITIONS> partitionsToDrop = ParenthesizedLiteralOptionCommaList()
+        |
+      <PRIMARY> <KEY> { dropPrimaryKey = true; }
     ) ]
 
     [ <ADD> (
       <COLUMNS> columnsToAdd = FieldListParens()
         |
       <PARTITIONS> partitionsToAdd = ParenthesizedLiteralOptionCommaList()
+        |
+      <PRIMARY> { primaryKeySpan = span(); } <KEY>
+      primaryKeyToAdd = ParenthesizedSimpleIdentifierList()
+      NotEnforced(primaryKeySpan)
     ) ]
 
     [ (<RESET> | <UNSET>) <LPAREN> resetProps = ArgList() <RPAREN> ]
@@ -651,7 +707,9 @@ SqlCall SqlAlterTable(Span s, String scope) :
             partitionsToAdd,
             partitionsToDrop,
             setProps,
-            resetProps);
+            resetProps,
+            primaryKeyToAdd,
+            dropPrimaryKey);
     }
 }