Support renaming table in MongoDB

Cherry-pick of trinodb/trino#11423 Support renaming table in MongoDB Co-authored-by: Yuya Ebihara <ebyhry@gmail.com>
obinnaokechukwu · Jun 25, 2022 · c61c2e4 · c61c2e4
1 parent e53d846
commit c61c2e4
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 1 deletion.
diff --git a/presto-docs/src/main/sphinx/connector/mongodb.rst b/presto-docs/src/main/sphinx/connector/mongodb.rst
@@ -261,3 +261,11 @@ MongoDB collection has the special field ``_id``. The connector tries to follow
 .. note::
 
     Unfortunately, there is no way to represent ``_id`` fields more fancy like ``55b151633864d6438c61a9ce``.
+
+SQL support
+-----------
+
+ALTER TABLE
+^^^^^^^^^^^
+
+The connector supports ``ALTER TABLE RENAME TO`` operation. Other uses of ``ALTER TABLE`` are not supported.
diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
@@ -212,7 +212,8 @@ public void dropTable(ConnectorSession session, ConnectorTableHandle tableHandle
     @Override
     public void renameTable(ConnectorSession session, ConnectorTableHandle tableHandle, SchemaTableName newTableName)
     {
-        throw new UnsupportedOperationException();
+        MongoTableHandle table = (MongoTableHandle) tableHandle;
+        mongoSession.renameTable(table.getSchemaTableName(), newTableName);
     }
 
     @Override

diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
@@ -37,6 +37,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.mongodb.MongoClient;
+import com.mongodb.MongoNamespace;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoCursor;
@@ -239,6 +240,29 @@ public MongoCursor<Document> execute(MongoSplit split, List<MongoColumnHandle> c
         return iterable.iterator();
     }
 
+    public void renameTable(SchemaTableName oldName, SchemaTableName newName)
+    {
+        String oldSchemaName = oldName.getSchemaName();
+        String oldTableName = oldName.getTableName();
+        String newSchemaName = newName.getSchemaName();
+
+        // Schema collection should always have the source table definition
+        MongoCollection<Document> oldSchema = client.getDatabase(oldSchemaName).getCollection(schemaCollection);
+        Document tableDefinition = oldSchema.findOneAndDelete(new Document(TABLE_NAME_KEY, oldTableName));
+        requireNonNull(tableDefinition, "Table definition not found in schema collection: " + oldTableName);
+
+        MongoCollection<Document> newSchema = client.getDatabase(newSchemaName).getCollection(schemaCollection);
+        tableDefinition.append(TABLE_NAME_KEY, newName.getTableName());
+        newSchema.insertOne(tableDefinition);
+
+        // Need to check explicitly because the old collection may not exist when it doesn't have any data
+        if (collectionExists(client.getDatabase(oldSchemaName), oldTableName)) {
+            getCollection(oldName).renameCollection(new MongoNamespace(newSchemaName, newName.getTableName()));
+        }
+
+        tableCache.invalidate(oldName);
+    }
+
     @VisibleForTesting
     static Document buildQuery(TupleDomain<ColumnHandle> tupleDomain)
     {

diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
@@ -19,6 +19,7 @@
 import com.facebook.presto.tests.AbstractTestIntegrationSmokeTest;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.mongodb.client.MongoCollection;
 import org.bson.Document;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -253,4 +254,20 @@ private void assertOneNotNullResult(String query)
         assertEquals(results.getMaterializedRows().get(0).getFieldCount(), 1);
         assertNotNull(results.getMaterializedRows().get(0).getField(0));
     }
+
+    @Test
+    public void testRenameTable()
+    {
+        assertUpdate("CREATE TABLE test_rename.tmp_rename_table (value bigint)");
+        MongoCollection<Document> collection = mongoQueryRunner.getMongoClient().getDatabase("test_rename").getCollection("tmp_rename_table");
+        collection.insertOne(new Document(ImmutableMap.of("value", 1)));
+        assertQuery("SHOW TABLES IN test_rename", "SELECT 'tmp_rename_table'");
+        assertQuery("SELECT value FROM test_rename.tmp_rename_table", "SELECT 1");
+
+        assertUpdate("ALTER TABLE test_rename.tmp_rename_table RENAME TO test_rename.tmp_rename_new_table");
+
+        assertQuery("SHOW TABLES IN test_rename", "SELECT 'tmp_rename_new_table'");
+        assertQuery("SELECT value FROM test_rename.tmp_rename_new_table", "SELECT 1");
+        assertUpdate("DROP TABLE test_rename.tmp_rename_new_table");
+    }
 }