From a9ab48a9bce70e2cc01d4b8640ba4a8181b45fb4 Mon Sep 17 00:00:00 2001 From: yangyicheng Date: Thu, 12 Oct 2023 10:44:48 +0800 Subject: [PATCH] Ensure that jdbc temp tables can be dropped if the insertions fail --- .../src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 index a27197faef24..c4a208d0ea65 100644 --- 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 @@ -879,7 +879,7 @@ protected void renameTable(ConnectorSession session, Connection connection, Stri quoted(catalogName, newRemoteSchemaName, newRemoteTableName))); } - private RemoteTableName constructPageSinkIdsTable(ConnectorSession session, Connection connection, JdbcOutputTableHandle handle, Set pageSinkIds) + private RemoteTableName constructPageSinkIdsTable(ConnectorSession session, Connection connection, JdbcOutputTableHandle handle, Set pageSinkIds, Closer closer) throws SQLException { verify(handle.getPageSinkIdColumnName().isPresent(), "Output table handle's pageSinkIdColumn is empty"); @@ -903,6 +903,7 @@ private RemoteTableName constructPageSinkIdsTable(ConnectorSession session, Conn LongWriteFunction pageSinkIdWriter = (LongWriteFunction) toWriteMapping(session, TRINO_PAGE_SINK_ID_COLUMN_TYPE).getWriteFunction(); execute(session, connection, pageSinkTableSql); + closer.register(() -> dropTable(session, pageSinkTable, true)); try (PreparedStatement statement = connection.prepareStatement(pageSinkInsertSql)) { int batchSize = 0; @@ -959,8 +960,7 @@ public void finishInsertTable(ConnectorSession session, JdbcOutputTableHandle ha quoted(temporaryTable)); if (handle.getPageSinkIdColumnName().isPresent()) { - RemoteTableName pageSinkTable = constructPageSinkIdsTable(session, connection, handle, pageSinkIds); - closer.register(() -> dropTable(session, pageSinkTable, true)); + RemoteTableName pageSinkTable = constructPageSinkIdsTable(session, connection, handle, pageSinkIds, closer); insertSql += format(" WHERE EXISTS (SELECT 1 FROM %s page_sink_table WHERE page_sink_table.%s = temp_table.%s)", quoted(pageSinkTable),