Invalidate table handle cache when column changed

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

@@ -412,28 +412,28 @@ public void setTableComment(ConnectorSession session, JdbcTableHandle handle, Op
     public void setColumnComment(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column, Optional<String> comment)
     {
         delegate.setColumnComment(session, handle, column, comment);
-        invalidateColumnsCache(handle.asPlainTable().getSchemaTableName());
+        invalidateTableCaches(handle.asPlainTable().getSchemaTableName());
     }
 
     @Override
     public void addColumn(ConnectorSession session, JdbcTableHandle handle, ColumnMetadata column)
     {
         delegate.addColumn(session, handle, column);
-        invalidateColumnsCache(handle.asPlainTable().getSchemaTableName());
+        invalidateTableCaches(handle.asPlainTable().getSchemaTableName());
     }
 
     @Override
     public void dropColumn(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column)
     {
         delegate.dropColumn(session, handle, column);
-        invalidateColumnsCache(handle.asPlainTable().getSchemaTableName());
+        invalidateTableCaches(handle.asPlainTable().getSchemaTableName());
     }
 
     @Override
     public void renameColumn(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle jdbcColumn, String newColumnName)
     {
         delegate.renameColumn(session, handle, jdbcColumn, newColumnName);
-        invalidateColumnsCache(handle.asPlainTable().getSchemaTableName());
+        invalidateTableCaches(handle.asPlainTable().getSchemaTableName());
     }
 
     @Override
@@ -578,6 +578,12 @@ CacheStats getTableNamesCacheStats()
         return tableNamesCache.stats();
     }
 
+    @VisibleForTesting
+    CacheStats getTableHandlesByNameCacheStats()
+    {
+        return tableHandlesByNameCache.stats();
+    }
+
     @VisibleForTesting
     CacheStats getTableHandlesByQueryCacheStats()
     {

plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestCachingJdbcClient.java

@@ -55,6 +55,7 @@
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
 import static io.trino.plugin.jdbc.TestCachingJdbcClient.CachingJdbcCache.STATISTICS_CACHE;
+import static io.trino.plugin.jdbc.TestCachingJdbcClient.CachingJdbcCache.TABLE_HANDLES_BY_NAME_CACHE;
 import static io.trino.plugin.jdbc.TestCachingJdbcClient.CachingJdbcCache.TABLE_HANDLES_BY_QUERY_CACHE;
 import static io.trino.spi.session.PropertyMetadata.stringProperty;
 import static io.trino.spi.testing.InterfaceTestUtils.assertAllMethodsOverridden;
@@ -273,6 +274,38 @@ public void testTableHandleOfQueryCached()
         dropTable(phantomTable);
     }
 
+    @Test
+    public void testTableHandleInvalidatedOnColumnsModifications()
+    {
+        JdbcTableHandle table = createTable(new SchemaTableName(schema, "a_table"));
+        JdbcColumnHandle existingColumn = addColumn(table, "a_column");
+
+        // warm-up cache
+        assertTableHandlesByNameCacheIsInvalidated(table);
+        JdbcColumnHandle newColumn = addColumn(cachingJdbcClient, table, "new_column");
+        assertTableHandlesByNameCacheIsInvalidated(table);
+        cachingJdbcClient.setColumnComment(SESSION, table, newColumn, Optional.empty());
+        assertTableHandlesByNameCacheIsInvalidated(table);
+        cachingJdbcClient.renameColumn(SESSION, table, newColumn, "new_column_name");
+        assertTableHandlesByNameCacheIsInvalidated(table);
+        cachingJdbcClient.dropColumn(SESSION, table, existingColumn);
+        assertTableHandlesByNameCacheIsInvalidated(table);
+
+        dropTable(table);
+    }
+
+    private void assertTableHandlesByNameCacheIsInvalidated(JdbcTableHandle table)
+    {
+        SchemaTableName tableName = table.asPlainTable().getSchemaTableName();
+
+        assertCacheStats(cachingJdbcClient, TABLE_HANDLES_BY_NAME_CACHE).misses(1).loads(1).afterRunning(() -> {
+            assertThat(cachingJdbcClient.getTableHandle(SESSION, tableName).orElseThrow()).isEqualTo(table);
+        });
+        assertCacheStats(cachingJdbcClient, TABLE_HANDLES_BY_NAME_CACHE).hits(1).afterRunning(() -> {
+            assertThat(cachingJdbcClient.getTableHandle(SESSION, tableName).orElseThrow()).isEqualTo(table);
+        });
+    }
+
     @Test
     public void testEmptyTableHandleIsCachedWhenCacheMissingIsTrue()
     {
@@ -304,6 +337,11 @@ private JdbcTableHandle createTable(SchemaTableName phantomTable)
         return jdbcClient.getTableHandle(SESSION, phantomTable).orElseThrow();
     }
 
+    private void dropTable(JdbcTableHandle tableHandle)
+    {
+        jdbcClient.dropTable(SESSION, tableHandle);
+    }
+
     private void dropTable(SchemaTableName phantomTable)
     {
         JdbcTableHandle tableHandle = jdbcClient.getTableHandle(SESSION, phantomTable).orElseThrow();
@@ -819,10 +857,15 @@ private JdbcColumnHandle addColumn(JdbcTableHandle tableHandle)
     }
 
     private JdbcColumnHandle addColumn(JdbcTableHandle tableHandle, String columnName)
+    {
+        return addColumn(jdbcClient, tableHandle, columnName);
+    }
+
+    private JdbcColumnHandle addColumn(JdbcClient client, JdbcTableHandle tableHandle, String columnName)
     {
         ColumnMetadata columnMetadata = new ColumnMetadata(columnName, INTEGER);
-        jdbcClient.addColumn(SESSION, tableHandle, columnMetadata);
-        return jdbcClient.getColumns(SESSION, tableHandle)
+        client.addColumn(SESSION, tableHandle, columnMetadata);
+        return client.getColumns(SESSION, tableHandle)
                 .stream()
                 .filter(jdbcColumnHandle -> jdbcColumnHandle.getColumnMetadata().equals(columnMetadata))
                 .findAny()
@@ -994,6 +1037,7 @@ public <T> T calling(Callable<T> callable)
     enum CachingJdbcCache
     {
         TABLE_NAMES_CACHE(CachingJdbcClient::getTableNamesCacheStats),
+        TABLE_HANDLES_BY_NAME_CACHE(CachingJdbcClient::getTableHandlesByNameCacheStats),
         TABLE_HANDLES_BY_QUERY_CACHE(CachingJdbcClient::getTableHandlesByQueryCacheStats),
         COLUMNS_CACHE(CachingJdbcClient::getColumnsCacheStats),
         STATISTICS_CACHE(CachingJdbcClient::getStatisticsCacheStats),

plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/TestingH2JdbcClient.java

@@ -109,6 +109,16 @@ public Optional<String> getTableComment(ResultSet resultSet)
         return Optional.empty();
     }
 
+    @Override
+    public void setColumnComment(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column, Optional<String> comment)
+    {
+        // Ignore (not fail) when comment is empty for testing purposes.
+        // however do not allow to set non-empty comment, not to have increased expectations from the invoking test
+        if (comment.isPresent()) {
+            throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting column comments");
+        }
+    }
+
     @Override
     public boolean supportsAggregationPushdown(ConnectorSession session, JdbcTableHandle table, List<AggregateFunction> aggregates, Map<String, ColumnHandle> assignments, List<List<ColumnHandle>> groupingSets)
     {