Retry JdbcClient::execute method

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

@@ -1609,7 +1609,8 @@ protected void renameSchema(ConnectorSession session, Connection connection, Str
         execute(session, connection, "ALTER SCHEMA " + quoted(remoteSchemaName) + " RENAME TO " + quoted(newRemoteSchemaName));
     }
 
-    protected void execute(ConnectorSession session, String query)
+    @Override
+    public void execute(ConnectorSession session, String query)
     {
         try (Connection connection = connectionFactory.openConnection(session)) {
             execute(session, connection, query);

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

@@ -273,6 +273,12 @@ public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcP
         return delegate.getConnection(session, split, procedureHandle);
     }
 
+    @Override
+    public void execute(ConnectorSession session, String query)
+    {
+        delegate.execute(session, query);
+    }
+
     @Override
     public void abortReadConnection(Connection connection, ResultSet resultSet)
             throws SQLException

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

@@ -197,6 +197,12 @@ public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcP
         return delegate().getConnection(session, split, procedureHandle);
     }
 
+    @Override
+    public void execute(ConnectorSession session, String query)
+    {
+        delegate().execute(session, query);
+    }
+
     @Override
     public void abortReadConnection(Connection connection, ResultSet resultSet)
             throws SQLException

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

@@ -119,6 +119,8 @@ Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcTableHan
     Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcProcedureHandle procedureHandle)
             throws SQLException;
 
+    void execute(ConnectorSession session, String query);
+
     default void abortReadConnection(Connection connection, ResultSet resultSet)
             throws SQLException
     {

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

@@ -202,6 +202,13 @@ public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcP
         return delegate.getConnection(session, split, procedureHandle);
     }
 
+    @Override
+    public void execute(ConnectorSession session, String query)
+    {
+        // we do a nested retry as opening a connection is already retried, however it is better to retry on intermittent issue than fail
+        retry(policy, () -> delegate.execute(session, query));
+    }
+
     @Override
     public void abortReadConnection(Connection connection, ResultSet resultSet)
             throws SQLException

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

@@ -42,9 +42,11 @@ public final class JdbcClientStats
     private final JdbcApiStats finishMergeTable = new JdbcApiStats();
     private final JdbcApiStats getColumns = new JdbcApiStats();
     private final JdbcApiStats getAllTableComments = new JdbcApiStats();
+    private final JdbcApiStats getConnection = new JdbcApiStats();
     private final JdbcApiStats getConnectionWithHandle = new JdbcApiStats();
     private final JdbcApiStats getConnectionWithSplit = new JdbcApiStats();
     private final JdbcApiStats getConnectionWithProcedure = new JdbcApiStats();
+    private final JdbcApiStats execute = new JdbcApiStats();
     private final JdbcApiStats getPreparedStatement = new JdbcApiStats();
     private final JdbcApiStats getSchemaNames = new JdbcApiStats();
     private final JdbcApiStats getSplits = new JdbcApiStats();
@@ -241,6 +243,13 @@ public JdbcApiStats getGetAllTableComments()
         return getAllTableComments;
     }
 
+    @Managed
+    @Nested
+    public JdbcApiStats getGetConnection()
+    {
+        return getConnection;
+    }
+
     @Managed
     @Nested
     public JdbcApiStats getGetConnectionWithHandle()
@@ -262,6 +271,13 @@ public JdbcApiStats getGetConnectionWithProcedure()
         return getConnectionWithProcedure;
     }
 
+    @Managed
+    @Nested
+    public JdbcApiStats getExecute()
+    {
+        return execute;
+    }
+
     @Managed
     @Nested
     public JdbcApiStats getGetPreparedStatement()

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

@@ -220,6 +220,12 @@ public Connection getConnection(ConnectorSession session, JdbcSplit split, JdbcP
         return stats.getGetConnectionWithProcedure().wrap(() -> delegate().getConnection(session, split, procedureHandle));
     }
 
+    @Override
+    public void execute(ConnectorSession session, String query)
+    {
+        stats.getExecute().wrap(() -> delegate().execute(session, query));
+    }
+
     @Override
     public void abortReadConnection(Connection connection, ResultSet resultSet)
             throws SQLException
@@ -414,7 +420,7 @@ public String buildInsertSql(JdbcOutputTableHandle handle, List<WriteFunction> c
     public Connection getConnection(ConnectorSession session)
             throws SQLException
     {
-        return stats.getGetConnectionWithHandle().wrap(() -> delegate().getConnection(session));
+        return stats.getGetConnection().wrap(() -> delegate().getConnection(session));
     }
 
     @Override