Skip to content

refactor(plugin-iceberg): Fix thread safety in Iceberg procedures #27341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
agrawalreetika merged 1 commit into from
Mar 16, 2026
Merged

refactor(plugin-iceberg): Fix thread safety in Iceberg procedures #27341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions ...eberg/src/main/java/com/facebook/presto/iceberg/procedure/FastForwardBranchProcedure.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import com.facebook.presto.iceberg.transaction.IcebergTransactionMetadata;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.SchemaTableName;
import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
import com.facebook.presto.spi.procedure.Procedure;
import com.facebook.presto.spi.procedure.Procedure.Argument;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -68,10 +69,12 @@ public Procedure get()

public void fastForwardToBranch(ConnectorSession clientSession, String schemaName, String tableName, String fromBranch, String targetBranch)
{
SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName);
IcebergTransactionMetadata metadata = metadataFactory.create();
Table icebergTable = getIcebergTable(metadata, clientSession, schemaTableName);
icebergTable.manageSnapshots().fastForwardBranch(fromBranch, targetBranch).commit();
metadata.commit();
try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(getClass().getClassLoader())) {
SchemaTableName schemaTableName = new SchemaTableName(schemaName, tableName);
IcebergTransactionMetadata metadata = metadataFactory.create();
Table icebergTable = getIcebergTable(metadata, clientSession, schemaTableName);
icebergTable.manageSnapshots().fastForwardBranch(fromBranch, targetBranch).commit();
metadata.commit();
}
}
}
13 changes: 8 additions & 5 deletions ...berg/src/main/java/com/facebook/presto/iceberg/procedure/RollbackToSnapshotProcedure.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import com.facebook.presto.iceberg.transaction.IcebergTransactionMetadata;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.SchemaTableName;
import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
import com.facebook.presto.spi.procedure.Procedure;
import com.facebook.presto.spi.procedure.Procedure.Argument;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -66,10 +67,12 @@ public Procedure get()

public void rollbackToSnapshot(ConnectorSession clientSession, String schema, String table, Long snapshotId)
{
SchemaTableName schemaTableName = new SchemaTableName(schema, table);
IcebergTransactionMetadata metadata = metadataFactory.create();
getIcebergTable(metadata, clientSession, schemaTableName)
.manageSnapshots().rollbackTo(snapshotId).commit();
metadata.commit();
try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(getClass().getClassLoader())) {
SchemaTableName schemaTableName = new SchemaTableName(schema, table);
IcebergTransactionMetadata metadata = metadataFactory.create();
getIcebergTable(metadata, clientSession, schemaTableName)
.manageSnapshots().rollbackTo(snapshotId).commit();
metadata.commit();
}
}
}
19 changes: 11 additions & 8 deletions ...berg/src/main/java/com/facebook/presto/iceberg/procedure/SetCurrentSnapshotProcedure.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import com.facebook.presto.iceberg.transaction.IcebergTransactionMetadata;
import com.facebook.presto.spi.ConnectorSession;
import com.facebook.presto.spi.SchemaTableName;
import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
import com.facebook.presto.spi.procedure.Procedure;
import com.facebook.presto.spi.procedure.Procedure.Argument;
import com.google.common.collect.ImmutableList;
Expand Down Expand Up @@ -71,14 +72,16 @@ public Procedure get()

public void setCurrentSnapshot(ConnectorSession clientSession, String schema, String table, Long snapshotId, String reference)
{
checkState((snapshotId != null && reference == null) || (snapshotId == null && reference != null),
"Either snapshot_id or reference must be provided, not both");
SchemaTableName schemaTableName = new SchemaTableName(schema, table);
IcebergTransactionMetadata metadata = metadataFactory.create();
Table icebergTable = getIcebergTable(metadata, clientSession, schemaTableName);
long targetSnapshotId = snapshotId != null ? snapshotId : getSnapshotIdFromReference(icebergTable, reference);
icebergTable.manageSnapshots().setCurrentSnapshot(targetSnapshotId).commit();
metadata.commit();
try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(getClass().getClassLoader())) {
checkState((snapshotId != null && reference == null) || (snapshotId == null && reference != null),
"Either snapshot_id or reference must be provided, not both");
SchemaTableName schemaTableName = new SchemaTableName(schema, table);
IcebergTransactionMetadata metadata = metadataFactory.create();
Table icebergTable = getIcebergTable(metadata, clientSession, schemaTableName);
long targetSnapshotId = snapshotId != null ? snapshotId : getSnapshotIdFromReference(icebergTable, reference);
icebergTable.manageSnapshots().setCurrentSnapshot(targetSnapshotId).commit();
metadata.commit();
}
}

private long getSnapshotIdFromReference(Table table, String refName)
Expand Down
Loading