Add rollback action for destination table for CTAS statements

findinpath · findinpath · commit 88664e3bbb3c · 2025-12-23T16:13:32.000+01:00
While CREATE TABLE is an atomic operation,
CREATE TABLE AS SELECT statement happens in two phases
for JDBC connectors from Trino:

- creating the table
- inserting content into the table

If the insertion of the data fails, the CTAS may leave
residual empty tables behind.

The JDBC connectors, in dealing
with CREATE TABLE AS SELECT scenarios,
remove the destination table in case of dealing with exceptions
while inserting the content.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
@@ -829,12 +829,13 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
         try {
             if (shouldUseFaultTolerantExecution(session)) {
                 // Create the target table
-                createTable(session, tableMetadata);
+                JdbcOutputTableHandle destinationTableHandle = createTable(session, tableMetadata, tableMetadata.getTable().getTableName());
+                rollbackActionConsumer.accept(() -> rollbackCreateDestinationTable(session, destinationTableHandle));
                 // Create the temporary table
                 ColumnMetadata pageSinkIdColumn = getPageSinkIdColumn(
                         tableMetadata.getColumns().stream().map(ColumnMetadata::getName).toList());
@@ -849,6 +850,11 @@ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, Connecto
         }
     }
 
+    protected void rollbackCreateDestinationTable(ConnectorSession session, JdbcOutputTableHandle handle)
+    {
+        dropTable(session, handle.getRemoteTableName(), false);
+    }
+
     protected JdbcOutputTableHandle createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, String targetTableName)
             throws SQLException
     {
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java
@@ -600,9 +600,9 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
-        return delegate.beginCreateTable(session, tableMetadata);
+        return delegate.beginCreateTable(session, tableMetadata, rollbackActionConsumer);
     }
 
     @Override
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DefaultJdbcMetadata.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/DefaultJdbcMetadata.java
@@ -137,7 +137,7 @@ public class DefaultJdbcMetadata
     private final boolean precalculateStatisticsForPushdown;
     private final Set<JdbcQueryEventListener> jdbcQueryEventListeners;
 
-    private final List<Runnable> rollbackActions = new ArrayList<>();
+    protected final List<Runnable> rollbackActions = new ArrayList<>();
 
     public DefaultJdbcMetadata(
             JdbcClient jdbcClient,
@@ -1199,7 +1199,7 @@ public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, Con
         if (replace) {
             throw new TrinoException(NOT_SUPPORTED, "This connector does not support replacing tables");
         }
-        JdbcOutputTableHandle handle = jdbcClient.beginCreateTable(session, tableMetadata);
+        JdbcOutputTableHandle handle = jdbcClient.beginCreateTable(session, tableMetadata, rollbackActions::add);
         rollbackActions.add(() -> jdbcClient.rollbackCreateTable(session, handle));
         return handle;
     }
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java
@@ -265,9 +265,9 @@ public Optional<PreparedQuery> legacyImplementJoin(
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
-        return delegate().beginCreateTable(session, tableMetadata);
+        return delegate().beginCreateTable(session, tableMetadata, rollbackActionConsumer);
     }
 
     @Override
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java
@@ -210,7 +210,7 @@ default void setTableProperties(ConnectorSession session, JdbcTableHandle handle
 
     void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata);
 
-    JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata);
+    JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer);
 
     void commitCreateTable(ConnectorSession session, JdbcOutputTableHandle handle, Set<Long> pageSinkIds);
 
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingJdbcClient.java
@@ -368,10 +368,10 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
         // no retrying as it could be not idempotent operation
-        return delegate.beginCreateTable(session, tableMetadata);
+        return delegate.beginCreateTable(session, tableMetadata, rollbackActionConsumer);
     }
 
     @Override
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java
@@ -353,9 +353,9 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
-        return stats.getBeginCreateTable().wrap(() -> delegate().beginCreateTable(session, tableMetadata));
+        return stats.getBeginCreateTable().wrap(() -> delegate().beginCreateTable(session, tableMetadata, rollbackActionConsumer));
     }
 
     @Override
diff --git a/plugin/trino-druid/src/main/java/io/trino/plugin/druid/DruidJdbcClient.java b/plugin/trino-druid/src/main/java/io/trino/plugin/druid/DruidJdbcClient.java
@@ -74,6 +74,7 @@
 import java.util.Set;
 import java.util.TimeZone;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Iterables.getOnlyElement;
@@ -476,7 +477,7 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
         throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with data");
     }
diff --git a/plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ExasolClient.java b/plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ExasolClient.java
@@ -56,6 +56,7 @@
 import java.util.OptionalLong;
 import java.util.Set;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 
 import static io.trino.plugin.jdbc.StandardColumnMappings.bigintColumnMapping;
 import static io.trino.plugin.jdbc.StandardColumnMappings.booleanColumnMapping;
@@ -314,7 +315,7 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
         throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with data");
     }
diff --git a/plugin/trino-ignite/src/main/java/io/trino/plugin/ignite/IgniteClient.java b/plugin/trino-ignite/src/main/java/io/trino/plugin/ignite/IgniteClient.java
@@ -83,6 +83,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.isNullOrEmpty;
@@ -381,7 +382,7 @@ private static void checkDateValue(long epochDay)
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
         if (tableMetadata.getComment().isPresent()) {
             throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with table comment");
@@ -417,13 +418,15 @@ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, Connecto
 
         try (Connection connection = connectionFactory.openConnection(session)) {
             execute(session, connection, sql);
-
-            return new IgniteOutputTableHandle(
+            IgniteOutputTableHandle destinationTableHandle = new IgniteOutputTableHandle(
                     new RemoteTableName(Optional.empty(), Optional.of(schemaTableName.getSchemaName()), schemaTableName.getTableName()),
                     columnNames,
                     columnTypes.build(),
                     Optional.empty(),
                     primaryKeys.isEmpty() ? Optional.of(IGNITE_DUMMY_ID) : Optional.empty());
+            rollbackActionConsumer.accept(() -> rollbackCreateDestinationTable(session, destinationTableHandle));
+            return destinationTableHandle;
+
         }
         catch (SQLException e) {
             throw new TrinoException(JDBC_ERROR, e);
diff --git a/plugin/trino-ignite/src/main/java/io/trino/plugin/ignite/IgniteMetadata.java b/plugin/trino-ignite/src/main/java/io/trino/plugin/ignite/IgniteMetadata.java
@@ -220,7 +220,7 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe
         if (saveMode == REPLACE) {
             throw new TrinoException(NOT_SUPPORTED, "This connector does not support replacing tables");
         }
-        igniteClient.beginCreateTable(session, tableMetadata);
+        igniteClient.beginCreateTable(session, tableMetadata, _ -> {});
     }
 
     @Override
@@ -232,7 +232,7 @@ public ConnectorOutputTableHandle beginCreateTable(ConnectorSession session, Con
         if (replace) {
             throw new TrinoException(NOT_SUPPORTED, "This connector does not support replacing tables");
         }
-        return igniteClient.beginCreateTable(session, tableMetadata);
+        return igniteClient.beginCreateTable(session, tableMetadata, rollbackActions::add);
     }
 
     @Override
diff --git a/plugin/trino-sqlserver/src/main/java/io/trino/plugin/sqlserver/SqlServerClient.java b/plugin/trino-sqlserver/src/main/java/io/trino/plugin/sqlserver/SqlServerClient.java
@@ -129,6 +129,7 @@
 import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.function.BiFunction;
+import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
@@ -363,9 +364,9 @@ protected void dropSchema(ConnectorSession session, Connection connection, Strin
     }
 
     @Override
-    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)
     {
-        JdbcOutputTableHandle table = super.beginCreateTable(session, tableMetadata);
+        JdbcOutputTableHandle table = super.beginCreateTable(session, tableMetadata, rollbackActionConsumer);
         enableTableLockOnBulkLoadTableOption(session, table);
         return table;
     }

Original file line number	Diff line number	Diff line change
`@@ -829,12 +829,13 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe`
`829`	`829`	`}`
`830`	`830`
`831`	`831`	`@Override`
`832`		`- public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)`
	`832`	`+ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)`
`833`	`833`	`{`
`834`	`834`	`try {`
`835`	`835`	`if (shouldUseFaultTolerantExecution(session)) {`
`836`	`836`	`// Create the target table`
`837`		`- createTable(session, tableMetadata);`
	`837`	`+ JdbcOutputTableHandle destinationTableHandle = createTable(session, tableMetadata, tableMetadata.getTable().getTableName());`
	`838`	`+ rollbackActionConsumer.accept(() -> rollbackCreateDestinationTable(session, destinationTableHandle));`
`838`	`839`	`// Create the temporary table`
`839`	`840`	`ColumnMetadata pageSinkIdColumn = getPageSinkIdColumn(`
`840`	`841`	`tableMetadata.getColumns().stream().map(ColumnMetadata::getName).toList());`
`@@ -849,6 +850,11 @@ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, Connecto`
`849`	`850`	`}`
`850`	`851`	`}`
`851`	`852`
	`853`	`+ protected void rollbackCreateDestinationTable(ConnectorSession session, JdbcOutputTableHandle handle)`
	`854`	`+ {`
	`855`	`+ dropTable(session, handle.getRemoteTableName(), false);`
	`856`	`+ }`
	`857`	`+`
`852`	`858`	`protected JdbcOutputTableHandle createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, String targetTableName)`
`853`	`859`	`throws SQLException`
`854`	`860`	`{`
Original file line number	Diff line number	Diff line change
`@@ -600,9 +600,9 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe`
`600`	`600`	`}`
`601`	`601`
`602`	`602`	`@Override`
`603`		`- public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)`
	`603`	`+ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)`
`604`	`604`	`{`
`605`		`- return delegate.beginCreateTable(session, tableMetadata);`
	`605`	`+ return delegate.beginCreateTable(session, tableMetadata, rollbackActionConsumer);`
`606`	`606`	`}`
`607`	`607`
`608`	`608`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -265,9 +265,9 @@ public Optional<PreparedQuery> legacyImplementJoin(`
`265`	`265`	`}`
`266`	`266`
`267`	`267`	`@Override`
`268`		`- public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)`
	`268`	`+ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)`
`269`	`269`	`{`
`270`		`- return delegate().beginCreateTable(session, tableMetadata);`
	`270`	`+ return delegate().beginCreateTable(session, tableMetadata, rollbackActionConsumer);`
`271`	`271`	`}`
`272`	`272`
`273`	`273`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -368,10 +368,10 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe`
`368`	`368`	`}`
`369`	`369`
`370`	`370`	`@Override`
`371`		`- public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)`
	`371`	`+ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)`
`372`	`372`	`{`
`373`	`373`	`// no retrying as it could be not idempotent operation`
`374`		`- return delegate.beginCreateTable(session, tableMetadata);`
	`374`	`+ return delegate.beginCreateTable(session, tableMetadata, rollbackActionConsumer);`
`375`	`375`	`}`
`376`	`376`
`377`	`377`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -353,9 +353,9 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe`
`353`	`353`	`}`
`354`	`354`
`355`	`355`	`@Override`
`356`		`- public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)`
	`356`	`+ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)`
`357`	`357`	`{`
`358`		`- return stats.getBeginCreateTable().wrap(() -> delegate().beginCreateTable(session, tableMetadata));`
	`358`	`+ return stats.getBeginCreateTable().wrap(() -> delegate().beginCreateTable(session, tableMetadata, rollbackActionConsumer));`
`359`	`359`	`}`
`360`	`360`
`361`	`361`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,7 @@`
`74`	`74`	`import java.util.Set;`
`75`	`75`	`import java.util.TimeZone;`
`76`	`76`	`import java.util.function.BiFunction;`
	`77`	`+import java.util.function.Consumer;`
`77`	`78`
`78`	`79`	`import static com.google.common.base.Preconditions.checkArgument;`
`79`	`80`	`import static com.google.common.collect.Iterables.getOnlyElement;`
`@@ -476,7 +477,7 @@ public void createTable(ConnectorSession session, ConnectorTableMetadata tableMe`
`476`	`477`	`}`
`477`	`478`
`478`	`479`	`@Override`
`479`		`- public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)`
	`480`	`+ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)`
`480`	`481`	`{`
`481`	`482`	`throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with data");`
`482`	`483`	`}`
Original file line number	Diff line number	Diff line change
`@@ -56,6 +56,7 @@`
`56`	`56`	`import java.util.OptionalLong;`
`57`	`57`	`import java.util.Set;`
`58`	`58`	`import java.util.function.BiFunction;`
	`59`	`+import java.util.function.Consumer;`
`59`	`60`
`60`	`61`	`import static io.trino.plugin.jdbc.StandardColumnMappings.bigintColumnMapping;`
`61`	`62`	`import static io.trino.plugin.jdbc.StandardColumnMappings.booleanColumnMapping;`
`@@ -314,7 +315,7 @@ public WriteMapping toWriteMapping(ConnectorSession session, Type type)`
`314`	`315`	`}`
`315`	`316`
`316`	`317`	`@Override`
`317`		`- public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata)`
	`318`	`+ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, ConnectorTableMetadata tableMetadata, Consumer<Runnable> rollbackActionConsumer)`
`318`	`319`	`{`
`319`	`320`	`throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables with data");`
`320`	`321`	`}`