Expose rollbacks for the creation tables in StatisticsAwareJdbcClient

Author: findinpath

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java

@@ -851,7 +851,8 @@ public JdbcOutputTableHandle beginCreateTable(ConnectorSession session, Connecto
         }
     }
 
-    protected void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
+    @Override
+    public void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
     {
         dropTable(session, remoteTableName, false);
     }

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/CachingJdbcClient.java

@@ -447,6 +447,12 @@ public void dropTable(ConnectorSession session, JdbcTableHandle jdbcTableHandle)
         invalidateTableCaches(jdbcTableHandle.asPlainTable().getSchemaTableName());
     }
 
+    @Override
+    public void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
+    {
+        delegate.rollbackDestinationTableCreation(session, remoteTableName);
+    }
+
     @Override
     public void rollbackTemporaryTableCreation(ConnectorSession session, JdbcOutputTableHandle handle)
     {

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/ForwardingJdbcClient.java

@@ -311,6 +311,12 @@ public void dropTable(ConnectorSession session, JdbcTableHandle jdbcTableHandle)
         delegate().dropTable(session, jdbcTableHandle);
     }
 
+    @Override
+    public void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
+    {
+        delegate().rollbackDestinationTableCreation(session, remoteTableName);
+    }
+
     @Override
     public void rollbackTemporaryTableCreation(ConnectorSession session, JdbcOutputTableHandle handle)
     {

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/JdbcClient.java

@@ -235,6 +235,8 @@ default void finishMerge(ConnectorSession session, JdbcMergeTableHandle handle,
 
     void dropTable(ConnectorSession session, JdbcTableHandle jdbcTableHandle);
 
+    void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName);
+
     void rollbackTemporaryTableCreation(ConnectorSession session, JdbcOutputTableHandle handle);
 
     boolean supportsRetries();

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/RetryingJdbcClient.java

@@ -421,6 +421,13 @@ public void dropTable(ConnectorSession session, JdbcTableHandle jdbcTableHandle)
         delegate.dropTable(session, jdbcTableHandle);
     }
 
+    @Override
+    public void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
+    {
+        // no retrying as it could be not idempotent operation
+        delegate.rollbackDestinationTableCreation(session, remoteTableName);
+    }
+
     @Override
     public void rollbackTemporaryTableCreation(ConnectorSession session, JdbcOutputTableHandle handle)
     {

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/JdbcClientStats.java

@@ -61,6 +61,7 @@ public final class JdbcClientStats
     private final JdbcApiStats dropNotNullConstraint = new JdbcApiStats();
     private final JdbcApiStats renameTable = new JdbcApiStats();
     private final JdbcApiStats setTableProperties = new JdbcApiStats();
+    private final JdbcApiStats rollbackDestinationTableCreation = new JdbcApiStats();
     private final JdbcApiStats rollbackTemporaryTableCreation = new JdbcApiStats();
     private final JdbcApiStats schemaExists = new JdbcApiStats();
     private final JdbcApiStats toTrinoType = new JdbcApiStats();
@@ -376,6 +377,13 @@ public JdbcApiStats getSetTableProperties()
         return setTableProperties;
     }
 
+    @Managed
+    @Nested
+    public JdbcApiStats getRollbackDestinationTableCreation()
+    {
+        return rollbackDestinationTableCreation;
+    }
+
     @Managed
     @Nested
     public JdbcApiStats getRollbackTemporaryTableCreation()

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/jmx/StatisticsAwareJdbcClient.java

@@ -399,6 +399,12 @@ public void dropTable(ConnectorSession session, JdbcTableHandle jdbcTableHandle)
         stats.getDropTable().wrap(() -> delegate().dropTable(session, jdbcTableHandle));
     }
 
+    @Override
+    public void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
+    {
+        stats.getRollbackDestinationTableCreation().wrap(() -> delegate().rollbackDestinationTableCreation(session, remoteTableName));
+    }
+
     @Override
     public void rollbackTemporaryTableCreation(ConnectorSession session, JdbcOutputTableHandle handle)
     {

plugin/trino-druid/src/main/java/io/trino/plugin/druid/DruidJdbcClient.java

@@ -542,6 +542,12 @@ public void truncateTable(ConnectorSession session, JdbcTableHandle handle)
         throw new TrinoException(NOT_SUPPORTED, "This connector does not support truncating tables");
     }
 
+    @Override
+    public void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
+    {
+        throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables");
+    }
+
     @Override
     public void rollbackTemporaryTableCreation(ConnectorSession session, JdbcOutputTableHandle handle)
     {

plugin/trino-exasol/src/main/java/io/trino/plugin/exasol/ExasolClient.java

@@ -31,6 +31,7 @@
 import io.trino.plugin.jdbc.LongReadFunction;
 import io.trino.plugin.jdbc.LongWriteFunction;
 import io.trino.plugin.jdbc.QueryBuilder;
+import io.trino.plugin.jdbc.RemoteTableName;
 import io.trino.plugin.jdbc.SliceReadFunction;
 import io.trino.plugin.jdbc.SliceWriteFunction;
 import io.trino.plugin.jdbc.WriteFunction;
@@ -159,6 +160,12 @@ public void dropTable(ConnectorSession session, JdbcTableHandle handle)
         throw new TrinoException(NOT_SUPPORTED, "This connector does not support dropping tables");
     }
 
+    @Override
+    public void rollbackDestinationTableCreation(ConnectorSession session, RemoteTableName remoteTableName)
+    {
+        throw new TrinoException(NOT_SUPPORTED, "This connector does not support creating tables");
+    }
+
     @Override
     public void rollbackTemporaryTableCreation(ConnectorSession session, JdbcOutputTableHandle handle)
     {