diff --git a/docs/docs/spark/procedures.md b/docs/docs/spark/procedures.md index 2048c49b0718..947de28348bb 100644 --- a/docs/docs/spark/procedures.md +++ b/docs/docs/spark/procedures.md @@ -89,6 +89,7 @@ Choose a group, then use the page contents to jump to a procedure: [`fast_forward`](./procedures/versions#fast_forward), [`merge_branch`](./procedures/versions#merge_branch), [`rollback`](./procedures/versions#rollback), +[`rollback_to_as_latest`](./procedures/versions#rollback_to_as_latest), [`rollback_to_timestamp`](./procedures/versions#rollback_to_timestamp), [`rollback_to_watermark`](./procedures/versions#rollback_to_watermark) diff --git a/docs/docs/spark/procedures/versions.md b/docs/docs/spark/procedures/versions.md index d9b96729722c..f4d4a76f086b 100644 --- a/docs/docs/spark/procedures/versions.md +++ b/docs/docs/spark/procedures/versions.md @@ -240,6 +240,29 @@ CALL sys.rollback(table => 'default.T', tag => 'tag1'); CALL sys.rollback(table => 'default.T', snapshot => 2); ``` +## rollback_to_as_latest + +Roll back to a retained snapshot or tag and materialize it as the latest snapshot, without deleting +later snapshots or tags. Specify exactly one of `snapshot_id` or `tag`. Batch and time-travel reads +are correct; for deletion-vector tables, a rollback whose only difference is a deletion-vector +change is not guaranteed to be observed by streaming overwrite readers. + +When rolling back to a snapshot, the procedure creates a tag named +`rollback-to-as-latest--` that protects the restored data files from snapshot +expiration. + +**Arguments** + +- `table` (`STRING`, required): the target table identifier. +- `tag` (`STRING`, optional): tag that will roll back to. +- `snapshot_id` (`BIGINT`, optional): snapshot that will roll back to. + +```sql +CALL sys.rollback_to_as_latest(table => 'default.T', snapshot_id => 10); + +CALL sys.rollback_to_as_latest(table => 'default.T', tag => 'tag1'); +``` + ## rollback_to_timestamp Rollback to the snapshot which earlier or equal than timestamp. diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java index f534773959f6..18bb32837533 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/SparkProcedures.java @@ -67,6 +67,7 @@ import org.apache.paimon.spark.procedure.RevokePermissionProcedure; import org.apache.paimon.spark.procedure.RewriteFileIndexProcedure; import org.apache.paimon.spark.procedure.RollbackProcedure; +import org.apache.paimon.spark.procedure.RollbackToAsLatestProcedure; import org.apache.paimon.spark.procedure.RollbackToTimestampProcedure; import org.apache.paimon.spark.procedure.RollbackToWatermarkProcedure; import org.apache.paimon.spark.procedure.TriggerTagAutomaticCreationProcedure; @@ -98,6 +99,7 @@ private static Map> initProcedureBuilders() { ImmutableMap.Builder> procedureBuilders = ImmutableMap.builder(); procedureBuilders.put("rollback", RollbackProcedure::builder); + procedureBuilders.put("rollback_to_as_latest", RollbackToAsLatestProcedure::builder); procedureBuilders.put("rollback_to_timestamp", RollbackToTimestampProcedure::builder); procedureBuilders.put("rollback_to_watermark", RollbackToWatermarkProcedure::builder); procedureBuilders.put("purge_files", PurgeFilesProcedure::builder); diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java new file mode 100644 index 000000000000..69366598eee0 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedure.java @@ -0,0 +1,242 @@ +/* + * 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.paimon.spark.procedure; + +import org.apache.paimon.FileStore; +import org.apache.paimon.Snapshot; +import org.apache.paimon.spark.SparkTable; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.sink.TableCommitImpl; +import org.apache.paimon.tag.Tag; +import org.apache.paimon.utils.Preconditions; +import org.apache.paimon.utils.SnapshotManager; +import org.apache.paimon.utils.StringUtils; +import org.apache.paimon.utils.TagManager; + +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.SortedMap; +import java.util.UUID; +import java.util.function.Function; + +import static org.apache.spark.sql.types.DataTypes.LongType; +import static org.apache.spark.sql.types.DataTypes.StringType; + +/** + * Rollback to a snapshot or tag as the latest snapshot, without dropping the snapshots and tags + * created after it (unlike {@link RollbackProcedure}). Mirrors Flink's {@code + * rollback_to_as_latest} procedure. + */ +public class RollbackToAsLatestProcedure extends BaseProcedure { + + private static final String ROLLBACK_TO_AS_LATEST_TAG_PREFIX = "rollback-to-as-latest-"; + + private static final ProcedureParameter[] PARAMETERS = + new ProcedureParameter[] { + ProcedureParameter.required("table", StringType), + ProcedureParameter.optional("tag", StringType), + ProcedureParameter.optional("snapshot_id", LongType) + }; + + private static final StructType OUTPUT_TYPE = + new StructType( + new StructField[] { + new StructField( + "previous_snapshot_id", + DataTypes.LongType, + false, + Metadata.empty()), + new StructField( + "rolled_back_snapshot_id", + DataTypes.LongType, + false, + Metadata.empty()), + new StructField( + "current_snapshot_id", DataTypes.LongType, false, Metadata.empty()) + }); + + protected RollbackToAsLatestProcedure(TableCatalog tableCatalog) { + super(tableCatalog); + } + + @Override + public ProcedureParameter[] parameters() { + return PARAMETERS; + } + + @Override + public StructType outputType() { + return OUTPUT_TYPE; + } + + @Override + public InternalRow[] call(InternalRow args) { + Identifier tableIdent = toIdentifier(args.getString(0), PARAMETERS[0].name()); + String tagName = args.isNullAt(1) ? null : args.getString(1); + Long snapshotId = args.isNullAt(2) ? null : args.getLong(2); + + return modifyPaimonTableRefreshingCacheOnFailure( + tableIdent, + table -> { + FileStoreTable fileStoreTable = (FileStoreTable) table; + FileStore store = fileStoreTable.store(); + Snapshot latestSnapshot = store.snapshotManager().latestSnapshot(); + Preconditions.checkNotNull( + latestSnapshot, "Latest snapshot is null, can not roll back."); + + boolean hasTag = !StringUtils.isNullOrWhitespaceOnly(tagName); + boolean hasSnapshot = snapshotId != null; + Preconditions.checkArgument( + hasTag != hasSnapshot, + "Must specify exactly one of tag and snapshot_id."); + + TagManager tagManager = store.newTagManager(); + Tag targetTag; + Snapshot targetSnapshot; + if (hasTag) { + targetTag = tagManager.getOrThrow(tagName); + targetSnapshot = targetTag.trimToSnapshot(); + } else { + targetTag = null; + targetSnapshot = findSnapshot(store, tagManager, snapshotId); + } + + String createdRollbackTag = null; + boolean canDeleteCreatedTag = true; + String commitUser = ROLLBACK_TO_AS_LATEST_TAG_PREFIX + UUID.randomUUID(); + try { + if (!hasTag) { + createdRollbackTag = + createRollbackToAsLatestTag(tagManager, targetSnapshot); + targetTag = tagManager.getOrThrow(createdRollbackTag); + } + try (TableCommitImpl commit = fileStoreTable.newCommit(commitUser)) { + // The core rollback reads latest again, so another writer may change + // the rollback snapshot ID. An exception may also come after a + // successful commit (for example, from a callback). Keep the + // protection tag unless the commit returns false. + canDeleteCreatedTag = false; + boolean success = commit.rollbackToAsLatest(targetTag); + canDeleteCreatedTag = !success; + Preconditions.checkState( + success, + "Failed to roll back to snapshot %s as latest.", + targetSnapshot.id()); + } + } catch (Exception e) { + try { + if (createdRollbackTag != null && canDeleteCreatedTag) { + tagManager.deleteTag( + createdRollbackTag, + store.newTagDeletion(), + store.snapshotManager(), + Collections.emptyList()); + } + } catch (Exception cleanupException) { + e.addSuppressed(cleanupException); + } + throw new RuntimeException( + String.format( + "Failed to roll back to snapshot %s as latest.", + targetSnapshot.id()), + e); + } + + InternalRow outputRow = + newInternalRow( + latestSnapshot.id(), + targetSnapshot.id(), + store.snapshotManager().latestSnapshotId()); + return new InternalRow[] {outputRow}; + }); + } + + /** + * Like {@link #modifyPaimonTable} but also refreshes Spark's cached plans when {@code func} + * throws. {@code rollback_to_as_latest} can publish the rollback snapshot and then fail (for + * example a post-commit callback throws), so the table state is already durable; the shared + * success-only refresh would otherwise leave {@code CACHE TABLE} serving the pre-rollback data. + * The original failure is preserved; a refresh error is only added as suppressed. + */ + private InternalRow[] modifyPaimonTableRefreshingCacheOnFailure( + Identifier ident, Function func) { + SparkTable sparkTable = loadSparkTable(ident); + try { + InternalRow[] result = func.apply(sparkTable.getTable()); + refreshSparkCache(ident, sparkTable); + return result; + } catch (RuntimeException e) { + try { + refreshSparkCache(ident, sparkTable); + } catch (RuntimeException refreshError) { + e.addSuppressed(refreshError); + } + throw e; + } + } + + private String createRollbackToAsLatestTag(TagManager tagManager, Snapshot targetSnapshot) { + String tagName = + ROLLBACK_TO_AS_LATEST_TAG_PREFIX + targetSnapshot.id() + "-" + UUID.randomUUID(); + tagManager.createTag(targetSnapshot, tagName, null, Collections.emptyList(), false); + return tagName; + } + + private Snapshot findSnapshot(FileStore store, TagManager tagManager, long snapshotId) { + SnapshotManager snapshotManager = store.snapshotManager(); + if (snapshotManager.snapshotExists(snapshotId)) { + return snapshotManager.snapshot(snapshotId); + } + + SortedMap> tags = tagManager.tags(); + for (Map.Entry> entry : tags.entrySet()) { + if (entry.getKey().id() == snapshotId) { + return entry.getKey(); + } else if (entry.getKey().id() > snapshotId) { + break; + } + } + + throw new IllegalArgumentException( + String.format("Snapshot '%s' to roll back to doesn't exist.", snapshotId)); + } + + public static ProcedureBuilder builder() { + return new BaseProcedure.Builder() { + @Override + public RollbackToAsLatestProcedure doBuild() { + return new RollbackToAsLatestProcedure(tableCatalog()); + } + }; + } + + @Override + public String description() { + return "RollbackToAsLatestProcedure"; + } +} diff --git a/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala new file mode 100644 index 000000000000..e8b01d611b40 --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/scala/org/apache/paimon/spark/procedure/RollbackToAsLatestProcedureTest.scala @@ -0,0 +1,186 @@ +/* + * 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.paimon.spark.procedure + +import org.apache.paimon.data.{BinaryString, GenericRow} +import org.apache.paimon.manifest.ManifestCommittable +import org.apache.paimon.spark.PaimonSparkTestBase +import org.apache.paimon.table.{ExpireSnapshotsImpl, FileStoreTable} +import org.apache.paimon.table.sink.CommitCallback + +import org.apache.spark.sql.Row +import org.assertj.core.api.Assertions.assertThatThrownBy + +class RollbackToAsLatestProcedureTest extends PaimonSparkTestBase { + + test("Paimon Procedure: rollback to a snapshot as the latest snapshot") { + spark.sql("CREATE TABLE T (id INT, name STRING)") + spark.sql("INSERT INTO T VALUES (1, 'a')") + spark.sql("INSERT INTO T VALUES (2, 'b')") + spark.sql("INSERT INTO T VALUES (3, 'c')") + + val snapshotManager = loadTable("T").snapshotManager + + // Non-destructive: roll snapshot 1 forward as the new latest (snapshot 4); + // snapshots 2 and 3 stay, unlike the destructive `rollback`. + checkAnswer( + spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 1)"), + Row(3L, 1L, 4L) :: Nil) + assert(snapshotManager.snapshotExists(2)) + assert(snapshotManager.snapshotExists(3)) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "a") :: Nil) + + // Roll forward to snapshot 3 as the latest (snapshot 5). + checkAnswer( + spark.sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 3)"), + Row(4L, 3L, 5L) :: Nil) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "a") :: Row(2, "b") :: Row(3, "c") :: Nil) + } + + test("Paimon Procedure: rollback keeps the protection tag after a post-commit failure") { + Seq(false, true).foreach { + concurrentCommit => + withTable("T") { + createTableWithFailingCallback() + spark.sql("INSERT INTO T VALUES (1, 'original')") + // The restored file must not also belong to the snapshots retained before the rollback. + spark.sql("INSERT OVERWRITE T VALUES (2, 'replacement')") + spark.sql("INSERT INTO T VALUES (3, 'retained')") + + FailingRollbackCallback.concurrentCommit = concurrentCommit + FailingRollbackCallback.failRollbackCommit = true + try { + assertThatThrownBy(() => rollbackToSnapshot1()) + .hasStackTraceContaining("Injected post-commit callback failure") + } finally { + FailingRollbackCallback.reset() + } + + val table = loadTable("T") + val latest = table.snapshotManager().latestSnapshot() + assert(latest.id() == (if (concurrentCommit) 5L else 4L)) + assert(latest.commitUser().startsWith("rollback-to-as-latest-")) + val tags = table.tagManager().allTagNames() + assert(tags.size() == 1 && tags.get(0).startsWith("rollback-to-as-latest-1-")) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "original") :: Nil) + val expire = table.newExpireSnapshots().asInstanceOf[ExpireSnapshotsImpl] + expire.expireUntil(1, latest.id()) + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "original") :: Nil) + } + } + } + + test("Paimon Procedure: rollback removes the protection tag when it never started") { + createTableWithFailingCallback() + spark.sql("INSERT INTO T VALUES (1, 'original')") + + FailingRollbackCallback.failCommitCreation = true + try { + assertThatThrownBy(() => rollbackToSnapshot1()) + .hasStackTraceContaining("Injected commit creation failure") + } finally { + FailingRollbackCallback.reset() + } + + val table = loadTable("T") + assert(table.tagManager().allTagNames().isEmpty) + assert(table.snapshotManager().latestSnapshotId() == 1L) + } + + test("Paimon Procedure: rollback refreshes a cached table after a post-commit failure") { + createTableWithFailingCallback() + spark.sql("INSERT INTO T VALUES (1, 'original')") + spark.sql("INSERT OVERWRITE T VALUES (2, 'replacement')") + spark.sql("CACHE TABLE T") + checkAnswer(spark.sql("SELECT * FROM T"), Row(2, "replacement") :: Nil) + + FailingRollbackCallback.failRollbackCommit = true + try { + assertThatThrownBy(() => rollbackToSnapshot1()) + .hasStackTraceContaining("Injected post-commit callback failure") + } finally { + FailingRollbackCallback.reset() + } + + // Snapshot 3 (the rollback) is durable despite the callback failure, so the cached read must + // reflect snapshot 1's data rather than the stale pre-rollback replacement row. + checkAnswer(spark.sql("SELECT * FROM T"), Row(1, "original") :: Nil) + } + + private def createTableWithFailingCallback(): Unit = { + val callback = classOf[FailingRollbackCallback].getName + spark.sql( + s"CREATE TABLE T (id INT, name STRING) TBLPROPERTIES ('commit.callbacks' = '$callback')") + } + + private def rollbackToSnapshot1(): Unit = { + spark + .sql("CALL paimon.sys.rollback_to_as_latest(table => 'test.T', snapshot_id => 1)") + .collect() + } +} + +/** Injects rollback failures, optionally after another writer commits a snapshot first. */ +class FailingRollbackCallback extends CommitCallback { + + override def setTable(table: FileStoreTable): Unit = { + if (FailingRollbackCallback.failCommitCreation) { + throw new RuntimeException("Injected commit creation failure") + } + if (FailingRollbackCallback.concurrentCommit) { + // The procedure already read latest; the core rollback has not read it yet. + FailingRollbackCallback.concurrentCommit = false + val builder = table.newBatchWriteBuilder() + val write = builder.newWrite() + val commit = builder.newCommit() + try { + write.write(GenericRow.of(Int.box(4), BinaryString.fromString("concurrent"))) + commit.commit(write.prepareCommit()) + } finally { + write.close() + commit.close() + } + } + } + + override def call(context: CommitCallback.Context): Unit = { + if ( + FailingRollbackCallback.failRollbackCommit && + context.snapshot.commitUser().startsWith("rollback-to-as-latest-") + ) { + throw new RuntimeException("Injected post-commit callback failure") + } + } + + override def retry(committable: ManifestCommittable): Unit = {} + + override def close(): Unit = {} +} + +object FailingRollbackCallback { + @volatile var failRollbackCommit = false + @volatile var concurrentCommit = false + @volatile var failCommitCreation = false + + def reset(): Unit = { + failRollbackCommit = false + concurrentCommit = false + failCommitCreation = false + } +}