Skip to content

Commit

Permalink
Adding alter support for mongodb
Browse files Browse the repository at this point in the history
  • Loading branch information
@NaveenNitturu @tdcmeehan
NaveenNitturu authored and tdcmeehan committed Sep 3, 2024
1 parent 87db58c commit c408975
Show file tree
Hide file tree
Showing 4 changed files with 133 additions and 1 deletion.
17 changes: 16 additions & 1 deletion presto-docs/src/main/sphinx/connector/mongodb.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -268,4 +268,19 @@ SQL support
ALTER TABLE
^^^^^^^^^^^

The connector supports ``ALTER TABLE RENAME TO`` operation. Other uses of ``ALTER TABLE`` are not supported.
.. code-block:: sql
ALTER TABLE mongodb.admin.sample_table ADD COLUMN new_col INT;
ALTER TABLE mongodb.admin.sample_table DROP COLUMN new_col;
ALTER TABLE mongodb.admin.sample_table RENAME COLUMN is_active TO is_enabled;
ALTER TABLE mongodb.admin.sample_table RENAME TO renamed_table;
.. note:: Presto does not support altering the data type of a column directly with the ALTER TABLE command.

.. code-block:: sql
ALTER TABLE mongodb.admin.users ALTER COLUMN age TYPE BIGINT;
returns an error similar to the following:

``Query 20240720_123348_00014_v7vrn failed: line 1:55: mismatched input 'int'. Expecting: 'FUNCTION', 'SCHEMA', 'TABLE'``
18 changes: 18 additions & 0 deletions presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoMetadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -301,4 +301,22 @@ private static List<MongoColumnHandle> buildColumnHandles(ConnectorTableMetadata
.map(m -> new MongoColumnHandle(m.getName(), m.getType(), m.isHidden()))
.collect(toList());
}

@Override
public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnMetadata column)
{
mongoSession.addColumn(((MongoTableHandle) tableHandle), column);
}

@Override
public void renameColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle source, String target)
{
mongoSession.renameColumn(((MongoTableHandle) tableHandle), ((MongoColumnHandle) source).getName(), target);
}

@Override
public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column)
{
mongoSession.dropColumn(((MongoTableHandle) tableHandle), ((MongoColumnHandle) column).getName());
}
}
84 changes: 84 additions & 0 deletions presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import com.facebook.presto.common.type.TypeSignatureParameter;
import com.facebook.presto.common.type.VarbinaryType;
import com.facebook.presto.spi.ColumnHandle;
import com.facebook.presto.spi.ColumnMetadata;
import com.facebook.presto.spi.PrestoException;
import com.facebook.presto.spi.SchemaNotFoundException;
import com.facebook.presto.spi.SchemaTableName;
Expand All @@ -43,7 +44,9 @@
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.Filters;
import com.mongodb.client.model.IndexOptions;
import com.mongodb.client.model.Updates;
import com.mongodb.client.result.DeleteResult;
import io.airlift.slice.Slice;
import org.bson.Document;
Expand All @@ -70,6 +73,7 @@
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Throwables.throwIfInstanceOf;
import static com.google.common.base.Verify.verify;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.HOURS;
import static java.util.concurrent.TimeUnit.MINUTES;
Expand All @@ -82,6 +86,7 @@ public class MongoSession
private static final List<String> SYSTEM_TABLES = Arrays.asList("system.indexes", "system.users", "system.version");

private static final String TABLE_NAME_KEY = "table";
private static final String COMMENT_KEY = "comment";
private static final String FIELDS_KEY = "fields";
private static final String FIELDS_NAME_KEY = "name";
private static final String FIELDS_TYPE_KEY = "type";
Expand Down Expand Up @@ -564,4 +569,83 @@ else if (value instanceof Document) {

return Optional.ofNullable(typeSignature);
}

public void addColumn(MongoTableHandle table, ColumnMetadata columnMetadata)
{
String remoteSchemaName = table.getSchemaTableName().getSchemaName();
String remoteTableName = table.getSchemaTableName().getTableName();
SchemaTableName schemaTableName = new SchemaTableName(remoteSchemaName, remoteTableName);

Document metadata = getTableMetadata(schemaTableName);

List<Document> columns = new ArrayList<>(getColumnMetadata(metadata));

Document newColumn = new Document()
.append(FIELDS_NAME_KEY, columnMetadata.getName())
.append(FIELDS_TYPE_KEY, columnMetadata.getType().getTypeSignature().toString())
.append(COMMENT_KEY, columnMetadata.getComment())
.append(FIELDS_HIDDEN_KEY, false);
columns.add(newColumn);

metadata.append(FIELDS_KEY, columns);

MongoCollection<Document> schema = client.getDatabase(remoteSchemaName)
.getCollection(schemaCollection);
schema.findOneAndReplace(new Document(TABLE_NAME_KEY, remoteTableName), metadata);

tableCache.invalidate(table.getSchemaTableName());
}

public void renameColumn(MongoTableHandle table, String source, String target)
{
String remoteSchemaName = table.getSchemaTableName().getSchemaName();
String remoteTableName = table.getSchemaTableName().getTableName();
SchemaTableName schemaTableName = new SchemaTableName(remoteSchemaName, remoteTableName);

Document metadata = getTableMetadata(schemaTableName);

List<Document> columns = getColumnMetadata(metadata).stream()
.map(document -> {
if (document.getString(FIELDS_NAME_KEY).equals(source)) {
document.put(FIELDS_NAME_KEY, target);
}
return document;
})
.collect(toImmutableList());

metadata.append(FIELDS_KEY, columns);

MongoDatabase database = client.getDatabase(remoteSchemaName);
MongoCollection<Document> schema = database.getCollection(schemaCollection);
schema.findOneAndReplace(new Document(TABLE_NAME_KEY, remoteTableName), metadata);

database.getCollection(remoteTableName)
.updateMany(new Document(), Updates.rename(source, target));

tableCache.invalidate(table.getSchemaTableName());
}

public void dropColumn(MongoTableHandle table, String columnName)
{
String remoteSchemaName = table.getSchemaTableName().getSchemaName();
String remoteTableName = table.getSchemaTableName().getTableName();
SchemaTableName schemaTableName = new SchemaTableName(remoteSchemaName, remoteTableName);

Document metadata = getTableMetadata(schemaTableName);

List<Document> columns = getColumnMetadata(metadata).stream()
.filter(document -> !document.getString(FIELDS_NAME_KEY).equals(columnName))
.collect(toImmutableList());

metadata.append(FIELDS_KEY, columns);

MongoDatabase database = client.getDatabase(remoteSchemaName);
MongoCollection<Document> schema = database.getCollection(schemaCollection);
schema.findOneAndReplace(new Document(TABLE_NAME_KEY, remoteTableName), metadata);

database.getCollection(remoteTableName)
.updateMany(Filters.exists(columnName), Updates.unset(columnName));

tableCache.invalidate(table.getSchemaTableName());
}
}
15 changes: 15 additions & 0 deletions presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -283,4 +283,19 @@ public void testRenameTable()
assertQuery("SELECT value FROM test_rename.tmp_rename_new_table", "SELECT 1");
assertUpdate("DROP TABLE test_rename.tmp_rename_new_table");
}

@Test
public void testAlterTable()
{
assertUpdate("CREATE TABLE test_alter.tmp_alter_table (value bigint)");
MongoCollection<Document> collection = mongoQueryRunner.getMongoClient().getDatabase("test_alter").getCollection("tmp_alter_table");
collection.insertOne(new Document(ImmutableMap.of("value", 1)));

assertUpdate("ALTER TABLE test_alter.tmp_alter_table ADD COLUMN email varchar");
collection.insertOne(new Document(ImmutableMap.of("value", 2, "email", "[email protected]")));
assertQuery("SELECT email from test_alter.tmp_alter_table WHERE email IS NOT NULL", "SELECT '[email protected]'");
assertUpdate("ALTER TABLE test_alter.tmp_alter_table RENAME COLUMN email TO email_id");
assertQuery("SELECT email_id from test_alter.tmp_alter_table WHERE email_id IS NOT NULL", "SELECT '[email protected]'");
assertUpdate("ALTER TABLE test_alter.tmp_alter_table DROP COLUMN email_id");
}
}

0 comments on commit c408975

Please sign in to comment.