forked from opensearch-project/data-prepper
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
ENH: docdb credential auto refreshment (opensearch-project#4399)
* ENH: docdb credential refreshment Signed-off-by: George Chen <[email protected]>
- Loading branch information
1 parent
ad1b173
commit 9f01409
Showing
14 changed files
with
715 additions
and
107 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
103 changes: 103 additions & 0 deletions
103
...rc/main/java/org/opensearch/dataprepper/plugins/mongo/documentdb/MongoTasksRefresher.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.dataprepper.plugins.mongo.documentdb; | ||
|
|
||
| import io.micrometer.core.instrument.Counter; | ||
| import org.opensearch.dataprepper.metrics.PluginMetrics; | ||
| import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; | ||
| import org.opensearch.dataprepper.model.buffer.Buffer; | ||
| import org.opensearch.dataprepper.model.event.Event; | ||
| import org.opensearch.dataprepper.model.plugin.PluginConfigObserver; | ||
| import org.opensearch.dataprepper.model.record.Record; | ||
| import org.opensearch.dataprepper.model.source.coordinator.enhanced.EnhancedSourceCoordinator; | ||
| import org.opensearch.dataprepper.plugins.mongo.configuration.CollectionConfig; | ||
| import org.opensearch.dataprepper.plugins.mongo.configuration.MongoDBSourceConfig; | ||
| import org.opensearch.dataprepper.plugins.mongo.export.ExportScheduler; | ||
| import org.opensearch.dataprepper.plugins.mongo.export.ExportWorker; | ||
| import org.opensearch.dataprepper.plugins.mongo.export.MongoDBExportPartitionSupplier; | ||
| import org.opensearch.dataprepper.plugins.mongo.stream.StreamScheduler; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.Objects; | ||
| import java.util.concurrent.ExecutorService; | ||
| import java.util.function.Function; | ||
|
|
||
| public class MongoTasksRefresher implements PluginConfigObserver<MongoDBSourceConfig> { | ||
| private static final Logger LOG = LoggerFactory.getLogger(MongoTasksRefresher.class); | ||
| static final String CREDENTIALS_CHANGED = "credentialsChanged"; | ||
| static final String EXECUTOR_REFRESH_ERRORS = "executorRefreshErrors"; | ||
| private final EnhancedSourceCoordinator sourceCoordinator; | ||
| private final PluginMetrics pluginMetrics; | ||
| private final AcknowledgementSetManager acknowledgementSetManager; | ||
| private final Buffer<Record<Event>> buffer; | ||
| private final Function<Integer, ExecutorService> executorServiceFunction; | ||
| private final Counter credentialsChangeCounter; | ||
| private final Counter executorRefreshErrorsCounter; | ||
| private MongoDBExportPartitionSupplier currentMongoDBExportPartitionSupplier; | ||
| private MongoDBSourceConfig currentMongoDBSourceConfig; | ||
| private ExecutorService currentExecutor; | ||
|
|
||
| public MongoTasksRefresher(final Buffer<Record<Event>> buffer, | ||
| final EnhancedSourceCoordinator sourceCoordinator, | ||
| final PluginMetrics pluginMetrics, | ||
| final AcknowledgementSetManager acknowledgementSetManager, | ||
| final Function<Integer, ExecutorService> executorServiceFunction) { | ||
| this.sourceCoordinator = sourceCoordinator; | ||
| this.pluginMetrics = pluginMetrics; | ||
| this.acknowledgementSetManager = acknowledgementSetManager; | ||
| this.buffer = buffer; | ||
| this.executorServiceFunction = executorServiceFunction; | ||
| credentialsChangeCounter = pluginMetrics.counter(CREDENTIALS_CHANGED); | ||
| executorRefreshErrorsCounter = pluginMetrics.counter(EXECUTOR_REFRESH_ERRORS); | ||
| } | ||
|
|
||
| public void initialize(final MongoDBSourceConfig sourceConfig) { | ||
| this.currentMongoDBSourceConfig = sourceConfig; | ||
| refreshJobs(sourceConfig); | ||
| } | ||
|
|
||
| @Override | ||
| public void update(MongoDBSourceConfig pluginConfig) { | ||
| final MongoDBSourceConfig.AuthenticationConfig newAuthConfig = pluginConfig.getCredentialsConfig(); | ||
| if (basicAuthChanged(newAuthConfig)) { | ||
| credentialsChangeCounter.increment(); | ||
| try { | ||
| currentExecutor.shutdownNow(); | ||
| refreshJobs(pluginConfig); | ||
| currentMongoDBSourceConfig = pluginConfig; | ||
| } catch (Exception e) { | ||
| executorRefreshErrorsCounter.increment(); | ||
| LOG.error("Refreshing executor failed.", e); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| private void refreshJobs(MongoDBSourceConfig pluginConfig) { | ||
| final List<Runnable> runnables = new ArrayList<>(); | ||
| if (pluginConfig.getCollections().stream().anyMatch(CollectionConfig::isExport)) { | ||
| currentMongoDBExportPartitionSupplier = new MongoDBExportPartitionSupplier(pluginConfig); | ||
| runnables.add(new ExportScheduler(sourceCoordinator, currentMongoDBExportPartitionSupplier, pluginMetrics)); | ||
| runnables.add(new ExportWorker( | ||
| sourceCoordinator, buffer, pluginMetrics, acknowledgementSetManager, pluginConfig)); | ||
| } | ||
| if (pluginConfig.getCollections().stream().anyMatch(CollectionConfig::isStream)) { | ||
| runnables.add(new StreamScheduler( | ||
| sourceCoordinator, buffer, acknowledgementSetManager, pluginConfig, pluginMetrics)); | ||
| } | ||
| this.currentExecutor = executorServiceFunction.apply(runnables.size()); | ||
| runnables.forEach(currentExecutor::submit); | ||
| } | ||
|
|
||
| private boolean basicAuthChanged(final MongoDBSourceConfig.AuthenticationConfig newAuthConfig) { | ||
| final MongoDBSourceConfig.AuthenticationConfig currentAuthConfig = currentMongoDBSourceConfig | ||
| .getCredentialsConfig(); | ||
| return !Objects.equals(currentAuthConfig.getUsername(), newAuthConfig.getUsername()) || | ||
| !Objects.equals(currentAuthConfig.getPassword(), newAuthConfig.getPassword()); | ||
| } | ||
| } |
Oops, something went wrong.