|
5 | 5 |
|
6 | 6 | package com.microsoft.azure.toolkit.intellij.cosmos.dbtools; |
7 | 7 |
|
8 | | -import com.intellij.database.dataSource.DatabaseDriver; |
9 | 8 | import com.intellij.database.dataSource.DatabaseDriverImpl; |
10 | 9 | import com.intellij.database.dataSource.DatabaseDriverManager; |
11 | 10 | import com.intellij.database.dataSource.DatabaseDriverManagerImpl; |
|
17 | 16 | import com.microsoft.azure.toolkit.intellij.common.IntelliJAzureIcons; |
18 | 17 | import com.microsoft.azure.toolkit.lib.common.task.AzureTaskManager; |
19 | 18 | import lombok.SneakyThrows; |
| 19 | +import org.apache.commons.lang3.StringUtils; |
20 | 20 | import org.apache.commons.lang3.reflect.FieldUtils; |
21 | 21 | import org.jdom.Element; |
22 | 22 |
|
@@ -47,24 +47,28 @@ public void preload(@Nonnull ProgressIndicator indicator) { |
47 | 47 |
|
48 | 48 | private static void loadMongoDriver() { |
49 | 49 | final DatabaseDriverManager manager = DatabaseDriverManager.getInstance(); |
50 | | - final DatabaseDriver oldDriver = manager.getDriver(COSMOS_MONGO_DRIVER_ID); |
51 | | - if (Objects.isNull(oldDriver) || !"Azure Cosmos DB API for MongoDB".equals(oldDriver.getName())) { // remove if old driver is not user defined. |
| 50 | + final DatabaseDriverImpl oldDriver = (DatabaseDriverImpl) manager.getDriver(COSMOS_MONGO_DRIVER_ID); |
| 51 | + final boolean legacy = Optional.ofNullable(oldDriver).map(DatabaseDriverImpl::getUrlTemplates) |
| 52 | + .flatMap(ts -> ts.stream().findFirst()) |
| 53 | + .filter(t -> StringUtils.containsIgnoreCase(t.getTemplate(), "retrywrites")) |
| 54 | + .isEmpty(); |
| 55 | + if (legacy) { // remove if old driver is not user defined. |
52 | 56 | Optional.ofNullable(oldDriver).ifPresent(manager::removeDriver); |
53 | 57 | addAsUserDriver(COSMOS_MONGO_DRIVER_CONFIG); |
54 | 58 | } else { |
55 | | - ((DatabaseDriverImpl) oldDriver).setIcon(IntelliJAzureIcons.getIcon(COSMOS_MONGO_ICON)); |
| 59 | + oldDriver.setIcon(IntelliJAzureIcons.getIcon(COSMOS_MONGO_ICON)); |
56 | 60 | } |
57 | 61 | } |
58 | 62 |
|
59 | 63 | private static void loadCassandraDriver() { |
60 | 64 | final DatabaseDriverManager manager = DatabaseDriverManager.getInstance(); |
61 | | - final DatabaseDriver oldDriver = manager.getDriver(COSMOS_CASSANDRA_DRIVER_ID); |
| 65 | + final DatabaseDriverImpl oldDriver = (DatabaseDriverImpl) manager.getDriver(COSMOS_CASSANDRA_DRIVER_ID); |
62 | 66 | if (Registry.is("azure.toolkit.cosmos_cassandra.dbtools.enabled")) { |
63 | 67 | if (Objects.isNull(oldDriver) || !"Azure Cosmos DB API for Cassandra".equals(oldDriver.getName())) { // remove if old driver is not user defined. |
64 | 68 | Optional.ofNullable(oldDriver).ifPresent(manager::removeDriver); |
65 | 69 | addAsUserDriver(COSMOS_CASSANDRA_DRIVER_CONFIG); |
66 | 70 | } else { |
67 | | - ((DatabaseDriverImpl) oldDriver).setIcon(IntelliJAzureIcons.getIcon(COSMOS_CASSANDRA_ICON)); |
| 71 | + oldDriver.setIcon(IntelliJAzureIcons.getIcon(COSMOS_CASSANDRA_ICON)); |
68 | 72 | } |
69 | 73 | } else { |
70 | 74 | Optional.ofNullable(oldDriver).ifPresent(manager::removeDriver); |
|
0 commit comments