From dee1bc31539483e6a8b2b4d8802c87a0987bc94f Mon Sep 17 00:00:00 2001 From: Heemin Kim Date: Mon, 24 Apr 2023 21:44:30 -0700 Subject: [PATCH] Added unit tests with some refactoring of codes * Add Unit tests * Set cache true for search query * Remove in memory cache implementation (Two way door decision) * Relying on search cache without custom cache * Renamed datasource state from FAILED to CREATE_FAILED * Renamed class name from *Helper to *Facade * Changed updateIntervalInDays to updateInterval * Changed value type of default update_interval from TimeValue to Long * Read setting value from cluster settings directly Signed-off-by: Heemin Kim --- lombok.config | 0 .../annotation/VisibleForTesting.java | 12 + .../ip2geo/action/PutDatasourceRequest.java | 26 +- .../action/PutDatasourceTransportAction.java | 163 +++------ .../action/RestPutDatasourceHandler.java | 17 +- ...ourceHelper.java => DatasourceFacade.java} | 29 +- .../ip2geo/common/DatasourceState.java | 16 +- ...IpDataHelper.java => GeoIpDataFacade.java} | 111 +++--- ...xecutorHelper.java => Ip2GeoExecutor.java} | 10 +- .../ip2geo/common/Ip2GeoSettings.java | 23 +- .../ip2geo/jobscheduler/Datasource.java | 80 +++- .../jobscheduler/DatasourceExtension.java | 2 +- .../ip2geo/jobscheduler/DatasourceRunner.java | 225 +++--------- .../jobscheduler/DatasourceUpdateService.java | 206 +++++++++++ .../ip2geo/processor/Ip2GeoCache.java | 112 ------ .../ip2geo/processor/Ip2GeoProcessor.java | 267 +++++++------- .../geospatial/plugin/GeospatialPlugin.java | 36 +- .../geospatial/ip2geo/Ip2GeoTestCase.java | 227 ++++++++++++ .../action/PutDatasourceRequestTests.java | 96 ++++- .../PutDatasourceTransportActionTests.java | 118 ++++++ .../action/RestPutDatasourceHandlerTests.java | 51 ++- .../ip2geo/common/DatasourceFacadeTests.java | 130 +++++++ .../ip2geo/common/DatasourceHelperTests.java | 78 ---- .../ip2geo/common/GeoIpDataFacadeTests.java | 324 ++++++++++++++++ .../ip2geo/common/GeoIpDataHelperTests.java | 22 -- .../DatasourceExtensionTests.java | 47 +++ .../jobscheduler/DatasourceRunnerTests.java | 110 ++++++ .../ip2geo/jobscheduler/DatasourceTests.java | 54 ++- .../DatasourceUpdateServiceTests.java | 152 ++++++++ .../ip2geo/processor/Ip2GeoCacheTests.java | 58 --- .../processor/Ip2GeoProcessorTests.java | 346 ++++++++++++++++++ .../plugin/GeospatialPluginTests.java | 118 +++++- src/test/resources/ip2geo/manifest.json | 8 + .../ip2geo/manifest_invalid_url.json | 8 + .../resources/ip2geo/manifest_template.json | 8 + .../sample_invalid_less_than_two_fields.csv | 2 + src/test/resources/ip2geo/sample_valid.csv | 3 + src/test/resources/ip2geo/sample_valid.zip | Bin 0 -> 250 bytes 38 files changed, 2424 insertions(+), 871 deletions(-) create mode 100644 lombok.config create mode 100644 src/main/java/org/opensearch/geospatial/annotation/VisibleForTesting.java rename src/main/java/org/opensearch/geospatial/ip2geo/common/{DatasourceHelper.java => DatasourceFacade.java} (83%) rename src/main/java/org/opensearch/geospatial/ip2geo/common/{GeoIpDataHelper.java => GeoIpDataFacade.java} (77%) rename src/main/java/org/opensearch/geospatial/ip2geo/common/{Ip2GeoExecutorHelper.java => Ip2GeoExecutor.java} (85%) create mode 100644 src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateService.java delete mode 100644 src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCache.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoTestCase.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportActionTests.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceFacadeTests.java delete mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceHelperTests.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataFacadeTests.java delete mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataHelperTests.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtensionTests.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunnerTests.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateServiceTests.java delete mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCacheTests.java create mode 100644 src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorTests.java create mode 100644 src/test/resources/ip2geo/manifest.json create mode 100644 src/test/resources/ip2geo/manifest_invalid_url.json create mode 100644 src/test/resources/ip2geo/manifest_template.json create mode 100644 src/test/resources/ip2geo/sample_invalid_less_than_two_fields.csv create mode 100644 src/test/resources/ip2geo/sample_valid.csv create mode 100644 src/test/resources/ip2geo/sample_valid.zip diff --git a/lombok.config b/lombok.config new file mode 100644 index 00000000..e69de29b diff --git a/src/main/java/org/opensearch/geospatial/annotation/VisibleForTesting.java b/src/main/java/org/opensearch/geospatial/annotation/VisibleForTesting.java new file mode 100644 index 00000000..d48c6dc2 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/annotation/VisibleForTesting.java @@ -0,0 +1,12 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.annotation; + +public @interface VisibleForTesting { +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequest.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequest.java index 24266b0d..82e513f0 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequest.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequest.java @@ -14,6 +14,7 @@ import java.net.URL; import java.util.Locale; +import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.Setter; import lombok.extern.log4j.Log4j2; @@ -28,11 +29,12 @@ import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; /** - * GeoIP datasource creation request + * Ip2Geo datasource creation request */ @Getter @Setter @Log4j2 +@EqualsAndHashCode public class PutDatasourceRequest extends AcknowledgedRequest { private static final ParseField ENDPOINT_FIELD = new ParseField("endpoint"); private static final ParseField UPDATE_INTERVAL_IN_DAYS_FIELD = new ParseField("update_interval_in_days"); @@ -47,10 +49,10 @@ public class PutDatasourceRequest extends AcknowledgedRequest("put_datasource"); PARSER.declareString((request, val) -> request.setEndpoint(val), ENDPOINT_FIELD); - PARSER.declareLong((request, val) -> request.setUpdateIntervalInDays(TimeValue.timeValueDays(val)), UPDATE_INTERVAL_IN_DAYS_FIELD); + PARSER.declareLong((request, val) -> request.setUpdateInterval(TimeValue.timeValueDays(val)), UPDATE_INTERVAL_IN_DAYS_FIELD); } /** @@ -79,7 +81,7 @@ public PutDatasourceRequest(final StreamInput in) throws IOException { super(in); this.datasourceName = in.readString(); this.endpoint = in.readString(); - this.updateIntervalInDays = in.readTimeValue(); + this.updateInterval = in.readTimeValue(); } @Override @@ -87,7 +89,7 @@ public void writeTo(final StreamOutput out) throws IOException { super.writeTo(out); out.writeString(datasourceName); out.writeString(endpoint); - out.writeTimeValue(updateIntervalInDays); + out.writeTimeValue(updateInterval); } @Override @@ -120,7 +122,7 @@ private void validateEndpoint(final ActionRequestValidationException errors) { * Conduct following validation on url * 1. can read manifest file from the endpoint * 2. the url in the manifest file complies with RFC-2396 - * 3. updateIntervalInDays is less than validForInDays value in the manifest file + * 3. updateInterval is less than validForInDays value in the manifest file * * @param url the url to validate * @param errors the errors to add error messages @@ -143,12 +145,12 @@ private void validateManifestFile(final URL url, final ActionRequestValidationEx return; } - if (updateIntervalInDays.days() >= manifest.getValidForInDays()) { + if (updateInterval.days() >= manifest.getValidForInDays()) { errors.addValidationError( String.format( Locale.ROOT, - "updateInterval %d is should be smaller than %d", - updateIntervalInDays.days(), + "updateInterval %d should be smaller than %d", + updateInterval.days(), manifest.getValidForInDays() ) ); @@ -156,12 +158,12 @@ private void validateManifestFile(final URL url, final ActionRequestValidationEx } /** - * Validate updateIntervalInDays is larger than 0 + * Validate updateInterval is equal or larger than 1 * * @param errors the errors to add error messages */ private void validateUpdateInterval(final ActionRequestValidationException errors) { - if (updateIntervalInDays.compareTo(TimeValue.timeValueDays(1)) > 0) { + if (updateInterval.compareTo(TimeValue.timeValueDays(1)) < 0) { errors.addValidationError("Update interval should be equal to or larger than 1 day"); } } diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportAction.java index 291e1087..96d15026 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportAction.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportAction.java @@ -8,16 +8,10 @@ package org.opensearch.geospatial.ip2geo.action; -import java.io.IOException; -import java.net.URL; -import java.time.Duration; import java.time.Instant; -import java.util.Iterator; import lombok.extern.log4j.Log4j2; -import org.apache.commons.csv.CSVParser; -import org.apache.commons.csv.CSVRecord; import org.opensearch.ResourceAlreadyExistsException; import org.opensearch.action.ActionListener; import org.opensearch.action.DocWriteRequest; @@ -28,19 +22,14 @@ import org.opensearch.action.support.WriteRequest; import org.opensearch.action.support.master.AcknowledgedResponse; import org.opensearch.client.Client; -import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; -import org.opensearch.common.settings.ClusterSettings; -import org.opensearch.common.settings.Settings; -import org.opensearch.common.unit.TimeValue; import org.opensearch.common.xcontent.json.JsonXContent; -import org.opensearch.geospatial.ip2geo.common.DatasourceHelper; -import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.ip2geo.common.DatasourceFacade; import org.opensearch.geospatial.ip2geo.common.DatasourceState; -import org.opensearch.geospatial.ip2geo.common.GeoIpDataHelper; -import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; import org.opensearch.index.engine.VersionConflictEngineException; import org.opensearch.tasks.Task; import org.opensearch.threadpool.ThreadPool; @@ -52,138 +41,96 @@ @Log4j2 public class PutDatasourceTransportAction extends HandledTransportAction { private final Client client; - private final ClusterService clusterService; private final ThreadPool threadPool; - - private TimeValue timeout; - private int indexingBulkSize; + private final DatasourceFacade datasourceFacade; + private final DatasourceUpdateService datasourceUpdateService; /** * Default constructor * @param transportService the transport service * @param actionFilters the action filters * @param client the client - * @param clusterService the cluster service * @param threadPool the thread pool - * @param settings the settings - * @param clusterSettings the cluster settings */ @Inject public PutDatasourceTransportAction( final TransportService transportService, final ActionFilters actionFilters, final Client client, - final ClusterService clusterService, final ThreadPool threadPool, - final Settings settings, - final ClusterSettings clusterSettings + final DatasourceFacade datasourceFacade, + final DatasourceUpdateService datasourceUpdateService ) { super(PutDatasourceAction.NAME, transportService, actionFilters, PutDatasourceRequest::new); this.client = client; - this.clusterService = clusterService; this.threadPool = threadPool; - timeout = Ip2GeoSettings.TIMEOUT_IN_SECONDS.get(settings); - clusterSettings.addSettingsUpdateConsumer(Ip2GeoSettings.TIMEOUT_IN_SECONDS, newValue -> timeout = newValue); - indexingBulkSize = Ip2GeoSettings.INDEXING_BULK_SIZE.get(settings); - clusterSettings.addSettingsUpdateConsumer(Ip2GeoSettings.INDEXING_BULK_SIZE, newValue -> indexingBulkSize = newValue); + this.datasourceFacade = datasourceFacade; + this.datasourceUpdateService = datasourceUpdateService; } @Override protected void doExecute(final Task task, final PutDatasourceRequest request, final ActionListener listener) { try { - Datasource jobParameter = Datasource.Builder.build(request); + Datasource datasource = Datasource.Builder.build(request); IndexRequest indexRequest = new IndexRequest().index(DatasourceExtension.JOB_INDEX_NAME) - .id(jobParameter.getId()) - .source(jobParameter.toXContent(JsonXContent.contentBuilder(), null)) + .id(datasource.getId()) + .source(datasource.toXContent(JsonXContent.contentBuilder(), null)) .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) .opType(DocWriteRequest.OpType.CREATE); - client.index(indexRequest, new ActionListener<>() { - @Override - public void onResponse(final IndexResponse indexResponse) { - // This is user initiated request. Therefore, we want to handle the first datasource update task in a generic thread - // pool. - threadPool.generic().submit(() -> { - try { - createDatasource(jobParameter); - } catch (Exception e) { - log.error("Failed to create datasource for {}", jobParameter.getId(), e); - jobParameter.getUpdateStats().setLastFailedAt(Instant.now()); - jobParameter.setState(DatasourceState.FAILED); - try { - DatasourceHelper.updateDatasource(client, jobParameter, timeout); - } catch (Exception ex) { - log.error("Failed to mark datasource state as FAILED for {}", jobParameter.getId(), ex); - } - } - }); - listener.onResponse(new AcknowledgedResponse(true)); - } - - @Override - public void onFailure(final Exception e) { - if (e instanceof VersionConflictEngineException) { - listener.onFailure( - new ResourceAlreadyExistsException("datasource [{}] already exists", request.getDatasourceName()) - ); - } else { - listener.onFailure(e); - } - } - }); + client.index(indexRequest, getIndexResponseListener(datasource, listener)); } catch (Exception e) { listener.onFailure(e); } } - private void createDatasource(final Datasource jobParameter) throws Exception { - if (!DatasourceState.PREPARING.equals(jobParameter.getState())) { - log.error("Invalid datasource state. Expecting {} but received {}", DatasourceState.AVAILABLE, jobParameter.getState()); - jobParameter.setState(DatasourceState.FAILED); - jobParameter.getUpdateStats().setLastFailedAt(Instant.now()); - DatasourceHelper.updateDatasource(client, jobParameter, timeout); - return; - } + @VisibleForTesting + protected ActionListener getIndexResponseListener( + final Datasource datasource, + final ActionListener listener + ) { + return new ActionListener<>() { + @Override + public void onResponse(final IndexResponse indexResponse) { + // This is user initiated request. Therefore, we want to handle the first datasource update task in a generic thread + // pool. + threadPool.generic().submit(() -> { createDatasource(datasource); }); + listener.onResponse(new AcknowledgedResponse(true)); + } - URL url = new URL(jobParameter.getEndpoint()); - DatasourceManifest manifest = DatasourceManifest.Builder.build(url); - String indexName = setupIndex(manifest, jobParameter); - Instant startTime = Instant.now(); - String[] fields = putIp2GeoData(indexName, manifest); - Instant endTime = Instant.now(); - updateJobParameterAsSucceeded(jobParameter, manifest, fields, startTime, endTime); - log.info("GeoIP database[{}] creation succeeded after {} seconds", jobParameter.getId(), Duration.between(startTime, endTime)); + @Override + public void onFailure(final Exception e) { + if (e instanceof VersionConflictEngineException) { + listener.onFailure(new ResourceAlreadyExistsException("datasource [{}] already exists", datasource.getId())); + } else { + listener.onFailure(e); + } + } + }; } - private void updateJobParameterAsSucceeded( - final Datasource jobParameter, - final DatasourceManifest manifest, - final String[] fields, - final Instant startTime, - final Instant endTime - ) throws IOException { - jobParameter.setDatabase(manifest, fields); - jobParameter.getUpdateStats().setLastSucceededAt(endTime); - jobParameter.getUpdateStats().setLastProcessingTimeInMillis(endTime.toEpochMilli() - startTime.toEpochMilli()); - jobParameter.enable(); - jobParameter.setState(DatasourceState.AVAILABLE); - DatasourceHelper.updateDatasource(client, jobParameter, timeout); - } + @VisibleForTesting + protected void createDatasource(final Datasource datasource) { + if (!DatasourceState.CREATING.equals(datasource.getState())) { + log.error("Invalid datasource state. Expecting {} but received {}", DatasourceState.CREATING, datasource.getState()); + markDatasourceAsCreateFailed(datasource); + return; + } - private String setupIndex(final DatasourceManifest manifest, final Datasource jobParameter) throws IOException { - String indexName = jobParameter.indexNameFor(manifest); - jobParameter.getIndices().add(indexName); - DatasourceHelper.updateDatasource(client, jobParameter, timeout); - GeoIpDataHelper.createIndexIfNotExists(clusterService, client, indexName, timeout); - return indexName; + try { + datasourceUpdateService.updateOrCreateGeoIpData(datasource); + } catch (Exception e) { + log.error("Failed to create datasource for {}", datasource.getId(), e); + markDatasourceAsCreateFailed(datasource); + } } - private String[] putIp2GeoData(final String indexName, final DatasourceManifest manifest) throws IOException { - String[] fields; - try (CSVParser reader = GeoIpDataHelper.getDatabaseReader(manifest)) { - Iterator iter = reader.iterator(); - fields = iter.next().values(); - GeoIpDataHelper.putGeoData(client, indexName, fields, iter, indexingBulkSize, timeout); + private void markDatasourceAsCreateFailed(final Datasource datasource) { + datasource.getUpdateStats().setLastFailedAt(Instant.now()); + datasource.setState(DatasourceState.CREATE_FAILED); + try { + datasourceFacade.updateDatasource(datasource); + } catch (Exception e) { + log.error("Failed to mark datasource state as CREATE_FAILED for {}", datasource.getId(), e); } - return fields; } } diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandler.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandler.java index 00c70f19..3ccffa06 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandler.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandler.java @@ -17,7 +17,6 @@ import org.opensearch.client.node.NodeClient; import org.opensearch.common.settings.ClusterSettings; -import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; @@ -41,14 +40,10 @@ */ public class RestPutDatasourceHandler extends BaseRestHandler { private static final String ACTION_NAME = "ip2geo_datasource"; - private String defaultDatasourceEndpoint; - private TimeValue defaultUpdateInterval; + private final ClusterSettings clusterSettings; - public RestPutDatasourceHandler(final Settings settings, final ClusterSettings clusterSettings) { - defaultDatasourceEndpoint = Ip2GeoSettings.DATASOURCE_ENDPOINT.get(settings); - clusterSettings.addSettingsUpdateConsumer(Ip2GeoSettings.DATASOURCE_ENDPOINT, newValue -> defaultDatasourceEndpoint = newValue); - defaultUpdateInterval = Ip2GeoSettings.DATASOURCE_UPDATE_INTERVAL.get(settings); - clusterSettings.addSettingsUpdateConsumer(Ip2GeoSettings.DATASOURCE_UPDATE_INTERVAL, newValue -> defaultUpdateInterval = newValue); + public RestPutDatasourceHandler(final ClusterSettings clusterSettings) { + this.clusterSettings = clusterSettings; } @Override @@ -65,10 +60,10 @@ protected RestChannelConsumer prepareRequest(final RestRequest request, final No } } if (putDatasourceRequest.getEndpoint() == null) { - putDatasourceRequest.setEndpoint(defaultDatasourceEndpoint); + putDatasourceRequest.setEndpoint(clusterSettings.get(Ip2GeoSettings.DATASOURCE_ENDPOINT)); } - if (putDatasourceRequest.getUpdateIntervalInDays() == null) { - putDatasourceRequest.setUpdateIntervalInDays(defaultUpdateInterval); + if (putDatasourceRequest.getUpdateInterval() == null) { + putDatasourceRequest.setUpdateInterval(TimeValue.timeValueDays(clusterSettings.get(Ip2GeoSettings.DATASOURCE_UPDATE_INTERVAL))); } return channel -> client.executeLocally(PutDatasourceAction.INSTANCE, putDatasourceRequest, new RestToXContentListener<>(channel)); } diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceHelper.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceFacade.java similarity index 83% rename from src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceHelper.java rename to src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceFacade.java index ea89e585..db8d3cff 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceHelper.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceFacade.java @@ -20,7 +20,7 @@ import org.opensearch.action.index.IndexRequestBuilder; import org.opensearch.action.index.IndexResponse; import org.opensearch.client.Client; -import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.xcontent.LoggingDeprecationHandler; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.common.xcontent.XContentHelper; @@ -32,42 +32,44 @@ import org.opensearch.index.IndexNotFoundException; /** - * Helper class for datasource + * Facade class for datasource */ @Log4j2 -public class DatasourceHelper { +public class DatasourceFacade { + private final Client client; + private final ClusterSettings clusterSettings; + + public DatasourceFacade(final Client client, final ClusterSettings clusterSettings) { + this.client = client; + this.clusterSettings = clusterSettings; + } /** * Update datasource in an index {@code DatasourceExtension.JOB_INDEX_NAME} - * @param client the client * @param datasource the datasource - * @param timeout the timeout * @return index response * @throws IOException exception */ - public static IndexResponse updateDatasource(final Client client, final Datasource datasource, final TimeValue timeout) - throws IOException { + public IndexResponse updateDatasource(final Datasource datasource) throws IOException { datasource.setLastUpdateTime(Instant.now()); IndexRequestBuilder requestBuilder = client.prepareIndex(DatasourceExtension.JOB_INDEX_NAME); requestBuilder.setId(datasource.getId()); requestBuilder.setOpType(DocWriteRequest.OpType.INDEX); requestBuilder.setSource(datasource.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)); - return client.index(requestBuilder.request()).actionGet(timeout); + return client.index(requestBuilder.request()).actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)); } /** * Get datasource from an index {@code DatasourceExtension.JOB_INDEX_NAME} - * @param client the client * @param id the name of a datasource - * @param timeout the timeout * @return datasource * @throws IOException exception */ - public static Datasource getDatasource(final Client client, final String id, final TimeValue timeout) throws IOException { + public Datasource getDatasource(final String id) throws IOException { GetRequest request = new GetRequest(DatasourceExtension.JOB_INDEX_NAME, id); GetResponse response; try { - response = client.get(request).actionGet(timeout); + response = client.get(request).actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)); if (!response.isExists()) { log.error("Datasource[{}] does not exist in an index[{}]", id, DatasourceExtension.JOB_INDEX_NAME); return null; @@ -87,11 +89,10 @@ public static Datasource getDatasource(final Client client, final String id, fin /** * Get datasource from an index {@code DatasourceExtension.JOB_INDEX_NAME} - * @param client the client * @param id the name of a datasource * @param actionListener the action listener */ - public static void getDatasource(final Client client, final String id, final ActionListener actionListener) { + public void getDatasource(final String id, final ActionListener actionListener) { GetRequest request = new GetRequest(DatasourceExtension.JOB_INDEX_NAME, id); client.get(request, new ActionListener() { @Override diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceState.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceState.java index 27523bda..85b0aecf 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceState.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceState.java @@ -11,28 +11,28 @@ /** * Ip2Geo datasource state * - * When data source is created, it starts with PREPARING state. Once the first GeoIP data is generated, the state changes to AVAILABLE. - * Only when the first GeoIP data generation failed, the state changes to FAILED. - * Subsequent GeoIP data failure won't change data source state from AVAILABLE to FAILED. + * When data source is created, it starts with CREATING state. Once the first GeoIP data is generated, the state changes to AVAILABLE. + * Only when the first GeoIP data generation failed, the state changes to CREATE_FAILED. + * Subsequent GeoIP data failure won't change data source state from AVAILABLE to CREATE_FAILED. * When delete request is received, the data source state changes to DELETING. * * State changed from left to right for the entire lifecycle of a datasource - * (PREPARING) to (FAILED or AVAILABLE) to (DELETING) + * (CREATING) to (CREATE_FAILED or AVAILABLE) to (DELETING) * */ public enum DatasourceState { /** - * Data source is being prepared + * Data source is being created */ - PREPARING, + CREATING, /** * Data source is ready to be used */ AVAILABLE, /** - * Data source preparation failed + * Data source creation failed */ - FAILED, + CREATE_FAILED, /** * Data source is being deleted */ diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataHelper.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataFacade.java similarity index 77% rename from src/main/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataHelper.java rename to src/main/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataFacade.java index 4543ce7f..264cfef4 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataHelper.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataFacade.java @@ -8,6 +8,8 @@ package org.opensearch.geospatial.ip2geo.common; +import static org.opensearch.geospatial.ip2geo.jobscheduler.Datasource.IP2GEO_DATA_INDEX_NAME_PREFIX; + import java.io.BufferedReader; import java.io.IOException; import java.io.InputStream; @@ -41,41 +43,45 @@ import org.opensearch.action.search.MultiSearchRequestBuilder; import org.opensearch.action.search.MultiSearchResponse; import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.IndicesOptions; import org.opensearch.action.support.WriteRequest; +import org.opensearch.action.support.master.AcknowledgedResponse; import org.opensearch.client.Client; import org.opensearch.client.Requests; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.SuppressForbidden; import org.opensearch.common.collect.Tuple; +import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.common.xcontent.XContentType; import org.opensearch.index.query.QueryBuilders; /** - * Helper class for GeoIp data + * Facade class for GeoIp data */ @Log4j2 -public class GeoIpDataHelper { +public class GeoIpDataFacade { private static final String IP_RANGE_FIELD_NAME = "_cidr"; private static final String DATA_FIELD_NAME = "_data"; private static final Tuple INDEX_SETTING_NUM_OF_SHARDS = new Tuple<>("index.number_of_shards", 1); private static final Tuple INDEX_SETTING_AUTO_EXPAND_REPLICAS = new Tuple<>("index.auto_expand_replicas", "0-all"); + private final ClusterService clusterService; + private final ClusterSettings clusterSettings; + private final Client client; + + public GeoIpDataFacade(final ClusterService clusterService, final Client client) { + this.clusterService = clusterService; + this.clusterSettings = clusterService.getClusterSettings(); + this.client = client; + } /** * Create an index of single shard with auto expand replicas to all nodes * - * @param clusterService cluster service - * @param client client * @param indexName index name - * @param timeout timeout */ - public static void createIndexIfNotExists( - final ClusterService clusterService, - final Client client, - final String indexName, - final TimeValue timeout - ) { + public void createIndexIfNotExists(final String indexName) { if (clusterService.state().metadata().hasIndex(indexName) == true) { return; } @@ -83,7 +89,7 @@ public static void createIndexIfNotExists( indexSettings.put(INDEX_SETTING_NUM_OF_SHARDS.v1(), INDEX_SETTING_NUM_OF_SHARDS.v2()); indexSettings.put(INDEX_SETTING_AUTO_EXPAND_REPLICAS.v1(), INDEX_SETTING_AUTO_EXPAND_REPLICAS.v2()); final CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName).settings(indexSettings).mapping(getIndexMapping()); - client.admin().indices().create(createIndexRequest).actionGet(timeout); + client.admin().indices().create(createIndexRequest).actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)); } /** @@ -101,11 +107,11 @@ public static void createIndexIfNotExists( * * @return String representing datasource database index mapping */ - private static String getIndexMapping() { + private String getIndexMapping() { try { - try (InputStream is = DatasourceHelper.class.getResourceAsStream("/mappings/ip2geo_datasource.json")) { + try (InputStream is = DatasourceFacade.class.getResourceAsStream("/mappings/ip2geo_datasource.json")) { try (BufferedReader reader = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) { - return reader.lines().collect(Collectors.joining()); + return reader.lines().map(line -> line.trim()).collect(Collectors.joining()); } } } catch (IOException e) { @@ -120,7 +126,7 @@ private static String getIndexMapping() { * @return CSVParser for GeoIP data */ @SuppressForbidden(reason = "Need to connect to http endpoint to read GeoIP database file") - public static CSVParser getDatabaseReader(final DatasourceManifest manifest) { + public CSVParser getDatabaseReader(final DatasourceManifest manifest) { SpecialPermission.check(); return AccessController.doPrivileged((PrivilegedAction) () -> { try { @@ -164,7 +170,10 @@ public static CSVParser getDatabaseReader(final DatasourceManifest manifest) { * @param values a list of values * @return Document in json string format */ - public static String createDocument(final String[] fields, final String[] values) { + public String createDocument(final String[] fields, final String[] values) { + if (fields.length != values.length) { + throw new OpenSearchException("header[{}] and record[{}] length does not match", fields, values); + } StringBuilder sb = new StringBuilder(); sb.append("{\""); sb.append(IP_RANGE_FIELD_NAME); @@ -188,35 +197,30 @@ public static String createDocument(final String[] fields, final String[] values } /** - * Query a given index using a given ip address to get geo data + * Query a given index using a given ip address to get geoip data * - * @param client client * @param indexName index * @param ip ip address * @param actionListener action listener */ - public static void getGeoData( - final Client client, - final String indexName, - final String ip, - final ActionListener> actionListener - ) { + public void getGeoIpData(final String indexName, final String ip, final ActionListener> actionListener) { client.prepareSearch(indexName) .setSize(1) .setQuery(QueryBuilders.termQuery(IP_RANGE_FIELD_NAME, ip)) .setPreference("_local") + .setRequestCache(true) .execute(new ActionListener<>() { @Override public void onResponse(final SearchResponse searchResponse) { if (searchResponse.getHits().getHits().length == 0) { actionListener.onResponse(Collections.emptyMap()); } else { - Map geoData = (Map) XContentHelper.convertToMap( + Map geoIpData = (Map) XContentHelper.convertToMap( searchResponse.getHits().getAt(0).getSourceRef(), false, XContentType.JSON ).v2().get(DATA_FIELD_NAME); - actionListener.onResponse(geoData); + actionListener.onResponse(geoIpData); } } @@ -228,28 +232,26 @@ public void onFailure(final Exception e) { } /** - * Query a given index using a given ip address iterator to get geo data + * Query a given index using a given ip address iterator to get geoip data * * This method calls itself recursively until it processes all ip addresses in bulk of {@code bulkSize}. * - * @param client the client * @param indexName the index name * @param ipIterator the iterator of ip addresses * @param maxBundleSize number of ip address to pass in multi search * @param maxConcurrentSearches the max concurrent search requests * @param firstOnly return only the first matching result if true - * @param geoData collected geo data + * @param geoIpData collected geo data * @param actionListener the action listener */ - public static void getGeoData( - final Client client, + public void getGeoIpData( final String indexName, final Iterator ipIterator, final Integer maxBundleSize, final Integer maxConcurrentSearches, final boolean firstOnly, - final Map> geoData, - final ActionListener actionListener + final Map> geoIpData, + final ActionListener>> actionListener ) { MultiSearchRequestBuilder mRequestBuilder = client.prepareMultiSearch(); if (maxConcurrentSearches != 0) { @@ -259,19 +261,20 @@ public static void getGeoData( List ipsToSearch = new ArrayList<>(maxBundleSize); while (ipIterator.hasNext() && ipsToSearch.size() < maxBundleSize) { String ip = ipIterator.next(); - if (geoData.get(ip) == null) { + if (geoIpData.get(ip) == null) { mRequestBuilder.add( client.prepareSearch(indexName) .setSize(1) .setQuery(QueryBuilders.termQuery(IP_RANGE_FIELD_NAME, ip)) .setPreference("_local") + .setRequestCache(true) ); ipsToSearch.add(ip); } } if (ipsToSearch.isEmpty()) { - actionListener.onResponse(null); + actionListener.onResponse(geoIpData); return; } @@ -285,7 +288,7 @@ public void onResponse(final MultiSearchResponse items) { } if (items.getResponses()[i].getResponse().getHits().getHits().length == 0) { - geoData.put(ipsToSearch.get(i), Collections.emptyMap()); + geoIpData.put(ipsToSearch.get(i), Collections.emptyMap()); continue; } @@ -295,14 +298,14 @@ public void onResponse(final MultiSearchResponse items) { XContentType.JSON ).v2().get(DATA_FIELD_NAME); - geoData.put(ipsToSearch.get(i), data); + geoIpData.put(ipsToSearch.get(i), data); if (firstOnly) { - actionListener.onResponse(null); + actionListener.onResponse(geoIpData); return; } } - getGeoData(client, indexName, ipIterator, maxBundleSize, maxConcurrentSearches, firstOnly, geoData, actionListener); + getGeoIpData(indexName, ipIterator, maxBundleSize, maxConcurrentSearches, firstOnly, geoIpData, actionListener); } @Override @@ -315,21 +318,13 @@ public void onFailure(final Exception e) { /** * Puts GeoIP data from CSVRecord iterator into a given index in bulk * - * @param client OpenSearch client * @param indexName Index name to puts the GeoIP data * @param fields Field name matching with data in CSVRecord in order * @param iterator GeoIP data to insert * @param bulkSize Bulk size of data to process - * @param timeout Timeout */ - public static void putGeoData( - final Client client, - final String indexName, - final String[] fields, - final Iterator iterator, - final int bulkSize, - final TimeValue timeout - ) { + public void putGeoIpData(final String indexName, final String[] fields, final Iterator iterator, final int bulkSize) { + TimeValue timeout = clusterSettings.get(Ip2GeoSettings.TIMEOUT); BulkRequest bulkRequest = new BulkRequest().setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL); while (iterator.hasNext()) { CSVRecord record = iterator.next(); @@ -351,4 +346,20 @@ public static void putGeoData( client.admin().indices().prepareRefresh(indexName).execute().actionGet(timeout); client.admin().indices().prepareForceMerge(indexName).setMaxNumSegments(1).execute().actionGet(timeout); } + + public AcknowledgedResponse deleteIp2GeoDataIndex(final String index) { + if (index.startsWith(IP2GEO_DATA_INDEX_NAME_PREFIX) == false) { + throw new OpenSearchException( + "the index[{}] is not ip2geo data index which should start with {}", + index, + IP2GEO_DATA_INDEX_NAME_PREFIX + ); + } + return client.admin() + .indices() + .prepareDelete(index) + .setIndicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN) + .execute() + .actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)); + } } diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutorHelper.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutor.java similarity index 85% rename from src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutorHelper.java rename to src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutor.java index 7c5d77f9..c2127482 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutorHelper.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutor.java @@ -18,8 +18,13 @@ /** * Provide a list of static methods related with executors for Ip2Geo */ -public class Ip2GeoExecutorHelper { +public class Ip2GeoExecutor { private static final String THREAD_POOL_NAME = "_plugin_geospatial_ip2geo_datasource_update"; + private final ThreadPool threadPool; + + public Ip2GeoExecutor(final ThreadPool threadPool) { + this.threadPool = threadPool; + } /** * We use fixed thread count of 1 for updating datasource as updating datasource is running background @@ -35,10 +40,9 @@ public static ExecutorBuilder executorBuilder(final Settings settings) { /** * Return an executor service for datasource update task * - * @param threadPool the thread pool * @return the executor service */ - public static ExecutorService forDatasourceUpdate(final ThreadPool threadPool) { + public ExecutorService forDatasourceUpdate() { return threadPool.executor(THREAD_POOL_NAME); } } diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettings.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettings.java index 5df30c04..f8348435 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettings.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettings.java @@ -36,10 +36,10 @@ public class Ip2GeoSettings { /** * Default update interval to be used in Ip2Geo datasource creation API */ - public static final Setting DATASOURCE_UPDATE_INTERVAL = Setting.timeSetting( + public static final Setting DATASOURCE_UPDATE_INTERVAL = Setting.longSetting( "plugins.geospatial.ip2geo.datasource.update_interval_in_days", - TimeValue.timeValueDays(3), - TimeValue.timeValueDays(1), + 3l, + 1l, Setting.Property.NodeScope, Setting.Property.Dynamic ); @@ -47,8 +47,8 @@ public class Ip2GeoSettings { /** * Timeout value for Ip2Geo processor */ - public static final Setting TIMEOUT_IN_SECONDS = Setting.timeSetting( - "plugins.geospatial.ip2geo.timeout_in_seconds", + public static final Setting TIMEOUT = Setting.timeSetting( + "plugins.geospatial.ip2geo.timeout", TimeValue.timeValueSeconds(30), TimeValue.timeValueSeconds(1), Setting.Property.NodeScope, @@ -66,16 +66,6 @@ public class Ip2GeoSettings { Setting.Property.Dynamic ); - /** - * Cache size for GeoIP data - */ - public static final Setting CACHE_SIZE = Setting.intSetting( - "plugins.geospatial.ip2geo.processor.cache_size", - 1000, - 0, - Setting.Property.NodeScope - ); - /** * Multi search bundle size for GeoIP data * @@ -113,9 +103,8 @@ public static final List> settings() { return List.of( DATASOURCE_ENDPOINT, DATASOURCE_UPDATE_INTERVAL, - TIMEOUT_IN_SECONDS, + TIMEOUT, INDEXING_BULK_SIZE, - CACHE_SIZE, MAX_BUNDLE_SIZE, MAX_CONCURRENT_SEARCHES ); diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/Datasource.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/Datasource.java index 4d42f9ad..242cf143 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/Datasource.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/Datasource.java @@ -12,20 +12,24 @@ import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.ArrayList; -import java.util.Arrays; +import java.util.HashSet; import java.util.List; import java.util.Locale; +import java.util.Set; import lombok.AccessLevel; import lombok.AllArgsConstructor; +import lombok.EqualsAndHashCode; import lombok.Getter; import lombok.NoArgsConstructor; import lombok.Setter; +import lombok.ToString; import org.opensearch.core.ParseField; import org.opensearch.core.xcontent.ConstructingObjectParser; import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.geospatial.annotation.VisibleForTesting; import org.opensearch.geospatial.ip2geo.action.PutDatasourceRequest; import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; import org.opensearch.geospatial.ip2geo.common.DatasourceState; @@ -38,6 +42,8 @@ */ @Getter @Setter +@ToString +@EqualsAndHashCode @AllArgsConstructor public class Datasource implements ScheduledJobParameter { /** @@ -45,6 +51,9 @@ public class Datasource implements ScheduledJobParameter { */ public static final String IP2GEO_DATA_INDEX_NAME_PREFIX = ".ip2geo-data"; private static final long LOCK_DURATION_IN_SECONDS = 60 * 60; + private static final long MAX_JITTER_IN_MINUTES = 5; + private static final long ONE_DAY_IN_HOURS = 24; + private static final long ONE_HOUR_IN_MINUTES = 60; /** * Default fields for job scheduling @@ -173,22 +182,20 @@ public class Datasource implements ScheduledJobParameter { } - /** - * Visible for testing - */ - protected Datasource() { + @VisibleForTesting + public Datasource() { this(null, null, null); } public Datasource(final String id, final IntervalSchedule schedule, final String endpoint) { this( id, - Instant.now(), + Instant.now().truncatedTo(ChronoUnit.MILLIS), null, false, schedule, endpoint, - DatasourceState.PREPARING, + DatasourceState.CREATING, new ArrayList<>(), new Database(), new UpdateStats() @@ -264,14 +271,17 @@ public Long getLockDurationSeconds() { */ @Override public Double getJitter() { - return 5.0 / (schedule.getInterval() * 24 * 60); + return MAX_JITTER_IN_MINUTES / ((double) schedule.getInterval() * ONE_DAY_IN_HOURS * ONE_HOUR_IN_MINUTES); } /** * Enable auto update of GeoIP data */ public void enable() { - enabledTime = Instant.now(); + if (isEnabled == true) { + return; + } + enabledTime = Instant.now().truncatedTo(ChronoUnit.MILLIS); isEnabled = true; } @@ -306,6 +316,11 @@ private String indexNameFor(final long suffix) { return String.format(Locale.ROOT, "%s.%s.%d", IP2GEO_DATA_INDEX_NAME_PREFIX, id, suffix); } + /** + * Tell if current datasource is expired or not + * + * @return true if datasource is expired false otherwise + */ public boolean isExpired() { if (database.validForInDays == null) { return false; @@ -322,12 +337,45 @@ public boolean isExpired() { return Instant.now().isAfter(lastCheckedAt.plus(database.validForInDays, ChronoUnit.DAYS)); } - public void setDatabase(final DatasourceManifest datasourceManifest, final String[] fields) { + /** + * Set database attributes with given input + * + * @param datasourceManifest the datasource manifest + * @param fields the fields + */ + public void setDatabase(final DatasourceManifest datasourceManifest, final List fields) { this.database.setProvider(datasourceManifest.getProvider()); this.database.setMd5Hash(datasourceManifest.getMd5Hash()); this.database.setUpdatedAt(Instant.ofEpochMilli(datasourceManifest.getUpdatedAt())); - this.database.setValidForInDays(database.validForInDays); - this.database.setFields(Arrays.asList(fields)); + this.database.setValidForInDays(datasourceManifest.getValidForInDays()); + this.database.setFields(fields); + } + + /** + * Checks if the database fields are compatible with the given set of fields. + * + * If database fields are null, it is compatible with any input fields + * as it hasn't been generated before. + * + * @param fields The set of input fields to check for compatibility. + * @return true if the database fields are compatible with the given input fields, false otherwise. + */ + public boolean isCompatible(final List fields) { + if (database.fields == null) { + return true; + } + + if (fields.size() < database.fields.size()) { + return false; + } + + Set fieldsSet = new HashSet<>(fields); + for (String field : database.fields) { + if (!fieldsSet.contains(field)) { + return false; + } + } + return true; } /** @@ -335,6 +383,8 @@ public void setDatabase(final DatasourceManifest datasourceManifest, final Strin */ @Getter @Setter + @ToString + @EqualsAndHashCode @AllArgsConstructor(access = AccessLevel.PRIVATE) @NoArgsConstructor(access = AccessLevel.PRIVATE) public static class Database implements ToXContent { @@ -427,6 +477,8 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa */ @Getter @Setter + @ToString + @EqualsAndHashCode @AllArgsConstructor(access = AccessLevel.PRIVATE) @NoArgsConstructor(access = AccessLevel.PRIVATE) public static class UpdateStats implements ToXContent { @@ -517,8 +569,8 @@ public static class Builder { public static Datasource build(final PutDatasourceRequest request) { String id = request.getDatasourceName(); IntervalSchedule schedule = new IntervalSchedule( - Instant.now(), - (int) request.getUpdateIntervalInDays().days(), + Instant.now().truncatedTo(ChronoUnit.MILLIS), + (int) request.getUpdateInterval().days(), ChronoUnit.DAYS ); String endpoint = request.getEndpoint(); diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtension.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtension.java index ed94f6fe..6def5295 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtension.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtension.java @@ -15,7 +15,7 @@ /** * Datasource job scheduler extension * - * This extension is responsible for scheduling Ip2Geo data update task + * This extension is responsible for scheduling GeoIp data update task * * See https://github.com/opensearch-project/job-scheduler/blob/main/README.md#getting-started */ diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunner.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunner.java index c408ff3b..c0f4795c 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunner.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunner.java @@ -8,40 +8,27 @@ package org.opensearch.geospatial.ip2geo.jobscheduler; -import java.net.URL; -import java.time.Duration; +import java.io.IOException; import java.time.Instant; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Iterator; -import java.util.List; import lombok.extern.log4j.Log4j2; -import org.apache.commons.csv.CSVParser; -import org.apache.commons.csv.CSVRecord; -import org.opensearch.OpenSearchException; import org.opensearch.action.ActionListener; -import org.opensearch.action.support.IndicesOptions; import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; -import org.opensearch.common.unit.TimeValue; -import org.opensearch.geospatial.ip2geo.common.DatasourceHelper; -import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.ip2geo.common.DatasourceFacade; import org.opensearch.geospatial.ip2geo.common.DatasourceState; -import org.opensearch.geospatial.ip2geo.common.GeoIpDataHelper; -import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutorHelper; -import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; import org.opensearch.jobscheduler.spi.JobExecutionContext; import org.opensearch.jobscheduler.spi.ScheduledJobParameter; import org.opensearch.jobscheduler.spi.ScheduledJobRunner; import org.opensearch.jobscheduler.spi.utils.LockService; -import org.opensearch.threadpool.ThreadPool; /** * Datasource update task * - * This is a background task which is responsible for updating Ip2Geo datasource + * This is a background task which is responsible for updating GeoIp data */ @Log4j2 public class DatasourceRunner implements ScheduledJobRunner { @@ -66,10 +53,10 @@ public static DatasourceRunner getJobRunnerInstance() { } private ClusterService clusterService; - private ThreadPool threadPool; private Client client; - private TimeValue timeout; - private Integer indexingBulkSize; + private DatasourceUpdateService datasourceUpdateService; + private Ip2GeoExecutor ip2GeoExecutor; + private DatasourceFacade datasourceFacade; private boolean initialized; private DatasourceRunner() { @@ -79,17 +66,18 @@ private DatasourceRunner() { /** * Initialize timeout and indexingBulkSize from settings */ - public void initialize(final ClusterService clusterService, final ThreadPool threadPool, final Client client) { + public void initialize( + final ClusterService clusterService, + final Client client, + final DatasourceUpdateService datasourceUpdateService, + final Ip2GeoExecutor ip2GeoExecutor, + final DatasourceFacade datasourceFacade + ) { this.clusterService = clusterService; - this.threadPool = threadPool; this.client = client; - - this.timeout = Ip2GeoSettings.TIMEOUT_IN_SECONDS.get(clusterService.getSettings()); - clusterService.getClusterSettings() - .addSettingsUpdateConsumer(Ip2GeoSettings.TIMEOUT_IN_SECONDS, newValue -> this.timeout = newValue); - this.indexingBulkSize = Ip2GeoSettings.INDEXING_BULK_SIZE.get(clusterService.getSettings()); - clusterService.getClusterSettings() - .addSettingsUpdateConsumer(Ip2GeoSettings.INDEXING_BULK_SIZE, newValue -> this.indexingBulkSize = newValue); + this.datasourceUpdateService = datasourceUpdateService; + this.ip2GeoExecutor = ip2GeoExecutor; + this.datasourceFacade = datasourceFacade; this.initialized = true; } @@ -102,11 +90,11 @@ public void runJob(final ScheduledJobParameter jobParameter, final JobExecutionC log.info("Update job started for a datasource[{}]", jobParameter.getName()); if (jobParameter instanceof Datasource == false) { throw new IllegalStateException( - "job parameter is not instance of DatasourceUpdateJobParameter, type: " + jobParameter.getClass().getCanonicalName() + "job parameter is not instance of Datasource, type: " + jobParameter.getClass().getCanonicalName() ); } - Ip2GeoExecutorHelper.forDatasourceUpdate(threadPool).submit(updateDatasourceRunner(jobParameter, context)); + ip2GeoExecutor.forDatasourceUpdate().submit(updateDatasourceRunner(jobParameter, context)); } /** @@ -118,161 +106,50 @@ public void runJob(final ScheduledJobParameter jobParameter, final JobExecutionC * @param jobParameter job parameter * @param context context */ - private Runnable updateDatasourceRunner(final ScheduledJobParameter jobParameter, final JobExecutionContext context) { + @VisibleForTesting + protected Runnable updateDatasourceRunner(final ScheduledJobParameter jobParameter, final JobExecutionContext context) { final LockService lockService = context.getLockService(); return () -> { - if (jobParameter.getLockDurationSeconds() != null) { - lockService.acquireLock(jobParameter, context, ActionListener.wrap(lock -> { - if (lock == null) { - return; - } - try { - Datasource datasource = DatasourceHelper.getDatasource(client, jobParameter.getName(), timeout); - if (datasource == null) { - log.info("Datasource[{}] is already deleted", jobParameter.getName()); - return; - } - - try { - deleteUnusedIndices(datasource); - updateDatasource(datasource); - deleteUnusedIndices(datasource); - } catch (Exception e) { - log.error("Failed to update datasource for {}", datasource.getId(), e); - datasource.getUpdateStats().setLastFailedAt(Instant.now()); - DatasourceHelper.updateDatasource(client, datasource, timeout); - } - } finally { - lockService.release( - lock, - ActionListener.wrap(released -> {}, exception -> { log.error("Failed to release lock [{}]", lock); }) - ); - } - }, exception -> { log.error("Failed to acquire lock for job [{}]", jobParameter.getName()); })); - } - }; - - } - - /** - * Delete all indices except the one which are being used - * - * @param parameter - */ - private void deleteUnusedIndices(final Datasource parameter) { - try { - List deletedIndices = new ArrayList<>(); - for (String index : parameter.getIndices()) { - if (index.equals(parameter.currentIndexName())) { - continue; + lockService.acquireLock(jobParameter, context, ActionListener.wrap(lock -> { + if (lock == null) { + return; } - - if (!clusterService.state().metadata().hasIndex(index)) { - deletedIndices.add(index); - continue; - } - try { - if (client.admin() - .indices() - .prepareDelete(index) - .setIndicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN) - .execute() - .actionGet(timeout) - .isAcknowledged()) { - deletedIndices.add(index); - } else { - log.error("Failed to delete an index [{}]", index); - } - } catch (Exception e) { - log.error("Failed to delete an index [{}]", index, e); + updateDatasource(jobParameter); + } finally { + lockService.release( + lock, + ActionListener.wrap(released -> {}, exception -> { log.error("Failed to release lock [{}]", lock, exception); }) + ); } - } - if (!deletedIndices.isEmpty()) { - parameter.getIndices().removeAll(deletedIndices); - DatasourceHelper.updateDatasource(client, parameter, timeout); - } - } catch (Exception e) { - log.error("Failed to delete old indices for {}", parameter.getId(), e); - } + }, exception -> { log.error("Failed to acquire lock for job [{}]", jobParameter.getName(), exception); })); + }; } - /** - * Update GeoIP data internal - * - * @param jobParameter - * @throws Exception - */ - private void updateDatasource(final Datasource jobParameter) throws Exception { - if (!DatasourceState.AVAILABLE.equals(jobParameter.getState())) { - log.error("Invalid datasource state. Expecting {} but received {}", DatasourceState.AVAILABLE, jobParameter.getState()); - jobParameter.disable(); - jobParameter.getUpdateStats().setLastFailedAt(Instant.now()); - DatasourceHelper.updateDatasource(client, jobParameter, timeout); + @VisibleForTesting + protected void updateDatasource(final ScheduledJobParameter jobParameter) throws IOException { + Datasource datasource = datasourceFacade.getDatasource(jobParameter.getName()); + if (datasource == null) { + log.info("Datasource[{}] does not exist", jobParameter.getName()); return; } - URL url = new URL(jobParameter.getEndpoint()); - DatasourceManifest manifest = DatasourceManifest.Builder.build(url); - - if (skipUpdate(jobParameter, manifest)) { - log.info("Skipping GeoIP database update. Update is not required for {}", jobParameter.getId()); - jobParameter.getUpdateStats().setLastSkippedAt(Instant.now()); - DatasourceHelper.updateDatasource(client, jobParameter, timeout); + if (!DatasourceState.AVAILABLE.equals(datasource.getState())) { + log.error("Invalid datasource state. Expecting {} but received {}", DatasourceState.AVAILABLE, datasource.getState()); + datasource.disable(); + datasource.getUpdateStats().setLastFailedAt(Instant.now()); + datasourceFacade.updateDatasource(datasource); return; } - Instant startTime = Instant.now(); - String indexName = jobParameter.indexNameFor(manifest); - jobParameter.getIndices().add(indexName); - DatasourceHelper.updateDatasource(client, jobParameter, timeout); - GeoIpDataHelper.createIndexIfNotExists(clusterService, client, indexName, timeout); - String[] fields; - try (CSVParser reader = GeoIpDataHelper.getDatabaseReader(manifest)) { - Iterator iter = reader.iterator(); - fields = iter.next().values(); - if (!jobParameter.getDatabase().getFields().equals(Arrays.asList(fields))) { - throw new OpenSearchException( - "fields does not match between old [{}] and new [{}]", - jobParameter.getDatabase().getFields().toString(), - Arrays.asList(fields).toString() - ); - } - GeoIpDataHelper.putGeoData(client, indexName, fields, iter, indexingBulkSize, timeout); - } - - Instant endTime = Instant.now(); - jobParameter.getDatabase().setProvider(manifest.getProvider()); - jobParameter.getDatabase().setMd5Hash(manifest.getMd5Hash()); - jobParameter.getDatabase().setUpdatedAt(Instant.ofEpochMilli(manifest.getUpdatedAt())); - jobParameter.getDatabase().setValidForInDays(manifest.getValidForInDays()); - jobParameter.getDatabase().setFields(Arrays.asList(fields)); - jobParameter.getUpdateStats().setLastSucceededAt(endTime); - jobParameter.getUpdateStats().setLastProcessingTimeInMillis(endTime.toEpochMilli() - startTime.toEpochMilli()); - DatasourceHelper.updateDatasource(client, jobParameter, timeout); - log.info( - "GeoIP database creation succeeded for {} and took {} seconds", - jobParameter.getId(), - Duration.between(startTime, endTime) - ); - } - - /** - * Determine if update is needed or not - * - * Update is needed when all following conditions are met - * 1. MD5 hash value in datasource is different with MD5 hash value in manifest - * 2. updatedAt value in datasource is before updateAt value in manifest - * - * @param parameter - * @param manifest - * @return - */ - private boolean skipUpdate(final Datasource parameter, final DatasourceManifest manifest) { - if (manifest.getMd5Hash().equals(parameter.getDatabase().getMd5Hash())) { - return true; + try { + datasourceUpdateService.deleteUnusedIndices(datasource); + datasourceUpdateService.updateOrCreateGeoIpData(datasource); + datasourceUpdateService.deleteUnusedIndices(datasource); + } catch (Exception e) { + log.error("Failed to update datasource for {}", datasource.getId(), e); + datasource.getUpdateStats().setLastFailedAt(Instant.now()); + datasourceFacade.updateDatasource(datasource); } - - return parameter.getDatabase().getUpdatedAt().toEpochMilli() >= manifest.getUpdatedAt(); } } diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateService.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateService.java new file mode 100644 index 00000000..e76ea388 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateService.java @@ -0,0 +1,206 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import java.io.IOException; +import java.net.URL; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; + +import lombok.extern.log4j.Log4j2; + +import org.apache.commons.csv.CSVParser; +import org.apache.commons.csv.CSVRecord; +import org.opensearch.OpenSearchException; +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.geospatial.ip2geo.common.DatasourceFacade; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.GeoIpDataFacade; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; + +@Log4j2 +public class DatasourceUpdateService { + private final ClusterService clusterService; + private final ClusterSettings clusterSettings; + private final Client client; + private final DatasourceFacade datasourceFacade; + private final GeoIpDataFacade geoIpDataFacade; + + public DatasourceUpdateService( + final ClusterService clusterService, + final Client client, + final DatasourceFacade datasourceFacade, + final GeoIpDataFacade geoIpDataFacade + ) { + this.clusterService = clusterService; + this.clusterSettings = clusterService.getClusterSettings(); + this.client = client; + this.datasourceFacade = datasourceFacade; + this.geoIpDataFacade = geoIpDataFacade; + } + + /** + * Update GeoIp data + * + * @param datasource + * @throws Exception + */ + public void updateOrCreateGeoIpData(final Datasource datasource) throws Exception { + URL url = new URL(datasource.getEndpoint()); + DatasourceManifest manifest = DatasourceManifest.Builder.build(url); + + if (skipUpdate(datasource, manifest)) { + log.info("Skipping GeoIP database update. Update is not required for {}", datasource.getId()); + datasource.getUpdateStats().setLastSkippedAt(Instant.now()); + datasourceFacade.updateDatasource(datasource); + return; + } + + Instant startTime = Instant.now(); + String indexName = setupIndex(manifest, datasource); + String[] header; + List fieldsToStore; + try (CSVParser reader = geoIpDataFacade.getDatabaseReader(manifest)) { + header = validateAndReturnHeader(reader.iterator()).values(); + fieldsToStore = Arrays.asList(header).subList(1, header.length); + if (!datasource.isCompatible(fieldsToStore)) { + throw new OpenSearchException( + "new fields [{}] does not contain all old fields [{}]", + fieldsToStore.toString(), + datasource.getDatabase().getFields().toString() + ); + } + geoIpDataFacade.putGeoIpData(indexName, header, reader.iterator(), clusterSettings.get(Ip2GeoSettings.INDEXING_BULK_SIZE)); + } + + Instant endTime = Instant.now(); + updateDatasourceAsSucceeded(datasource, manifest, fieldsToStore, startTime, endTime); + } + + /** + * Delete all indices except the one which are being used + * + * @param parameter + */ + public void deleteUnusedIndices(final Datasource parameter) { + try { + List deletedIndices = new ArrayList<>(); + for (String index : parameter.getIndices()) { + if (index.equals(parameter.currentIndexName())) { + continue; + } + + if (!clusterService.state().metadata().hasIndex(index)) { + deletedIndices.add(index); + continue; + } + + try { + if (geoIpDataFacade.deleteIp2GeoDataIndex(index).isAcknowledged()) { + deletedIndices.add(index); + } else { + log.error("Failed to delete an index [{}]", index); + } + } catch (Exception e) { + log.error("Failed to delete an index [{}]", index, e); + } + } + if (!deletedIndices.isEmpty()) { + parameter.getIndices().removeAll(deletedIndices); + datasourceFacade.updateDatasource(parameter); + } + } catch (Exception e) { + log.error("Failed to delete old indices for {}", parameter.getId(), e); + } + } + + /** + * Validate the next csv record as header + * This method move iterator one step forward + * + * @param iterator the csv record iterator + * @return next CSVRecord + */ + private CSVRecord validateAndReturnHeader(Iterator iterator) { + if (iterator.hasNext() == false) { + throw new OpenSearchException("geoip database is empty"); + } + CSVRecord header = iterator.next(); + if (header.values().length < 2) { + throw new OpenSearchException("geoip database should have at least two fields"); + } + return header; + } + + /*** + * Update datasource as succeeded + * + * @param manifest the manifest + * @param datasource the datasource + * @return + * @throws IOException + */ + private void updateDatasourceAsSucceeded( + final Datasource datasource, + final DatasourceManifest manifest, + final List fields, + final Instant startTime, + final Instant endTime + ) throws IOException { + datasource.setDatabase(manifest, fields); + datasource.getUpdateStats().setLastSucceededAt(endTime); + datasource.getUpdateStats().setLastProcessingTimeInMillis(endTime.toEpochMilli() - startTime.toEpochMilli()); + datasource.enable(); + datasource.setState(DatasourceState.AVAILABLE); + datasourceFacade.updateDatasource(datasource); + log.info("GeoIP database creation succeeded for {} and took {} seconds", datasource.getId(), Duration.between(startTime, endTime)); + } + + /*** + * Setup index to add a new geoip data + * + * @param manifest the manifest + * @param datasource the datasource + * @return + * @throws IOException + */ + private String setupIndex(final DatasourceManifest manifest, final Datasource datasource) throws IOException { + String indexName = datasource.indexNameFor(manifest); + datasource.getIndices().add(indexName); + datasourceFacade.updateDatasource(datasource); + geoIpDataFacade.createIndexIfNotExists(indexName); + return indexName; + } + + /** + * Determine if update is needed or not + * + * Update is needed when all following conditions are met + * 1. MD5 hash value in datasource is different with MD5 hash value in manifest + * 2. updatedAt value in datasource is before updateAt value in manifest + * + * @param datasource + * @param manifest + * @return + */ + private boolean skipUpdate(final Datasource datasource, final DatasourceManifest manifest) { + if (manifest.getMd5Hash().equals(datasource.getDatabase().getMd5Hash())) { + return true; + } + + return datasource.getDatabase().getUpdatedAt().toEpochMilli() >= manifest.getUpdatedAt(); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCache.java b/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCache.java deleted file mode 100644 index 2bcc1ed5..00000000 --- a/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCache.java +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.geospatial.ip2geo.processor; - -import java.util.Collections; -import java.util.Map; -import java.util.Objects; -import java.util.function.Function; - -import org.opensearch.common.cache.Cache; -import org.opensearch.common.cache.CacheBuilder; -import org.opensearch.common.unit.TimeValue; - -/** - * The in-memory cache for the ip2geo data. There should only be 1 instance of this class. - */ -public class Ip2GeoCache { - private static final TimeValue CACHING_PERIOD = TimeValue.timeValueMinutes(1); - private final Cache> cache; - - /** - * Default constructor - * - * @param maxSize size of a cache - */ - public Ip2GeoCache(final long maxSize) { - if (maxSize < 0) { - throw new IllegalArgumentException("ip2geo max cache size must be 0 or greater"); - } - this.cache = CacheBuilder.>builder() - .setMaximumWeight(maxSize) - .setExpireAfterWrite(CACHING_PERIOD) - .build(); - } - - /** - * Put data in a cache if it is absent and return the data - * - * @param ip the first part of a key - * @param datasourceName the second part of a key - * @param retrieveFunction function to retrieve a data to be stored in a cache - * @return data in a cache - */ - public Map putIfAbsent( - final String ip, - final String datasourceName, - final Function> retrieveFunction - ) { - CacheKey cacheKey = new CacheKey(ip, datasourceName); - Map response = cache.get(cacheKey); - if (response == null) { - response = retrieveFunction.apply(ip); - response = response == null ? Collections.emptyMap() : response; - cache.put(cacheKey, response); - } - return response; - } - - /** - * Put data in a cache - * - * @param ip the first part of a key - * @param datasourceName the second part of a key - * @param data the data - */ - public void put(final String ip, final String datasourceName, final Map data) { - CacheKey cacheKey = new CacheKey(ip, datasourceName); - cache.put(cacheKey, data); - } - - protected Map get(final String ip, final String datasourceName) { - CacheKey cacheKey = new CacheKey(ip, datasourceName); - return cache.get(cacheKey); - } - - /** - * The key to use for the cache. Since this cache can span multiple ip2geo processors that all use different datasource, the datasource - * name is needed to be included in the cache key. For example, if we only used the IP address as the key the same IP may be in multiple - * datasource with different values. The datasource name scopes the IP to the correct datasource - */ - private static class CacheKey { - - private final String ip; - private final String datasourceName; - - private CacheKey(final String ip, final String datasourceName) { - this.ip = ip; - this.datasourceName = datasourceName; - } - - // generated - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; - CacheKey cacheKey = (CacheKey) o; - return Objects.equals(ip, cacheKey.ip) && Objects.equals(datasourceName, cacheKey.datasourceName); - } - - // generated - @Override - public int hashCode() { - return Objects.hash(ip, datasourceName); - } - } -} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessor.java b/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessor.java index 0d6ddd8c..4385186d 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessor.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessor.java @@ -21,17 +21,17 @@ import java.util.Map; import java.util.Set; import java.util.function.BiConsumer; +import java.util.stream.Collectors; import lombok.extern.log4j.Log4j2; import org.opensearch.action.ActionListener; import org.opensearch.client.Client; -import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; -import org.opensearch.common.unit.TimeValue; -import org.opensearch.geospatial.ip2geo.common.DatasourceHelper; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.ip2geo.common.DatasourceFacade; import org.opensearch.geospatial.ip2geo.common.DatasourceState; -import org.opensearch.geospatial.ip2geo.common.GeoIpDataHelper; +import org.opensearch.geospatial.ip2geo.common.GeoIpDataFacade; import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; import org.opensearch.ingest.AbstractProcessor; @@ -52,12 +52,10 @@ public final class Ip2GeoProcessor extends AbstractProcessor { private final Set properties; private final boolean ignoreMissing; private final boolean firstOnly; - private final Ip2GeoCache cache; private final Client client; - private final ClusterService clusterService; - - private int maxBundleSize; - private int maxConcurrentSearches; + private final ClusterSettings clusterSettings; + private final DatasourceFacade datasourceFacade; + private final GeoIpDataFacade geoIpDataFacade; /** * Ip2Geo processor type @@ -74,9 +72,8 @@ public final class Ip2GeoProcessor extends AbstractProcessor { * @param properties the properties * @param ignoreMissing true if documents with a missing value for the field should be ignored * @param firstOnly true if only first result should be returned in case of array - * @param cache the Ip2Geo cache * @param client the client - * @param clusterService the cluster service + * @param clusterSettings the cluster settings */ public Ip2GeoProcessor( final String tag, @@ -87,9 +84,10 @@ public Ip2GeoProcessor( final Set properties, final boolean ignoreMissing, final boolean firstOnly, - final Ip2GeoCache cache, final Client client, - final ClusterService clusterService + final ClusterSettings clusterSettings, + final DatasourceFacade datasourceFacade, + final GeoIpDataFacade geoIpDataFacade ) { super(tag, description); this.field = field; @@ -98,15 +96,10 @@ public Ip2GeoProcessor( this.properties = properties; this.ignoreMissing = ignoreMissing; this.firstOnly = firstOnly; - this.cache = cache; this.client = client; - this.clusterService = clusterService; - - maxBundleSize = clusterService.getClusterSettings().get(Ip2GeoSettings.MAX_BUNDLE_SIZE); - clusterService.getClusterSettings().addSettingsUpdateConsumer(Ip2GeoSettings.MAX_BUNDLE_SIZE, newValue -> maxBundleSize = newValue); - maxConcurrentSearches = clusterService.getClusterSettings().get(Ip2GeoSettings.MAX_CONCURRENT_SEARCHES); - clusterService.getClusterSettings() - .addSettingsUpdateConsumer(Ip2GeoSettings.MAX_CONCURRENT_SEARCHES, newValue -> maxConcurrentSearches = newValue); + this.clusterSettings = clusterSettings; + this.datasourceFacade = datasourceFacade; + this.geoIpDataFacade = geoIpDataFacade; } /** @@ -119,12 +112,9 @@ public Ip2GeoProcessor( public void execute(IngestDocument ingestDocument, BiConsumer handler) { Object ip = ingestDocument.getFieldValue(field, Object.class, ignoreMissing); - if (ip == null && ignoreMissing) { + if (ip == null) { handler.accept(ingestDocument, null); return; - } else if (ip == null) { - handler.accept(null, new IllegalArgumentException("field [" + field + "] is null, cannot extract geo information.")); - return; } if (ip instanceof String) { @@ -147,49 +137,23 @@ public IngestDocument execute(IngestDocument ingestDocument) { throw new IllegalStateException("Not implemented"); } - /** - * Handle single ip - * - * @param ingestDocument the document - * @param handler the handler - * @param ip the ip - */ - private void executeInternal( + @VisibleForTesting + protected void executeInternal( final IngestDocument ingestDocument, final BiConsumer handler, final String ip ) { - Map geoData = cache.get(ip, datasourceName); - if (geoData != null) { - if (!geoData.isEmpty()) { - ingestDocument.setFieldValue(targetField, filteredGeoData(geoData, ip)); - } - handler.accept(ingestDocument, null); - return; - } - - DatasourceHelper.getDatasource(client, datasourceName, new ActionListener<>() { + datasourceFacade.getDatasource(datasourceName, new ActionListener<>() { @Override public void onResponse(final Datasource datasource) { - if (datasource == null) { - handler.accept(null, new IllegalStateException("datasource does not exist")); + if (handleInvalidDatasource(ingestDocument, datasource, handler)) { return; } - if (datasource.isExpired()) { - ingestDocument.setFieldValue(targetField, DATA_EXPIRED); - handler.accept(ingestDocument, null); - return; - } - - GeoIpDataHelper.getGeoData(client, datasource.currentIndexName(), ip, new ActionListener<>() { + geoIpDataFacade.getGeoIpData(datasource.currentIndexName(), ip, new ActionListener<>() { @Override - public void onResponse(final Map stringObjectMap) { - cache.put(ip, datasourceName, stringObjectMap); - if (!stringObjectMap.isEmpty()) { - ingestDocument.setFieldValue(targetField, filteredGeoData(stringObjectMap, ip)); - } - handler.accept(ingestDocument, null); + public void onResponse(final Map ipToGeoData) { + handleSingleIp(ip, ipToGeoData, ingestDocument, handler); } @Override @@ -206,6 +170,32 @@ public void onFailure(final Exception e) { }); } + @VisibleForTesting + protected void handleSingleIp( + final String ip, + final Map ipToGeoData, + final IngestDocument ingestDocument, + final BiConsumer handler + ) { + if (!ipToGeoData.isEmpty()) { + ingestDocument.setFieldValue(targetField, filteredGeoData(ipToGeoData, ip)); + } + handler.accept(ingestDocument, null); + } + + private Map filteredGeoData(final Map geoData, final String ip) { + Map filteredGeoData; + if (properties == null) { + return geoData; + } + + filteredGeoData = properties.stream().filter(p -> !p.equals(PROPERTY_IP)).collect(Collectors.toMap(p -> p, p -> geoData.get(p))); + if (properties.contains(PROPERTY_IP)) { + filteredGeoData.put(PROPERTY_IP, ip); + } + return filteredGeoData; + } + /** * Handle multiple ips * @@ -213,7 +203,8 @@ public void onFailure(final Exception e) { * @param handler the handler * @param ips the ip list */ - private void executeInternal( + @VisibleForTesting + protected void executeInternal( final IngestDocument ingestDocument, final BiConsumer handler, final List ips @@ -223,72 +214,23 @@ private void executeInternal( if (ip instanceof String == false) { throw new IllegalArgumentException("array in field [" + field + "] should only contain strings"); } - String ipAddr = (String) ip; - data.put(ipAddr, cache.get(ipAddr, datasourceName)); } List ipList = (List) ips; - DatasourceHelper.getDatasource(client, datasourceName, new ActionListener<>() { + datasourceFacade.getDatasource(datasourceName, new ActionListener<>() { @Override public void onResponse(final Datasource datasource) { - if (datasource == null) { - handler.accept(null, new IllegalStateException("datasource does not exist")); + if (handleInvalidDatasource(ingestDocument, datasource, handler)) { return; } - if (datasource.isExpired()) { - ingestDocument.setFieldValue(targetField, DATA_EXPIRED); - handler.accept(ingestDocument, null); - return; - } - GeoIpDataHelper.getGeoData( - client, + geoIpDataFacade.getGeoIpData( datasource.currentIndexName(), ipList.iterator(), - maxBundleSize, - maxConcurrentSearches, + clusterSettings.get(Ip2GeoSettings.MAX_BUNDLE_SIZE), + clusterSettings.get(Ip2GeoSettings.MAX_CONCURRENT_SEARCHES), firstOnly, data, - new ActionListener<>() { - @Override - public void onResponse(final Object obj) { - for (Map.Entry> entry : data.entrySet()) { - cache.put(entry.getKey(), datasourceName, entry.getValue()); - } - - if (firstOnly) { - for (String ipAddr : ipList) { - Map geoData = data.get(ipAddr); - // GeoData for ipAddr won't be null - if (!geoData.isEmpty()) { - ingestDocument.setFieldValue(targetField, geoData); - handler.accept(ingestDocument, null); - return; - } - } - handler.accept(ingestDocument, null); - } else { - boolean match = false; - List> geoDataList = new ArrayList<>(ipList.size()); - for (String ipAddr : ipList) { - Map geoData = data.get(ipAddr); - // GeoData for ipAddr won't be null - geoDataList.add(geoData.isEmpty() ? null : geoData); - if (!geoData.isEmpty()) { - match = true; - } - } - if (match) { - ingestDocument.setFieldValue(targetField, geoDataList); - } - handler.accept(ingestDocument, null); - } - } - - @Override - public void onFailure(final Exception e) { - handler.accept(null, e); - } - } + listenerToAppendDataToDocument(data, ipList, ingestDocument, handler) ); } @@ -299,21 +241,70 @@ public void onFailure(final Exception e) { }); } - private Map filteredGeoData(final Map geoData, final String ip) { - Map filteredGeoData; - if (properties == null) { - filteredGeoData = geoData; - } else { - filteredGeoData = new HashMap<>(); - for (String property : this.properties) { - if (property.equals(PROPERTY_IP)) { - filteredGeoData.put(PROPERTY_IP, ip); + @VisibleForTesting + protected ActionListener>> listenerToAppendDataToDocument( + final Map> data, + final List ipList, + final IngestDocument ingestDocument, + final BiConsumer handler + ) { + return new ActionListener<>() { + @Override + public void onResponse(final Map> response) { + if (firstOnly) { + for (String ipAddr : ipList) { + Map geoData = data.get(ipAddr); + // GeoData for ipAddr won't be null + if (!geoData.isEmpty()) { + ingestDocument.setFieldValue(targetField, filteredGeoData(geoData, ipAddr)); + handler.accept(ingestDocument, null); + return; + } + } } else { - filteredGeoData.put(property, geoData.get(property)); + boolean match = false; + List> geoDataList = new ArrayList<>(ipList.size()); + for (String ipAddr : ipList) { + Map geoData = data.get(ipAddr); + // GeoData for ipAddr won't be null + geoDataList.add(geoData.isEmpty() ? null : filteredGeoData(geoData, ipAddr)); + if (!geoData.isEmpty()) { + match = true; + } + } + if (match) { + ingestDocument.setFieldValue(targetField, geoDataList); + handler.accept(ingestDocument, null); + return; + } } + handler.accept(ingestDocument, null); + } + + @Override + public void onFailure(final Exception e) { + handler.accept(null, e); } + }; + } + + @VisibleForTesting + protected boolean handleInvalidDatasource( + final IngestDocument ingestDocument, + final Datasource datasource, + final BiConsumer handler + ) { + if (datasource == null) { + handler.accept(null, new IllegalStateException("datasource does not exist")); + return true; } - return filteredGeoData; + + if (datasource.isExpired()) { + ingestDocument.setFieldValue(targetField, DATA_EXPIRED); + handler.accept(ingestDocument, null); + return true; + } + return false; } @Override @@ -325,26 +316,27 @@ public String getType() { * Ip2Geo processor factory */ public static final class Factory implements Processor.Factory { - private final Ip2GeoCache cache; private final Client client; private final IngestService ingestService; - private TimeValue timeout; + private final DatasourceFacade datasourceFacade; + private final GeoIpDataFacade geoIpDataFacade; /** * Default constructor * - * @param cache the cache * @param client the client * @param ingestService the ingest service */ - public Factory(final Ip2GeoCache cache, final Client client, final IngestService ingestService) { - this.cache = cache; + public Factory( + final Client client, + final IngestService ingestService, + final DatasourceFacade datasourceFacade, + final GeoIpDataFacade geoIpDataFacade + ) { this.client = client; this.ingestService = ingestService; - - timeout = Ip2GeoSettings.TIMEOUT_IN_SECONDS.get(client.settings()); - ClusterSettings clusterSettings = ingestService.getClusterService().getClusterSettings(); - clusterSettings.addSettingsUpdateConsumer(Ip2GeoSettings.TIMEOUT_IN_SECONDS, newValue -> timeout = newValue); + this.datasourceFacade = datasourceFacade; + this.geoIpDataFacade = geoIpDataFacade; } /** @@ -386,14 +378,15 @@ public Ip2GeoProcessor create( propertyNames == null ? null : new HashSet<>(propertyNames), ignoreMissing, firstOnly, - cache, client, - ingestService.getClusterService() + ingestService.getClusterService().getClusterSettings(), + datasourceFacade, + geoIpDataFacade ); } private void validate(final String processorTag, final String datasourceName, final List propertyNames) throws IOException { - Datasource datasource = DatasourceHelper.getDatasource(client, datasourceName, timeout); + Datasource datasource = datasourceFacade.getDatasource(datasourceName); if (datasource == null) { throw newConfigurationException(TYPE, processorTag, "datasource", "datasource [" + datasourceName + "] doesn't exist"); diff --git a/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java b/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java index bc6ce32d..66752fc5 100644 --- a/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java +++ b/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java @@ -39,10 +39,12 @@ import org.opensearch.geospatial.ip2geo.action.PutDatasourceAction; import org.opensearch.geospatial.ip2geo.action.PutDatasourceTransportAction; import org.opensearch.geospatial.ip2geo.action.RestPutDatasourceHandler; -import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutorHelper; +import org.opensearch.geospatial.ip2geo.common.DatasourceFacade; +import org.opensearch.geospatial.ip2geo.common.GeoIpDataFacade; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceRunner; -import org.opensearch.geospatial.ip2geo.processor.Ip2GeoCache; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; import org.opensearch.geospatial.ip2geo.processor.Ip2GeoProcessor; import org.opensearch.geospatial.processor.FeatureProcessor; import org.opensearch.geospatial.rest.action.upload.geojson.RestUploadGeoJSONAction; @@ -87,9 +89,10 @@ public Map getProcessors(Processor.Parameters paramet .put( Ip2GeoProcessor.TYPE, new Ip2GeoProcessor.Factory( - new Ip2GeoCache(Ip2GeoSettings.CACHE_SIZE.get(parameters.client.settings())), parameters.client, - parameters.ingestService + parameters.ingestService, + new DatasourceFacade(parameters.client, parameters.ingestService.getClusterService().getClusterSettings()), + new GeoIpDataFacade(parameters.ingestService.getClusterService(), parameters.client) ) ) .immutableMap(); @@ -98,7 +101,7 @@ public Map getProcessors(Processor.Parameters paramet @Override public List> getExecutorBuilders(Settings settings) { List> executorBuilders = new ArrayList<>(); - executorBuilders.add(Ip2GeoExecutorHelper.executorBuilder(settings)); + executorBuilders.add(Ip2GeoExecutor.executorBuilder(settings)); return executorBuilders; } @@ -121,10 +124,23 @@ public Collection createComponents( IndexNameExpressionResolver indexNameExpressionResolver, Supplier repositoriesServiceSupplier ) { - // Initialize DatasourceUpdateRunner - DatasourceRunner.getJobRunnerInstance().initialize(clusterService, threadPool, client); - - return List.of(UploadStats.getInstance()); + GeoIpDataFacade geoIpDataFacade = new GeoIpDataFacade(clusterService, client); + DatasourceFacade datasourceFacade = new DatasourceFacade(client, clusterService.getClusterSettings()); + DatasourceUpdateService datasourceUpdateService = new DatasourceUpdateService( + clusterService, + client, + datasourceFacade, + geoIpDataFacade + ); + Ip2GeoExecutor ip2GeoExecutor = new Ip2GeoExecutor(threadPool); + /** + * We don't need to return datasource runner because it is used only by job scheduler and job scheduler + * does not use DI but it calls DatasourceExtension#getJobRunner to get DatasourceRunner instance. + */ + DatasourceRunner.getJobRunnerInstance() + .initialize(clusterService, client, datasourceUpdateService, ip2GeoExecutor, datasourceFacade); + + return List.of(UploadStats.getInstance(), datasourceUpdateService, datasourceFacade, ip2GeoExecutor, geoIpDataFacade); } @Override @@ -137,7 +153,7 @@ public List getRestHandlers( IndexNameExpressionResolver indexNameExpressionResolver, Supplier nodesInCluster ) { - return List.of(new RestUploadStatsAction(), new RestUploadGeoJSONAction(), new RestPutDatasourceHandler(settings, clusterSettings)); + return List.of(new RestUploadStatsAction(), new RestUploadGeoJSONAction(), new RestPutDatasourceHandler(clusterSettings)); } @Override diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoTestCase.java b/src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoTestCase.java new file mode 100644 index 00000000..c1aec21e --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoTestCase.java @@ -0,0 +1,227 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo; + +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import java.io.File; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.HashSet; +import java.util.Locale; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; +import java.util.stream.Collectors; + +import org.junit.After; +import org.junit.Before; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.action.ActionListener; +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionResponse; +import org.opensearch.action.ActionType; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.routing.RoutingTable; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Randomness; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.OpenSearchExecutors; +import org.opensearch.geospatial.ip2geo.common.DatasourceFacade; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.GeoIpDataFacade; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; +import org.opensearch.ingest.IngestService; +import org.opensearch.jobscheduler.spi.utils.LockService; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskListener; +import org.opensearch.test.client.NoOpNodeClient; +import org.opensearch.test.rest.RestActionTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +public abstract class Ip2GeoTestCase extends RestActionTestCase { + @Mock + protected ClusterService clusterService; + @Mock + protected DatasourceUpdateService datasourceUpdateService; + @Mock + protected DatasourceFacade datasourceFacade; + @Mock + protected Ip2GeoExecutor ip2GeoExecutor; + @Mock + protected ExecutorService executorService; + @Mock + protected GeoIpDataFacade geoIpDataFacade; + @Mock + protected ClusterState clusterState; + @Mock + protected Metadata metadata; + @Mock + protected IngestService ingestService; + @Mock + protected ActionFilters actionFilters; + @Mock + protected ThreadPool threadPool; + @Mock + protected TransportService transportService; + protected NoOpNodeClient client; + protected VerifyingClient verifyingClient; + protected LockService lockService; + protected ClusterSettings clusterSettings; + protected Settings settings; + private AutoCloseable openMocks; + + @Before + public void prepareIp2GeoTestCase() { + openMocks = MockitoAnnotations.openMocks(this); + settings = Settings.EMPTY; + client = new NoOpNodeClient(this.getTestName()); + verifyingClient = spy(new VerifyingClient(this.getTestName())); + clusterSettings = new ClusterSettings(settings, new HashSet<>(Ip2GeoSettings.settings())); + lockService = new LockService(client, clusterService); + when(clusterService.getSettings()).thenReturn(Settings.EMPTY); + when(clusterService.getClusterSettings()).thenReturn(clusterSettings); + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(clusterState.routingTable()).thenReturn(RoutingTable.EMPTY_ROUTING_TABLE); + when(ip2GeoExecutor.forDatasourceUpdate()).thenReturn(executorService); + when(ingestService.getClusterService()).thenReturn(clusterService); + when(threadPool.generic()).thenReturn(OpenSearchExecutors.newDirectExecutorService()); + } + + @After + public void clean() throws Exception { + openMocks.close(); + client.close(); + verifyingClient.close(); + } + + public DatasourceState randomStateExcept(DatasourceState state) { + assertNotNull(state); + return Arrays.stream(DatasourceState.values()) + .sequential() + .filter(s -> !s.equals(state)) + .collect(Collectors.toList()) + .get(Randomness.createSecure().nextInt(DatasourceState.values().length - 2)); + } + + public String randomIpAddress() { + return String.format( + Locale.ROOT, + "%d.%d.%d.%d", + Randomness.get().nextInt(255), + Randomness.get().nextInt(255), + Randomness.get().nextInt(255), + Randomness.get().nextInt(255) + ); + } + + @SuppressForbidden(reason = "unit test") + public String sampleManifestUrl() throws Exception { + return Paths.get(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").toURI()).toUri().toURL().toExternalForm(); + } + + @SuppressForbidden(reason = "unit test") + public String sampleManifestUrlWithInvalidUrl() throws Exception { + return Paths.get(this.getClass().getClassLoader().getResource("ip2geo/manifest_invalid_url.json").toURI()) + .toUri() + .toURL() + .toExternalForm(); + } + + @SuppressForbidden(reason = "unit test") + public File sampleIp2GeoFile() { + return new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.csv").getFile()); + } + + /** + * Temporary class of VerifyingClient until this PR(https://github.com/opensearch-project/OpenSearch/pull/7167) + * is merged in OpenSearch core + */ + public static class VerifyingClient extends NoOpNodeClient { + AtomicReference executeVerifier = new AtomicReference<>(); + AtomicReference executeLocallyVerifier = new AtomicReference<>(); + + public VerifyingClient(String testName) { + super(testName); + reset(); + } + + /** + * Clears any previously set verifier functions set by {@link #setExecuteVerifier(BiFunction)} and/or + * {@link #setExecuteLocallyVerifier(BiFunction)}. These functions are replaced with functions which will throw an + * {@link AssertionError} if called. + */ + public void reset() { + executeVerifier.set((arg1, arg2) -> { throw new AssertionError(); }); + executeLocallyVerifier.set((arg1, arg2) -> { throw new AssertionError(); }); + } + + /** + * Sets the function that will be called when {@link #doExecute(ActionType, ActionRequest, ActionListener)} is called. The given + * function should return either a subclass of {@link ActionResponse} or {@code null}. + * @param verifier A function which is called in place of {@link #doExecute(ActionType, ActionRequest, ActionListener)} + */ + public void setExecuteVerifier( + BiFunction, Request, Response> verifier + ) { + executeVerifier.set(verifier); + } + + @Override + public void doExecute( + ActionType action, + Request request, + ActionListener listener + ) { + listener.onResponse((Response) executeVerifier.get().apply(action, request)); + } + + /** + * Sets the function that will be called when {@link #executeLocally(ActionType, ActionRequest, TaskListener)}is called. The given + * function should return either a subclass of {@link ActionResponse} or {@code null}. + * @param verifier A function which is called in place of {@link #executeLocally(ActionType, ActionRequest, TaskListener)} + */ + public void setExecuteLocallyVerifier( + BiFunction, Request, Response> verifier + ) { + executeLocallyVerifier.set(verifier); + } + + @Override + public Task executeLocally( + ActionType action, + Request request, + ActionListener listener + ) { + listener.onResponse((Response) executeLocallyVerifier.get().apply(action, request)); + return null; + } + + @Override + public Task executeLocally( + ActionType action, + Request request, + TaskListener listener + ) { + listener.onResponse(null, (Response) executeLocallyVerifier.get().apply(action, request)); + return null; + } + + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequestTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequestTests.java index 3198f26a..383d832a 100644 --- a/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequestTests.java +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequestTests.java @@ -11,25 +11,31 @@ import java.util.Locale; import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.Randomness; +import org.opensearch.common.io.stream.BytesStreamInput; +import org.opensearch.common.io.stream.BytesStreamOutput; import org.opensearch.common.unit.TimeValue; -import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; -public class PutDatasourceRequestTests extends OpenSearchTestCase { +public class PutDatasourceRequestTests extends Ip2GeoTestCase { - public void testValidateInvalidUrl() { - PutDatasourceRequest request = new PutDatasourceRequest("test"); + public void testValidateWithInvalidUrl() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); request.setEndpoint("invalidUrl"); - request.setUpdateIntervalInDays(TimeValue.ZERO); + request.setUpdateInterval(TimeValue.timeValueDays(1)); ActionRequestValidationException exception = request.validate(); assertEquals(1, exception.validationErrors().size()); assertEquals("Invalid URL format is provided", exception.validationErrors().get(0)); } - public void testValidateInvalidManifestFile() { - PutDatasourceRequest request = new PutDatasourceRequest("test"); - request.setDatasourceName("test"); - request.setEndpoint("https://hi.com"); - request.setUpdateIntervalInDays(TimeValue.ZERO); + public void testValidateWithInvalidManifestFile() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String domain = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint(String.format(Locale.ROOT, "https://%s.com", domain)); + request.setUpdateInterval(TimeValue.timeValueDays(1)); ActionRequestValidationException exception = request.validate(); assertEquals(1, exception.validationErrors().size()); assertEquals( @@ -37,4 +43,74 @@ public void testValidateInvalidManifestFile() { exception.validationErrors().get(0) ); } + + public void testValidate() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(1)); + assertNull(request.validate()); + } + + public void testValidateWithZeroUpdateInterval() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(0)); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertEquals( + String.format(Locale.ROOT, "Update interval should be equal to or larger than 1 day"), + exception.validationErrors().get(0) + ); + } + + public void testValidateWithLargeUpdateInterval() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(30)); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertTrue(exception.validationErrors().get(0).contains("should be smaller")); + } + + public void testValidateWithInvalidUrlInsideManifest() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint(sampleManifestUrlWithInvalidUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(1)); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertTrue(exception.validationErrors().get(0).contains("Invalid URL format")); + } + + public void testStreamInOut() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String domain = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint(String.format(Locale.ROOT, "https://%s.com", domain)); + request.setUpdateInterval(TimeValue.timeValueDays(Randomness.get().nextInt(30) + 1)); + + // Run + BytesStreamOutput output = new BytesStreamOutput(); + request.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + PutDatasourceRequest copiedRequest = new PutDatasourceRequest(input); + + // Verify + assertEquals(request, copiedRequest); + } } diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportActionTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportActionTests.java new file mode 100644 index 00000000..38226c7f --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportActionTests.java @@ -0,0 +1,118 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import org.junit.Before; +import org.opensearch.ResourceAlreadyExistsException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.DocWriteRequest; +import org.opensearch.action.index.IndexRequest; +import org.opensearch.action.support.WriteRequest; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension; +import org.opensearch.index.engine.VersionConflictEngineException; +import org.opensearch.tasks.Task; + +public class PutDatasourceTransportActionTests extends Ip2GeoTestCase { + private PutDatasourceTransportAction action; + + @Before + public void init() { + action = new PutDatasourceTransportAction( + transportService, + actionFilters, + verifyingClient, + threadPool, + datasourceFacade, + datasourceUpdateService + ); + } + + public void testDoExecute() throws Exception { + Task task = mock(Task.class); + PutDatasourceRequest request = new PutDatasourceRequest("test"); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(1)); + ActionListener listener = mock(ActionListener.class); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof IndexRequest); + IndexRequest indexRequest = (IndexRequest) actionRequest; + assertEquals(DatasourceExtension.JOB_INDEX_NAME, indexRequest.index()); + assertEquals(request.getDatasourceName(), indexRequest.id()); + assertEquals(WriteRequest.RefreshPolicy.IMMEDIATE, indexRequest.getRefreshPolicy()); + assertEquals(DocWriteRequest.OpType.CREATE, indexRequest.opType()); + return null; + }); + action.doExecute(task, request, listener); + verify(verifyingClient).index(any(IndexRequest.class), any(ActionListener.class)); + verify(listener).onResponse(new AcknowledgedResponse(true)); + } + + public void testIndexResponseListenerFailure() { + Datasource datasource = new Datasource(); + ActionListener listener = mock(ActionListener.class); + action.getIndexResponseListener(datasource, listener) + .onFailure( + new VersionConflictEngineException( + null, + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString() + ) + ); + verify(listener).onFailure(any(ResourceAlreadyExistsException.class)); + } + + public void testCreateDatasourceInvalidState() throws Exception { + Datasource datasource = new Datasource(); + datasource.setState(randomStateExcept(DatasourceState.CREATING)); + datasource.getUpdateStats().setLastFailedAt(null); + + // Run + action.createDatasource(datasource); + + // Verify + assertEquals(DatasourceState.CREATE_FAILED, datasource.getState()); + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceFacade).updateDatasource(datasource); + } + + public void testCreateDatasourceWithException() throws Exception { + Datasource datasource = new Datasource(); + doThrow(new RuntimeException()).when(datasourceUpdateService).updateOrCreateGeoIpData(datasource); + + // Run + action.createDatasource(datasource); + + // Verify + assertEquals(DatasourceState.CREATE_FAILED, datasource.getState()); + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceFacade).updateDatasource(datasource); + } + + public void testCreateDatasource() throws Exception { + Datasource datasource = new Datasource(); + + // Run + action.createDatasource(datasource); + + // Verify + verify(datasourceUpdateService).updateOrCreateGeoIpData(datasource); + assertEquals(DatasourceState.CREATING, datasource.getState()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandlerTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandlerTests.java index 82cae4ad..b812cdfa 100644 --- a/src/test/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandlerTests.java +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandlerTests.java @@ -8,67 +8,78 @@ package org.opensearch.geospatial.ip2geo.action; +import static org.opensearch.geospatial.shared.URLBuilder.URL_DELIMITER; +import static org.opensearch.geospatial.shared.URLBuilder.getPluginURLPrefix; + import java.util.HashSet; +import java.util.concurrent.atomic.AtomicBoolean; import org.junit.Before; +import org.opensearch.common.SuppressForbidden; import org.opensearch.common.bytes.BytesArray; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.common.unit.TimeValue; import org.opensearch.common.xcontent.XContentType; +import org.opensearch.geospatial.GeospatialTestHelper; import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; import org.opensearch.rest.RestRequest; import org.opensearch.test.rest.FakeRestRequest; import org.opensearch.test.rest.RestActionTestCase; +@SuppressForbidden(reason = "unit test") public class RestPutDatasourceHandlerTests extends RestActionTestCase { + private String path; private RestPutDatasourceHandler action; @Before public void setupAction() { - action = new RestPutDatasourceHandler(Settings.EMPTY, new ClusterSettings(Settings.EMPTY, new HashSet(Ip2GeoSettings.settings()))); + action = new RestPutDatasourceHandler(new ClusterSettings(Settings.EMPTY, new HashSet(Ip2GeoSettings.settings()))); controller().registerHandler(action); + path = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/%s"); } public void testPrepareRequest() { - String content = "{\"endpoint\":\"https://test.com\", \"update_interval\":1}"; - RestRequest restRequest = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) - .withPath("/_geoip/datasource/test") + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String content = "{\"endpoint\":\"https://test.com\", \"update_interval_in_days\":1}"; + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) + .withPath(String.format(path, datasourceName)) .withContent(new BytesArray(content), XContentType.JSON) .build(); + AtomicBoolean isExecuted = new AtomicBoolean(false); - verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { assertTrue(actionRequest instanceof PutDatasourceRequest); PutDatasourceRequest putDatasourceRequest = (PutDatasourceRequest) actionRequest; assertEquals("https://test.com", putDatasourceRequest.getEndpoint()); - assertEquals(TimeValue.timeValueDays(1), putDatasourceRequest.getUpdateIntervalInDays()); - assertEquals("test", putDatasourceRequest.getDatasourceName()); + assertEquals(TimeValue.timeValueDays(1), putDatasourceRequest.getUpdateInterval()); + assertEquals(datasourceName, putDatasourceRequest.getDatasourceName()); + isExecuted.set(true); return null; }); - dispatchRequest(restRequest); + dispatchRequest(request); + assertTrue(isExecuted.get()); } public void testPrepareRequestDefaultValue() { - RestRequest restRequestWithEmptyContent = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) - .withPath("/_geoip/datasource/test") + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) + .withPath(String.format(path, datasourceName)) .withContent(new BytesArray("{}"), XContentType.JSON) .build(); - - RestRequest restRequestWithoutContent = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) - .withPath("/_geoip/datasource/test") - .build(); - - verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + AtomicBoolean isExecuted = new AtomicBoolean(false); + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { assertTrue(actionRequest instanceof PutDatasourceRequest); PutDatasourceRequest putDatasourceRequest = (PutDatasourceRequest) actionRequest; assertEquals("https://geoip.maps.opensearch.org/v1/geolite-2/manifest.json", putDatasourceRequest.getEndpoint()); - assertEquals(TimeValue.timeValueDays(3), putDatasourceRequest.getUpdateIntervalInDays()); - assertEquals("test", putDatasourceRequest.getDatasourceName()); + assertEquals(TimeValue.timeValueDays(3), putDatasourceRequest.getUpdateInterval()); + assertEquals(datasourceName, putDatasourceRequest.getDatasourceName()); + isExecuted.set(true); return null; }); - dispatchRequest(restRequestWithEmptyContent); - dispatchRequest(restRequestWithoutContent); + dispatchRequest(request); + assertTrue(isExecuted.get()); } } diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceFacadeTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceFacadeTests.java new file mode 100644 index 00000000..066ea27b --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceFacadeTests.java @@ -0,0 +1,130 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.common; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.HashSet; + +import org.junit.Before; +import org.opensearch.action.ActionListener; +import org.opensearch.action.DocWriteRequest; +import org.opensearch.action.get.GetRequest; +import org.opensearch.action.get.GetResponse; +import org.opensearch.action.index.IndexRequest; +import org.opensearch.common.bytes.BytesReference; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; + +public class DatasourceFacadeTests extends Ip2GeoTestCase { + private DatasourceFacade datasourceFacade; + + @Before + public void init() { + datasourceFacade = new DatasourceFacade( + verifyingClient, + new ClusterSettings(Settings.EMPTY, new HashSet<>(Ip2GeoSettings.settings())) + ); + } + + public void testUpdateDatasource() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Datasource datasource = new Datasource( + datasourceName, + new IntervalSchedule(Instant.now().truncatedTo(ChronoUnit.MILLIS), 1, ChronoUnit.DAYS), + "https://test.com" + ); + Instant previousTime = Instant.now().minusMillis(1); + datasource.setLastUpdateTime(previousTime); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof IndexRequest); + IndexRequest request = (IndexRequest) actionRequest; + assertEquals(datasource.getId(), request.id()); + assertEquals(DocWriteRequest.OpType.INDEX, request.opType()); + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); + return null; + }); + + datasourceFacade.updateDatasource(datasource); + assertTrue(previousTime.isBefore(datasource.getLastUpdateTime())); + } + + public void testGetDatasourceException() throws Exception { + Datasource datasource = setupClientForGetRequest(true, new IndexNotFoundException(DatasourceExtension.JOB_INDEX_NAME)); + assertNull(datasourceFacade.getDatasource(datasource.getId())); + } + + public void testGetDatasourceExist() throws Exception { + Datasource datasource = setupClientForGetRequest(true, null); + assertEquals(datasource, datasourceFacade.getDatasource(datasource.getId())); + } + + public void testGetDatasourceNotExist() throws Exception { + Datasource datasource = setupClientForGetRequest(false, null); + assertNull(datasourceFacade.getDatasource(datasource.getId())); + } + + public void testGetDatasourceExistWithListener() { + Datasource datasource = setupClientForGetRequest(true, null); + ActionListener listener = mock(ActionListener.class); + datasourceFacade.getDatasource(datasource.getId(), listener); + verify(listener).onResponse(eq(datasource)); + } + + public void testGetDatasourceNotExistWithListener() { + Datasource datasource = setupClientForGetRequest(false, null); + ActionListener listener = mock(ActionListener.class); + datasourceFacade.getDatasource(datasource.getId(), listener); + verify(listener).onResponse(null); + } + + private Datasource setupClientForGetRequest(final boolean isExist, final RuntimeException exception) { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Datasource datasource = new Datasource( + datasourceName, + new IntervalSchedule(Instant.now().truncatedTo(ChronoUnit.MILLIS), 1, ChronoUnit.DAYS), + "https://test.com" + ); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof GetRequest); + GetRequest request = (GetRequest) actionRequest; + assertEquals(datasource.getId(), request.id()); + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); + GetResponse response = mock(GetResponse.class); + when(response.isExists()).thenReturn(isExist); + try { + when(response.getSourceAsBytesRef()).thenReturn( + BytesReference.bytes(datasource.toXContent(JsonXContent.contentBuilder(), null)) + ); + } catch (IOException e) { + throw new RuntimeException(e); + } + if (exception != null) { + throw exception; + } + return response; + }); + return datasource; + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceHelperTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceHelperTests.java deleted file mode 100644 index 02ff8d58..00000000 --- a/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceHelperTests.java +++ /dev/null @@ -1,78 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.geospatial.ip2geo.common; - -import java.time.Instant; -import java.time.temporal.ChronoUnit; - -import org.opensearch.action.DocWriteRequest; -import org.opensearch.action.get.GetRequest; -import org.opensearch.action.index.IndexRequest; -import org.opensearch.common.unit.TimeValue; -import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; -import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension; -import org.opensearch.index.IndexNotFoundException; -import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; -import org.opensearch.test.rest.RestActionTestCase; - -public class DatasourceHelperTests extends RestActionTestCase { - - public void testUpdateDatasource() throws Exception { - Instant previousTime = Instant.now().minusMillis(1); - Datasource datasource = new Datasource( - "testId", - previousTime, - null, - false, - null, - null, - DatasourceState.PREPARING, - null, - null, - null - ); - - verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { - assertTrue(actionRequest instanceof IndexRequest); - IndexRequest request = (IndexRequest) actionRequest; - assertEquals(datasource.getId(), request.id()); - assertEquals(DocWriteRequest.OpType.INDEX, request.opType()); - assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); - return null; - }); - - DatasourceHelper.updateDatasource(verifyingClient, datasource, TimeValue.timeValueSeconds(30)); - assertTrue(previousTime.isBefore(datasource.getLastUpdateTime())); - } - - public void testGetDatasourceException() throws Exception { - Datasource datasource = new Datasource( - "testId", - Instant.now(), - null, - false, - new IntervalSchedule(Instant.now(), 1, ChronoUnit.DAYS), - "https://test.com", - DatasourceState.PREPARING, - null, - null, - null - ); - - verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { - assertTrue(actionRequest instanceof GetRequest); - GetRequest request = (GetRequest) actionRequest; - assertEquals(datasource.getId(), request.id()); - assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); - throw new IndexNotFoundException(DatasourceExtension.JOB_INDEX_NAME); - }); - - assertNull(DatasourceHelper.getDatasource(verifyingClient, datasource.getId(), TimeValue.timeValueSeconds(30))); - } -} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataFacadeTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataFacadeTests.java new file mode 100644 index 00000000..4a61e612 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataFacadeTests.java @@ -0,0 +1,324 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.common; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.opensearch.geospatial.ip2geo.jobscheduler.Datasource.IP2GEO_DATA_INDEX_NAME_PREFIX; + +import java.io.File; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +import org.apache.commons.csv.CSVFormat; +import org.apache.commons.csv.CSVParser; +import org.apache.commons.csv.CSVRecord; +import org.apache.lucene.search.TotalHits; +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.opensearch.OpenSearchException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionType; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.delete.DeleteIndexRequest; +import org.opensearch.action.admin.indices.forcemerge.ForceMergeRequest; +import org.opensearch.action.admin.indices.refresh.RefreshRequest; +import org.opensearch.action.bulk.BulkRequest; +import org.opensearch.action.bulk.BulkResponse; +import org.opensearch.action.search.MultiSearchRequest; +import org.opensearch.action.search.MultiSearchResponse; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.IndicesOptions; +import org.opensearch.action.support.WriteRequest; +import org.opensearch.common.Randomness; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.common.bytes.BytesReference; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; + +@SuppressForbidden(reason = "unit test") +public class GeoIpDataFacadeTests extends Ip2GeoTestCase { + private static final String IP_RANGE_FIELD_NAME = "_cidr"; + private static final String DATA_FIELD_NAME = "_data"; + private GeoIpDataFacade noOpsGeoIpDataFacade; + private GeoIpDataFacade verifyingGeoIpDataFacade; + + @Before + public void init() { + noOpsGeoIpDataFacade = new GeoIpDataFacade(clusterService, client); + verifyingGeoIpDataFacade = new GeoIpDataFacade(clusterService, verifyingClient); + } + + public void testCreateIndexIfNotExistsWithExistingIndex() { + String index = GeospatialTestHelper.randomLowerCaseString(); + when(metadata.hasIndex(index)).thenReturn(true); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { throw new RuntimeException("Shouldn't get called"); }); + verifyingGeoIpDataFacade.createIndexIfNotExists(index); + } + + public void testCreateIndexIfNotExistsWithoutExistingIndex() { + String index = GeospatialTestHelper.randomLowerCaseString(); + when(metadata.hasIndex(index)).thenReturn(false); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof CreateIndexRequest); + CreateIndexRequest request = (CreateIndexRequest) actionRequest; + assertEquals(index, request.index()); + assertEquals("1", request.settings().get("index.number_of_shards")); + assertEquals("0-all", request.settings().get("index.auto_expand_replicas")); + assertEquals( + "{\"dynamic\": false,\"properties\": {\"_cidr\": {\"type\": \"ip_range\",\"doc_values\": false}}}", + request.mappings() + ); + return null; + }); + verifyingGeoIpDataFacade.createIndexIfNotExists(index); + } + + public void testCreateDocument() { + String[] names = { "ip", "country", "city" }; + String[] values = { "1.0.0.0/25", "USA", "Seattle" }; + assertEquals( + "{\"_cidr\":\"1.0.0.0/25\",\"_data\":{\"country\":\"USA\",\"city\":\"Seattle\"}}", + noOpsGeoIpDataFacade.createDocument(names, values) + ); + } + + public void testGetDatabaseReader() throws Exception { + File zipFile = new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.zip").getFile()); + DatasourceManifest manifest = new DatasourceManifest( + zipFile.toURI().toURL().toExternalForm(), + "sample_valid.csv", + "fake_md5", + 1l, + Instant.now().toEpochMilli(), + "tester" + ); + CSVParser parser = noOpsGeoIpDataFacade.getDatabaseReader(manifest); + String[] expectedHeader = { "network", "country_name" }; + assertArrayEquals(expectedHeader, parser.iterator().next().values()); + String[] expectedValues = { "1.0.0.0/24", "Australia" }; + assertArrayEquals(expectedValues, parser.iterator().next().values()); + } + + public void testGetDatabaseReaderNoFile() throws Exception { + File zipFile = new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.zip").getFile()); + DatasourceManifest manifest = new DatasourceManifest( + zipFile.toURI().toURL().toExternalForm(), + "no_file.csv", + "fake_md5", + 1l, + Instant.now().toEpochMilli(), + "tester" + ); + OpenSearchException exception = expectThrows(OpenSearchException.class, () -> noOpsGeoIpDataFacade.getDatabaseReader(manifest)); + assertTrue(exception.getMessage().contains("does not exist")); + } + + public void testDeleteIp2GeoDataIndex() { + String index = String.format(Locale.ROOT, "%s.%s", IP2GEO_DATA_INDEX_NAME_PREFIX, GeospatialTestHelper.randomLowerCaseString()); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof DeleteIndexRequest); + DeleteIndexRequest request = (DeleteIndexRequest) actionRequest; + assertEquals(1, request.indices().length); + assertEquals(index, request.indices()[0]); + assertEquals(IndicesOptions.LENIENT_EXPAND_OPEN, request.indicesOptions()); + return null; + }); + verifyingGeoIpDataFacade.deleteIp2GeoDataIndex(index); + } + + public void testDeleteIp2GeoDataIndexWithNonIp2GeoDataIndex() { + String index = GeospatialTestHelper.randomLowerCaseString(); + Exception e = expectThrows(OpenSearchException.class, () -> verifyingGeoIpDataFacade.deleteIp2GeoDataIndex(index)); + assertTrue(e.getMessage().contains("not ip2geo data index")); + verify(verifyingClient, never()).index(any()); + } + + public void testPutGeoIpData() throws Exception { + String index = GeospatialTestHelper.randomLowerCaseString(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + if (actionRequest instanceof BulkRequest) { + BulkRequest request = (BulkRequest) actionRequest; + assertEquals(1, request.numberOfActions()); + assertEquals(WriteRequest.RefreshPolicy.WAIT_UNTIL, request.getRefreshPolicy()); + BulkResponse response = mock(BulkResponse.class); + when(response.hasFailures()).thenReturn(false); + return response; + } else if (actionRequest instanceof RefreshRequest) { + RefreshRequest request = (RefreshRequest) actionRequest; + assertEquals(1, request.indices().length); + assertEquals(index, request.indices()[0]); + return null; + } else if (actionRequest instanceof ForceMergeRequest) { + ForceMergeRequest request = (ForceMergeRequest) actionRequest; + assertEquals(1, request.indices().length); + assertEquals(index, request.indices()[0]); + assertEquals(1, request.maxNumSegments()); + return null; + } else { + throw new RuntimeException("invalid request is called"); + } + }); + try (CSVParser csvParser = CSVParser.parse(sampleIp2GeoFile(), StandardCharsets.UTF_8, CSVFormat.RFC4180)) { + Iterator iterator = csvParser.iterator(); + String[] fields = iterator.next().values(); + verifyingGeoIpDataFacade.putGeoIpData(index, fields, iterator, 1); + } + } + + public void testGetSingleGeoIpData() { + String indexName = GeospatialTestHelper.randomLowerCaseString(); + String ip = randomIpAddress(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assert actionRequest instanceof SearchRequest; + SearchRequest request = (SearchRequest) actionRequest; + assertEquals("_local", request.preference()); + assertEquals(1, request.source().size()); + assertEquals(QueryBuilders.termQuery(IP_RANGE_FIELD_NAME, ip), request.source().query()); + + String data = String.format( + Locale.ROOT, + "{\"%s\":\"1.0.0.1/16\",\"%s\":{\"city\":\"seattle\"}}", + IP_RANGE_FIELD_NAME, + DATA_FIELD_NAME + ); + SearchHit searchHit = new SearchHit(1); + searchHit.sourceRef(BytesReference.fromByteBuffer(ByteBuffer.wrap(data.getBytes(StandardCharsets.UTF_8)))); + SearchHit[] searchHitArray = { searchHit }; + SearchHits searchHits = new SearchHits(searchHitArray, new TotalHits(1l, TotalHits.Relation.EQUAL_TO), 1); + + SearchResponse response = mock(SearchResponse.class); + when(response.getHits()).thenReturn(searchHits); + return response; + }); + ActionListener> listener = mock(ActionListener.class); + verifyingGeoIpDataFacade.getGeoIpData(indexName, ip, listener); + ArgumentCaptor> captor = ArgumentCaptor.forClass(Map.class); + verify(listener).onResponse(captor.capture()); + assertEquals("seattle", captor.getValue().get("city")); + } + + public void testGetMultipleGeoIpDataNoSearchRequired() { + String indexName = GeospatialTestHelper.randomLowerCaseString(); + String ip1 = randomIpAddress(); + String ip2 = randomIpAddress(); + Iterator ipIterator = Arrays.asList(ip1, ip2).iterator(); + int maxBundleSize = 1; + int maxConcurrentSearches = 1; + boolean firstOnly = true; + Map> geoData = new HashMap<>(); + geoData.put(ip1, Map.of("city", "Seattle")); + geoData.put(ip2, Map.of("city", "Hawaii")); + ActionListener>> actionListener = mock(ActionListener.class); + + // Run + verifyingGeoIpDataFacade.getGeoIpData( + indexName, + ipIterator, + maxBundleSize, + maxConcurrentSearches, + firstOnly, + geoData, + actionListener + ); + + // Verify + verify(actionListener).onResponse(geoData); + } + + public void testGetMultipleGeoIpData() { + String indexName = GeospatialTestHelper.randomLowerCaseString(); + int dataSize = Randomness.get().nextInt(10) + 1; + List ips = new ArrayList<>(); + for (int i = 0; i < dataSize; i++) { + ips.add(randomIpAddress()); + } + int maxBundleSize = Randomness.get().nextInt(11) + 1; + int maxConcurrentSearches = 1; + boolean firstOnly = false; + Map> geoData = new HashMap<>(); + ActionListener>> actionListener = mock(ActionListener.class); + + List cities = new ArrayList<>(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assert actionRequest instanceof MultiSearchRequest; + MultiSearchRequest request = (MultiSearchRequest) actionRequest; + assertEquals(maxConcurrentSearches, request.maxConcurrentSearchRequests()); + assertTrue(request.requests().size() == maxBundleSize || request.requests().size() == dataSize % maxBundleSize); + for (SearchRequest searchRequest : request.requests()) { + assertEquals("_local", searchRequest.preference()); + assertEquals(1, searchRequest.source().size()); + } + + MultiSearchResponse.Item[] items = new MultiSearchResponse.Item[request.requests().size()]; + for (int i = 0; i < request.requests().size(); i++) { + String city = GeospatialTestHelper.randomLowerCaseString(); + cities.add(city); + String data = String.format( + Locale.ROOT, + "{\"%s\":\"1.0.0.1/16\",\"%s\":{\"city\":\"%s\"}}", + IP_RANGE_FIELD_NAME, + DATA_FIELD_NAME, + city + ); + SearchHit searchHit = new SearchHit(1); + searchHit.sourceRef(BytesReference.fromByteBuffer(ByteBuffer.wrap(data.getBytes(StandardCharsets.UTF_8)))); + SearchHit[] searchHitArray = { searchHit }; + SearchHits searchHits = new SearchHits(searchHitArray, new TotalHits(1l, TotalHits.Relation.EQUAL_TO), 1); + SearchResponse searchResponse = mock(SearchResponse.class); + when(searchResponse.getHits()).thenReturn(searchHits); + MultiSearchResponse.Item item = mock(MultiSearchResponse.Item.class); + when(item.isFailure()).thenReturn(false); + when(item.getResponse()).thenReturn(searchResponse); + items[i] = item; + } + MultiSearchResponse response = mock(MultiSearchResponse.class); + when(response.getResponses()).thenReturn(items); + return response; + }); + + // Run + verifyingGeoIpDataFacade.getGeoIpData( + indexName, + ips.iterator(), + maxBundleSize, + maxConcurrentSearches, + firstOnly, + geoData, + actionListener + ); + + // Verify + verify(verifyingClient, times((dataSize + maxBundleSize - 1) / maxBundleSize)).execute( + any(ActionType.class), + any(ActionRequest.class), + any(ActionListener.class) + ); + for (int i = 0; i < dataSize; i++) { + assertEquals(cities.get(i), geoData.get(ips.get(i)).get("city")); + } + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataHelperTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataHelperTests.java deleted file mode 100644 index 6b65026d..00000000 --- a/src/test/java/org/opensearch/geospatial/ip2geo/common/GeoIpDataHelperTests.java +++ /dev/null @@ -1,22 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.geospatial.ip2geo.common; - -import org.opensearch.test.OpenSearchTestCase; - -public class GeoIpDataHelperTests extends OpenSearchTestCase { - public void testCreateDocument() { - String[] names = { "ip", "country", "city" }; - String[] values = { "1.0.0.0/25", "USA", "Seattle" }; - assertEquals( - "{\"_cidr\":\"1.0.0.0/25\",\"_data\":{\"country\":\"USA\",\"city\":\"Seattle\"}}", - GeoIpDataHelper.createDocument(names, values) - ); - } -} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtensionTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtensionTests.java new file mode 100644 index 00000000..3632d9e9 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtensionTests.java @@ -0,0 +1,47 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import static org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension.JOB_INDEX_NAME; + +import java.time.Instant; +import java.time.temporal.ChronoUnit; + +import org.opensearch.common.Randomness; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.jobscheduler.spi.JobDocVersion; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; + +public class DatasourceExtensionTests extends Ip2GeoTestCase { + public void testBasic() throws Exception { + DatasourceExtension extension = new DatasourceExtension(); + assertEquals("scheduler_geospatial_ip2geo_datasource", extension.getJobType()); + assertEquals(JOB_INDEX_NAME, extension.getJobIndex()); + assertEquals(DatasourceRunner.getJobRunnerInstance(), extension.getJobRunner()); + } + + public void testParser() throws Exception { + DatasourceExtension extension = new DatasourceExtension(); + String id = GeospatialTestHelper.randomLowerCaseString(); + IntervalSchedule schedule = new IntervalSchedule(Instant.now().truncatedTo(ChronoUnit.MILLIS), 1, ChronoUnit.DAYS); + String endpoint = GeospatialTestHelper.randomLowerCaseString(); + Datasource datasource = new Datasource(id, schedule, endpoint); + + Datasource anotherDatasource = (Datasource) extension.getJobParser() + .parse( + createParser(datasource.toXContent(XContentFactory.jsonBuilder(), null)), + GeospatialTestHelper.randomLowerCaseString(), + new JobDocVersion(Randomness.get().nextLong(), Randomness.get().nextLong(), Randomness.get().nextLong()) + ); + + assertTrue(datasource.equals(anotherDatasource)); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunnerTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunnerTests.java new file mode 100644 index 00000000..0e33e232 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunnerTests.java @@ -0,0 +1,110 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.mockito.internal.verification.VerificationModeFactory.times; +import static org.opensearch.geospatial.GeospatialTestHelper.randomLowerCaseString; + +import java.time.Instant; + +import org.junit.Before; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.jobscheduler.spi.JobDocVersion; +import org.opensearch.jobscheduler.spi.JobExecutionContext; +import org.opensearch.jobscheduler.spi.ScheduledJobParameter; + +public class DatasourceRunnerTests extends Ip2GeoTestCase { + @Before + public void init() { + DatasourceRunner.getJobRunnerInstance() + .initialize(clusterService, client, datasourceUpdateService, ip2GeoExecutor, datasourceFacade); + } + + public void testRunJobInvalidClass() { + JobExecutionContext jobExecutionContext = mock(JobExecutionContext.class); + ScheduledJobParameter jobParameter = mock(ScheduledJobParameter.class); + expectThrows(IllegalStateException.class, () -> DatasourceRunner.getJobRunnerInstance().runJob(jobParameter, jobExecutionContext)); + } + + public void testRunJob() { + JobDocVersion jobDocVersion = new JobDocVersion(randomInt(), randomInt(), randomInt()); + String jobIndexName = randomLowerCaseString(); + String jobId = randomLowerCaseString(); + JobExecutionContext jobExecutionContext = new JobExecutionContext(Instant.now(), jobDocVersion, lockService, jobIndexName, jobId); + Datasource datasource = new Datasource(); + + // Run + DatasourceRunner.getJobRunnerInstance().runJob(datasource, jobExecutionContext); + + // Verify + verify(executorService).submit(any(Runnable.class)); + } + + public void testUpdateDatasourceNull() throws Exception { + Datasource datasource = new Datasource(); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource); + + // Verify + verify(datasourceUpdateService, never()).deleteUnusedIndices(any()); + } + + public void testUpdateDatasourceInvalidState() throws Exception { + Datasource datasource = new Datasource(); + datasource.enable(); + datasource.getUpdateStats().setLastFailedAt(null); + datasource.setState(randomStateExcept(DatasourceState.AVAILABLE)); + when(datasourceFacade.getDatasource(datasource.getId())).thenReturn(datasource); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource); + + // Verify + assertFalse(datasource.isEnabled()); + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceFacade).updateDatasource(datasource); + } + + public void testUpdateDatasource() throws Exception { + Datasource datasource = new Datasource(); + datasource.setState(DatasourceState.AVAILABLE); + datasource.setId(randomLowerCaseString()); + when(datasourceFacade.getDatasource(datasource.getId())).thenReturn(datasource); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource); + + // Verify + verify(datasourceUpdateService, times(2)).deleteUnusedIndices(datasource); + verify(datasourceUpdateService).updateOrCreateGeoIpData(datasource); + } + + public void testUpdateDatasourceExceptionHandling() throws Exception { + Datasource datasource = new Datasource(); + datasource.setId(randomLowerCaseString()); + datasource.getUpdateStats().setLastFailedAt(null); + when(datasourceFacade.getDatasource(datasource.getId())).thenReturn(datasource); + doThrow(new RuntimeException("test failure")).when(datasourceUpdateService).deleteUnusedIndices(any()); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource); + + // Verify + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceFacade).updateDatasource(datasource); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTests.java index 0faaa3e2..b4f76ef6 100644 --- a/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTests.java +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTests.java @@ -15,15 +15,41 @@ import java.time.Instant; import java.time.temporal.ChronoUnit; import java.util.ArrayList; +import java.util.Arrays; import java.util.Locale; import org.opensearch.common.Randomness; +import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.geospatial.GeospatialTestHelper; import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; import org.opensearch.test.OpenSearchTestCase; public class DatasourceTests extends OpenSearchTestCase { + + public void testParser() throws Exception { + String id = GeospatialTestHelper.randomLowerCaseString(); + IntervalSchedule schedule = new IntervalSchedule(Instant.now().truncatedTo(ChronoUnit.MILLIS), 1, ChronoUnit.DAYS); + String endpoint = GeospatialTestHelper.randomLowerCaseString(); + Datasource datasource = new Datasource(id, schedule, endpoint); + datasource.enable(); + datasource.getDatabase().setFields(Arrays.asList("field1", "field2")); + datasource.getDatabase().setProvider("test_provider"); + datasource.getDatabase().setUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MILLIS)); + datasource.getDatabase().setMd5Hash(GeospatialTestHelper.randomLowerCaseString()); + datasource.getDatabase().setValidForInDays(1l); + datasource.getUpdateStats().setLastProcessingTimeInMillis(Randomness.get().nextLong()); + datasource.getUpdateStats().setLastSucceededAt(Instant.now().truncatedTo(ChronoUnit.MILLIS)); + datasource.getUpdateStats().setLastSkippedAt(Instant.now().truncatedTo(ChronoUnit.MILLIS)); + datasource.getUpdateStats().setLastFailedAt(Instant.now().truncatedTo(ChronoUnit.MILLIS)); + + Datasource anotherDatasource = Datasource.PARSER.parse( + createParser(datasource.toXContent(XContentFactory.jsonBuilder(), null)), + null + ); + assertTrue(datasource.equals(anotherDatasource)); + } + public void testCurrentIndexName() { String id = GeospatialTestHelper.randomLowerCaseString(); Instant now = Instant.now(); @@ -56,9 +82,33 @@ public void testGetIndexNameFor() { public void testGetJitter() { Datasource datasource = new Datasource(); - datasource.setSchedule(new IntervalSchedule(Instant.now(), Randomness.get().nextInt(31), ChronoUnit.DAYS)); - long intervalInMinutes = datasource.getSchedule().getInterval() * 60 * 24; + datasource.setSchedule(new IntervalSchedule(Instant.now(), Randomness.get().ints(1, 31).findFirst().getAsInt(), ChronoUnit.DAYS)); + long intervalInMinutes = datasource.getSchedule().getInterval() * 60l * 24l; double sixMinutes = 6; assertTrue(datasource.getJitter() * intervalInMinutes <= sixMinutes); } + + public void testIsExpired() { + Datasource datasource = new Datasource(); + // never expire if validForInDays is null + assertFalse(datasource.isExpired()); + + datasource.getDatabase().setValidForInDays(1l); + + // if last skipped date is null, use only last succeeded date to determine + datasource.getUpdateStats().setLastSucceededAt(Instant.now().minus(25, ChronoUnit.HOURS)); + assertTrue(datasource.isExpired()); + + // use the latest date between last skipped date and last succeeded date to determine + datasource.getUpdateStats().setLastSkippedAt(Instant.now()); + assertFalse(datasource.isExpired()); + datasource.getUpdateStats().setLastSkippedAt(Instant.now().minus(25, ChronoUnit.HOURS)); + datasource.getUpdateStats().setLastSucceededAt(Instant.now()); + assertFalse(datasource.isExpired()); + } + + public void testLockDurationSeconds() { + Datasource datasource = new Datasource(); + assertNotNull(datasource.getLockDurationSeconds()); + } } diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateServiceTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateServiceTests.java new file mode 100644 index 00000000..3c9ec15d --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateServiceTests.java @@ -0,0 +1,152 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.File; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.Arrays; + +import org.apache.commons.csv.CSVFormat; +import org.apache.commons.csv.CSVParser; +import org.junit.Before; +import org.opensearch.OpenSearchException; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; + +@SuppressForbidden(reason = "unit test") +public class DatasourceUpdateServiceTests extends Ip2GeoTestCase { + private DatasourceUpdateService datasourceUpdateService; + + @Before + public void init() { + datasourceUpdateService = new DatasourceUpdateService(clusterService, client, datasourceFacade, geoIpDataFacade); + } + + public void testUpdateDatasourceSkip() throws Exception { + File manifestFile = new File(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").getFile()); + DatasourceManifest manifest = DatasourceManifest.Builder.build(manifestFile.toURI().toURL()); + + Datasource datasource = new Datasource(); + datasource.setState(DatasourceState.AVAILABLE); + datasource.getUpdateStats().setLastSkippedAt(null); + datasource.getDatabase().setUpdatedAt(Instant.ofEpochMilli(manifest.getUpdatedAt())); + datasource.getDatabase().setMd5Hash(manifest.getMd5Hash()); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + + // Run + datasourceUpdateService.updateOrCreateGeoIpData(datasource); + + // Verify + assertNotNull(datasource.getUpdateStats().getLastSkippedAt()); + verify(datasourceFacade).updateDatasource(datasource); + } + + public void testUpdateDatasourceInvalidFile() throws Exception { + File manifestFile = new File(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").getFile()); + DatasourceManifest manifest = DatasourceManifest.Builder.build(manifestFile.toURI().toURL()); + + File sampleFile = new File( + this.getClass().getClassLoader().getResource("ip2geo/sample_invalid_less_than_two_fields.csv").getFile() + ); + when(geoIpDataFacade.getDatabaseReader(any())).thenReturn(CSVParser.parse(sampleFile, StandardCharsets.UTF_8, CSVFormat.RFC4180)); + + Datasource datasource = new Datasource(); + datasource.setState(DatasourceState.AVAILABLE); + datasource.getDatabase().setUpdatedAt(Instant.ofEpochMilli(manifest.getUpdatedAt() - 1)); + datasource.getDatabase().setMd5Hash(manifest.getMd5Hash().substring(1)); + datasource.getDatabase().setFields(Arrays.asList("country_name")); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + + // Run + expectThrows(OpenSearchException.class, () -> datasourceUpdateService.updateOrCreateGeoIpData(datasource)); + } + + public void testUpdateDatasourceIncompatibleFields() throws Exception { + File manifestFile = new File(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").getFile()); + DatasourceManifest manifest = DatasourceManifest.Builder.build(manifestFile.toURI().toURL()); + + File sampleFile = new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.csv").getFile()); + when(geoIpDataFacade.getDatabaseReader(any())).thenReturn(CSVParser.parse(sampleFile, StandardCharsets.UTF_8, CSVFormat.RFC4180)); + + Datasource datasource = new Datasource(); + datasource.setState(DatasourceState.AVAILABLE); + datasource.getDatabase().setUpdatedAt(Instant.ofEpochMilli(manifest.getUpdatedAt() - 1)); + datasource.getDatabase().setMd5Hash(manifest.getMd5Hash().substring(1)); + datasource.getDatabase().setFields(Arrays.asList("country_name", "additional_field")); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + + // Run + expectThrows(OpenSearchException.class, () -> datasourceUpdateService.updateOrCreateGeoIpData(datasource)); + } + + public void testUpdateDatasource() throws Exception { + File manifestFile = new File(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").getFile()); + DatasourceManifest manifest = DatasourceManifest.Builder.build(manifestFile.toURI().toURL()); + + File sampleFile = new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.csv").getFile()); + when(geoIpDataFacade.getDatabaseReader(any())).thenReturn(CSVParser.parse(sampleFile, StandardCharsets.UTF_8, CSVFormat.RFC4180)); + + Datasource datasource = new Datasource(); + datasource.setState(DatasourceState.AVAILABLE); + datasource.getDatabase().setUpdatedAt(Instant.ofEpochMilli(manifest.getUpdatedAt() - 1)); + datasource.getDatabase().setMd5Hash(manifest.getMd5Hash().substring(1)); + datasource.getDatabase().setFields(Arrays.asList("country_name")); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + datasource.getUpdateStats().setLastSucceededAt(null); + datasource.getUpdateStats().setLastProcessingTimeInMillis(null); + + // Run + datasourceUpdateService.updateOrCreateGeoIpData(datasource); + + // Verify + assertEquals(manifest.getProvider(), datasource.getDatabase().getProvider()); + assertEquals(manifest.getMd5Hash(), datasource.getDatabase().getMd5Hash()); + assertEquals(Instant.ofEpochMilli(manifest.getUpdatedAt()), datasource.getDatabase().getUpdatedAt()); + assertEquals(manifest.getValidForInDays(), datasource.getDatabase().getValidForInDays()); + assertNotNull(datasource.getUpdateStats().getLastSucceededAt()); + assertNotNull(datasource.getUpdateStats().getLastProcessingTimeInMillis()); + verify(datasourceFacade, times(2)).updateDatasource(datasource); + } + + public void testDeleteUnusedIndices() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String indexPrefix = String.format(".ip2geo-data.%s.", datasourceName); + Instant now = Instant.now(); + String currentIndex = indexPrefix + now.toEpochMilli(); + String oldIndex = indexPrefix + now.minusMillis(1).toEpochMilli(); + String lingeringIndex = indexPrefix + now.minusMillis(2).toEpochMilli(); + Datasource datasource = new Datasource(); + datasource.setId(datasourceName); + datasource.getIndices().add(currentIndex); + datasource.getIndices().add(oldIndex); + datasource.getIndices().add(lingeringIndex); + datasource.getDatabase().setUpdatedAt(now); + + when(metadata.hasIndex(currentIndex)).thenReturn(true); + when(metadata.hasIndex(oldIndex)).thenReturn(true); + when(metadata.hasIndex(lingeringIndex)).thenReturn(false); + when(geoIpDataFacade.deleteIp2GeoDataIndex(any())).thenReturn(new AcknowledgedResponse(true)); + + datasourceUpdateService.deleteUnusedIndices(datasource); + + assertEquals(1, datasource.getIndices().size()); + assertEquals(currentIndex, datasource.getIndices().get(0)); + verify(datasourceFacade).updateDatasource(datasource); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCacheTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCacheTests.java deleted file mode 100644 index ad340f47..00000000 --- a/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoCacheTests.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * SPDX-License-Identifier: Apache-2.0 - * - * The OpenSearch Contributors require contributions made to - * this file be licensed under the Apache-2.0 license or a - * compatible open source license. - */ - -package org.opensearch.geospatial.ip2geo.processor; - -import java.util.HashMap; -import java.util.Map; - -import org.opensearch.OpenSearchException; -import org.opensearch.test.OpenSearchTestCase; - -public class Ip2GeoCacheTests extends OpenSearchTestCase { - public void testCachesAndEvictsResults() { - Ip2GeoCache cache = new Ip2GeoCache(1); - String datasource = "datasource"; - Map response1 = new HashMap<>(); - Map response2 = new HashMap<>(); - assertNotSame(response1, response2); - - // add a key - Map cachedResponse = cache.putIfAbsent("127.0.0.1", datasource, key -> response1); - assertSame(cachedResponse, response1); - assertSame(cachedResponse, cache.putIfAbsent("127.0.0.1", datasource, key -> response2)); - assertSame(cachedResponse, cache.get("127.0.0.1", datasource)); - - // evict old key by adding another value - cachedResponse = cache.putIfAbsent("127.0.0.2", datasource, key -> response2); - assertSame(cachedResponse, response2); - assertSame(cachedResponse, cache.putIfAbsent("127.0.0.2", datasource, ip -> response2)); - assertSame(cachedResponse, cache.get("127.0.0.2", datasource)); - - assertNotSame(response1, cache.get("127.0.0.1", datasource)); - } - - public void testThrowsFunctionsException() { - Ip2GeoCache cache = new Ip2GeoCache(1); - expectThrows( - OpenSearchException.class, - () -> cache.putIfAbsent("127.0.0.1", "datasource", ip -> { throw new OpenSearchException("bad"); }) - ); - } - - public void testNoExceptionForNullValue() { - Ip2GeoCache cache = new Ip2GeoCache(1); - Map response = cache.putIfAbsent("127.0.0.1", "datasource", ip -> null); - assertTrue(response.isEmpty()); - } - - public void testInvalidInit() { - IllegalArgumentException ex = expectThrows(IllegalArgumentException.class, () -> new Ip2GeoCache(-1)); - assertEquals("ip2geo max cache size must be 0 or greater", ex.getMessage()); - } -} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorTests.java new file mode 100644 index 00000000..bd6857ed --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorTests.java @@ -0,0 +1,346 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.geospatial.ip2geo.processor; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.function.BiConsumer; + +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.opensearch.OpenSearchException; +import org.opensearch.action.ActionListener; +import org.opensearch.common.Randomness; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.ingest.IngestDocument; + +public class Ip2GeoProcessorTests extends Ip2GeoTestCase { + private static final String DEFAULT_TARGET_FIELD = "ip2geo"; + private static final String CONFIG_DATASOURCE_KEY = "datasource"; + private static final String CONFIG_FIELD_KEY = "field"; + private static final List SUPPORTED_FIELDS = Arrays.asList("city", "country"); + private Ip2GeoProcessor.Factory factory; + + @Before + public void init() { + factory = new Ip2GeoProcessor.Factory(client, ingestService, datasourceFacade, geoIpDataFacade); + } + + public void testCreateWithNoDatasource() { + Map config = new HashMap<>(); + config.put("field", "ip"); + config.put(CONFIG_DATASOURCE_KEY, "no_datasource"); + OpenSearchException exception = expectThrows( + OpenSearchException.class, + () -> factory.create( + Collections.emptyMap(), + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString(), + config + ) + ); + assertTrue(exception.getDetailedMessage().contains("doesn't exist")); + } + + public void testCreateWithInvalidDatasourceState() { + Datasource datasource = new Datasource(); + datasource.setId(GeospatialTestHelper.randomLowerCaseString()); + datasource.setState(randomStateExcept(DatasourceState.AVAILABLE)); + OpenSearchException exception = expectThrows(OpenSearchException.class, () -> createProcessor(datasource, Collections.emptyMap())); + assertTrue(exception.getDetailedMessage().contains("available state")); + } + + public void testCreateWithInvalidProperties() { + Map config = new HashMap<>(); + config.put("properties", Arrays.asList("ip", "invalid_property")); + OpenSearchException exception = expectThrows( + OpenSearchException.class, + () -> createProcessor(GeospatialTestHelper.randomLowerCaseString(), config) + ); + assertTrue(exception.getDetailedMessage().contains("property")); + } + + public void testExecuteWithNoIpAndIgnoreMissing() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Map config = new HashMap<>(); + config.put("ignore_missing", true); + Ip2GeoProcessor processor = createProcessor(datasourceName, config); + IngestDocument document = new IngestDocument(new HashMap<>(), new HashMap<>()); + BiConsumer handler = (doc, e) -> { + assertEquals(document, doc); + assertNull(e); + }; + processor.execute(document, handler); + } + + public void testExecuteWithNoIp() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Map config = new HashMap<>(); + Ip2GeoProcessor processor = createProcessor(datasourceName, config); + IngestDocument document = new IngestDocument(new HashMap<>(), new HashMap<>()); + BiConsumer handler = (doc, e) -> {}; + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> processor.execute(document, handler)); + assertTrue(exception.getMessage().contains("not present")); + } + + public void testExecuteWithNonStringValue() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + Map source = new HashMap<>(); + source.put("ip", Randomness.get().nextInt()); + IngestDocument document = new IngestDocument(source, new HashMap<>()); + BiConsumer handler = (doc, e) -> {}; + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> processor.execute(document, handler)); + assertTrue(exception.getMessage().contains("string")); + } + + public void testExecuteWithNullDatasource() throws Exception { + BiConsumer handler = (doc, e) -> { + assertNull(doc); + assertTrue(e.getMessage().contains("datasource does not exist")); + }; + getActionListener(Collections.emptyMap(), handler).onResponse(null); + } + + public void testExecuteWithExpiredDatasource() throws Exception { + Datasource datasource = mock(Datasource.class); + when(datasource.isExpired()).thenReturn(true); + BiConsumer handler = (doc, e) -> { + assertEquals("ip2geo_data_expired", doc.getFieldValue(DEFAULT_TARGET_FIELD + ".error", String.class)); + assertNull(e); + }; + getActionListener(Collections.emptyMap(), handler).onResponse(datasource); + } + + public void testExecute() throws Exception { + Map ip2geoData = new HashMap<>(); + for (String field : SUPPORTED_FIELDS) { + ip2geoData.put(field, GeospatialTestHelper.randomLowerCaseString()); + } + + Datasource datasource = mock(Datasource.class); + when(datasource.isExpired()).thenReturn(false); + when(datasource.currentIndexName()).thenReturn(GeospatialTestHelper.randomLowerCaseString()); + BiConsumer handler = (doc, e) -> { + assertEquals( + ip2geoData.get(SUPPORTED_FIELDS.get(0)), + doc.getFieldValue(DEFAULT_TARGET_FIELD + "." + SUPPORTED_FIELDS.get(0), String.class) + ); + for (int i = 1; i < SUPPORTED_FIELDS.size(); i++) { + assertNull(doc.getFieldValue(DEFAULT_TARGET_FIELD + "." + SUPPORTED_FIELDS.get(i), String.class, true)); + } + assertNull(e); + }; + Map config = Map.of("properties", Arrays.asList(SUPPORTED_FIELDS.get(0))); + getActionListener(config, handler).onResponse(datasource); + + ArgumentCaptor>> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(geoIpDataFacade).getGeoIpData(anyString(), anyString(), captor.capture()); + captor.getValue().onResponse(ip2geoData); + } + + private ActionListener getActionListener( + final Map config, + final BiConsumer handler + ) throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, config); + + Map source = new HashMap<>(); + String ip = randomIpAddress(); + source.put("ip", ip); + IngestDocument document = new IngestDocument(source, new HashMap<>()); + + processor.execute(document, handler); + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(datasourceFacade).getDatasource(eq(datasourceName), captor.capture()); + return captor.getValue(); + } + + public void testExecuteNotImplemented() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + IngestDocument document = new IngestDocument(Collections.emptyMap(), Collections.emptyMap()); + Exception e = expectThrows(IllegalStateException.class, () -> processor.execute(document)); + assertTrue(e.getMessage().contains("Not implemented")); + } + + public void testGenerateDataToAppendWithFirstOnlyOption() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor( + datasourceName, + Map.of("first_only", true, "properties", Arrays.asList(SUPPORTED_FIELDS.get(0))) + ); + List ips = new ArrayList<>(); + Map> data = new HashMap<>(); + for (int i = 0; i < 3; i++) { + String ip = randomIpAddress(); + ips.add(ip); + Map geoData = new HashMap<>(); + for (String field : SUPPORTED_FIELDS) { + geoData.put(field, GeospatialTestHelper.randomLowerCaseString()); + } + data.put(ip, i == 0 ? Collections.emptyMap() : geoData); + } + IngestDocument document = new IngestDocument(new HashMap<>(), new HashMap<>()); + BiConsumer handler = mock(BiConsumer.class); + + // Run + processor.listenerToAppendDataToDocument(data, ips, document, handler).onResponse(data); + + // Verify + verify(handler).accept(document, null); + assertEquals(1, document.getFieldValue(DEFAULT_TARGET_FIELD, Map.class).size()); + assertEquals( + data.get(ips.get(1)).get(SUPPORTED_FIELDS.get(0)), + document.getFieldValue(DEFAULT_TARGET_FIELD, Map.class).get(SUPPORTED_FIELDS.get(0)) + ); + assertNull(document.getFieldValue(DEFAULT_TARGET_FIELD, Map.class).get(SUPPORTED_FIELDS.get(1))); + } + + public void testGenerateDataToAppendWithOutFirstOnlyOption() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor( + datasourceName, + Map.of("first_only", false, "properties", Arrays.asList(SUPPORTED_FIELDS.get(0))) + ); + List ips = new ArrayList<>(); + Map> data = new HashMap<>(); + for (int i = 0; i < 3; i++) { + String ip = randomIpAddress(); + ips.add(ip); + Map geoData = new HashMap<>(); + for (String field : SUPPORTED_FIELDS) { + geoData.put(field, GeospatialTestHelper.randomLowerCaseString()); + } + data.put(ip, i == 0 ? Collections.emptyMap() : geoData); + } + IngestDocument document = new IngestDocument(new HashMap<>(), new HashMap<>()); + BiConsumer handler = mock(BiConsumer.class); + + // Run + processor.listenerToAppendDataToDocument(data, ips, document, handler).onResponse(data); + + // Verify + verify(handler).accept(document, null); + assertEquals(ips.size(), document.getFieldValue(DEFAULT_TARGET_FIELD, List.class).size()); + for (int i = 0; i < ips.size(); i++) { + if (data.get(ips.get(i)).isEmpty()) { + assertNull(document.getFieldValue(DEFAULT_TARGET_FIELD, List.class).get(i)); + } else { + Map documentData = (Map) document.getFieldValue(DEFAULT_TARGET_FIELD, List.class).get(i); + assertEquals(1, documentData.size()); + assertEquals(data.get(ips.get(i)).get(SUPPORTED_FIELDS.get(0)), documentData.get(SUPPORTED_FIELDS.get(0))); + assertNull(documentData.get(SUPPORTED_FIELDS.get(1))); + } + } + } + + public void testGenerateDataToAppendWithNoData() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Map.of("first_only", Randomness.get().nextInt() % 2 == 0)); + List ips = new ArrayList<>(); + Map> data = new HashMap<>(); + for (int i = 0; i < 3; i++) { + String ip = randomIpAddress(); + ips.add(ip); + data.put(ip, Collections.emptyMap()); + } + IngestDocument document = new IngestDocument(new HashMap<>(), new HashMap<>()); + BiConsumer handler = mock(BiConsumer.class); + processor.listenerToAppendDataToDocument(data, ips, document, handler).onResponse(data); + verify(handler).accept(document, null); + Exception e = expectThrows(IllegalArgumentException.class, () -> document.getFieldValue(DEFAULT_TARGET_FIELD, Map.class)); + assertTrue(e.getMessage().contains("not present")); + } + + public void testExecuteInternalNonStringIp() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + List ips = Arrays.asList(randomIpAddress(), 1); + Map source = new HashMap<>(); + String ip = randomIpAddress(); + source.put("ip", ip); + IngestDocument document = new IngestDocument(source, new HashMap<>()); + + BiConsumer handler = mock(BiConsumer.class); + Exception e = expectThrows(IllegalArgumentException.class, () -> processor.executeInternal(document, handler, ips)); + assertTrue(e.getMessage().contains("should only contain strings")); + } + + public void testExecuteInternal() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + List ips = Arrays.asList(randomIpAddress(), randomIpAddress()); + Map source = new HashMap<>(); + String ip = randomIpAddress(); + source.put("ip", ip); + IngestDocument document = new IngestDocument(source, new HashMap<>()); + + BiConsumer handler = mock(BiConsumer.class); + processor.executeInternal(document, handler, ips); + + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(datasourceFacade).getDatasource(eq(datasourceName), captor.capture()); + Datasource datasource = mock(Datasource.class); + when(datasource.isExpired()).thenReturn(false); + when(datasource.currentIndexName()).thenReturn(GeospatialTestHelper.randomLowerCaseString()); + captor.getValue().onResponse(datasource); + verify(geoIpDataFacade).getGeoIpData( + anyString(), + any(Iterator.class), + anyInt(), + anyInt(), + anyBoolean(), + anyMap(), + any(ActionListener.class) + ); + } + + private Ip2GeoProcessor createProcessor(final String datasourceName, final Map config) throws Exception { + Datasource datasource = new Datasource(); + datasource.setId(datasourceName); + datasource.setState(DatasourceState.AVAILABLE); + datasource.getDatabase().setFields(SUPPORTED_FIELDS); + return createProcessor(datasource, config); + } + + private Ip2GeoProcessor createProcessor(final Datasource datasource, final Map config) throws Exception { + when(datasourceFacade.getDatasource(datasource.getId())).thenReturn(datasource); + Map baseConfig = new HashMap<>(); + baseConfig.put(CONFIG_FIELD_KEY, "ip"); + baseConfig.put(CONFIG_DATASOURCE_KEY, datasource.getName()); + baseConfig.putAll(config); + + return factory.create( + Collections.emptyMap(), + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString(), + baseConfig + ); + } +} diff --git a/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java b/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java index d7ce5594..6d9430bc 100644 --- a/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java +++ b/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java @@ -7,50 +7,148 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import static org.opensearch.geospatial.ip2geo.jobscheduler.Datasource.IP2GEO_DATA_INDEX_NAME_PREFIX; +import java.util.Collection; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.function.Supplier; +import org.junit.After; +import org.junit.Before; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionResponse; import org.opensearch.client.Client; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.io.stream.NamedWriteableRegistry; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; +import org.opensearch.core.xcontent.NamedXContentRegistry; import org.opensearch.env.Environment; +import org.opensearch.env.NodeEnvironment; import org.opensearch.geospatial.action.upload.geojson.UploadGeoJSONAction; import org.opensearch.geospatial.ip2geo.action.RestPutDatasourceHandler; +import org.opensearch.geospatial.ip2geo.common.DatasourceFacade; +import org.opensearch.geospatial.ip2geo.common.GeoIpDataFacade; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; import org.opensearch.geospatial.processor.FeatureProcessor; import org.opensearch.geospatial.rest.action.upload.geojson.RestUploadGeoJSONAction; import org.opensearch.geospatial.stats.upload.RestUploadStatsAction; +import org.opensearch.geospatial.stats.upload.UploadStats; +import org.opensearch.indices.SystemIndexDescriptor; import org.opensearch.ingest.IngestService; import org.opensearch.ingest.Processor; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.IngestPlugin; +import org.opensearch.repositories.RepositoriesService; import org.opensearch.rest.RestHandler; import org.opensearch.script.ScriptService; import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.watcher.ResourceWatcherService; public class GeospatialPluginTests extends OpenSearchTestCase { private final ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, new HashSet(Ip2GeoSettings.settings())); private final List SUPPORTED_REST_HANDLERS = List.of( new RestUploadGeoJSONAction(), new RestUploadStatsAction(), - new RestPutDatasourceHandler(Settings.EMPTY, clusterSettings) + new RestPutDatasourceHandler(clusterSettings) ); - private final Client client; - private final ClusterService clusterService; - private final IngestService ingestService; - - public GeospatialPluginTests() { - client = mock(Client.class); - when(client.settings()).thenReturn(Settings.EMPTY); - clusterService = mock(ClusterService.class); + + private final Set SUPPORTED_SYSTEM_INDEX_PATTERN = Set.of(IP2GEO_DATA_INDEX_NAME_PREFIX); + + private final Set SUPPORTED_COMPONENTS = Set.of( + UploadStats.class, + DatasourceUpdateService.class, + DatasourceFacade.class, + Ip2GeoExecutor.class, + GeoIpDataFacade.class + ); + + @Mock + private Client client; + @Mock + private ClusterService clusterService; + @Mock + private IngestService ingestService; + @Mock + private ThreadPool threadPool; + @Mock + private ResourceWatcherService resourceWatcherService; + @Mock + private ScriptService scriptService; + @Mock + private NamedXContentRegistry xContentRegistry; + @Mock + private Environment environment; + @Mock + private NamedWriteableRegistry namedWriteableRegistry; + @Mock + private IndexNameExpressionResolver indexNameExpressionResolver; + @Mock + private Supplier repositoriesServiceSupplier; + private NodeEnvironment nodeEnvironment; + private Settings settings; + private AutoCloseable openMocks; + + @Before + public void init() { + openMocks = MockitoAnnotations.openMocks(this); + settings = Settings.EMPTY; + when(client.settings()).thenReturn(settings); when(clusterService.getClusterSettings()).thenReturn(clusterSettings); - ingestService = mock(IngestService.class); + when(clusterService.getSettings()).thenReturn(settings); when(ingestService.getClusterService()).thenReturn(clusterService); + nodeEnvironment = null; + } + + @After + public void close() throws Exception { + openMocks.close(); + } + + public void testSystemIndexDescriptors() { + GeospatialPlugin plugin = new GeospatialPlugin(); + Set registeredSystemIndexPatterns = new HashSet<>(); + for (SystemIndexDescriptor descriptor : plugin.getSystemIndexDescriptors(Settings.EMPTY)) { + registeredSystemIndexPatterns.add(descriptor.getIndexPattern()); + } + assertEquals(SUPPORTED_SYSTEM_INDEX_PATTERN, registeredSystemIndexPatterns); + + } + + public void testExecutorBuilders() { + GeospatialPlugin plugin = new GeospatialPlugin(); + assertEquals(1, plugin.getExecutorBuilders(Settings.EMPTY).size()); + } + + public void testCreateComponents() { + GeospatialPlugin plugin = new GeospatialPlugin(); + Set registeredComponents = new HashSet<>(); + Collection components = plugin.createComponents( + client, + clusterService, + threadPool, + resourceWatcherService, + scriptService, + xContentRegistry, + environment, + nodeEnvironment, + namedWriteableRegistry, + indexNameExpressionResolver, + repositoriesServiceSupplier + ); + for (Object component : components) { + registeredComponents.add(component.getClass()); + } + assertEquals(SUPPORTED_COMPONENTS, registeredComponents); } public void testIsAnIngestPlugin() { diff --git a/src/test/resources/ip2geo/manifest.json b/src/test/resources/ip2geo/manifest.json new file mode 100644 index 00000000..652bc9d8 --- /dev/null +++ b/src/test/resources/ip2geo/manifest.json @@ -0,0 +1,8 @@ +{ + "url": "https://test.com/db.zip", + "db_name": "sample_valid.csv", + "md5_hash": "safasdfaskkkesadfasdf", + "valid_for_in_days": 30, + "updated_at": 3134012341236, + "provider": "sample_provider" +} \ No newline at end of file diff --git a/src/test/resources/ip2geo/manifest_invalid_url.json b/src/test/resources/ip2geo/manifest_invalid_url.json new file mode 100644 index 00000000..77d68aaf --- /dev/null +++ b/src/test/resources/ip2geo/manifest_invalid_url.json @@ -0,0 +1,8 @@ +{ + "url": "invalid://test.com/db.zip", + "db_name": "sample_valid.csv", + "md5_hash": "safasdfaskkkesadfasdf", + "valid_for_in_days": 30, + "updated_at": 3134012341236, + "provider": "sample_provider" +} \ No newline at end of file diff --git a/src/test/resources/ip2geo/manifest_template.json b/src/test/resources/ip2geo/manifest_template.json new file mode 100644 index 00000000..4c273fa4 --- /dev/null +++ b/src/test/resources/ip2geo/manifest_template.json @@ -0,0 +1,8 @@ +{ + "url": "URL", + "db_name": "sample_valid.csv", + "md5_hash": "safasdfaskkkesadfasdf", + "valid_for_in_days": 30, + "updated_at": 3134012341236, + "provider": "maxmind" +} \ No newline at end of file diff --git a/src/test/resources/ip2geo/sample_invalid_less_than_two_fields.csv b/src/test/resources/ip2geo/sample_invalid_less_than_two_fields.csv new file mode 100644 index 00000000..08670061 --- /dev/null +++ b/src/test/resources/ip2geo/sample_invalid_less_than_two_fields.csv @@ -0,0 +1,2 @@ +network +1.0.0.0/24 \ No newline at end of file diff --git a/src/test/resources/ip2geo/sample_valid.csv b/src/test/resources/ip2geo/sample_valid.csv new file mode 100644 index 00000000..a6d08935 --- /dev/null +++ b/src/test/resources/ip2geo/sample_valid.csv @@ -0,0 +1,3 @@ +network,country_name +1.0.0.0/24,Australia +10.0.0.0/24,USA \ No newline at end of file diff --git a/src/test/resources/ip2geo/sample_valid.zip b/src/test/resources/ip2geo/sample_valid.zip new file mode 100644 index 0000000000000000000000000000000000000000..0bdeeadbf1f9d2c9c7d542cde8694556b4e055fa GIT binary patch literal 250 zcmWIWW@Zs#-~d9W>KS1SP@oB<1sD_E%D&A`a=gOPy&XbJVQ4^ZKW*d7jfhtF5`J=e+kB-={`bj2Rqjm4#)G zUwSlY(UCc4To?kp**Vz%&wil+)C00Bz?+dtgc;!uWI2#KU|>ljh()Ta0=!w-K>8Sg LFb+t!fjA5RP5V9S literal 0 HcmV?d00001