Separate cosmos module in base and java parts

PluginsAndFeatures/azure-toolkit-for-intellij/azure-intellij-plugin-cosmos-java/build.gradle.kts

@@ -0,0 +1,28 @@
+fun properties(key: String) = providers.gradleProperty(key)
+
+sourceSets {
+    main {
+        resources.srcDirs("src/main/resources")
+    }
+}
+
+dependencies {
+    implementation(project(":azure-intellij-plugin-lib"))
+    // runtimeOnly project(path: ":azure-intellij-plugin-lib", configuration: "instrumentedJar")
+    implementation(project(":azure-intellij-resource-connector-lib"))
+    // runtimeOnly project(path: ":azure-intellij-resource-connector-lib", configuration: "instrumentedJar")
+    implementation(project(":azure-intellij-resource-connector-lib-java"))
+    implementation(project(":azure-intellij-plugin-cosmos"))
+    // runtimeOnly project(path: ":azure-intellij-resource-connector-lib-java", configuration: "instrumentedJar")
+    implementation("com.microsoft.azure:azure-toolkit-cosmos-lib")
+    implementation("com.microsoft.azure:azure-toolkit-ide-common-lib")
+    implementation("com.microsoft.azure:azure-toolkit-ide-cosmos-lib")
+    implementation("com.microsoft.azure:azure-toolkit-identity-lib")
+
+    intellijPlatform {
+        intellijIdeaUltimate(properties("platformVersion").get())
+        // Plugin Dependencies. Uses `platformBundledPlugins` property from the gradle.properties file for bundled IntelliJ Platform plugins.
+        bundledPlugin("com.intellij.java")
+        bundledPlugin("com.intellij.database")
+    }
+}

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

@@ -0,0 +1,72 @@
+/*
+ * 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.openapi.actionSystem.AnActionEvent;
+import com.intellij.openapi.project.Project;
+import com.microsoft.azure.toolkit.ide.common.IActionsContributor;
+import com.microsoft.azure.toolkit.ide.common.action.ResourceCommonActionsContributor;
+import com.microsoft.azure.toolkit.intellij.connector.AzureServiceResource;
+import com.microsoft.azure.toolkit.intellij.connector.ConnectorDialog;
+import com.microsoft.azure.toolkit.intellij.cosmos.connection.*;
+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.messager.AzureMessager;
+import com.microsoft.azure.toolkit.lib.common.model.AzResource;
+import com.microsoft.azure.toolkit.lib.common.task.AzureTaskManager;
+import com.microsoft.azure.toolkit.lib.cosmos.*;
+import com.microsoft.azure.toolkit.lib.cosmos.cassandra.CassandraCosmosDBAccount;
+import com.microsoft.azure.toolkit.lib.cosmos.cassandra.CassandraKeyspace;
+import com.microsoft.azure.toolkit.lib.cosmos.mongo.MongoCosmosDBAccount;
+import com.microsoft.azure.toolkit.lib.cosmos.mongo.MongoDatabase;
+import com.microsoft.azure.toolkit.lib.cosmos.sql.SqlCosmosDBAccount;
+import com.microsoft.azure.toolkit.lib.cosmos.sql.SqlDatabase;
+
+import javax.annotation.Nonnull;
+import java.util.Objects;
+import java.util.function.Function;
+
+public class IntellijJavaCosmosActionsContributor implements IActionsContributor {
+    @Override
+    public void registerHandlers(AzureActionManager am) {
+
+        final Function<MongoCosmosDBAccount, MongoDatabase> mongoFunction = account -> account.mongoDatabases().list().stream().findFirst().orElse(null);
+        am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof MongoCosmosDBAccount && r.getFormalStatus().isConnected(), (AzResource r, AnActionEvent e) ->
+            openResourceConnector((MongoCosmosDBAccount) r, mongoFunction, MongoCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+        am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof MongoDatabase && r.getFormalStatus().isConnected(), (AzResource r, AnActionEvent e) ->
+            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 && r.getFormalStatus().isConnected(), (AzResource r, AnActionEvent e) ->
+            openResourceConnector((SqlCosmosDBAccount) r, sqlFunction, SqlCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+        am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof SqlDatabase && r.getFormalStatus().isConnected(), (AzResource r, AnActionEvent e) ->
+            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 && r.getFormalStatus().isConnected(), (AzResource r, AnActionEvent e) ->
+            openResourceConnector((CassandraCosmosDBAccount) r, cassandraFunction, CassandraCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+        am.registerHandler(ResourceCommonActionsContributor.CONNECT, (r, e) -> r instanceof CassandraKeyspace && r.getFormalStatus().isConnected(), (AzResource r, AnActionEvent e) ->
+            openResourceConnector((CassandraKeyspace) r, CassandraCosmosDBAccountResourceDefinition.INSTANCE, e.getProject()));
+    }
+
+    private <T extends AzResource> void openResourceConnector(@Nonnull final T resource, @Nonnull final AzureServiceResource.Definition<T> definition, Project project) {
+        AzureTaskManager.getInstance().runLater(() -> {
+            final ConnectorDialog dialog = new ConnectorDialog(project);
+            dialog.setResource(new AzureServiceResource<>(resource, definition));
+            dialog.show();
+        });
+    }
+
+    private <T extends AzResource, R extends CosmosDBAccount> void openResourceConnector(@Nonnull final R account, @Nonnull Function<R, T> databaseFunction,
+                                                                                         @Nonnull final AzureServiceResource.Definition<T> definition, Project project) {
+        final T database = databaseFunction.apply(account);
+        if (Objects.isNull(database)) {
+            AzureMessager.getMessager().warning(AzureString.format("Can not connect to %s as there is no database in selected account", account.getName()));
+        } else {
+            openResourceConnector(database, definition, project);
+        }
+    }
+}

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

@@ -0,0 +1,33 @@
+/*
+ * 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.connection;
+
+import com.microsoft.azure.toolkit.intellij.connector.Connection;
+import com.microsoft.azure.toolkit.intellij.connector.spring.SpringSupported;
+import com.microsoft.azure.toolkit.lib.cosmos.cassandra.CassandraKeyspace;
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+import java.util.ArrayList;
+import java.util.List;
+
+public class CassandraCosmosDBAccountResourceDefinition extends BaseCassandraCosmosDBAccountResourceDefinition implements SpringSupported<CassandraKeyspace> {
+    public static final CassandraCosmosDBAccountResourceDefinition INSTANCE = new CassandraCosmosDBAccountResourceDefinition();
+
+    @Override
+    public List<Pair<String, String>> getSpringProperties(@Nullable final String key) {
+        final List<Pair<String, String>> properties = new ArrayList<>();
+        properties.add(Pair.of("spring.cassandra.contact-points", String.format("${%s_CONTACT_POINT}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cassandra.port", String.format("${%s_PORT}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cassandra.username", String.format("${%s_USERNAME}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cassandra.password", String.format("${%s_PASSWORD}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cassandra.keyspace-name", String.format("${%s_KEYSPACE}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cassandra.schema-action", "create_if_not_exists"));
+        properties.add(Pair.of("spring.cassandra.ssl.enabled", "true"));
+        properties.add(Pair.of("spring.cassandra.local-datacenter", String.format("${%s_REGION}", Connection.ENV_PREFIX)));
+        return properties;
+    }
+}

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

@@ -0,0 +1,27 @@
+/*
+ * 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.connection;
+
+import com.microsoft.azure.toolkit.intellij.connector.Connection;
+import com.microsoft.azure.toolkit.intellij.connector.spring.SpringSupported;
+import com.microsoft.azure.toolkit.lib.cosmos.mongo.MongoDatabase;
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+import java.util.ArrayList;
+import java.util.List;
+
+public class MongoCosmosDBAccountResourceDefinition extends BaseMongoCosmosDBAccountResourceDefinition implements SpringSupported<MongoDatabase> {
+    public static final MongoCosmosDBAccountResourceDefinition INSTANCE = new MongoCosmosDBAccountResourceDefinition();
+
+    @Override
+    public List<Pair<String, String>> getSpringProperties(@Nullable final String key) {
+        final List<Pair<String, String>> properties = new ArrayList<>();
+        properties.add(Pair.of("spring.data.mongodb.database", String.format("${%s_DATABASE}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.data.mongodb.uri", String.format("${%s_CONNECTION_STRING}", Connection.ENV_PREFIX)));
+        return properties;
+    }
+}

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

@@ -0,0 +1,79 @@
+/*
+ * 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.connection;
+
+import com.azure.resourcemanager.authorization.models.BuiltInRole;
+import com.intellij.openapi.project.Project;
+import com.microsoft.azure.toolkit.intellij.connector.AuthenticationType;
+import com.microsoft.azure.toolkit.intellij.connector.Connection;
+import com.microsoft.azure.toolkit.intellij.connector.Resource;
+import com.microsoft.azure.toolkit.intellij.connector.spring.SpringManagedIdentitySupported;
+import com.microsoft.azure.toolkit.intellij.connector.spring.SpringSupported;
+import com.microsoft.azure.toolkit.lib.cosmos.CosmosDBAccount;
+import com.microsoft.azure.toolkit.lib.cosmos.sql.SqlDatabase;
+import org.apache.commons.lang3.tuple.Pair;
+
+import javax.annotation.Nullable;
+import java.util.*;
+
+public class SqlCosmosDBAccountResourceDefinition extends BaseSqlCosmosDBAccountResourceDefinition
+        implements SpringSupported<SqlDatabase>, SpringManagedIdentitySupported<SqlDatabase> {
+
+    public static final SqlCosmosDBAccountResourceDefinition INSTANCE = new SqlCosmosDBAccountResourceDefinition();
+
+    @Override
+    public List<Pair<String, String>> getSpringProperties(@Nullable final String key) {
+        final List<Pair<String, String>> properties = new ArrayList<>();
+        properties.add(Pair.of("spring.cloud.azure.cosmos.endpoint", String.format("${%s_ENDPOINT}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cloud.azure.cosmos.key", String.format("${%s_KEY}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cloud.azure.cosmos.database", String.format("${%s_DATABASE}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cloud.azure.cosmos.populate-query-metrics", String.valueOf(true)));
+        return properties;
+    }
+
+
+
+    @Override
+    public Map<String, String> initIdentityEnv(Connection<SqlDatabase, ?> data, Project project) {
+        final SqlDatabase database = data.getResource().getData();
+        final CosmosDBAccount account = database.getParent();
+        final HashMap<String, String> env = new HashMap<>();
+        env.put(String.format("%s_ENDPOINT", Connection.ENV_PREFIX), account.getDocumentEndpoint());
+        env.put(String.format("%s_DATABASE", Connection.ENV_PREFIX), database.getName());
+        if (data.getAuthenticationType() == AuthenticationType.USER_ASSIGNED_MANAGED_IDENTITY) {
+            Optional.ofNullable(data.getUserAssignedManagedIdentity()).map(Resource::getData).ifPresent(identity -> {
+                env.put(String.format("%s_CLIENT_ID", Connection.ENV_PREFIX), identity.getClientId());
+                env.put("AZURE_CLIENT_ID", identity.getClientId());
+            });
+        }
+        return env;
+    }
+
+    @Override
+    public List<String> getRequiredPermissions() {
+        return List.of("Microsoft.DocumentDB/databaseAccounts/readMetadata",
+                "Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/items/*",
+                "Microsoft.DocumentDB/databaseAccounts/sqlDatabases/containers/*");
+    }
+
+    @Nullable
+    @Override
+    public Map<String, BuiltInRole> getBuiltInRoles() {
+        return null;
+    }
+
+    @Override
+    public List<Pair<String, String>> getSpringPropertiesForManagedIdentity(String key, Connection<?, ?> connection) {
+        final List<Pair<String, String>> properties = new ArrayList<>();
+        properties.add(Pair.of("spring.cloud.azure.cosmos.endpoint", String.format("${%s_ENDPOINT}", Connection.ENV_PREFIX)));
+        properties.add(Pair.of("spring.cloud.azure.cosmos.database", String.format("${%s_DATABASE}", Connection.ENV_PREFIX)));
+        // properties.add(Pair.of("spring.cloud.azure.cosmos.credential.managed-identity-enabled", String.valueOf(true)));
+        if (connection.getAuthenticationType() == AuthenticationType.USER_ASSIGNED_MANAGED_IDENTITY) {
+            properties.add(Pair.of("spring.cloud.azure.cosmos.credential.client-id", String.format("${%s_CLIENT_ID}", Connection.ENV_PREFIX)));
+        }
+        return properties;
+    }
+}

PluginsAndFeatures/azure-toolkit-for-intellij/azure-intellij-plugin-cosmos-java/src/main/resources/META-INF/azure-intellij-plugin-cosmos-java.xml

@@ -0,0 +1,13 @@
+<idea-plugin>
+    <extensions defaultExtensionNs="com.intellij">
+        <completion.contributor id="azCosmosDBCompletion" language="JAVA" order="after azStorageStringLiteral" implementationClass="com.microsoft.azure.toolkit.intellij.cosmos.code.function.AzureCosmosDBFunctionAnnotationCompletionContributor"/>
+        <psi.referenceContributor language="JAVA" implementation="com.microsoft.azure.toolkit.intellij.cosmos.code.function.AzureCosmosDBResourceReferenceContributor"/>
+        <annotator language="JAVA" implementationClass="com.microsoft.azure.toolkit.intellij.cosmos.code.function.CosmosDBFunctionPathAnnotator"/>
+    </extensions>
+    <extensions defaultExtensionNs="com.microsoft.tooling.msservices.intellij.azure">
+        <actions implementation="com.microsoft.azure.toolkit.intellij.cosmos.IntellijJavaCosmosActionsContributor"/>
+        <connectorResourceType implementation="com.microsoft.azure.toolkit.intellij.cosmos.connection.CassandraCosmosDBAccountResourceDefinition"/>
+        <connectorResourceType implementation="com.microsoft.azure.toolkit.intellij.cosmos.connection.MongoCosmosDBAccountResourceDefinition"/>
+        <connectorResourceType implementation="com.microsoft.azure.toolkit.intellij.cosmos.connection.SqlCosmosDBAccountResourceDefinition"/>
+    </extensions>
+</idea-plugin>

PluginsAndFeatures/azure-toolkit-for-intellij/azure-intellij-plugin-cosmos/build.gradle.kts

@@ -11,8 +11,6 @@ dependencies {
     // runtimeOnly project(path: ":azure-intellij-plugin-lib", configuration: "instrumentedJar")
     implementation(project(":azure-intellij-resource-connector-lib"))
     // runtimeOnly project(path: ":azure-intellij-resource-connector-lib", configuration: "instrumentedJar")
-    implementation(project(":azure-intellij-resource-connector-lib-java"))
-    // runtimeOnly project(path: ":azure-intellij-resource-connector-lib-java", configuration: "instrumentedJar")
     implementation("com.microsoft.azure:azure-toolkit-cosmos-lib")
     implementation("com.microsoft.azure:azure-toolkit-ide-common-lib")
     implementation("com.microsoft.azure:azure-toolkit-ide-cosmos-lib")
@@ -21,7 +19,6 @@ dependencies {
     intellijPlatform {
         intellijIdeaUltimate(properties("platformVersion").get())
         // Plugin Dependencies. Uses `platformBundledPlugins` property from the gradle.properties file for bundled IntelliJ Platform plugins.
-        bundledPlugin("com.intellij.java")
         bundledPlugin("com.intellij.database")
     }
 }