add feature toggle for Cassandra API

PluginsAndFeatures/azure-toolkit-for-intellij/azure-intellij-plugin-cosmos/src/main/java/com/microsoft/azure/toolkit/intellij/cosmos/IntelliJCosmosActionsContributor.java

@@ -5,14 +5,11 @@
 
 package com.microsoft.azure.toolkit.intellij.cosmos;
 
-import com.intellij.database.autoconfig.DataSourceDetector;
-import com.intellij.database.autoconfig.DataSourceRegistry;
-import com.intellij.database.psi.DbPsiFacade;
-import com.intellij.database.view.ui.DataSourceManagerDialog;
 import com.intellij.ide.plugins.PluginManagerCore;
 import com.intellij.openapi.actionSystem.AnActionEvent;
 import com.intellij.openapi.extensions.PluginId;
 import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.registry.Registry;
 import com.microsoft.azure.toolkit.ide.common.IActionsContributor;
 import com.microsoft.azure.toolkit.ide.common.action.ResourceCommonActionsContributor;
 import com.microsoft.azure.toolkit.ide.cosmos.CosmosActionsContributor;
@@ -24,7 +21,6 @@
 import com.microsoft.azure.toolkit.intellij.cosmos.connection.SqlCosmosDBAccountResourceDefinition;
 import com.microsoft.azure.toolkit.intellij.cosmos.creation.CreateCosmosDBAccountAction;
 import com.microsoft.azure.toolkit.intellij.cosmos.creation.CreateCosmosDatabaseAction;
-import com.microsoft.azure.toolkit.intellij.cosmos.dbtools.Dbms;
 import com.microsoft.azure.toolkit.lib.common.action.AzureActionManager;
 import com.microsoft.azure.toolkit.lib.common.bundle.AzureString;
 import com.microsoft.azure.toolkit.lib.common.exception.AzureToolkitRuntimeException;
@@ -52,7 +48,6 @@
 import java.util.function.BiPredicate;
 import java.util.function.Function;
 
-import static com.microsoft.azure.toolkit.intellij.cosmos.dbtools.AzureCosmosDbAccountParamEditor.KEY_COSMOS_ACCOUNT_ID;
 import static com.microsoft.azure.toolkit.lib.cosmos.CosmosDBAccountDraft.Config.getDefaultConfig;
 import static com.microsoft.azure.toolkit.lib.cosmos.model.DatabaseConfig.getDefaultDatabaseConfig;
 
@@ -64,66 +59,56 @@ public void registerHandlers(AzureActionManager am) {
         am.registerHandler(ResourceCommonActionsContributor.CREATE, serviceCondition, handler);
 
         final BiConsumer<ResourceGroup, AnActionEvent> groupCreateHandler = (r, e) ->
-                CreateCosmosDBAccountAction.create(e.getProject(), getDefaultConfig(r));
+            CreateCosmosDBAccountAction.create(e.getProject(), getDefaultConfig(r));
         am.registerHandler(CosmosActionsContributor.GROUP_CREATE_COSMOS_SERVICE, (r, e) -> true, groupCreateHandler);
 
         final Function<MongoCosmosDBAccount, MongoDatabase> mongoFunction = account -> account.mongoDatabases().list().stream().findFirst().orElse(null);
         am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof MongoCosmosDBAccount, (AzResource<?, ?, ?> r, AnActionEvent e) ->
-                openResourceConnector((MongoCosmosDBAccount) r, mongoFunction, MongoCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+            openResourceConnector((MongoCosmosDBAccount) r, mongoFunction, MongoCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
         am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof MongoDatabase, (AzResource<?, ?, ?> r, AnActionEvent e) ->
-                openResourceConnector((MongoDatabase) r, MongoCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+            openResourceConnector((MongoDatabase) r, MongoCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
 
         final Function<SqlCosmosDBAccount, SqlDatabase> sqlFunction = account -> account.sqlDatabases().list().stream().findFirst().orElse(null);
         am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof SqlCosmosDBAccount, (AzResource<?, ?, ?> r, AnActionEvent e) ->
-                openResourceConnector((SqlCosmosDBAccount) r, sqlFunction, SqlCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+            openResourceConnector((SqlCosmosDBAccount) r, sqlFunction, SqlCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
         am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof SqlDatabase, (AzResource<?, ?, ?> r, AnActionEvent e) ->
-                openResourceConnector((SqlDatabase) r, SqlCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+            openResourceConnector((SqlDatabase) r, SqlCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
 
         final Function<CassandraCosmosDBAccount, CassandraKeyspace> cassandraFunction = account -> account.keySpaces().list().stream().findFirst().orElse(null);
         am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof CassandraCosmosDBAccount, (AzResource<?, ?, ?> r, AnActionEvent e) ->
-                openResourceConnector((CassandraCosmosDBAccount) r, cassandraFunction, CassandraCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+            openResourceConnector((CassandraCosmosDBAccount) r, cassandraFunction, CassandraCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
         am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof CassandraKeyspace, (AzResource<?, ?, ?> r, AnActionEvent e) ->
-                openResourceConnector((CassandraKeyspace) r, CassandraCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+            openResourceConnector((CassandraKeyspace) r, CassandraCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
 
         final BiFunction<MongoCosmosDBAccount, DatabaseConfig, ICosmosDatabaseDraft<?, ?>> mongoDraftSupplier = (account, config) ->
-                (ICosmosDatabaseDraft<?, ?>) account.mongoDatabases().getOrDraft(config.getName(), account.getResourceGroupName());
+            (ICosmosDatabaseDraft<?, ?>) account.mongoDatabases().getOrDraft(config.getName(), account.getResourceGroupName());
         am.registerHandler(ResourceCommonActionsContributor.CREATE, (r, e) -> r instanceof MongoCosmosDBAccount, (Object r, AnActionEvent e) ->
-                CreateCosmosDatabaseAction.create(e.getProject(), (MongoCosmosDBAccount) r, mongoDraftSupplier, getDefaultDatabaseConfig()));
+            CreateCosmosDatabaseAction.create(e.getProject(), (MongoCosmosDBAccount) r, mongoDraftSupplier, getDefaultDatabaseConfig()));
 
         final BiFunction<SqlCosmosDBAccount, DatabaseConfig, ICosmosDatabaseDraft<?, ?>> sqlDraftSupplier = (account, config) ->
-                (ICosmosDatabaseDraft<?, ?>) account.sqlDatabases().getOrDraft(config.getName(), account.getResourceGroupName());
+            (ICosmosDatabaseDraft<?, ?>) account.sqlDatabases().getOrDraft(config.getName(), account.getResourceGroupName());
         am.registerHandler(ResourceCommonActionsContributor.CREATE, (r, e) -> r instanceof SqlCosmosDBAccount, (Object r, AnActionEvent e) ->
-                CreateCosmosDatabaseAction.create(e.getProject(), (SqlCosmosDBAccount) r, sqlDraftSupplier, getDefaultDatabaseConfig()));
+            CreateCosmosDatabaseAction.create(e.getProject(), (SqlCosmosDBAccount) r, sqlDraftSupplier, getDefaultDatabaseConfig()));
 
         final BiFunction<CassandraCosmosDBAccount, DatabaseConfig, ICosmosDatabaseDraft<?, ?>> cassandraDraftSupplier = (account, config) ->
-                (ICosmosDatabaseDraft<?, ?>) account.keySpaces().getOrDraft(config.getName(), account.getResourceGroupName());
+            (ICosmosDatabaseDraft<?, ?>) account.keySpaces().getOrDraft(config.getName(), account.getResourceGroupName());
         am.registerHandler(ResourceCommonActionsContributor.CREATE, (r, e) -> r instanceof CassandraCosmosDBAccount, (Object r, AnActionEvent e) ->
-                CreateCosmosDatabaseAction.create(e.getProject(), (CassandraCosmosDBAccount) r, cassandraDraftSupplier, getDefaultDatabaseConfig()));
+            CreateCosmosDatabaseAction.create(e.getProject(), (CassandraCosmosDBAccount) r, cassandraDraftSupplier, getDefaultDatabaseConfig()));
 
-        final BiConsumer<CosmosDBAccount, AnActionEvent> openDatabaseHandler = (c, e) -> openDatabaseTool(e.getProject(), c);
-        am.registerHandler(CosmosActionsContributor.OPEN_DATABASE_TOOL, (r, e) -> r instanceof MongoCosmosDBAccount || r instanceof CassandraCosmosDBAccount, openDatabaseHandler);
+        final String DATABASE_TOOLS_PLUGIN_ID = "com.intellij.database";
+        if (PluginManagerCore.getPlugin(PluginId.findId(DATABASE_TOOLS_PLUGIN_ID)) == null) {
+            final BiConsumer<CosmosDBAccount, AnActionEvent> openDatabaseHandler = (c, e) -> openDatabaseTool(e.getProject(), c);
+            final boolean cassandraOn = Registry.is("azure.toolkit.cosmos_cassandra.dbtools.enabled");
+            am.registerHandler(CosmosActionsContributor.OPEN_DATABASE_TOOL, (r, e) -> r instanceof MongoCosmosDBAccount || (r instanceof CassandraCosmosDBAccount && cassandraOn), openDatabaseHandler);
+        }
     }
 
     @AzureOperation(name = "cosmos.open_database_tools.account", params = {"account.getName()"}, type = AzureOperation.Type.ACTION)
     private void openDatabaseTool(Project project, CosmosDBAccount account) {
         final String DATABASE_TOOLS_PLUGIN_ID = "com.intellij.database";
         final String DATABASE_PLUGIN_NOT_INSTALLED = "\"Database tools and SQL\" plugin is not installed.";
         final String NOT_SUPPORT_ERROR_ACTION = "\"Database tools and SQL\" plugin is only provided in IntelliJ Ultimate edition.";
-        if (PluginManagerCore.getPlugin(PluginId.findId(DATABASE_TOOLS_PLUGIN_ID)) == null) {
-            throw new AzureToolkitRuntimeException(DATABASE_PLUGIN_NOT_INSTALLED, NOT_SUPPORT_ERROR_ACTION, IntellijActionsContributor.TRY_ULTIMATE);
-        }
-        AzureTaskManager.getInstance().runLater(() -> {
-            final DataSourceRegistry registry = new DataSourceRegistry(project);
-            final com.intellij.database.Dbms dbms =
-                account instanceof MongoCosmosDBAccount ? Dbms.AZ_COSMOS_MONGO :
-                    account instanceof CassandraCosmosDBAccount ? Dbms.AZ_COSMOS_CASSANDRA :
-                        com.intellij.database.Dbms.UNKNOWN;
-            final DataSourceDetector.Builder builder = registry.getBuilder()
-                .withDbms(dbms)
-                .withJdbcAdditionalProperty(KEY_COSMOS_ACCOUNT_ID, account.getId())
-                .commit();
-            DataSourceManagerDialog.showDialog(DbPsiFacade.getInstance(project), registry);
-        });
+        throw new AzureToolkitRuntimeException(DATABASE_PLUGIN_NOT_INSTALLED, NOT_SUPPORT_ERROR_ACTION, IntellijActionsContributor.TRY_ULTIMATE);
     }
 
     private <T extends AzResource<?, ?, ?>> void openResourceConnector(@Nonnull final T resource, @Nonnull final AzureServiceResource.Definition<T> definition, Project project) {

PluginsAndFeatures/azure-toolkit-for-intellij/azure-intellij-plugin-cosmos/src/main/java/com/microsoft/azure/toolkit/intellij/cosmos/IntelliJCosmosActionsContributorForUltimate.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License. See License.txt in the project root for license information.
+ */
+
+package com.microsoft.azure.toolkit.intellij.cosmos;
+
+import com.intellij.database.autoconfig.DataSourceRegistry;
+import com.intellij.database.dataSource.DatabaseDriverManager;
+import com.intellij.database.dataSource.LocalDataSource;
+import com.intellij.database.dataSource.LocalDataSourceManager;
+import com.intellij.database.psi.DbDataSource;
+import com.intellij.database.psi.DbPsiFacade;
+import com.intellij.database.psi.DbPsiFacadeImpl;
+import com.intellij.database.view.ui.DataSourceManagerDialog;
+import com.intellij.openapi.actionSystem.AnActionEvent;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.util.registry.Registry;
+import com.microsoft.azure.toolkit.ide.common.IActionsContributor;
+import com.microsoft.azure.toolkit.ide.cosmos.CosmosActionsContributor;
+import com.microsoft.azure.toolkit.lib.common.action.AzureActionManager;
+import com.microsoft.azure.toolkit.lib.common.operation.AzureOperation;
+import com.microsoft.azure.toolkit.lib.common.task.AzureTaskManager;
+import com.microsoft.azure.toolkit.lib.cosmos.CosmosDBAccount;
+import com.microsoft.azure.toolkit.lib.cosmos.cassandra.CassandraCosmosDBAccount;
+import com.microsoft.azure.toolkit.lib.cosmos.mongo.MongoCosmosDBAccount;
+
+import java.util.function.BiConsumer;
+
+import static com.microsoft.azure.toolkit.intellij.cosmos.dbtools.AzureCosmosDbAccountParamEditor.KEY_COSMOS_ACCOUNT_ID;
+
+public class IntelliJCosmosActionsContributorForUltimate implements IActionsContributor {
+    @Override
+    public void registerHandlers(AzureActionManager am) {
+        final BiConsumer<CosmosDBAccount, AnActionEvent> openDatabaseHandler = (c, e) -> openDatabaseTool(e.getProject(), c);
+        final boolean cassandraOn = Registry.is("azure.toolkit.cosmos_cassandra.dbtools.enabled");
+        am.registerHandler(CosmosActionsContributor.OPEN_DATABASE_TOOL, (r, e) -> r instanceof MongoCosmosDBAccount || (r instanceof CassandraCosmosDBAccount && cassandraOn), openDatabaseHandler);
+    }
+
+    @AzureOperation(name = "cosmos.open_database_tools.account", params = {"account.getName()"}, type = AzureOperation.Type.ACTION)
+    private void openDatabaseTool(Project project, CosmosDBAccount account) {
+        AzureTaskManager.getInstance().runLater(() -> {
+            final DataSourceRegistry registry = new DataSourceRegistry(project);
+            final String driver = account instanceof MongoCosmosDBAccount ? "az_cosmos_mongo" : "az_cosmos_cassandra";
+            final LocalDataSource ds = DatabaseDriverManager.getInstance().getDriver(driver).createDataSource(null);
+            final DbPsiFacade facade = DbPsiFacade.getInstance(project);
+            final LocalDataSourceManager manager = LocalDataSourceManager.getInstance(project);
+            final DbDataSource newElement = ((DbPsiFacadeImpl) facade).createDataSourceWrapperElement(ds, manager);
+            ds.setAdditionalProperty(KEY_COSMOS_ACCOUNT_ID, account.getId());
+            DataSourceManagerDialog.showDialog(facade, newElement, null, null, null);
+        });
+    }
+}

PluginsAndFeatures/azure-toolkit-for-intellij/azure-intellij-plugin-cosmos/src/main/java/com/microsoft/azure/toolkit/intellij/cosmos/dbtools/AzureCosmosDbAccountConnectionInterceptor.java

@@ -6,7 +6,6 @@
 package com.microsoft.azure.toolkit.intellij.cosmos.dbtools;
 
 import com.azure.resourcemanager.resources.fluentcore.arm.ResourceId;
-import com.intellij.database.Dbms;
 import com.intellij.database.dataSource.DatabaseConnectionInterceptor;
 import com.intellij.database.dataSource.DatabaseConnectionPoint;
 import com.microsoft.azure.toolkit.lib.common.operation.AzureOperation;
@@ -30,15 +29,13 @@ public class AzureCosmosDbAccountConnectionInterceptor implements DatabaseConnec
     @Nullable
     public CompletionStage<ProtoConnection> intercept(@NotNull DatabaseConnectionInterceptor.ProtoConnection proto, boolean silent) {
         final DatabaseConnectionPoint point = proto.getConnectionPoint();
-        final Dbms dbms = point.getDbms();
-        if (StringUtils.startsWithIgnoreCase(dbms.getName(), "AZ_COSMOS")) {
-            final String accountId = point.getAdditionalProperty(AzureCosmosDbAccountParamEditor.KEY_COSMOS_ACCOUNT_ID);
+        final String accountId = point.getAdditionalProperty(AzureCosmosDbAccountParamEditor.KEY_COSMOS_ACCOUNT_ID);
+        if (StringUtils.isNotBlank(accountId)) {
             final Map<String, String> properties = new HashMap<>();
-            if (StringUtils.isNotBlank(accountId) && !StringUtils.equals(accountId, AzureCosmosDbAccountParamEditor.NONE)) {
+            if (!StringUtils.equals(accountId, AzureCosmosDbAccountParamEditor.NONE)) {
                 final ResourceId id = ResourceId.fromString(accountId);
                 properties.put("subscriptionId", id.subscriptionId());
             }
-            properties.put("kind", dbms.getName());
             properties.put(SERVICE_NAME, "cosmos");
             properties.put(OPERATION_NAME, "connect_jdbc_from_dbtools");
             properties.put(OP_NAME, "cosmos.connect_jdbc_from_dbtools");