Merge branch 'trinodb:master' into hsqldb

Author: prrvchr

plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java

@@ -165,32 +165,40 @@ public synchronized void renameSchema(ConnectorSession session, String source, S
         }
         schemas.add(target);
 
-        for (Map.Entry<SchemaTableName, Long> table : tableIds.entrySet()) {
+        Map<SchemaTableName, Long> newTableIds = new HashMap<>();
+        for (Iterator<Map.Entry<SchemaTableName, Long>> iterator = tableIds.entrySet().iterator(); iterator.hasNext(); ) {
+            Map.Entry<SchemaTableName, Long> table = iterator.next();
             if (table.getKey().getSchemaName().equals(source)) {
-                tableIds.remove(table.getKey());
-                tableIds.put(new SchemaTableName(target, table.getKey().getTableName()), table.getValue());
+                iterator.remove();
+                newTableIds.put(new SchemaTableName(target, table.getKey().getTableName()), table.getValue());
             }
         }
+        tableIds.putAll(newTableIds);
 
-        for (TableInfo table : tables.values()) {
-            if (table.schemaName().equals(source)) {
-                tables.put(table.id(), new TableInfo(table.id(), target, table.tableName(), table.columns(), false, table.dataFragments(), table.comment()));
-            }
-        }
+        tables.replaceAll((tableId, table) ->
+                table.schemaName().equals(source)
+                        ? new TableInfo(tableId, target, table.tableName(), table.columns(), table.truncated(), table.dataFragments(), table.comment())
+                        : table);
 
-        for (Map.Entry<SchemaTableName, ConnectorViewDefinition> view : views.entrySet()) {
+        Map<SchemaTableName, ConnectorViewDefinition> newViews = new HashMap<>();
+        for (Iterator<Map.Entry<SchemaTableName, ConnectorViewDefinition>> iterator = views.entrySet().iterator(); iterator.hasNext(); ) {
+            Map.Entry<SchemaTableName, ConnectorViewDefinition> view = iterator.next();
             if (view.getKey().getSchemaName().equals(source)) {
-                views.remove(view.getKey());
-                views.put(new SchemaTableName(target, view.getKey().getTableName()), view.getValue());
+                iterator.remove();
+                newViews.put(new SchemaTableName(target, view.getKey().getTableName()), view.getValue());
             }
         }
+        views.putAll(newViews);
 
-        for (Map.Entry<SchemaFunctionName, Map<String, LanguageFunction>> function : functions.entrySet()) {
+        Map<SchemaFunctionName, Map<String, LanguageFunction>> newFunctions = new HashMap<>();
+        for (Iterator<Map.Entry<SchemaFunctionName, Map<String, LanguageFunction>>> iterator = functions.entrySet().iterator(); iterator.hasNext(); ) {
+            Map.Entry<SchemaFunctionName, Map<String, LanguageFunction>> function = iterator.next();
             if (function.getKey().getSchemaName().equals(source)) {
-                functions.remove(function.getKey());
-                functions.put(new SchemaFunctionName(target, function.getKey().getFunctionName()), function.getValue());
+                iterator.remove();
+                newFunctions.put(new SchemaFunctionName(target, function.getKey().getFunctionName()), function.getValue());
             }
         }
+        functions.putAll(newFunctions);
     }
 
     @GuardedBy("this")

plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryMetadata.java

@@ -32,7 +32,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
+import java.util.stream.IntStream;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.trino.spi.StandardErrorCode.ALREADY_EXISTS;
 import static io.trino.spi.StandardErrorCode.NOT_FOUND;
 import static io.trino.spi.connector.RetryMode.NO_RETRIES;
@@ -344,6 +347,33 @@ public void testRenameTable()
         assertThat(metadata.listTables(SESSION, Optional.of("test_different_schema"))).isEqualTo(ImmutableList.of(differentSchemaTableName));
     }
 
+    @Test
+    public void testRenameSchema()
+    {
+        Set<SchemaTableName> tableNames = IntStream.range(1, 10)
+                .mapToObj(idx -> new SchemaTableName("test_schema_to_be_renamed", "test_table_" + idx))
+                .collect(toImmutableSet());
+        MemoryMetadata metadata = createMetadata();
+        metadata.createSchema(SESSION, "test_schema_to_be_renamed", ImmutableMap.of(), new TrinoPrincipal(USER, SESSION.getUser()));
+        tableNames.forEach(tableName -> {
+            ConnectorOutputTableHandle table = metadata.beginCreateTable(
+                    SESSION,
+                    new ConnectorTableMetadata(tableName, ImmutableList.of(), ImmutableMap.of()),
+                    Optional.empty(),
+                    NO_RETRIES,
+                    false);
+            metadata.finishCreateTable(SESSION, table, ImmutableList.of(), ImmutableList.of());
+        });
+
+        // rename schema
+        Set<SchemaTableName> renamedTableNames = tableNames.stream()
+                .map(tableName -> new SchemaTableName("test_schema", tableName.getTableName()))
+                .collect(toImmutableSet());
+        metadata.renameSchema(SESSION, "test_schema_to_be_renamed", "test_schema");
+        assertThat(metadata.listTables(SESSION, Optional.of("test_schema_to_be_renamed"))).isEmpty();
+        assertThat(metadata.listTables(SESSION, Optional.of("test_schema"))).containsAll(renamedTableNames);
+    }
+
     private static void assertNoTables(MemoryMetadata metadata)
     {
         assertThat(metadata.listTables(SESSION, Optional.empty()))