diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a19cc0e5e..5b3ac85afa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Enhancements ### Bug Fixes ### Infrastructure +* Make jacoco report to be generated faster in local ([#267](https://github.com/opensearch-project/geospatial/pull/267)) +* Exclude lombok generated code from jacoco coverage report ([#268](https://github.com/opensearch-project/geospatial/pull/268)) ### Documentation ### Maintenance +* Change package for Strings.hasText ([#314](https://github.com/opensearch-project/geospatial/pull/314)) ### Refactoring diff --git a/build.gradle b/build.gradle index 92f1b13ee2..ca35448a2e 100644 --- a/build.gradle +++ b/build.gradle @@ -5,6 +5,8 @@ import org.opensearch.gradle.test.RestIntegTestTask +import java.util.concurrent.Callable + apply plugin: 'java' apply plugin: 'idea' apply plugin: 'opensearch.opensearchplugin' @@ -35,6 +37,7 @@ opensearchplugin { classname "${projectPath}.${pathToPlugin}.${pluginClassName}" licenseFile rootProject.file('LICENSE') noticeFile rootProject.file('NOTICE') + extendedPlugins = ['opensearch-job-scheduler'] } // This requires an additional Jar not published as part of build-tools @@ -142,6 +145,10 @@ publishing { } +configurations { + zipArchive +} + //****************************************************************************/ // Dependencies //****************************************************************************/ @@ -154,6 +161,9 @@ dependencies { implementation "org.apache.commons:commons-lang3:3.12.0" implementation "org.locationtech.spatial4j:spatial4j:${versions.spatial4j}" implementation "org.locationtech.jts:jts-core:${versions.jts}" + implementation "org.apache.commons:commons-csv:1.10.0" + zipArchive group: 'org.opensearch.plugin', name:'opensearch-job-scheduler', version: "${opensearch_build}" + compileOnly "org.opensearch:opensearch-job-scheduler-spi:${opensearch_build}" } licenseHeaders.enabled = true @@ -206,8 +216,6 @@ integTest { testClusters.integTest { testDistribution = "ARCHIVE" - // This installs our plugin into the testClusters - plugin(project.tasks.bundlePlugin.archiveFile) // Cluster shrink exception thrown if we try to set numberOfNodes to 1, so only apply if > 1 if (_numNodes > 1) numberOfNodes = _numNodes // When running integration tests it doesn't forward the --debug-jvm to the cluster anymore @@ -220,6 +228,49 @@ testClusters.integTest { debugPort += 1 } } + + // This installs our plugin into the testClusters + plugin(project.tasks.bundlePlugin.archiveFile) + plugin(provider(new Callable(){ + @Override + RegularFile call() throws Exception { + return new RegularFile() { + @Override + File getAsFile() { + return configurations.zipArchive.asFileTree.getSingleFile() + } + } + } + })) + + // opensearch-geospatial plugin is being added to the list of plugins for the testCluster during build before + // the opensearch-job-scheduler plugin, which is causing build failures. From the stack trace, this looks like a bug. + // + // Exception in thread "main" java.lang.IllegalArgumentException: Missing plugin [opensearch-job-scheduler], dependency of [opensearch-geospatial] + // at org.opensearch.plugins.PluginsService.addSortedBundle(PluginsService.java:515) + // + // A temporary hack is to reorder the plugins list after evaluation but prior to task execution when the plugins are installed. + // See https://github.com/opensearch-project/anomaly-detection/blob/fd547014fdde5114bbc9c8e49fe7aaa37eb6e793/build.gradle#L400-L422 + nodes.each { node -> + def plugins = node.plugins + def firstPlugin = plugins.get(0) + plugins.remove(0) + plugins.add(firstPlugin) + } +} + +testClusters.yamlRestTest { + plugin(provider(new Callable(){ + @Override + RegularFile call() throws Exception { + return new RegularFile() { + @Override + File getAsFile() { + return configurations.zipArchive.asFileTree.getSingleFile() + } + } + } + })) } run { 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 0000000000..d48c6dc2b7 --- /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/constants/IndexSetting.java b/src/main/java/org/opensearch/geospatial/constants/IndexSetting.java new file mode 100644 index 0000000000..0978c411d3 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/constants/IndexSetting.java @@ -0,0 +1,18 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.constants; + +/** + * Collection of keys for index setting + */ +public class IndexSetting { + public static final String NUMBER_OF_SHARDS = "index.number_of_shards"; + public static final String NUMBER_OF_REPLICAS = "index.number_of_replicas"; + public static final String REFRESH_INTERVAL = "index.refresh_interval"; + public static final String AUTO_EXPAND_REPLICAS = "index.auto_expand_replicas"; + public static final String HIDDEN = "index.hidden"; + public static final String BLOCKS_WRITE = "index.blocks.write"; +} diff --git a/src/main/java/org/opensearch/geospatial/exceptions/ConcurrentModificationException.java b/src/main/java/org/opensearch/geospatial/exceptions/ConcurrentModificationException.java new file mode 100644 index 0000000000..f3a2ae11db --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/exceptions/ConcurrentModificationException.java @@ -0,0 +1,37 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.exceptions; + +import java.io.IOException; + +import org.opensearch.OpenSearchException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; + +/** + * General ConcurrentModificationException corresponding to the {@link RestStatus#BAD_REQUEST} status code + * + * The exception is thrown when multiple mutation API is called for a same resource at the same time + */ +public class ConcurrentModificationException extends OpenSearchException { + + public ConcurrentModificationException(String msg, Object... args) { + super(msg, args); + } + + public ConcurrentModificationException(String msg, Throwable cause, Object... args) { + super(msg, cause, args); + } + + public ConcurrentModificationException(StreamInput in) throws IOException { + super(in); + } + + @Override + public final RestStatus status() { + return RestStatus.BAD_REQUEST; + } +} diff --git a/src/main/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceException.java b/src/main/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceException.java new file mode 100644 index 0000000000..ce30d5c0f7 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceException.java @@ -0,0 +1,38 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.exceptions; + +import java.io.IOException; + +import org.opensearch.OpenSearchException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; + +/** + * IncompatibleDatasourceException corresponding to the {@link RestStatus#BAD_REQUEST} status code + * + * The exception is thrown when a user tries to update datasource with new endpoint which is not compatible + * with current datasource + */ +public class IncompatibleDatasourceException extends OpenSearchException { + + public IncompatibleDatasourceException(String msg, Object... args) { + super(msg, args); + } + + public IncompatibleDatasourceException(String msg, Throwable cause, Object... args) { + super(msg, cause, args); + } + + public IncompatibleDatasourceException(StreamInput in) throws IOException { + super(in); + } + + @Override + public final RestStatus status() { + return RestStatus.BAD_REQUEST; + } +} diff --git a/src/main/java/org/opensearch/geospatial/exceptions/ResourceInUseException.java b/src/main/java/org/opensearch/geospatial/exceptions/ResourceInUseException.java new file mode 100644 index 0000000000..d102bb9d32 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/exceptions/ResourceInUseException.java @@ -0,0 +1,35 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.exceptions; + +import java.io.IOException; + +import org.opensearch.OpenSearchException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.rest.RestStatus; + +/** + * Generic ResourceInUseException corresponding to the {@link RestStatus#BAD_REQUEST} status code + */ +public class ResourceInUseException extends OpenSearchException { + + public ResourceInUseException(String msg, Object... args) { + super(msg, args); + } + + public ResourceInUseException(String msg, Throwable cause, Object... args) { + super(msg, cause, args); + } + + public ResourceInUseException(StreamInput in) throws IOException { + super(in); + } + + @Override + public final RestStatus status() { + return RestStatus.BAD_REQUEST; + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceAction.java new file mode 100644 index 0000000000..b08e0861e4 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceAction.java @@ -0,0 +1,27 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import org.opensearch.action.ActionType; +import org.opensearch.action.support.master.AcknowledgedResponse; + +/** + * Ip2Geo datasource delete action + */ +public class DeleteDatasourceAction extends ActionType { + /** + * Delete datasource action instance + */ + public static final DeleteDatasourceAction INSTANCE = new DeleteDatasourceAction(); + /** + * Delete datasource action name + */ + public static final String NAME = "cluster:admin/geospatial/datasource/delete"; + + private DeleteDatasourceAction() { + super(NAME, AcknowledgedResponse::new); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceRequest.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceRequest.java new file mode 100644 index 0000000000..60bc3f0c61 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceRequest.java @@ -0,0 +1,58 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; + +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.Setter; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +/** + * GeoIP datasource delete request + */ +@Getter +@Setter +@AllArgsConstructor +public class DeleteDatasourceRequest extends ActionRequest { + /** + * @param name the datasource name + * @return the datasource name + */ + private String name; + + /** + * Constructor + * + * @param in the stream input + * @throws IOException IOException + */ + public DeleteDatasourceRequest(final StreamInput in) throws IOException { + super(in); + this.name = in.readString(); + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException errors = null; + if (name == null || name.isBlank()) { + errors = new ActionRequestValidationException(); + errors.addValidationError("Datasource name should not be empty"); + } + return errors; + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(name); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceTransportAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceTransportAction.java new file mode 100644 index 0000000000..8046f6ff03 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceTransportAction.java @@ -0,0 +1,139 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; + +import lombok.extern.log4j.Log4j2; + +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.inject.Inject; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.exceptions.ConcurrentModificationException; +import org.opensearch.geospatial.exceptions.ResourceInUseException; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.dao.GeoIpDataDao; +import org.opensearch.geospatial.ip2geo.dao.Ip2GeoProcessorDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.ingest.IngestService; +import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +/** + * Transport action to delete datasource + */ +@Log4j2 +public class DeleteDatasourceTransportAction extends HandledTransportAction { + private static final long LOCK_DURATION_IN_SECONDS = 300l; + private final Ip2GeoLockService lockService; + private final IngestService ingestService; + private final DatasourceDao datasourceDao; + private final GeoIpDataDao geoIpDataDao; + private final Ip2GeoProcessorDao ip2GeoProcessorDao; + private final ThreadPool threadPool; + + /** + * Constructor + * @param transportService the transport service + * @param actionFilters the action filters + * @param lockService the lock service + * @param ingestService the ingest service + * @param datasourceDao the datasource facade + */ + @Inject + public DeleteDatasourceTransportAction( + final TransportService transportService, + final ActionFilters actionFilters, + final Ip2GeoLockService lockService, + final IngestService ingestService, + final DatasourceDao datasourceDao, + final GeoIpDataDao geoIpDataDao, + final Ip2GeoProcessorDao ip2GeoProcessorDao, + final ThreadPool threadPool + ) { + super(DeleteDatasourceAction.NAME, transportService, actionFilters, DeleteDatasourceRequest::new); + this.lockService = lockService; + this.ingestService = ingestService; + this.datasourceDao = datasourceDao; + this.geoIpDataDao = geoIpDataDao; + this.ip2GeoProcessorDao = ip2GeoProcessorDao; + this.threadPool = threadPool; + } + + /** + * We delete datasource regardless of its state as long as we can acquire a lock + * + * @param task the task + * @param request the request + * @param listener the listener + */ + @Override + protected void doExecute(final Task task, final DeleteDatasourceRequest request, final ActionListener listener) { + lockService.acquireLock(request.getName(), LOCK_DURATION_IN_SECONDS, ActionListener.wrap(lock -> { + if (lock == null) { + listener.onFailure( + new ConcurrentModificationException("another processor is holding a lock on the resource. Try again later") + ); + return; + } + try { + // TODO: makes every sub-methods as async call to avoid using a thread in generic pool + threadPool.generic().submit(() -> { + try { + deleteDatasource(request.getName()); + lockService.releaseLock(lock); + listener.onResponse(new AcknowledgedResponse(true)); + } catch (Exception e) { + lockService.releaseLock(lock); + listener.onFailure(e); + } + }); + } catch (Exception e) { + lockService.releaseLock(lock); + listener.onFailure(e); + } + }, exception -> { listener.onFailure(exception); })); + } + + @VisibleForTesting + protected void deleteDatasource(final String datasourceName) throws IOException { + Datasource datasource = datasourceDao.getDatasource(datasourceName); + if (datasource == null) { + throw new ResourceNotFoundException("no such datasource exist"); + } + + setDatasourceStateAsDeleting(datasource); + geoIpDataDao.deleteIp2GeoDataIndex(datasource.getIndices()); + datasourceDao.deleteDatasource(datasource); + } + + private void setDatasourceStateAsDeleting(final Datasource datasource) { + if (ip2GeoProcessorDao.getProcessors(datasource.getName()).isEmpty() == false) { + throw new ResourceInUseException("datasource is being used by one of processors"); + } + + DatasourceState previousState = datasource.getState(); + datasource.setState(DatasourceState.DELETING); + datasourceDao.updateDatasource(datasource); + + // Check again as processor might just have been created. + // If it fails to update the state back to the previous state, the new processor + // will fail to convert an ip to a geo data. + // In such case, user have to delete the processor and delete this datasource again. + if (ip2GeoProcessorDao.getProcessors(datasource.getName()).isEmpty() == false) { + datasource.setState(previousState); + datasourceDao.updateDatasource(datasource); + throw new ResourceInUseException("datasource is being used by one of processors"); + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceAction.java new file mode 100644 index 0000000000..039ab35bc4 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceAction.java @@ -0,0 +1,26 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import org.opensearch.action.ActionType; + +/** + * Ip2Geo datasource get action + */ +public class GetDatasourceAction extends ActionType { + /** + * Get datasource action instance + */ + public static final GetDatasourceAction INSTANCE = new GetDatasourceAction(); + /** + * Get datasource action name + */ + public static final String NAME = "cluster:admin/geospatial/datasource/get"; + + private GetDatasourceAction() { + super(NAME, GetDatasourceResponse::new); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceRequest.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceRequest.java new file mode 100644 index 0000000000..3cf94e557c --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceRequest.java @@ -0,0 +1,67 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; + +import lombok.Getter; +import lombok.Setter; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; + +/** + * Ip2Geo datasource get request + */ +@Getter +@Setter +public class GetDatasourceRequest extends ActionRequest { + /** + * @param names the datasource names + * @return the datasource names + */ + private String[] names; + + /** + * Constructs a new get datasource request with a list of datasources. + * + * If the list of datasources is empty or it contains a single element "_all", all registered datasources + * are returned. + * + * @param names list of datasource names + */ + public GetDatasourceRequest(final String[] names) { + this.names = names; + } + + /** + * Constructor with stream input + * @param in the stream input + * @throws IOException IOException + */ + public GetDatasourceRequest(final StreamInput in) throws IOException { + super(in); + this.names = in.readStringArray(); + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException errors = null; + if (names == null) { + errors = new ActionRequestValidationException(); + errors.addValidationError("names should not be null"); + } + return errors; + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + super.writeTo(out); + out.writeStringArray(names); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceResponse.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceResponse.java new file mode 100644 index 0000000000..c6afac0b9c --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceResponse.java @@ -0,0 +1,88 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; +import java.time.Instant; +import java.util.List; + +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.Setter; + +import org.opensearch.action.ActionResponse; +import org.opensearch.core.ParseField; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.ToXContentObject; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; + +/** + * Ip2Geo datasource get request + */ +@Getter +@Setter +@EqualsAndHashCode(callSuper = false) +public class GetDatasourceResponse extends ActionResponse implements ToXContentObject { + private static final ParseField FIELD_NAME_DATASOURCES = new ParseField("datasources"); + private static final ParseField FIELD_NAME_NAME = new ParseField("name"); + private static final ParseField FIELD_NAME_STATE = new ParseField("state"); + private static final ParseField FIELD_NAME_ENDPOINT = new ParseField("endpoint"); + private static final ParseField FIELD_NAME_UPDATE_INTERVAL = new ParseField("update_interval_in_days"); + private static final ParseField FIELD_NAME_NEXT_UPDATE_AT = new ParseField("next_update_at_in_epoch_millis"); + private static final ParseField FIELD_NAME_NEXT_UPDATE_AT_READABLE = new ParseField("next_update_at"); + private static final ParseField FIELD_NAME_DATABASE = new ParseField("database"); + private static final ParseField FIELD_NAME_UPDATE_STATS = new ParseField("update_stats"); + private List datasources; + + /** + * Default constructor + * + * @param datasources List of datasources + */ + public GetDatasourceResponse(final List datasources) { + this.datasources = datasources; + } + + /** + * Constructor with StreamInput + * + * @param in the stream input + */ + public GetDatasourceResponse(final StreamInput in) throws IOException { + datasources = in.readList(Datasource::new); + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + out.writeList(datasources); + } + + @Override + public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { + builder.startObject(); + builder.startArray(FIELD_NAME_DATASOURCES.getPreferredName()); + for (Datasource datasource : datasources) { + builder.startObject(); + builder.field(FIELD_NAME_NAME.getPreferredName(), datasource.getName()); + builder.field(FIELD_NAME_STATE.getPreferredName(), datasource.getState()); + builder.field(FIELD_NAME_ENDPOINT.getPreferredName(), datasource.getEndpoint()); + builder.field(FIELD_NAME_UPDATE_INTERVAL.getPreferredName(), datasource.getUserSchedule().getInterval()); + builder.timeField( + FIELD_NAME_NEXT_UPDATE_AT.getPreferredName(), + FIELD_NAME_NEXT_UPDATE_AT_READABLE.getPreferredName(), + datasource.getUserSchedule().getNextExecutionTime(Instant.now()).toEpochMilli() + ); + builder.field(FIELD_NAME_DATABASE.getPreferredName(), datasource.getDatabase()); + builder.field(FIELD_NAME_UPDATE_STATS.getPreferredName(), datasource.getUpdateStats()); + builder.endObject(); + } + builder.endArray(); + builder.endObject(); + return builder; + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceTransportAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceTransportAction.java new file mode 100644 index 0000000000..d8e1aba7ed --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceTransportAction.java @@ -0,0 +1,81 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.util.Collections; +import java.util.List; + +import org.opensearch.OpenSearchException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.common.inject.Inject; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.tasks.Task; +import org.opensearch.transport.TransportService; + +/** + * Transport action to get datasource + */ +public class GetDatasourceTransportAction extends HandledTransportAction { + private final DatasourceDao datasourceDao; + + /** + * Default constructor + * @param transportService the transport service + * @param actionFilters the action filters + * @param datasourceDao the datasource facade + */ + @Inject + public GetDatasourceTransportAction( + final TransportService transportService, + final ActionFilters actionFilters, + final DatasourceDao datasourceDao + ) { + super(GetDatasourceAction.NAME, transportService, actionFilters, GetDatasourceRequest::new); + this.datasourceDao = datasourceDao; + } + + @Override + protected void doExecute(final Task task, final GetDatasourceRequest request, final ActionListener listener) { + if (shouldGetAllDatasource(request)) { + // We don't expect too many data sources. Therefore, querying all data sources without pagination should be fine. + datasourceDao.getAllDatasources(newActionListener(listener)); + } else { + datasourceDao.getDatasources(request.getNames(), newActionListener(listener)); + } + } + + private boolean shouldGetAllDatasource(final GetDatasourceRequest request) { + if (request.getNames() == null) { + throw new OpenSearchException("names in a request should not be null"); + } + + return request.getNames().length == 0 || (request.getNames().length == 1 && "_all".equals(request.getNames()[0])); + } + + @VisibleForTesting + protected ActionListener> newActionListener(final ActionListener listener) { + return new ActionListener<>() { + @Override + public void onResponse(final List datasources) { + listener.onResponse(new GetDatasourceResponse(datasources)); + } + + @Override + public void onFailure(final Exception e) { + if (e instanceof IndexNotFoundException) { + listener.onResponse(new GetDatasourceResponse(Collections.emptyList())); + return; + } + listener.onFailure(e); + } + }; + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceAction.java new file mode 100644 index 0000000000..2554b4f5a3 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceAction.java @@ -0,0 +1,27 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import org.opensearch.action.ActionType; +import org.opensearch.action.support.master.AcknowledgedResponse; + +/** + * Ip2Geo datasource creation action + */ +public class PutDatasourceAction extends ActionType { + /** + * Put datasource action instance + */ + public static final PutDatasourceAction INSTANCE = new PutDatasourceAction(); + /** + * Put datasource action name + */ + public static final String NAME = "cluster:admin/geospatial/datasource/put"; + + private PutDatasourceAction() { + super(NAME, AcknowledgedResponse::new); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequest.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequest.java new file mode 100644 index 0000000000..e764f6c498 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequest.java @@ -0,0 +1,210 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; +import java.io.UnsupportedEncodingException; +import java.net.MalformedURLException; +import java.net.URISyntaxException; +import java.net.URL; +import java.util.Locale; + +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.Setter; +import lombok.extern.log4j.Log4j2; + +import org.opensearch.OpenSearchException; +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.Strings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.ParseField; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.ObjectParser; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; + +/** + * Ip2Geo datasource creation request + */ +@Getter +@Setter +@Log4j2 +@EqualsAndHashCode(callSuper = false) +public class PutDatasourceRequest extends ActionRequest { + private static final int MAX_DATASOURCE_NAME_BYTES = 255; + public static final ParseField ENDPOINT_FIELD = new ParseField("endpoint"); + public static final ParseField UPDATE_INTERVAL_IN_DAYS_FIELD = new ParseField("update_interval_in_days"); + /** + * @param name the datasource name + * @return the datasource name + */ + private String name; + /** + * @param endpoint url to a manifest file for a datasource + * @return url to a manifest file for a datasource + */ + private String endpoint; + /** + * @param updateInterval update interval of a datasource + * @return update interval of a datasource + */ + private TimeValue updateInterval; + + /** + * Parser of a datasource + */ + public static final ObjectParser PARSER; + static { + PARSER = new ObjectParser<>("put_datasource"); + PARSER.declareString((request, val) -> request.setEndpoint(val), ENDPOINT_FIELD); + PARSER.declareLong((request, val) -> request.setUpdateInterval(TimeValue.timeValueDays(val)), UPDATE_INTERVAL_IN_DAYS_FIELD); + } + + /** + * Default constructor + * @param name name of a datasource + */ + public PutDatasourceRequest(final String name) { + this.name = name; + } + + /** + * Constructor with stream input + * @param in the stream input + * @throws IOException IOException + */ + public PutDatasourceRequest(final StreamInput in) throws IOException { + super(in); + this.name = in.readString(); + this.endpoint = in.readString(); + this.updateInterval = in.readTimeValue(); + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(name); + out.writeString(endpoint); + out.writeTimeValue(updateInterval); + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException errors = new ActionRequestValidationException(); + validateDatasourceName(errors); + validateEndpoint(errors); + validateUpdateInterval(errors); + return errors.validationErrors().isEmpty() ? null : errors; + } + + private void validateDatasourceName(final ActionRequestValidationException errors) { + if (!Strings.validFileName(name)) { + errors.addValidationError("Datasource name must not contain the following characters " + Strings.INVALID_FILENAME_CHARS); + return; + } + if (name.isEmpty()) { + errors.addValidationError("Datasource name must not be empty"); + return; + } + if (name.contains("#")) { + errors.addValidationError("Datasource name must not contain '#'"); + return; + } + if (name.contains(":")) { + errors.addValidationError("Datasource name must not contain ':'"); + return; + } + if (name.charAt(0) == '_' || name.charAt(0) == '-' || name.charAt(0) == '+') { + errors.addValidationError("Datasource name must not start with '_', '-', or '+'"); + return; + } + int byteCount = 0; + try { + byteCount = name.getBytes("UTF-8").length; + } catch (UnsupportedEncodingException e) { + // UTF-8 should always be supported, but rethrow this if it is not for some reason + throw new OpenSearchException("Unable to determine length of datasource name", e); + } + if (byteCount > MAX_DATASOURCE_NAME_BYTES) { + errors.addValidationError("Datasource name is too long, (" + byteCount + " > " + MAX_DATASOURCE_NAME_BYTES + ")"); + return; + } + if (name.equals(".") || name.equals("..")) { + errors.addValidationError("Datasource name must not be '.' or '..'"); + return; + } + } + + /** + * Conduct following validation on endpoint + * 1. endpoint format complies with RFC-2396 + * 2. validate manifest file from the endpoint + * + * @param errors the errors to add error messages + */ + private void validateEndpoint(final ActionRequestValidationException errors) { + try { + URL url = new URL(endpoint); + url.toURI(); // Validate URL complies with RFC-2396 + validateManifestFile(url, errors); + } catch (MalformedURLException | URISyntaxException e) { + log.info("Invalid URL[{}] is provided", endpoint, e); + errors.addValidationError("Invalid URL format is provided"); + } + } + + /** + * 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. updateInterval is less than validForInDays value in the manifest file + * + * @param url the url to validate + * @param errors the errors to add error messages + */ + private void validateManifestFile(final URL url, final ActionRequestValidationException errors) { + DatasourceManifest manifest; + try { + manifest = DatasourceManifest.Builder.build(url); + } catch (Exception e) { + log.info("Error occurred while reading a file from {}", url, e); + errors.addValidationError(String.format(Locale.ROOT, "Error occurred while reading a file from %s: %s", url, e.getMessage())); + return; + } + + try { + new URL(manifest.getUrl()).toURI(); // Validate URL complies with RFC-2396 + } catch (MalformedURLException | URISyntaxException e) { + log.info("Invalid URL[{}] is provided for url field in the manifest file", manifest.getUrl(), e); + errors.addValidationError("Invalid URL format is provided for url field in the manifest file"); + return; + } + + if (manifest.getValidForInDays() != null && updateInterval.days() >= manifest.getValidForInDays()) { + errors.addValidationError( + String.format( + Locale.ROOT, + "updateInterval %d should be smaller than %d", + updateInterval.days(), + manifest.getValidForInDays() + ) + ); + } + } + + /** + * Validate updateInterval is equal or larger than 1 + * + * @param errors the errors to add error messages + */ + private void validateUpdateInterval(final ActionRequestValidationException errors) { + 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 new file mode 100644 index 0000000000..ed4d1ee7d1 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportAction.java @@ -0,0 +1,173 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService.LOCK_DURATION_IN_SECONDS; + +import java.time.Instant; +import java.util.concurrent.atomic.AtomicReference; + +import lombok.extern.log4j.Log4j2; + +import org.opensearch.ResourceAlreadyExistsException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.StepListener; +import org.opensearch.action.index.IndexResponse; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.inject.Inject; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.exceptions.ConcurrentModificationException; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; +import org.opensearch.index.engine.VersionConflictEngineException; +import org.opensearch.jobscheduler.spi.LockModel; +import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +/** + * Transport action to create datasource + */ +@Log4j2 +public class PutDatasourceTransportAction extends HandledTransportAction { + private final ThreadPool threadPool; + private final DatasourceDao datasourceDao; + private final DatasourceUpdateService datasourceUpdateService; + private final Ip2GeoLockService lockService; + + /** + * Default constructor + * @param transportService the transport service + * @param actionFilters the action filters + * @param threadPool the thread pool + * @param datasourceDao the datasource facade + * @param datasourceUpdateService the datasource update service + * @param lockService the lock service + */ + @Inject + public PutDatasourceTransportAction( + final TransportService transportService, + final ActionFilters actionFilters, + final ThreadPool threadPool, + final DatasourceDao datasourceDao, + final DatasourceUpdateService datasourceUpdateService, + final Ip2GeoLockService lockService + ) { + super(PutDatasourceAction.NAME, transportService, actionFilters, PutDatasourceRequest::new); + this.threadPool = threadPool; + this.datasourceDao = datasourceDao; + this.datasourceUpdateService = datasourceUpdateService; + this.lockService = lockService; + } + + @Override + protected void doExecute(final Task task, final PutDatasourceRequest request, final ActionListener listener) { + lockService.acquireLock(request.getName(), LOCK_DURATION_IN_SECONDS, ActionListener.wrap(lock -> { + if (lock == null) { + listener.onFailure( + new ConcurrentModificationException("another processor is holding a lock on the resource. Try again later") + ); + return; + } + try { + internalDoExecute(request, lock, listener); + } catch (Exception e) { + lockService.releaseLock(lock); + listener.onFailure(e); + } + }, exception -> { listener.onFailure(exception); })); + } + + /** + * This method takes lock as a parameter and is responsible for releasing lock + * unless exception is thrown + */ + @VisibleForTesting + protected void internalDoExecute( + final PutDatasourceRequest request, + final LockModel lock, + final ActionListener listener + ) { + StepListener createIndexStep = new StepListener<>(); + datasourceDao.createIndexIfNotExists(createIndexStep); + createIndexStep.whenComplete(v -> { + Datasource datasource = Datasource.Builder.build(request); + datasourceDao.putDatasource(datasource, getIndexResponseListener(datasource, lock, listener)); + }, exception -> { + lockService.releaseLock(lock); + listener.onFailure(exception); + }); + } + + /** + * This method takes lock as a parameter and is responsible for releasing lock + * unless exception is thrown + */ + @VisibleForTesting + protected ActionListener getIndexResponseListener( + final Datasource datasource, + final LockModel lock, + 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(() -> { + AtomicReference lockReference = new AtomicReference<>(lock); + try { + createDatasource(datasource, lockService.getRenewLockRunnable(lockReference)); + } finally { + lockService.releaseLock(lockReference.get()); + } + }); + listener.onResponse(new AcknowledgedResponse(true)); + } + + @Override + public void onFailure(final Exception e) { + lockService.releaseLock(lock); + if (e instanceof VersionConflictEngineException) { + listener.onFailure(new ResourceAlreadyExistsException("datasource [{}] already exists", datasource.getName())); + } else { + listener.onFailure(e); + } + } + }; + } + + @VisibleForTesting + protected void createDatasource(final Datasource datasource, final Runnable renewLock) { + if (DatasourceState.CREATING.equals(datasource.getState()) == false) { + log.error("Invalid datasource state. Expecting {} but received {}", DatasourceState.CREATING, datasource.getState()); + markDatasourceAsCreateFailed(datasource); + return; + } + + try { + datasourceUpdateService.updateOrCreateGeoIpData(datasource, renewLock); + } catch (Exception e) { + log.error("Failed to create datasource for {}", datasource.getName(), e); + markDatasourceAsCreateFailed(datasource); + } + } + + private void markDatasourceAsCreateFailed(final Datasource datasource) { + datasource.getUpdateStats().setLastFailedAt(Instant.now()); + datasource.setState(DatasourceState.CREATE_FAILED); + try { + datasourceDao.updateDatasource(datasource); + } catch (Exception e) { + log.error("Failed to mark datasource state as CREATE_FAILED for {}", datasource.getName(), e); + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/RestDeleteDatasourceHandler.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestDeleteDatasourceHandler.java new file mode 100644 index 0000000000..dc2dd1179b --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestDeleteDatasourceHandler.java @@ -0,0 +1,49 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.opensearch.geospatial.shared.URLBuilder.URL_DELIMITER; +import static org.opensearch.geospatial.shared.URLBuilder.getPluginURLPrefix; +import static org.opensearch.rest.RestRequest.Method.DELETE; + +import java.util.List; +import java.util.Locale; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +/** + * Rest handler for Ip2Geo datasource delete request + */ +public class RestDeleteDatasourceHandler extends BaseRestHandler { + private static final String ACTION_NAME = "ip2geo_datasource_delete"; + private static final String PARAMS_NAME = "name"; + + @Override + public String getName() { + return ACTION_NAME; + } + + @Override + protected RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) { + final String name = request.param(PARAMS_NAME); + final DeleteDatasourceRequest deleteDatasourceRequest = new DeleteDatasourceRequest(name); + + return channel -> client.executeLocally( + DeleteDatasourceAction.INSTANCE, + deleteDatasourceRequest, + new RestToXContentListener<>(channel) + ); + } + + @Override + public List routes() { + String path = String.join(URL_DELIMITER, getPluginURLPrefix(), String.format(Locale.ROOT, "ip2geo/datasource/{%s}", PARAMS_NAME)); + return List.of(new Route(DELETE, path)); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/RestGetDatasourceHandler.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestGetDatasourceHandler.java new file mode 100644 index 0000000000..2c9379b59b --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestGetDatasourceHandler.java @@ -0,0 +1,46 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.opensearch.geospatial.shared.URLBuilder.URL_DELIMITER; +import static org.opensearch.geospatial.shared.URLBuilder.getPluginURLPrefix; +import static org.opensearch.rest.RestRequest.Method.GET; + +import java.util.List; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.common.Strings; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +/** + * Rest handler for Ip2Geo datasource get request + */ +public class RestGetDatasourceHandler extends BaseRestHandler { + private static final String ACTION_NAME = "ip2geo_datasource_get"; + + @Override + public String getName() { + return ACTION_NAME; + } + + @Override + protected RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) { + final String[] names = request.paramAsStringArray("name", Strings.EMPTY_ARRAY); + final GetDatasourceRequest getDatasourceRequest = new GetDatasourceRequest(names); + + return channel -> client.executeLocally(GetDatasourceAction.INSTANCE, getDatasourceRequest, new RestToXContentListener<>(channel)); + } + + @Override + public List routes() { + return List.of( + new Route(GET, String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource")), + new Route(GET, String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/{name}")) + ); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandler.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandler.java new file mode 100644 index 0000000000..0fac4ec9af --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandler.java @@ -0,0 +1,73 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.opensearch.geospatial.shared.URLBuilder.URL_DELIMITER; +import static org.opensearch.geospatial.shared.URLBuilder.getPluginURLPrefix; +import static org.opensearch.rest.RestRequest.Method.PUT; + +import java.io.IOException; +import java.util.List; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +/** + * Rest handler for Ip2Geo datasource creation + * + * This handler handles a request of + * PUT /_plugins/geospatial/ip2geo/datasource/{id} + * { + * "endpoint": {endpoint}, + * "update_interval_in_days": 3 + * } + * + * When request is received, it will create a datasource by downloading GeoIp data from the endpoint. + * After the creation of datasource is completed, it will schedule the next update task after update_interval_in_days. + * + */ +public class RestPutDatasourceHandler extends BaseRestHandler { + private static final String ACTION_NAME = "ip2geo_datasource_put"; + private final ClusterSettings clusterSettings; + + public RestPutDatasourceHandler(final ClusterSettings clusterSettings) { + this.clusterSettings = clusterSettings; + } + + @Override + public String getName() { + return ACTION_NAME; + } + + @Override + protected RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) throws IOException { + final PutDatasourceRequest putDatasourceRequest = new PutDatasourceRequest(request.param("name")); + if (request.hasContentOrSourceParam()) { + try (XContentParser parser = request.contentOrSourceParamParser()) { + PutDatasourceRequest.PARSER.parse(parser, putDatasourceRequest, null); + } + } + if (putDatasourceRequest.getEndpoint() == null) { + putDatasourceRequest.setEndpoint(clusterSettings.get(Ip2GeoSettings.DATASOURCE_ENDPOINT)); + } + if (putDatasourceRequest.getUpdateInterval() == null) { + putDatasourceRequest.setUpdateInterval(TimeValue.timeValueDays(clusterSettings.get(Ip2GeoSettings.DATASOURCE_UPDATE_INTERVAL))); + } + return channel -> client.executeLocally(PutDatasourceAction.INSTANCE, putDatasourceRequest, new RestToXContentListener<>(channel)); + } + + @Override + public List routes() { + String path = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/{name}"); + return List.of(new Route(PUT, path)); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandler.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandler.java new file mode 100644 index 0000000000..f9ba73ecaa --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandler.java @@ -0,0 +1,52 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.opensearch.geospatial.shared.URLBuilder.URL_DELIMITER; +import static org.opensearch.geospatial.shared.URLBuilder.getPluginURLPrefix; +import static org.opensearch.rest.RestRequest.Method.PUT; + +import java.io.IOException; +import java.util.List; + +import org.opensearch.client.node.NodeClient; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +/** + * Rest handler for Ip2Geo datasource update request + */ +public class RestUpdateDatasourceHandler extends BaseRestHandler { + private static final String ACTION_NAME = "ip2geo_datasource_update"; + + @Override + public String getName() { + return ACTION_NAME; + } + + @Override + protected RestChannelConsumer prepareRequest(final RestRequest request, final NodeClient client) throws IOException { + final UpdateDatasourceRequest updateDatasourceRequest = new UpdateDatasourceRequest(request.param("name")); + if (request.hasContentOrSourceParam()) { + try (XContentParser parser = request.contentOrSourceParamParser()) { + UpdateDatasourceRequest.PARSER.parse(parser, updateDatasourceRequest, null); + } + } + return channel -> client.executeLocally( + UpdateDatasourceAction.INSTANCE, + updateDatasourceRequest, + new RestToXContentListener<>(channel) + ); + } + + @Override + public List routes() { + String path = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/{name}/_settings"); + return List.of(new Route(PUT, path)); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceAction.java new file mode 100644 index 0000000000..96cd00df26 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceAction.java @@ -0,0 +1,27 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import org.opensearch.action.ActionType; +import org.opensearch.action.support.master.AcknowledgedResponse; + +/** + * Ip2Geo datasource update action + */ +public class UpdateDatasourceAction extends ActionType { + /** + * Update datasource action instance + */ + public static final UpdateDatasourceAction INSTANCE = new UpdateDatasourceAction(); + /** + * Update datasource action name + */ + public static final String NAME = "cluster:admin/geospatial/datasource/update"; + + private UpdateDatasourceAction() { + super(NAME, AcknowledgedResponse::new); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceRequest.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceRequest.java new file mode 100644 index 0000000000..2114d6c1e7 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceRequest.java @@ -0,0 +1,168 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; +import java.net.MalformedURLException; +import java.net.URISyntaxException; +import java.net.URL; +import java.util.Locale; + +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.Setter; +import lombok.extern.log4j.Log4j2; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.ParseField; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.xcontent.ObjectParser; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; + +/** + * Ip2Geo datasource update request + */ +@Getter +@Setter +@Log4j2 +@EqualsAndHashCode(callSuper = false) +public class UpdateDatasourceRequest extends ActionRequest { + public static final ParseField ENDPOINT_FIELD = new ParseField("endpoint"); + public static final ParseField UPDATE_INTERVAL_IN_DAYS_FIELD = new ParseField("update_interval_in_days"); + private static final int MAX_DATASOURCE_NAME_BYTES = 255; + /** + * @param name the datasource name + * @return the datasource name + */ + private String name; + /** + * @param endpoint url to a manifest file for a datasource + * @return url to a manifest file for a datasource + */ + private String endpoint; + /** + * @param updateInterval update interval of a datasource + * @return update interval of a datasource + */ + private TimeValue updateInterval; + + /** + * Parser of a datasource + */ + public static final ObjectParser PARSER; + static { + PARSER = new ObjectParser<>("update_datasource"); + PARSER.declareString((request, val) -> request.setEndpoint(val), ENDPOINT_FIELD); + PARSER.declareLong((request, val) -> request.setUpdateInterval(TimeValue.timeValueDays(val)), UPDATE_INTERVAL_IN_DAYS_FIELD); + } + + /** + * Constructor + * @param name name of a datasource + */ + public UpdateDatasourceRequest(final String name) { + this.name = name; + } + + /** + * Constructor + * @param in the stream input + * @throws IOException IOException + */ + public UpdateDatasourceRequest(final StreamInput in) throws IOException { + super(in); + this.name = in.readString(); + this.endpoint = in.readOptionalString(); + this.updateInterval = in.readOptionalTimeValue(); + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(name); + out.writeOptionalString(endpoint); + out.writeOptionalTimeValue(updateInterval); + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException errors = new ActionRequestValidationException(); + if (endpoint == null && updateInterval == null) { + errors.addValidationError("no values to update"); + } + + validateEndpoint(errors); + validateUpdateInterval(errors); + + return errors.validationErrors().isEmpty() ? null : errors; + } + + /** + * Conduct following validation on endpoint + * 1. endpoint format complies with RFC-2396 + * 2. validate manifest file from the endpoint + * + * @param errors the errors to add error messages + */ + private void validateEndpoint(final ActionRequestValidationException errors) { + if (endpoint == null) { + return; + } + + try { + URL url = new URL(endpoint); + url.toURI(); // Validate URL complies with RFC-2396 + validateManifestFile(url, errors); + } catch (MalformedURLException | URISyntaxException e) { + log.info("Invalid URL[{}] is provided", endpoint, e); + errors.addValidationError("Invalid URL format is provided"); + } + } + + /** + * Conduct following validation on url + * 1. can read manifest file from the endpoint + * 2. the url in the manifest file complies with RFC-2396 + * + * @param url the url to validate + * @param errors the errors to add error messages + */ + private void validateManifestFile(final URL url, final ActionRequestValidationException errors) { + DatasourceManifest manifest; + try { + manifest = DatasourceManifest.Builder.build(url); + } catch (Exception e) { + log.info("Error occurred while reading a file from {}", url, e); + errors.addValidationError(String.format(Locale.ROOT, "Error occurred while reading a file from %s: %s", url, e.getMessage())); + return; + } + + try { + new URL(manifest.getUrl()).toURI(); // Validate URL complies with RFC-2396 + } catch (MalformedURLException | URISyntaxException e) { + log.info("Invalid URL[{}] is provided for url field in the manifest file", manifest.getUrl(), e); + errors.addValidationError("Invalid URL format is provided for url field in the manifest file"); + } + } + + /** + * Validate updateInterval is equal or larger than 1 + * + * @param errors the errors to add error messages + */ + private void validateUpdateInterval(final ActionRequestValidationException errors) { + if (updateInterval == null) { + return; + } + + 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/UpdateDatasourceTransportAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceTransportAction.java new file mode 100644 index 0000000000..c832dc8982 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceTransportAction.java @@ -0,0 +1,216 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; +import java.net.URL; +import java.security.InvalidParameterException; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.List; +import java.util.Locale; + +import lombok.extern.log4j.Log4j2; + +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.HandledTransportAction; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.inject.Inject; +import org.opensearch.geospatial.exceptions.ConcurrentModificationException; +import org.opensearch.geospatial.exceptions.IncompatibleDatasourceException; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceTask; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; +import org.opensearch.tasks.Task; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +/** + * Transport action to update datasource + */ +@Log4j2 +public class UpdateDatasourceTransportAction extends HandledTransportAction { + private static final long LOCK_DURATION_IN_SECONDS = 300l; + private final Ip2GeoLockService lockService; + private final DatasourceDao datasourceDao; + private final DatasourceUpdateService datasourceUpdateService; + private final ThreadPool threadPool; + + /** + * Constructor + * + * @param transportService the transport service + * @param actionFilters the action filters + * @param lockService the lock service + * @param datasourceDao the datasource facade + * @param datasourceUpdateService the datasource update service + */ + @Inject + public UpdateDatasourceTransportAction( + final TransportService transportService, + final ActionFilters actionFilters, + final Ip2GeoLockService lockService, + final DatasourceDao datasourceDao, + final DatasourceUpdateService datasourceUpdateService, + final ThreadPool threadPool + ) { + super(UpdateDatasourceAction.NAME, transportService, actionFilters, UpdateDatasourceRequest::new); + this.lockService = lockService; + this.datasourceUpdateService = datasourceUpdateService; + this.datasourceDao = datasourceDao; + this.threadPool = threadPool; + } + + /** + * Get a lock and update datasource + * + * @param task the task + * @param request the request + * @param listener the listener + */ + @Override + protected void doExecute(final Task task, final UpdateDatasourceRequest request, final ActionListener listener) { + lockService.acquireLock(request.getName(), LOCK_DURATION_IN_SECONDS, ActionListener.wrap(lock -> { + if (lock == null) { + listener.onFailure( + new ConcurrentModificationException("another processor is holding a lock on the resource. Try again later") + ); + return; + } + try { + // TODO: makes every sub-methods as async call to avoid using a thread in generic pool + threadPool.generic().submit(() -> { + try { + Datasource datasource = datasourceDao.getDatasource(request.getName()); + if (datasource == null) { + throw new ResourceNotFoundException("no such datasource exist"); + } + validate(request, datasource); + updateIfChanged(request, datasource); + lockService.releaseLock(lock); + listener.onResponse(new AcknowledgedResponse(true)); + } catch (Exception e) { + lockService.releaseLock(lock); + listener.onFailure(e); + } + }); + } catch (Exception e) { + lockService.releaseLock(lock); + listener.onFailure(e); + } + }, exception -> listener.onFailure(exception))); + } + + private void updateIfChanged(final UpdateDatasourceRequest request, final Datasource datasource) { + boolean isChanged = false; + if (isEndpointChanged(request, datasource)) { + datasource.setEndpoint(request.getEndpoint()); + isChanged = true; + } + if (isUpdateIntervalChanged(request)) { + datasource.setUserSchedule(new IntervalSchedule(Instant.now(), (int) request.getUpdateInterval().getDays(), ChronoUnit.DAYS)); + datasource.setSystemSchedule(datasource.getUserSchedule()); + datasource.setTask(DatasourceTask.ALL); + isChanged = true; + } + + if (isChanged) { + datasourceDao.updateDatasource(datasource); + } + } + + /** + * Additional validation based on an existing datasource + * + * Basic validation is done in UpdateDatasourceRequest#validate + * In this method we do additional validation based on an existing datasource + * + * 1. Check the compatibility of new fields and old fields + * 2. Check the updateInterval is less than validForInDays in datasource + * + * This method throws exception if one of validation fails. + * + * @param request the update request + * @param datasource the existing datasource + * @throws IOException the exception + */ + private void validate(final UpdateDatasourceRequest request, final Datasource datasource) throws IOException { + validateFieldsCompatibility(request, datasource); + validateUpdateIntervalIsLessThanValidForInDays(request, datasource); + validateNextUpdateScheduleIsBeforeExpirationDay(request, datasource); + } + + private void validateNextUpdateScheduleIsBeforeExpirationDay(final UpdateDatasourceRequest request, final Datasource datasource) { + if (request.getUpdateInterval() == null) { + return; + } + + IntervalSchedule newSchedule = new IntervalSchedule(Instant.now(), (int) request.getUpdateInterval().getDays(), ChronoUnit.DAYS); + + if (newSchedule.getNextExecutionTime(Instant.now()).isAfter(datasource.expirationDay())) { + throw new IllegalArgumentException( + String.format(Locale.ROOT, "datasource will expire at %s with the update interval", datasource.expirationDay().toString()) + ); + } + } + + private void validateFieldsCompatibility(final UpdateDatasourceRequest request, final Datasource datasource) throws IOException { + if (isEndpointChanged(request, datasource) == false) { + return; + } + + List fields = datasourceUpdateService.getHeaderFields(request.getEndpoint()); + if (datasource.isCompatible(fields) == false) { + throw new IncompatibleDatasourceException( + "new fields [{}] does not contain all old fields [{}]", + fields.toString(), + datasource.getDatabase().getFields().toString() + ); + } + } + + private void validateUpdateIntervalIsLessThanValidForInDays(final UpdateDatasourceRequest request, final Datasource datasource) + throws IOException { + if (isEndpointChanged(request, datasource) == false && isUpdateIntervalChanged(request) == false) { + return; + } + + long validForInDays = isEndpointChanged(request, datasource) + ? DatasourceManifest.Builder.build(new URL(request.getEndpoint())).getValidForInDays() + : datasource.getDatabase().getValidForInDays(); + + long updateInterval = isUpdateIntervalChanged(request) + ? request.getUpdateInterval().days() + : datasource.getUserSchedule().getInterval(); + + if (updateInterval >= validForInDays) { + throw new InvalidParameterException( + String.format(Locale.ROOT, "updateInterval %d should be smaller than %d", updateInterval, validForInDays) + ); + } + } + + private boolean isEndpointChanged(final UpdateDatasourceRequest request, final Datasource datasource) { + return request.getEndpoint() != null && request.getEndpoint().equals(datasource.getEndpoint()) == false; + } + + /** + * Update interval is changed as long as user provide one because + * start time will get updated even if the update interval is same as current one. + * + * @param request the update datasource request + * @return true if update interval is changed, and false otherwise + */ + private boolean isUpdateIntervalChanged(final UpdateDatasourceRequest request) { + return request.getUpdateInterval() != null; + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceManifest.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceManifest.java new file mode 100644 index 0000000000..5382aa56fa --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceManifest.java @@ -0,0 +1,147 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStreamReader; +import java.net.URL; +import java.net.URLConnection; +import java.nio.CharBuffer; +import java.security.AccessController; +import java.security.PrivilegedAction; + +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.Setter; + +import org.opensearch.SpecialPermission; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.ParseField; +import org.opensearch.core.xcontent.ConstructingObjectParser; +import org.opensearch.core.xcontent.DeprecationHandler; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.shared.Constants; + +/** + * Ip2Geo datasource manifest file object + * + * Manifest file is stored in an external endpoint. OpenSearch read the file and store values it in this object. + */ +@Setter +@Getter +@AllArgsConstructor +public class DatasourceManifest { + private static final ParseField URL_FIELD = new ParseField("url"); + private static final ParseField DB_NAME_FIELD = new ParseField("db_name"); + private static final ParseField SHA256_HASH_FIELD = new ParseField("sha256_hash"); + private static final ParseField VALID_FOR_IN_DAYS_FIELD = new ParseField("valid_for_in_days"); + private static final ParseField UPDATED_AT_FIELD = new ParseField("updated_at_in_epoch_milli"); + private static final ParseField PROVIDER_FIELD = new ParseField("provider"); + + /** + * @param url URL of a ZIP file containing a database + * @return URL of a ZIP file containing a database + */ + private String url; + /** + * @param dbName A database file name inside the ZIP file + * @return A database file name inside the ZIP file + */ + private String dbName; + /** + * @param sha256Hash SHA256 hash value of a database file + * @return SHA256 hash value of a database file + */ + private String sha256Hash; + /** + * @param validForInDays A duration in which the database file is valid to use + * @return A duration in which the database file is valid to use + */ + private Long validForInDays; + /** + * @param updatedAt A date when the database was updated + * @return A date when the database was updated + */ + private Long updatedAt; + /** + * @param provider A database provider name + * @return A database provider name + */ + private String provider; + + /** + * Ddatasource manifest parser + */ + public static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( + "datasource_manifest", + true, + args -> { + String url = (String) args[0]; + String dbName = (String) args[1]; + String sha256Hash = (String) args[2]; + Long validForInDays = (Long) args[3]; + Long updatedAt = (Long) args[4]; + String provider = (String) args[5]; + return new DatasourceManifest(url, dbName, sha256Hash, validForInDays, updatedAt, provider); + } + ); + static { + PARSER.declareString(ConstructingObjectParser.constructorArg(), URL_FIELD); + PARSER.declareString(ConstructingObjectParser.constructorArg(), DB_NAME_FIELD); + PARSER.declareString(ConstructingObjectParser.constructorArg(), SHA256_HASH_FIELD); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), VALID_FOR_IN_DAYS_FIELD); + PARSER.declareLong(ConstructingObjectParser.constructorArg(), UPDATED_AT_FIELD); + PARSER.declareString(ConstructingObjectParser.constructorArg(), PROVIDER_FIELD); + } + + /** + * Datasource manifest builder + */ + public static class Builder { + private static final int MANIFEST_FILE_MAX_BYTES = 1024 * 8; + + /** + * Build DatasourceManifest from a given url + * + * @param url url to downloads a manifest file + * @return DatasourceManifest representing the manifest file + */ + @SuppressForbidden(reason = "Need to connect to http endpoint to read manifest file") + public static DatasourceManifest build(final URL url) { + SpecialPermission.check(); + return AccessController.doPrivileged((PrivilegedAction) () -> { + try { + URLConnection connection = url.openConnection(); + return internalBuild(connection); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + } + + @VisibleForTesting + @SuppressForbidden(reason = "Need to connect to http endpoint to read manifest file") + protected static DatasourceManifest internalBuild(final URLConnection connection) throws IOException { + connection.addRequestProperty(Constants.USER_AGENT_KEY, Constants.USER_AGENT_VALUE); + InputStreamReader inputStreamReader = new InputStreamReader(connection.getInputStream()); + try (BufferedReader reader = new BufferedReader(inputStreamReader)) { + CharBuffer charBuffer = CharBuffer.allocate(MANIFEST_FILE_MAX_BYTES); + reader.read(charBuffer); + charBuffer.flip(); + XContentParser parser = JsonXContent.jsonXContent.createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.IGNORE_DEPRECATIONS, + charBuffer.toString() + ); + return PARSER.parse(parser, null); + } + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceState.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceState.java new file mode 100644 index 0000000000..3fbb064ceb --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/DatasourceState.java @@ -0,0 +1,37 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +/** + * Ip2Geo datasource state + * + * 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 + * (CREATING) to (CREATE_FAILED or AVAILABLE) to (DELETING) + * + */ +public enum DatasourceState { + /** + * Data source is being created + */ + CREATING, + /** + * Data source is ready to be used + */ + AVAILABLE, + /** + * Data source creation failed + */ + CREATE_FAILED, + /** + * Data source is being deleted + */ + DELETING +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutor.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutor.java new file mode 100644 index 0000000000..f6230e0430 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoExecutor.java @@ -0,0 +1,45 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +import java.util.concurrent.ExecutorService; + +import org.opensearch.common.settings.Settings; +import org.opensearch.threadpool.ExecutorBuilder; +import org.opensearch.threadpool.FixedExecutorBuilder; +import org.opensearch.threadpool.ThreadPool; + +/** + * Provide a list of static methods related with executors for Ip2Geo + */ +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 + * once a day at most and no need to expedite the task. + * + * @param settings the settings + * @return the executor builder + */ + public static ExecutorBuilder executorBuilder(final Settings settings) { + return new FixedExecutorBuilder(settings, THREAD_POOL_NAME, 1, 1000, THREAD_POOL_NAME, false); + } + + /** + * Return an executor service for datasource update task + * + * @return the executor service + */ + public ExecutorService forDatasourceUpdate() { + return threadPool.executor(THREAD_POOL_NAME); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoLockService.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoLockService.java new file mode 100644 index 0000000000..7bd0331628 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoLockService.java @@ -0,0 +1,157 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +import static org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension.JOB_INDEX_NAME; + +import java.time.Instant; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import lombok.extern.log4j.Log4j2; + +import org.opensearch.OpenSearchException; +import org.opensearch.action.ActionListener; +import org.opensearch.client.Client; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.jobscheduler.spi.LockModel; +import org.opensearch.jobscheduler.spi.utils.LockService; + +/** + * A wrapper of job scheduler's lock service for datasource + */ +@Log4j2 +public class Ip2GeoLockService { + public static final long LOCK_DURATION_IN_SECONDS = 300l; + public static final long RENEW_AFTER_IN_SECONDS = 120l; + private final ClusterService clusterService; + private final LockService lockService; + + /** + * Constructor + * + * @param clusterService the cluster service + * @param client the client + */ + public Ip2GeoLockService(final ClusterService clusterService, final Client client) { + this.clusterService = clusterService; + this.lockService = new LockService(client, clusterService); + } + + /** + * Wrapper method of LockService#acquireLockWithId + * + * Datasource use its name as doc id in job scheduler. Therefore, we can use datasource name to acquire + * a lock on a datasource. + * + * @param datasourceName datasourceName to acquire lock on + * @param lockDurationSeconds the lock duration in seconds + * @param listener the listener + */ + public void acquireLock(final String datasourceName, final Long lockDurationSeconds, final ActionListener listener) { + lockService.acquireLockWithId(JOB_INDEX_NAME, lockDurationSeconds, datasourceName, listener); + } + + /** + * Synchronous method of #acquireLock + * + * @param datasourceName datasourceName to acquire lock on + * @param lockDurationSeconds the lock duration in seconds + * @return lock model + */ + public Optional acquireLock(final String datasourceName, final Long lockDurationSeconds) { + AtomicReference lockReference = new AtomicReference(); + CountDownLatch countDownLatch = new CountDownLatch(1); + lockService.acquireLockWithId(JOB_INDEX_NAME, lockDurationSeconds, datasourceName, new ActionListener<>() { + @Override + public void onResponse(final LockModel lockModel) { + lockReference.set(lockModel); + countDownLatch.countDown(); + } + + @Override + public void onFailure(final Exception e) { + lockReference.set(null); + countDownLatch.countDown(); + } + }); + + try { + countDownLatch.await(clusterService.getClusterSettings().get(Ip2GeoSettings.TIMEOUT).getSeconds(), TimeUnit.SECONDS); + return Optional.ofNullable(lockReference.get()); + } catch (InterruptedException e) { + return Optional.empty(); + } + } + + /** + * Wrapper method of LockService#release + * + * @param lockModel the lock model + */ + public void releaseLock(final LockModel lockModel) { + lockService.release( + lockModel, + ActionListener.wrap(released -> {}, exception -> log.error("Failed to release the lock", exception)) + ); + } + + /** + * Synchronous method of LockService#renewLock + * + * @param lockModel lock to renew + * @return renewed lock if renew succeed and null otherwise + */ + public LockModel renewLock(final LockModel lockModel) { + AtomicReference lockReference = new AtomicReference(); + CountDownLatch countDownLatch = new CountDownLatch(1); + lockService.renewLock(lockModel, new ActionListener<>() { + @Override + public void onResponse(final LockModel lockModel) { + lockReference.set(lockModel); + countDownLatch.countDown(); + } + + @Override + public void onFailure(final Exception e) { + lockReference.set(null); + countDownLatch.countDown(); + } + }); + + try { + countDownLatch.await(clusterService.getClusterSettings().get(Ip2GeoSettings.TIMEOUT).getSeconds(), TimeUnit.SECONDS); + return lockReference.get(); + } catch (InterruptedException e) { + return null; + } + } + + /** + * Return a runnable which can renew the given lock model + * + * The runnable renews the lock and store the renewed lock in the AtomicReference. + * It only renews the lock when it passed {@code RENEW_AFTER_IN_SECONDS} since + * the last time the lock was renewed to avoid resource abuse. + * + * @param lockModel lock model to renew + * @return runnable which can renew the given lock for every call + */ + public Runnable getRenewLockRunnable(final AtomicReference lockModel) { + return () -> { + LockModel preLock = lockModel.get(); + if (Instant.now().isBefore(preLock.getLockTime().plusSeconds(RENEW_AFTER_IN_SECONDS))) { + return; + } + lockModel.set(renewLock(lockModel.get())); + if (lockModel.get() == null) { + new OpenSearchException("failed to renew a lock [{}]", preLock); + } + }; + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettings.java b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettings.java new file mode 100644 index 0000000000..16aba0e152 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettings.java @@ -0,0 +1,97 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +import java.net.MalformedURLException; +import java.net.URISyntaxException; +import java.net.URL; +import java.util.List; + +import org.opensearch.common.settings.Setting; +import org.opensearch.common.unit.TimeValue; + +/** + * Settings for Ip2Geo datasource operations + */ +public class Ip2GeoSettings { + + /** + * Default endpoint to be used in GeoIP datasource creation API + */ + public static final Setting DATASOURCE_ENDPOINT = Setting.simpleString( + "plugins.geospatial.ip2geo.datasource.endpoint", + "https://geoip.maps.opensearch.org/v1/geolite2-city/manifest.json", + new DatasourceEndpointValidator(), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Default update interval to be used in Ip2Geo datasource creation API + */ + public static final Setting DATASOURCE_UPDATE_INTERVAL = Setting.longSetting( + "plugins.geospatial.ip2geo.datasource.update_interval_in_days", + 3l, + 1l, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Bulk size for indexing GeoIP data + */ + public static final Setting BATCH_SIZE = Setting.intSetting( + "plugins.geospatial.ip2geo.datasource.batch_size", + 10000, + 1, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Timeout value for Ip2Geo processor + */ + public static final Setting TIMEOUT = Setting.timeSetting( + "plugins.geospatial.ip2geo.timeout", + TimeValue.timeValueSeconds(30), + TimeValue.timeValueSeconds(1), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Max size for geo data cache + */ + public static final Setting CACHE_SIZE = Setting.longSetting( + "plugins.geospatial.ip2geo.processor.cache_size", + 1000, + 0, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + /** + * Return all settings of Ip2Geo feature + * @return a list of all settings for Ip2Geo feature + */ + public static final List> settings() { + return List.of(DATASOURCE_ENDPOINT, DATASOURCE_UPDATE_INTERVAL, BATCH_SIZE, TIMEOUT, CACHE_SIZE); + } + + /** + * Visible for testing + */ + protected static class DatasourceEndpointValidator implements Setting.Validator { + @Override + public void validate(final String value) { + try { + new URL(value).toURI(); + } catch (MalformedURLException | URISyntaxException e) { + throw new IllegalArgumentException("Invalid URL format is provided"); + } + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/dao/DatasourceDao.java b/src/main/java/org/opensearch/geospatial/ip2geo/dao/DatasourceDao.java new file mode 100644 index 0000000000..144add66d0 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/dao/DatasourceDao.java @@ -0,0 +1,377 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + +import lombok.extern.log4j.Log4j2; + +import org.opensearch.OpenSearchException; +import org.opensearch.ResourceAlreadyExistsException; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.DocWriteRequest; +import org.opensearch.action.StepListener; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.action.bulk.BulkRequest; +import org.opensearch.action.bulk.BulkResponse; +import org.opensearch.action.delete.DeleteResponse; +import org.opensearch.action.get.GetRequest; +import org.opensearch.action.get.GetResponse; +import org.opensearch.action.get.MultiGetItemResponse; +import org.opensearch.action.get.MultiGetResponse; +import org.opensearch.action.index.IndexRequest; +import org.opensearch.action.index.IndexResponse; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.WriteRequest; +import org.opensearch.client.Client; +import org.opensearch.cluster.routing.Preference; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.xcontent.LoggingDeprecationHandler; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.common.xcontent.XContentHelper; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.ToXContent; +import org.opensearch.core.xcontent.XContentParser; +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.shared.StashedThreadContext; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; + +/** + * Data access object for datasource + */ +@Log4j2 +public class DatasourceDao { + private static final Integer MAX_SIZE = 1000; + private final Client client; + private final ClusterService clusterService; + private final ClusterSettings clusterSettings; + + public DatasourceDao(final Client client, final ClusterService clusterService) { + this.client = client; + this.clusterService = clusterService; + this.clusterSettings = clusterService.getClusterSettings(); + } + + /** + * Create datasource index + * + * @param stepListener setp listener + */ + public void createIndexIfNotExists(final StepListener stepListener) { + if (clusterService.state().metadata().hasIndex(DatasourceExtension.JOB_INDEX_NAME) == true) { + stepListener.onResponse(null); + return; + } + final CreateIndexRequest createIndexRequest = new CreateIndexRequest(DatasourceExtension.JOB_INDEX_NAME).mapping(getIndexMapping()) + .settings(DatasourceExtension.INDEX_SETTING); + StashedThreadContext.run(client, () -> client.admin().indices().create(createIndexRequest, new ActionListener<>() { + @Override + public void onResponse(final CreateIndexResponse createIndexResponse) { + stepListener.onResponse(null); + } + + @Override + public void onFailure(final Exception e) { + if (e instanceof ResourceAlreadyExistsException) { + log.info("index[{}] already exist", DatasourceExtension.JOB_INDEX_NAME); + stepListener.onResponse(null); + return; + } + stepListener.onFailure(e); + } + })); + } + + private String getIndexMapping() { + try { + try (InputStream is = DatasourceDao.class.getResourceAsStream("/mappings/ip2geo_datasource.json")) { + try (BufferedReader reader = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) { + return reader.lines().map(String::trim).collect(Collectors.joining()); + } + } + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + /** + * Update datasource in an index {@code DatasourceExtension.JOB_INDEX_NAME} + * @param datasource the datasource + * @return index response + */ + public IndexResponse updateDatasource(final Datasource datasource) { + datasource.setLastUpdateTime(Instant.now()); + return StashedThreadContext.run(client, () -> { + try { + return client.prepareIndex(DatasourceExtension.JOB_INDEX_NAME) + .setId(datasource.getName()) + .setOpType(DocWriteRequest.OpType.INDEX) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .setSource(datasource.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)) + .execute() + .actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)); + } catch (IOException e) { + throw new RuntimeException(e); + } + }); + } + + /** + * Update datasources in an index {@code DatasourceExtension.JOB_INDEX_NAME} + * @param datasources the datasources + * @param listener action listener + */ + public void updateDatasource(final List datasources, final ActionListener listener) { + BulkRequest bulkRequest = new BulkRequest(); + datasources.stream().map(datasource -> { + datasource.setLastUpdateTime(Instant.now()); + return datasource; + }).map(this::toIndexRequest).forEach(indexRequest -> bulkRequest.add(indexRequest)); + StashedThreadContext.run(client, () -> client.bulk(bulkRequest, listener)); + } + + private IndexRequest toIndexRequest(Datasource datasource) { + try { + IndexRequest indexRequest = new IndexRequest(); + indexRequest.index(DatasourceExtension.JOB_INDEX_NAME); + indexRequest.id(datasource.getName()); + indexRequest.opType(DocWriteRequest.OpType.INDEX); + indexRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE); + indexRequest.source(datasource.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)); + return indexRequest; + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + /** + * Put datasource in an index {@code DatasourceExtension.JOB_INDEX_NAME} + * + * @param datasource the datasource + * @param listener the listener + */ + public void putDatasource(final Datasource datasource, final ActionListener listener) { + datasource.setLastUpdateTime(Instant.now()); + StashedThreadContext.run(client, () -> { + try { + client.prepareIndex(DatasourceExtension.JOB_INDEX_NAME) + .setId(datasource.getName()) + .setOpType(DocWriteRequest.OpType.CREATE) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .setSource(datasource.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)) + .execute(listener); + } catch (IOException e) { + new RuntimeException(e); + } + }); + } + + /** + * Delete datasource in an index {@code DatasourceExtension.JOB_INDEX_NAME} + * + * @param datasource the datasource + * + */ + public void deleteDatasource(final Datasource datasource) { + DeleteResponse response = client.prepareDelete() + .setIndex(DatasourceExtension.JOB_INDEX_NAME) + .setId(datasource.getName()) + .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE) + .execute() + .actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)); + + if (response.status().equals(RestStatus.OK)) { + log.info("deleted datasource[{}] successfully", datasource.getName()); + } else if (response.status().equals(RestStatus.NOT_FOUND)) { + throw new ResourceNotFoundException("datasource[{}] does not exist", datasource.getName()); + } else { + throw new OpenSearchException("failed to delete datasource[{}] with status[{}]", datasource.getName(), response.status()); + } + } + + /** + * Get datasource from an index {@code DatasourceExtension.JOB_INDEX_NAME} + * @param name the name of a datasource + * @return datasource + * @throws IOException exception + */ + public Datasource getDatasource(final String name) throws IOException { + GetRequest request = new GetRequest(DatasourceExtension.JOB_INDEX_NAME, name).preference(Preference.PRIMARY.type()); + GetResponse response; + try { + response = StashedThreadContext.run(client, () -> client.get(request).actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT))); + if (response.isExists() == false) { + log.error("Datasource[{}] does not exist in an index[{}]", name, DatasourceExtension.JOB_INDEX_NAME); + return null; + } + } catch (IndexNotFoundException e) { + log.error("Index[{}] is not found", DatasourceExtension.JOB_INDEX_NAME); + return null; + } + + XContentParser parser = XContentHelper.createParser( + NamedXContentRegistry.EMPTY, + LoggingDeprecationHandler.INSTANCE, + response.getSourceAsBytesRef() + ); + return Datasource.PARSER.parse(parser, null); + } + + /** + * Get datasource from an index {@code DatasourceExtension.JOB_INDEX_NAME} + * @param name the name of a datasource + * @param actionListener the action listener + */ + public void getDatasource(final String name, final ActionListener actionListener) { + GetRequest request = new GetRequest(DatasourceExtension.JOB_INDEX_NAME, name).preference(Preference.PRIMARY.type()); + StashedThreadContext.run(client, () -> client.get(request, new ActionListener<>() { + @Override + public void onResponse(final GetResponse response) { + if (response.isExists() == false) { + actionListener.onResponse(null); + return; + } + + try { + XContentParser parser = XContentHelper.createParser( + NamedXContentRegistry.EMPTY, + LoggingDeprecationHandler.INSTANCE, + response.getSourceAsBytesRef() + ); + actionListener.onResponse(Datasource.PARSER.parse(parser, null)); + } catch (IOException e) { + actionListener.onFailure(e); + } + } + + @Override + public void onFailure(final Exception e) { + actionListener.onFailure(e); + } + })); + } + + /** + * Get datasources from an index {@code DatasourceExtension.JOB_INDEX_NAME} + * @param names the array of datasource names + * @param actionListener the action listener + */ + public void getDatasources(final String[] names, final ActionListener> actionListener) { + StashedThreadContext.run( + client, + () -> client.prepareMultiGet() + .add(DatasourceExtension.JOB_INDEX_NAME, names) + .setPreference(Preference.PRIMARY.type()) + .execute(createGetDataSourceQueryActionLister(MultiGetResponse.class, actionListener)) + ); + } + + /** + * Get all datasources up to {@code MAX_SIZE} from an index {@code DatasourceExtension.JOB_INDEX_NAME} + * @param actionListener the action listener + */ + public void getAllDatasources(final ActionListener> actionListener) { + StashedThreadContext.run( + client, + () -> client.prepareSearch(DatasourceExtension.JOB_INDEX_NAME) + .setQuery(QueryBuilders.matchAllQuery()) + .setPreference(Preference.PRIMARY.type()) + .setSize(MAX_SIZE) + .execute(createGetDataSourceQueryActionLister(SearchResponse.class, actionListener)) + ); + } + + /** + * Get all datasources up to {@code MAX_SIZE} from an index {@code DatasourceExtension.JOB_INDEX_NAME} + */ + public List getAllDatasources() { + SearchResponse response = StashedThreadContext.run( + client, + () -> client.prepareSearch(DatasourceExtension.JOB_INDEX_NAME) + .setQuery(QueryBuilders.matchAllQuery()) + .setPreference(Preference.PRIMARY.type()) + .setSize(MAX_SIZE) + .execute() + .actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)) + ); + + List bytesReferences = toBytesReferences(response); + return bytesReferences.stream().map(bytesRef -> toDatasource(bytesRef)).collect(Collectors.toList()); + } + + private ActionListener createGetDataSourceQueryActionLister( + final Class response, + final ActionListener> actionListener + ) { + return new ActionListener() { + @Override + public void onResponse(final T response) { + try { + List bytesReferences = toBytesReferences(response); + List datasources = bytesReferences.stream() + .map(bytesRef -> toDatasource(bytesRef)) + .collect(Collectors.toList()); + actionListener.onResponse(datasources); + } catch (Exception e) { + actionListener.onFailure(e); + } + } + + @Override + public void onFailure(final Exception e) { + actionListener.onFailure(e); + } + }; + } + + private List toBytesReferences(final Object response) { + if (response instanceof SearchResponse) { + SearchResponse searchResponse = (SearchResponse) response; + return Arrays.stream(searchResponse.getHits().getHits()).map(SearchHit::getSourceRef).collect(Collectors.toList()); + } else if (response instanceof MultiGetResponse) { + MultiGetResponse multiGetResponse = (MultiGetResponse) response; + return Arrays.stream(multiGetResponse.getResponses()) + .map(MultiGetItemResponse::getResponse) + .filter(Objects::nonNull) + .filter(GetResponse::isExists) + .map(GetResponse::getSourceAsBytesRef) + .collect(Collectors.toList()); + } else { + throw new OpenSearchException("No supported instance type[{}] is provided", response.getClass()); + } + } + + private Datasource toDatasource(final BytesReference bytesReference) { + try { + XContentParser parser = XContentHelper.createParser( + NamedXContentRegistry.EMPTY, + LoggingDeprecationHandler.INSTANCE, + bytesReference + ); + return Datasource.PARSER.parse(parser, null); + } catch (IOException e) { + throw new RuntimeException(e); + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/dao/GeoIpDataDao.java b/src/main/java/org/opensearch/geospatial/ip2geo/dao/GeoIpDataDao.java new file mode 100644 index 0000000000..a538e813b0 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/dao/GeoIpDataDao.java @@ -0,0 +1,345 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +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; +import java.io.InputStreamReader; +import java.net.URL; +import java.net.URLConnection; +import java.nio.charset.StandardCharsets; +import java.security.AccessController; +import java.security.PrivilegedAction; +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.LinkedList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Queue; +import java.util.stream.Collectors; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; + +import lombok.NonNull; +import lombok.extern.log4j.Log4j2; + +import org.apache.commons.csv.CSVFormat; +import org.apache.commons.csv.CSVParser; +import org.apache.commons.csv.CSVRecord; +import org.apache.logging.log4j.util.Strings; +import org.opensearch.OpenSearchException; +import org.opensearch.SpecialPermission; +import org.opensearch.action.DocWriteRequest; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.bulk.BulkRequest; +import org.opensearch.action.bulk.BulkResponse; +import org.opensearch.action.index.IndexRequest; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.IndicesOptions; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.client.Client; +import org.opensearch.client.Requests; +import org.opensearch.cluster.routing.Preference; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.common.xcontent.XContentHelper; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.constants.IndexSetting; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.shared.Constants; +import org.opensearch.geospatial.shared.StashedThreadContext; +import org.opensearch.index.query.QueryBuilders; + +/** + * Data access object for GeoIp data + */ +@Log4j2 +public class GeoIpDataDao { + private static final String IP_RANGE_FIELD_NAME = "_cidr"; + private static final String DATA_FIELD_NAME = "_data"; + private static final Map INDEX_SETTING_TO_CREATE = Map.of( + IndexSetting.NUMBER_OF_SHARDS, + 1, + IndexSetting.NUMBER_OF_REPLICAS, + 0, + IndexSetting.REFRESH_INTERVAL, + -1, + IndexSetting.HIDDEN, + true + ); + private static final Map INDEX_SETTING_TO_FREEZE = Map.of( + IndexSetting.AUTO_EXPAND_REPLICAS, + "0-all", + IndexSetting.BLOCKS_WRITE, + true + ); + private final ClusterService clusterService; + private final ClusterSettings clusterSettings; + private final Client client; + + public GeoIpDataDao(final ClusterService clusterService, final Client client) { + this.clusterService = clusterService; + this.clusterSettings = clusterService.getClusterSettings(); + this.client = client; + } + + /** + * Create an index for GeoIP data + * + * Index setting start with single shard, zero replica, no refresh interval, and hidden. + * Once the GeoIP data is indexed, do refresh and force merge. + * Then, change the index setting to expand replica to all nodes, and read only allow delete. + * See {@link #freezeIndex} + * + * @param indexName index name + */ + public void createIndexIfNotExists(final String indexName) { + if (clusterService.state().metadata().hasIndex(indexName) == true) { + return; + } + final CreateIndexRequest createIndexRequest = new CreateIndexRequest(indexName).settings(INDEX_SETTING_TO_CREATE) + .mapping(getIndexMapping()); + StashedThreadContext.run( + client, + () -> client.admin().indices().create(createIndexRequest).actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)) + ); + } + + private void freezeIndex(final String indexName) { + TimeValue timeout = clusterSettings.get(Ip2GeoSettings.TIMEOUT); + StashedThreadContext.run(client, () -> { + client.admin().indices().prepareForceMerge(indexName).setMaxNumSegments(1).execute().actionGet(timeout); + client.admin().indices().prepareRefresh(indexName).execute().actionGet(timeout); + client.admin() + .indices() + .prepareUpdateSettings(indexName) + .setSettings(INDEX_SETTING_TO_FREEZE) + .execute() + .actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)); + }); + } + + /** + * Generate XContentBuilder representing datasource database index mapping + * + * { + * "dynamic": false, + * "properties": { + * "_cidr": { + * "type": "ip_range", + * "doc_values": false + * } + * } + * } + * + * @return String representing datasource database index mapping + */ + private String getIndexMapping() { + try { + try (InputStream is = DatasourceDao.class.getResourceAsStream("/mappings/ip2geo_geoip.json")) { + try (BufferedReader reader = new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) { + return reader.lines().map(String::trim).collect(Collectors.joining()); + } + } + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + /** + * Create CSVParser of a GeoIP data + * + * @param manifest Datasource manifest + * @return CSVParser for GeoIP data + */ + @SuppressForbidden(reason = "Need to connect to http endpoint to read GeoIP database file") + public CSVParser getDatabaseReader(final DatasourceManifest manifest) { + SpecialPermission.check(); + return AccessController.doPrivileged((PrivilegedAction) () -> { + try { + URL zipUrl = new URL(manifest.getUrl()); + return internalGetDatabaseReader(manifest, zipUrl.openConnection()); + } catch (IOException e) { + throw new OpenSearchException("failed to read geoip data from {}", manifest.getUrl(), e); + } + }); + } + + @VisibleForTesting + @SuppressForbidden(reason = "Need to connect to http endpoint to read GeoIP database file") + protected CSVParser internalGetDatabaseReader(final DatasourceManifest manifest, final URLConnection connection) throws IOException { + connection.addRequestProperty(Constants.USER_AGENT_KEY, Constants.USER_AGENT_VALUE); + ZipInputStream zipIn = new ZipInputStream(connection.getInputStream()); + ZipEntry zipEntry = zipIn.getNextEntry(); + while (zipEntry != null) { + if (zipEntry.getName().equalsIgnoreCase(manifest.getDbName()) == false) { + zipEntry = zipIn.getNextEntry(); + continue; + } + return new CSVParser(new BufferedReader(new InputStreamReader(zipIn)), CSVFormat.RFC4180); + } + throw new OpenSearchException("database file [{}] does not exist in the zip file [{}]", manifest.getDbName(), manifest.getUrl()); + } + + /** + * Create a document to ingest in datasource database index + * + * It assumes the first field as ip_range. The rest is added under data field. + * + * Document example + * { + * "_cidr":"1.0.0.1/25", + * "_data":{ + * "country": "USA", + * "city": "Seattle", + * "location":"13.23,42.12" + * } + * } + * + * @param fields a list of field name + * @param values a list of values + * @return Document in json string format + * @throws IOException the exception + */ + public XContentBuilder createDocument(final String[] fields, final String[] values) throws IOException { + if (fields.length != values.length) { + throw new OpenSearchException("header[{}] and record[{}] length does not match", fields, values); + } + XContentBuilder builder = XContentFactory.jsonBuilder(); + builder.startObject(); + builder.field(IP_RANGE_FIELD_NAME, values[0]); + builder.startObject(DATA_FIELD_NAME); + for (int i = 1; i < fields.length; i++) { + if (Strings.isBlank(values[i])) { + continue; + } + builder.field(fields[i], values[i]); + } + builder.endObject(); + builder.endObject(); + builder.close(); + return builder; + } + + /** + * Query a given index using a given ip address to get geoip data + * + * @param indexName index + * @param ip ip address + * @return geoIP data + */ + public Map getGeoIpData(final String indexName, final String ip) { + SearchResponse response = StashedThreadContext.run( + client, + () -> client.prepareSearch(indexName) + .setSize(1) + .setQuery(QueryBuilders.termQuery(IP_RANGE_FIELD_NAME, ip)) + .setPreference(Preference.LOCAL.type()) + .setRequestCache(true) + .get(clusterSettings.get(Ip2GeoSettings.TIMEOUT)) + ); + + if (response.getHits().getHits().length == 0) { + return Collections.emptyMap(); + } else { + return (Map) XContentHelper.convertToMap(response.getHits().getAt(0).getSourceRef(), false, XContentType.JSON) + .v2() + .get(DATA_FIELD_NAME); + } + } + + /** + * Puts GeoIP data from CSVRecord iterator into a given index in bulk + * + * @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 renewLock Runnable to renew lock + */ + public void putGeoIpData( + @NonNull final String indexName, + @NonNull final String[] fields, + @NonNull final Iterator iterator, + @NonNull final Runnable renewLock + ) throws IOException { + TimeValue timeout = clusterSettings.get(Ip2GeoSettings.TIMEOUT); + Integer batchSize = clusterSettings.get(Ip2GeoSettings.BATCH_SIZE); + final BulkRequest bulkRequest = new BulkRequest(); + Queue requests = new LinkedList<>(); + for (int i = 0; i < batchSize; i++) { + requests.add(Requests.indexRequest(indexName)); + } + while (iterator.hasNext()) { + CSVRecord record = iterator.next(); + XContentBuilder document = createDocument(fields, record.values()); + IndexRequest indexRequest = (IndexRequest) requests.poll(); + indexRequest.source(document); + indexRequest.id(record.get(0)); + bulkRequest.add(indexRequest); + if (iterator.hasNext() == false || bulkRequest.requests().size() == batchSize) { + BulkResponse response = StashedThreadContext.run(client, () -> client.bulk(bulkRequest).actionGet(timeout)); + if (response.hasFailures()) { + throw new OpenSearchException( + "error occurred while ingesting GeoIP data in {} with an error {}", + indexName, + response.buildFailureMessage() + ); + } + requests.addAll(bulkRequest.requests()); + bulkRequest.requests().clear(); + } + renewLock.run(); + } + freezeIndex(indexName); + + } + + public void deleteIp2GeoDataIndex(final String index) { + deleteIp2GeoDataIndex(Arrays.asList(index)); + } + + public void deleteIp2GeoDataIndex(final List indices) { + if (indices == null || indices.isEmpty()) { + return; + } + + Optional invalidIndex = indices.stream() + .filter(index -> index.startsWith(IP2GEO_DATA_INDEX_NAME_PREFIX) == false) + .findAny(); + if (invalidIndex.isPresent()) { + throw new OpenSearchException( + "the index[{}] is not ip2geo data index which should start with {}", + invalidIndex.get(), + IP2GEO_DATA_INDEX_NAME_PREFIX + ); + } + + AcknowledgedResponse response = StashedThreadContext.run( + client, + () -> client.admin() + .indices() + .prepareDelete(indices.toArray(new String[0])) + .setIndicesOptions(IndicesOptions.LENIENT_EXPAND_OPEN_CLOSED_HIDDEN) + .execute() + .actionGet(clusterSettings.get(Ip2GeoSettings.TIMEOUT)) + ); + + if (response.isAcknowledged() == false) { + throw new OpenSearchException("failed to delete data[{}] in datasource", String.join(",", indices)); + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoCachedDao.java b/src/main/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoCachedDao.java new file mode 100644 index 0000000000..8c2e686d13 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoCachedDao.java @@ -0,0 +1,218 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +import java.io.IOException; +import java.time.Instant; +import java.util.Iterator; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.function.Function; + +import lombok.AllArgsConstructor; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import lombok.extern.log4j.Log4j2; + +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.cache.Cache; +import org.opensearch.common.cache.CacheBuilder; +import org.opensearch.common.xcontent.LoggingDeprecationHandler; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.core.xcontent.NamedXContentRegistry; +import org.opensearch.core.xcontent.XContentParser; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.index.engine.Engine; +import org.opensearch.index.shard.IndexingOperationListener; + +/** + * Data access object for Datasource and GeoIP data with added caching layer + * + * Ip2GeoCachedDao has a memory cache to store Datasource and GeoIP data. To fully utilize the cache, + * do not create multiple Ip2GeoCachedDao. Ip2GeoCachedDao instance is bound to guice so that you can use + * it through injection. + * + * All IP2Geo processors share single Ip2GeoCachedDao instance. + */ +@Log4j2 +public class Ip2GeoCachedDao implements IndexingOperationListener { + private final DatasourceDao datasourceDao; + private final GeoIpDataDao geoIpDataDao; + private final GeoDataCache geoDataCache; + private Map metadata; + + public Ip2GeoCachedDao(final ClusterService clusterService, final DatasourceDao datasourceDao, final GeoIpDataDao geoIpDataDao) { + this.datasourceDao = datasourceDao; + this.geoIpDataDao = geoIpDataDao; + this.geoDataCache = new GeoDataCache(clusterService.getClusterSettings().get(Ip2GeoSettings.CACHE_SIZE)); + clusterService.getClusterSettings() + .addSettingsUpdateConsumer(Ip2GeoSettings.CACHE_SIZE, setting -> this.geoDataCache.updateMaxSize(setting.longValue())); + } + + public String getIndexName(final String datasourceName) { + return getMetadata().getOrDefault(datasourceName, DatasourceMetadata.EMPTY_METADATA).getIndexName(); + } + + public boolean isExpired(final String datasourceName) { + return getMetadata().getOrDefault(datasourceName, DatasourceMetadata.EMPTY_METADATA).getExpirationDate().isBefore(Instant.now()); + } + + public boolean has(final String datasourceName) { + return getMetadata().containsKey(datasourceName); + } + + public DatasourceState getState(final String datasourceName) { + return getMetadata().getOrDefault(datasourceName, DatasourceMetadata.EMPTY_METADATA).getState(); + } + + public Map getGeoData(final String indexName, final String ip) { + try { + return geoDataCache.putIfAbsent(indexName, ip, addr -> geoIpDataDao.getGeoIpData(indexName, ip)); + } catch (ExecutionException e) { + throw new RuntimeException(e); + } + } + + private Map getMetadata() { + if (metadata != null) { + return metadata; + } + synchronized (this) { + if (metadata != null) { + return metadata; + } + Map tempData = new ConcurrentHashMap<>(); + try { + datasourceDao.getAllDatasources() + .stream() + .forEach(datasource -> tempData.put(datasource.getName(), new DatasourceMetadata(datasource))); + } catch (IndexNotFoundException e) { + log.debug("Datasource has never been created"); + } + metadata = tempData; + return metadata; + } + } + + private void put(final Datasource datasource) { + DatasourceMetadata metadata = new DatasourceMetadata(datasource); + getMetadata().put(datasource.getName(), metadata); + } + + private void remove(final String datasourceName) { + getMetadata().remove(datasourceName); + } + + @Override + public void postIndex(ShardId shardId, Engine.Index index, Engine.IndexResult result) { + if (Engine.Result.Type.FAILURE.equals(result.getResultType())) { + return; + } + + try { + XContentParser parser = XContentType.JSON.xContent() + .createParser(NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, index.source().utf8ToString()); + parser.nextToken(); + Datasource datasource = Datasource.PARSER.parse(parser, null); + put(datasource); + } catch (IOException e) { + log.error("IOException occurred updating datasource metadata for datasource {} ", index.id(), e); + } + } + + @Override + public void postDelete(ShardId shardId, Engine.Delete delete, Engine.DeleteResult result) { + if (result.getResultType().equals(Engine.Result.Type.FAILURE)) { + return; + } + remove(delete.id()); + } + + @Getter + private static class DatasourceMetadata { + private static DatasourceMetadata EMPTY_METADATA = new DatasourceMetadata(); + private String indexName; + private Instant expirationDate; + private DatasourceState state; + + private DatasourceMetadata() { + expirationDate = Instant.MIN; + } + + public DatasourceMetadata(final Datasource datasource) { + this.indexName = datasource.currentIndexName(); + this.expirationDate = datasource.expirationDay(); + this.state = datasource.getState(); + } + } + + /** + * Cache to hold geo data + * + * GeoData in an index in immutable. Therefore, invalidation is not needed. + */ + @VisibleForTesting + protected static class GeoDataCache { + private Cache> cache; + + public GeoDataCache(final long maxSize) { + if (maxSize < 0) { + throw new IllegalArgumentException("ip2geo max cache size must be 0 or greater"); + } + this.cache = CacheBuilder.>builder().setMaximumWeight(maxSize).build(); + } + + public Map putIfAbsent( + final String indexName, + final String ip, + final Function> retrieveFunction + ) throws ExecutionException { + CacheKey cacheKey = new CacheKey(indexName, ip); + return cache.computeIfAbsent(cacheKey, key -> retrieveFunction.apply(key.ip)); + } + + public Map get(final String indexName, final String ip) { + return cache.get(new CacheKey(indexName, ip)); + } + + /** + * Create a new cache with give size and replace existing cache + * + * Try to populate the existing value from previous cache to the new cache in best effort + * + * @param maxSize + */ + public void updateMaxSize(final long maxSize) { + if (maxSize < 0) { + throw new IllegalArgumentException("ip2geo max cache size must be 0 or greater"); + } + Cache> temp = CacheBuilder.>builder() + .setMaximumWeight(maxSize) + .build(); + int count = 0; + Iterator it = cache.keys().iterator(); + while (it.hasNext() && count < maxSize) { + CacheKey key = it.next(); + temp.put(key, cache.get(key)); + count++; + } + cache = temp; + } + + @AllArgsConstructor + @EqualsAndHashCode + private static class CacheKey { + private final String indexName; + private final String ip; + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoProcessorDao.java b/src/main/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoProcessorDao.java new file mode 100644 index 0000000000..55e1152da4 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoProcessorDao.java @@ -0,0 +1,40 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import org.opensearch.common.inject.Inject; +import org.opensearch.geospatial.ip2geo.processor.Ip2GeoProcessor; +import org.opensearch.ingest.IngestMetadata; +import org.opensearch.ingest.IngestService; + +/** + * Data access object for Ip2Geo processors + */ +public class Ip2GeoProcessorDao { + private final IngestService ingestService; + + @Inject + public Ip2GeoProcessorDao(final IngestService ingestService) { + this.ingestService = ingestService; + } + + public List getProcessors(final String datasourceName) { + IngestMetadata ingestMetadata = ingestService.getClusterService().state().getMetadata().custom(IngestMetadata.TYPE); + if (ingestMetadata == null) { + return Collections.emptyList(); + } + return ingestMetadata.getPipelines() + .keySet() + .stream() + .flatMap(pipelineId -> ingestService.getProcessorsInPipeline(pipelineId, Ip2GeoProcessor.class).stream()) + .filter(ip2GeoProcessor -> ip2GeoProcessor.getDatasourceName().equals(datasourceName)) + .collect(Collectors.toList()); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/Datasource.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/Datasource.java new file mode 100644 index 0000000000..a256aa27cb --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/Datasource.java @@ -0,0 +1,720 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import java.io.IOException; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.ArrayList; +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.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; +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; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.jobscheduler.spi.ScheduledJobParameter; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; +import org.opensearch.jobscheduler.spi.schedule.ScheduleParser; + +/** + * Ip2Geo datasource job parameter + */ +@Getter +@Setter +@ToString +@EqualsAndHashCode +@AllArgsConstructor +public class Datasource implements Writeable, ScheduledJobParameter { + /** + * Prefix of indices having Ip2Geo data + */ + public static final String IP2GEO_DATA_INDEX_NAME_PREFIX = ".geospatial.ip2geo.data"; + + /** + * Default fields for job scheduling + */ + private static final ParseField NAME_FIELD = new ParseField("name"); + private static final ParseField ENABLED_FIELD = new ParseField("update_enabled"); + private static final ParseField LAST_UPDATE_TIME_FIELD = new ParseField("last_update_time"); + private static final ParseField LAST_UPDATE_TIME_FIELD_READABLE = new ParseField("last_update_time_field"); + /** + * Schedule that user set + */ + private static final ParseField USER_SCHEDULE_FIELD = new ParseField("user_schedule"); + /** + * System schedule which will be used by job scheduler + * + * If datasource is going to get expired before next update, we want to run clean up task before the next update + * by changing system schedule. + * + * If datasource is restored from snapshot, we want to run clean up task immediately to handle expired datasource + * by changing system schedule. + * + * For every task run, we revert system schedule back to user schedule. + */ + private static final ParseField SYSTEM_SCHEDULE_FIELD = new ParseField("system_schedule"); + /** + * {@link DatasourceTask} that DatasourceRunner will execute in next run + * + * For every task run, we revert task back to {@link DatasourceTask#ALL} + */ + private static final ParseField TASK_FIELD = new ParseField("task"); + private static final ParseField ENABLED_TIME_FIELD = new ParseField("enabled_time"); + private static final ParseField ENABLED_TIME_FIELD_READABLE = new ParseField("enabled_time_field"); + + /** + * Additional fields for datasource + */ + private static final ParseField ENDPOINT_FIELD = new ParseField("endpoint"); + private static final ParseField STATE_FIELD = new ParseField("state"); + private static final ParseField CURRENT_INDEX_FIELD = new ParseField("current_index"); + private static final ParseField INDICES_FIELD = new ParseField("indices"); + private static final ParseField DATABASE_FIELD = new ParseField("database"); + private static final ParseField UPDATE_STATS_FIELD = new ParseField("update_stats"); + + /** + * Default variables for job scheduling + */ + + /** + * @param name name of a datasource + * @return name of a datasource + */ + private String name; + /** + * @param lastUpdateTime Last update time of a datasource + * @return Last update time of a datasource + */ + private Instant lastUpdateTime; + /** + * @param enabledTime Last time when a scheduling is enabled for a GeoIP data update + * @return Last time when a scheduling is enabled for the job scheduler + */ + private Instant enabledTime; + /** + * @param isEnabled Indicate if GeoIP data update is scheduled or not + * @return Indicate if scheduling is enabled or not + */ + private boolean isEnabled; + /** + * @param userSchedule Schedule that user provided + * @return Schedule that user provided + */ + private IntervalSchedule userSchedule; + + /** + * @param systemSchedule Schedule that job scheduler use + * @return Schedule that job scheduler use + */ + private IntervalSchedule systemSchedule; + + /** + * @param task Task that {@link DatasourceRunner} will execute + * @return Task that {@link DatasourceRunner} will execute + */ + private DatasourceTask task; + + /** + * Additional variables for datasource + */ + + /** + * @param endpoint URL of a manifest file + * @return URL of a manifest file + */ + private String endpoint; + /** + * @param state State of a datasource + * @return State of a datasource + */ + private DatasourceState state; + /** + * @param currentIndex the current index name having GeoIP data + * @return the current index name having GeoIP data + */ + @Getter(AccessLevel.NONE) + private String currentIndex; + /** + * @param indices A list of indices having GeoIP data including currentIndex + * @return A list of indices having GeoIP data including currentIndex + */ + private List indices; + /** + * @param database GeoIP database information + * @return GeoIP database information + */ + private Database database; + /** + * @param updateStats GeoIP database update statistics + * @return GeoIP database update statistics + */ + private UpdateStats updateStats; + + /** + * Datasource parser + */ + public static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( + "datasource_metadata", + true, + args -> { + String name = (String) args[0]; + Instant lastUpdateTime = Instant.ofEpochMilli((long) args[1]); + Instant enabledTime = args[2] == null ? null : Instant.ofEpochMilli((long) args[2]); + boolean isEnabled = (boolean) args[3]; + IntervalSchedule userSchedule = (IntervalSchedule) args[4]; + IntervalSchedule systemSchedule = (IntervalSchedule) args[5]; + DatasourceTask task = DatasourceTask.valueOf((String) args[6]); + String endpoint = (String) args[7]; + DatasourceState state = DatasourceState.valueOf((String) args[8]); + String currentIndex = (String) args[9]; + List indices = (List) args[10]; + Database database = (Database) args[11]; + UpdateStats updateStats = (UpdateStats) args[12]; + Datasource parameter = new Datasource( + name, + lastUpdateTime, + enabledTime, + isEnabled, + userSchedule, + systemSchedule, + task, + endpoint, + state, + currentIndex, + indices, + database, + updateStats + ); + + return parameter; + } + ); + static { + PARSER.declareString(ConstructingObjectParser.constructorArg(), NAME_FIELD); + PARSER.declareLong(ConstructingObjectParser.constructorArg(), LAST_UPDATE_TIME_FIELD); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), ENABLED_TIME_FIELD); + PARSER.declareBoolean(ConstructingObjectParser.constructorArg(), ENABLED_FIELD); + PARSER.declareObject(ConstructingObjectParser.constructorArg(), (p, c) -> ScheduleParser.parse(p), USER_SCHEDULE_FIELD); + PARSER.declareObject(ConstructingObjectParser.constructorArg(), (p, c) -> ScheduleParser.parse(p), SYSTEM_SCHEDULE_FIELD); + PARSER.declareString(ConstructingObjectParser.constructorArg(), TASK_FIELD); + PARSER.declareString(ConstructingObjectParser.constructorArg(), ENDPOINT_FIELD); + PARSER.declareString(ConstructingObjectParser.constructorArg(), STATE_FIELD); + PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), CURRENT_INDEX_FIELD); + PARSER.declareStringArray(ConstructingObjectParser.constructorArg(), INDICES_FIELD); + PARSER.declareObject(ConstructingObjectParser.constructorArg(), Database.PARSER, DATABASE_FIELD); + PARSER.declareObject(ConstructingObjectParser.constructorArg(), UpdateStats.PARSER, UPDATE_STATS_FIELD); + } + + @VisibleForTesting + public Datasource() { + this(null, null, null); + } + + public Datasource(final String name, final IntervalSchedule schedule, final String endpoint) { + this( + name, + Instant.now().truncatedTo(ChronoUnit.MILLIS), + null, + false, + schedule, + schedule, + DatasourceTask.ALL, + endpoint, + DatasourceState.CREATING, + null, + new ArrayList<>(), + new Database(), + new UpdateStats() + ); + } + + public Datasource(final StreamInput in) throws IOException { + name = in.readString(); + lastUpdateTime = toInstant(in.readVLong()); + enabledTime = toInstant(in.readOptionalVLong()); + isEnabled = in.readBoolean(); + userSchedule = new IntervalSchedule(in); + systemSchedule = new IntervalSchedule(in); + task = DatasourceTask.valueOf(in.readString()); + endpoint = in.readString(); + state = DatasourceState.valueOf(in.readString()); + currentIndex = in.readOptionalString(); + indices = in.readStringList(); + database = new Database(in); + updateStats = new UpdateStats(in); + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + out.writeString(name); + out.writeVLong(lastUpdateTime.toEpochMilli()); + out.writeOptionalVLong(enabledTime == null ? null : enabledTime.toEpochMilli()); + out.writeBoolean(isEnabled); + userSchedule.writeTo(out); + systemSchedule.writeTo(out); + out.writeString(task.name()); + out.writeString(endpoint); + out.writeString(state.name()); + out.writeOptionalString(currentIndex); + out.writeStringCollection(indices); + database.writeTo(out); + updateStats.writeTo(out); + } + + @Override + public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { + builder.startObject(); + builder.field(NAME_FIELD.getPreferredName(), name); + builder.timeField( + LAST_UPDATE_TIME_FIELD.getPreferredName(), + LAST_UPDATE_TIME_FIELD_READABLE.getPreferredName(), + lastUpdateTime.toEpochMilli() + ); + if (enabledTime != null) { + builder.timeField( + ENABLED_TIME_FIELD.getPreferredName(), + ENABLED_TIME_FIELD_READABLE.getPreferredName(), + enabledTime.toEpochMilli() + ); + } + builder.field(ENABLED_FIELD.getPreferredName(), isEnabled); + builder.field(USER_SCHEDULE_FIELD.getPreferredName(), userSchedule); + builder.field(SYSTEM_SCHEDULE_FIELD.getPreferredName(), systemSchedule); + builder.field(TASK_FIELD.getPreferredName(), task.name()); + builder.field(ENDPOINT_FIELD.getPreferredName(), endpoint); + builder.field(STATE_FIELD.getPreferredName(), state.name()); + if (currentIndex != null) { + builder.field(CURRENT_INDEX_FIELD.getPreferredName(), currentIndex); + } + builder.field(INDICES_FIELD.getPreferredName(), indices); + builder.field(DATABASE_FIELD.getPreferredName(), database); + builder.field(UPDATE_STATS_FIELD.getPreferredName(), updateStats); + builder.endObject(); + return builder; + } + + @Override + public String getName() { + return name; + } + + @Override + public Instant getLastUpdateTime() { + return lastUpdateTime; + } + + @Override + public Instant getEnabledTime() { + return enabledTime; + } + + @Override + public IntervalSchedule getSchedule() { + return systemSchedule; + } + + @Override + public boolean isEnabled() { + return isEnabled; + } + + @Override + public Long getLockDurationSeconds() { + return Ip2GeoLockService.LOCK_DURATION_IN_SECONDS; + } + + /** + * Enable auto update of GeoIP data + */ + public void enable() { + if (isEnabled == true) { + return; + } + enabledTime = Instant.now().truncatedTo(ChronoUnit.MILLIS); + isEnabled = true; + } + + /** + * Disable auto update of GeoIP data + */ + public void disable() { + enabledTime = null; + isEnabled = false; + } + + /** + * Current index name of a datasource + * + * @return Current index name of a datasource + */ + public String currentIndexName() { + if (isExpired()) { + return null; + } + + return currentIndex; + } + + /** + * Index name for a datasource with given suffix + * + * @param suffix the suffix of a index name + * @return index name for a datasource with given suffix + */ + public String newIndexName(final String suffix) { + return String.format(Locale.ROOT, "%s.%s.%s", IP2GEO_DATA_INDEX_NAME_PREFIX, name, suffix); + } + + /** + * Reset database so that it can be updated in next run regardless there is new update or not + */ + public void resetDatabase() { + database.setUpdatedAt(null); + database.setSha256Hash(null); + } + + /** + * Checks if datasource is expired or not + * + * @return true if datasource is expired, and false otherwise + */ + public boolean isExpired() { + return willExpire(Instant.now()); + } + + /** + * Checks if datasource will expire at given time + * + * @return true if datasource will expired at given time, and false otherwise + */ + public boolean willExpire(Instant instant) { + if (database.validForInDays == null) { + return false; + } + + return instant.isAfter(expirationDay()); + } + + /** + * Day when datasource will expire + * + * @return Day when datasource will expire + */ + public Instant expirationDay() { + if (database.validForInDays == null) { + return Instant.MAX; + } + return lastCheckedAt().plus(database.validForInDays, ChronoUnit.DAYS); + } + + private Instant lastCheckedAt() { + Instant lastCheckedAt; + if (updateStats.lastSkippedAt == null) { + lastCheckedAt = updateStats.lastSucceededAt; + } else { + lastCheckedAt = updateStats.lastSucceededAt.isBefore(updateStats.lastSkippedAt) + ? updateStats.lastSkippedAt + : updateStats.lastSucceededAt; + } + return lastCheckedAt; + } + + /** + * 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.setSha256Hash(datasourceManifest.getSha256Hash()); + this.database.setUpdatedAt(Instant.ofEpochMilli(datasourceManifest.getUpdatedAt())); + 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) == false) { + return false; + } + } + return true; + } + + private static Instant toInstant(final Long epochMilli) { + return epochMilli == null ? null : Instant.ofEpochMilli(epochMilli); + } + + /** + * Database of a datasource + */ + @Getter + @Setter + @ToString + @EqualsAndHashCode + @AllArgsConstructor(access = AccessLevel.PRIVATE) + @NoArgsConstructor(access = AccessLevel.PRIVATE) + public static class Database implements Writeable, ToXContent { + private static final ParseField PROVIDER_FIELD = new ParseField("provider"); + private static final ParseField SHA256_HASH_FIELD = new ParseField("sha256_hash"); + private static final ParseField UPDATED_AT_FIELD = new ParseField("updated_at_in_epoch_millis"); + private static final ParseField UPDATED_AT_FIELD_READABLE = new ParseField("updated_at"); + private static final ParseField FIELDS_FIELD = new ParseField("fields"); + private static final ParseField VALID_FOR_IN_DAYS_FIELD = new ParseField("valid_for_in_days"); + + /** + * @param provider A database provider name + * @return A database provider name + */ + private String provider; + /** + * @param sha256Hash SHA256 hash value of a database file + * @return SHA256 hash value of a database file + */ + private String sha256Hash; + /** + * @param updatedAt A date when the database was updated + * @return A date when the database was updated + */ + private Instant updatedAt; + /** + * @param validForInDays A duration in which the database file is valid to use + * @return A duration in which the database file is valid to use + */ + private Long validForInDays; + /** + * @param fields A list of available fields in the database + * @return A list of available fields in the database + */ + private List fields; + + private static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( + "datasource_metadata_database", + true, + args -> { + String provider = (String) args[0]; + String sha256Hash = (String) args[1]; + Instant updatedAt = args[2] == null ? null : Instant.ofEpochMilli((Long) args[2]); + Long validForInDays = (Long) args[3]; + List fields = (List) args[4]; + return new Database(provider, sha256Hash, updatedAt, validForInDays, fields); + } + ); + static { + PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), PROVIDER_FIELD); + PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), SHA256_HASH_FIELD); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), UPDATED_AT_FIELD); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), VALID_FOR_IN_DAYS_FIELD); + PARSER.declareStringArray(ConstructingObjectParser.optionalConstructorArg(), FIELDS_FIELD); + } + + public Database(final StreamInput in) throws IOException { + provider = in.readOptionalString(); + sha256Hash = in.readOptionalString(); + updatedAt = toInstant(in.readOptionalVLong()); + validForInDays = in.readOptionalVLong(); + fields = in.readOptionalStringList(); + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + out.writeOptionalString(provider); + out.writeOptionalString(sha256Hash); + out.writeOptionalVLong(updatedAt == null ? null : updatedAt.toEpochMilli()); + out.writeOptionalVLong(validForInDays); + out.writeOptionalStringCollection(fields); + } + + @Override + public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { + builder.startObject(); + if (provider != null) { + builder.field(PROVIDER_FIELD.getPreferredName(), provider); + } + if (sha256Hash != null) { + builder.field(SHA256_HASH_FIELD.getPreferredName(), sha256Hash); + } + if (updatedAt != null) { + builder.timeField( + UPDATED_AT_FIELD.getPreferredName(), + UPDATED_AT_FIELD_READABLE.getPreferredName(), + updatedAt.toEpochMilli() + ); + } + if (validForInDays != null) { + builder.field(VALID_FOR_IN_DAYS_FIELD.getPreferredName(), validForInDays); + } + if (fields != null) { + builder.startArray(FIELDS_FIELD.getPreferredName()); + for (String field : fields) { + builder.value(field); + } + builder.endArray(); + } + builder.endObject(); + return builder; + } + } + + /** + * Update stats of a datasource + */ + @Getter + @Setter + @ToString + @EqualsAndHashCode + @AllArgsConstructor(access = AccessLevel.PRIVATE) + @NoArgsConstructor(access = AccessLevel.PRIVATE) + public static class UpdateStats implements Writeable, ToXContent { + private static final ParseField LAST_SUCCEEDED_AT_FIELD = new ParseField("last_succeeded_at_in_epoch_millis"); + private static final ParseField LAST_SUCCEEDED_AT_FIELD_READABLE = new ParseField("last_succeeded_at"); + private static final ParseField LAST_PROCESSING_TIME_IN_MILLIS_FIELD = new ParseField("last_processing_time_in_millis"); + private static final ParseField LAST_FAILED_AT_FIELD = new ParseField("last_failed_at_in_epoch_millis"); + private static final ParseField LAST_FAILED_AT_FIELD_READABLE = new ParseField("last_failed_at"); + private static final ParseField LAST_SKIPPED_AT = new ParseField("last_skipped_at_in_epoch_millis"); + private static final ParseField LAST_SKIPPED_AT_READABLE = new ParseField("last_skipped_at"); + + /** + * @param lastSucceededAt The last time when GeoIP data update was succeeded + * @return The last time when GeoIP data update was succeeded + */ + private Instant lastSucceededAt; + /** + * @param lastProcessingTimeInMillis The last processing time when GeoIP data update was succeeded + * @return The last processing time when GeoIP data update was succeeded + */ + private Long lastProcessingTimeInMillis; + /** + * @param lastFailedAt The last time when GeoIP data update was failed + * @return The last time when GeoIP data update was failed + */ + private Instant lastFailedAt; + /** + * @param lastSkippedAt The last time when GeoIP data update was skipped as there was no new update from an endpoint + * @return The last time when GeoIP data update was skipped as there was no new update from an endpoint + */ + private Instant lastSkippedAt; + + private static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( + "datasource_metadata_update_stats", + true, + args -> { + Instant lastSucceededAt = args[0] == null ? null : Instant.ofEpochMilli((long) args[0]); + Long lastProcessingTimeInMillis = (Long) args[1]; + Instant lastFailedAt = args[2] == null ? null : Instant.ofEpochMilli((long) args[2]); + Instant lastSkippedAt = args[3] == null ? null : Instant.ofEpochMilli((long) args[3]); + return new UpdateStats(lastSucceededAt, lastProcessingTimeInMillis, lastFailedAt, lastSkippedAt); + } + ); + + static { + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), LAST_SUCCEEDED_AT_FIELD); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), LAST_PROCESSING_TIME_IN_MILLIS_FIELD); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), LAST_FAILED_AT_FIELD); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), LAST_SKIPPED_AT); + } + + public UpdateStats(final StreamInput in) throws IOException { + lastSucceededAt = toInstant(in.readOptionalVLong()); + lastProcessingTimeInMillis = in.readOptionalVLong(); + lastFailedAt = toInstant(in.readOptionalVLong()); + lastSkippedAt = toInstant(in.readOptionalVLong()); + } + + @Override + public void writeTo(final StreamOutput out) throws IOException { + out.writeOptionalVLong(lastSucceededAt == null ? null : lastSucceededAt.toEpochMilli()); + out.writeOptionalVLong(lastProcessingTimeInMillis); + out.writeOptionalVLong(lastFailedAt == null ? null : lastFailedAt.toEpochMilli()); + out.writeOptionalVLong(lastSkippedAt == null ? null : lastSkippedAt.toEpochMilli()); + } + + @Override + public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { + builder.startObject(); + if (lastSucceededAt != null) { + builder.timeField( + LAST_SUCCEEDED_AT_FIELD.getPreferredName(), + LAST_SUCCEEDED_AT_FIELD_READABLE.getPreferredName(), + lastSucceededAt.toEpochMilli() + ); + } + if (lastProcessingTimeInMillis != null) { + builder.field(LAST_PROCESSING_TIME_IN_MILLIS_FIELD.getPreferredName(), lastProcessingTimeInMillis); + } + if (lastFailedAt != null) { + builder.timeField( + LAST_FAILED_AT_FIELD.getPreferredName(), + LAST_FAILED_AT_FIELD_READABLE.getPreferredName(), + lastFailedAt.toEpochMilli() + ); + } + if (lastSkippedAt != null) { + builder.timeField( + LAST_SKIPPED_AT.getPreferredName(), + LAST_SKIPPED_AT_READABLE.getPreferredName(), + lastSkippedAt.toEpochMilli() + ); + } + builder.endObject(); + return builder; + } + } + + /** + * Builder class for Datasource + */ + public static class Builder { + public static Datasource build(final PutDatasourceRequest request) { + String id = request.getName(); + IntervalSchedule schedule = new IntervalSchedule( + Instant.now().truncatedTo(ChronoUnit.MILLIS), + (int) request.getUpdateInterval().days(), + ChronoUnit.DAYS + ); + String endpoint = request.getEndpoint(); + return new Datasource(id, schedule, endpoint); + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtension.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtension.java new file mode 100644 index 0000000000..ce610ebe41 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtension.java @@ -0,0 +1,57 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import static org.opensearch.geospatial.constants.IndexSetting.AUTO_EXPAND_REPLICAS; +import static org.opensearch.geospatial.constants.IndexSetting.HIDDEN; +import static org.opensearch.geospatial.constants.IndexSetting.NUMBER_OF_SHARDS; + +import java.util.Map; + +import org.opensearch.jobscheduler.spi.JobSchedulerExtension; +import org.opensearch.jobscheduler.spi.ScheduledJobParser; +import org.opensearch.jobscheduler.spi.ScheduledJobRunner; + +/** + * Datasource job scheduler extension + * + * This extension is responsible for scheduling GeoIp data update task + * + * See https://github.com/opensearch-project/job-scheduler/blob/main/README.md#getting-started + */ +public class DatasourceExtension implements JobSchedulerExtension { + /** + * Job index name for a datasource + */ + public static final String JOB_INDEX_NAME = ".scheduler_geospatial_ip2geo_datasource"; + /** + * Job index setting + * + * We want it to be single shard so that job can be run only in a single node by job scheduler. + * We want it to expand to all replicas so that querying to this index can be done locally to reduce latency. + */ + public static final Map INDEX_SETTING = Map.of(NUMBER_OF_SHARDS, 1, AUTO_EXPAND_REPLICAS, "0-all", HIDDEN, true); + + @Override + public String getJobType() { + return "scheduler_geospatial_ip2geo_datasource"; + } + + @Override + public String getJobIndex() { + return JOB_INDEX_NAME; + } + + @Override + public ScheduledJobRunner getJobRunner() { + return DatasourceRunner.getJobRunnerInstance(); + } + + @Override + public ScheduledJobParser getJobParser() { + return (parser, id, jobDocVersion) -> Datasource.PARSER.parse(parser, null); + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunner.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunner.java new file mode 100644 index 0000000000..fe4aa36bfa --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunner.java @@ -0,0 +1,189 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import java.io.IOException; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; + +import lombok.extern.log4j.Log4j2; + +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.jobscheduler.spi.JobExecutionContext; +import org.opensearch.jobscheduler.spi.LockModel; +import org.opensearch.jobscheduler.spi.ScheduledJobParameter; +import org.opensearch.jobscheduler.spi.ScheduledJobRunner; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; + +/** + * Datasource update task + * + * This is a background task which is responsible for updating GeoIp data + */ +@Log4j2 +public class DatasourceRunner implements ScheduledJobRunner { + private static final int DELETE_INDEX_RETRY_IN_MIN = 15; + private static final int DELETE_INDEX_DELAY_IN_MILLIS = 10000; + + private static DatasourceRunner INSTANCE; + + /** + * Return a singleton job runner instance + * @return job runner + */ + public static DatasourceRunner getJobRunnerInstance() { + if (INSTANCE != null) { + return INSTANCE; + } + synchronized (DatasourceRunner.class) { + if (INSTANCE != null) { + return INSTANCE; + } + INSTANCE = new DatasourceRunner(); + return INSTANCE; + } + } + + private ClusterService clusterService; + private DatasourceUpdateService datasourceUpdateService; + private Ip2GeoExecutor ip2GeoExecutor; + private DatasourceDao datasourceDao; + private Ip2GeoLockService ip2GeoLockService; + private boolean initialized; + + private DatasourceRunner() { + // Singleton class, use getJobRunner method instead of constructor + } + + /** + * Initialize timeout and indexingBulkSize from settings + */ + public void initialize( + final ClusterService clusterService, + final DatasourceUpdateService datasourceUpdateService, + final Ip2GeoExecutor ip2GeoExecutor, + final DatasourceDao datasourceDao, + final Ip2GeoLockService ip2GeoLockService + ) { + this.clusterService = clusterService; + this.datasourceUpdateService = datasourceUpdateService; + this.ip2GeoExecutor = ip2GeoExecutor; + this.datasourceDao = datasourceDao; + this.ip2GeoLockService = ip2GeoLockService; + this.initialized = true; + } + + @Override + public void runJob(final ScheduledJobParameter jobParameter, final JobExecutionContext context) { + if (initialized == false) { + throw new AssertionError("this instance is not initialized"); + } + + log.info("Update job started for a datasource[{}]", jobParameter.getName()); + if (jobParameter instanceof Datasource == false) { + throw new IllegalStateException( + "job parameter is not instance of Datasource, type: " + jobParameter.getClass().getCanonicalName() + ); + } + + ip2GeoExecutor.forDatasourceUpdate().submit(updateDatasourceRunner(jobParameter)); + } + + /** + * Update GeoIP data + * + * Lock is used so that only one of nodes run this task. + * + * @param jobParameter job parameter + */ + @VisibleForTesting + protected Runnable updateDatasourceRunner(final ScheduledJobParameter jobParameter) { + return () -> { + Optional lockModel = ip2GeoLockService.acquireLock( + jobParameter.getName(), + Ip2GeoLockService.LOCK_DURATION_IN_SECONDS + ); + if (lockModel.isEmpty()) { + log.error("Failed to update. Another processor is holding a lock for datasource[{}]", jobParameter.getName()); + return; + } + + LockModel lock = lockModel.get(); + try { + updateDatasource(jobParameter, ip2GeoLockService.getRenewLockRunnable(new AtomicReference<>(lock))); + } catch (Exception e) { + log.error("Failed to update datasource[{}]", jobParameter.getName(), e); + } finally { + ip2GeoLockService.releaseLock(lock); + } + }; + } + + @VisibleForTesting + protected void updateDatasource(final ScheduledJobParameter jobParameter, final Runnable renewLock) throws IOException { + Datasource datasource = datasourceDao.getDatasource(jobParameter.getName()); + /** + * If delete request comes while update task is waiting on a queue for other update tasks to complete, + * because update task for this datasource didn't acquire a lock yet, delete request is processed. + * When it is this datasource's turn to run, it will find that the datasource is deleted already. + * Therefore, we stop the update process when data source does not exist. + */ + if (datasource == null) { + log.info("Datasource[{}] does not exist", jobParameter.getName()); + return; + } + + if (DatasourceState.AVAILABLE.equals(datasource.getState()) == false) { + log.error("Invalid datasource state. Expecting {} but received {}", DatasourceState.AVAILABLE, datasource.getState()); + datasource.disable(); + datasource.getUpdateStats().setLastFailedAt(Instant.now()); + datasourceDao.updateDatasource(datasource); + return; + } + + try { + datasourceUpdateService.deleteUnusedIndices(datasource); + if (DatasourceTask.DELETE_UNUSED_INDICES.equals(datasource.getTask()) == false) { + datasourceUpdateService.updateOrCreateGeoIpData(datasource, renewLock); + } + datasourceUpdateService.deleteUnusedIndices(datasource); + } catch (Exception e) { + log.error("Failed to update datasource for {}", datasource.getName(), e); + datasource.getUpdateStats().setLastFailedAt(Instant.now()); + datasourceDao.updateDatasource(datasource); + } finally { + postProcessing(datasource); + } + } + + private void postProcessing(final Datasource datasource) { + if (datasource.isExpired()) { + // Try to delete again as it could have just been expired + datasourceUpdateService.deleteUnusedIndices(datasource); + datasourceUpdateService.updateDatasource(datasource, datasource.getUserSchedule(), DatasourceTask.ALL); + return; + } + + if (datasource.willExpire(datasource.getUserSchedule().getNextExecutionTime(Instant.now()))) { + IntervalSchedule intervalSchedule = new IntervalSchedule( + datasource.expirationDay(), + DELETE_INDEX_RETRY_IN_MIN, + ChronoUnit.MINUTES, + DELETE_INDEX_DELAY_IN_MILLIS + ); + datasourceUpdateService.updateDatasource(datasource, intervalSchedule, DatasourceTask.DELETE_UNUSED_INDICES); + } else { + datasourceUpdateService.updateDatasource(datasource, datasource.getUserSchedule(), DatasourceTask.ALL); + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTask.java b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTask.java new file mode 100644 index 0000000000..bfbcb1d2d9 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTask.java @@ -0,0 +1,21 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +/** + * Task that {@link DatasourceRunner} will run + */ +public enum DatasourceTask { + /** + * Do everything + */ + ALL, + + /** + * Only delete unused indices + */ + DELETE_UNUSED_INDICES +} 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 0000000000..2abf8a7983 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateService.java @@ -0,0 +1,288 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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.List; +import java.util.UUID; +import java.util.stream.Collectors; + +import lombok.extern.log4j.Log4j2; + +import org.apache.commons.csv.CSVParser; +import org.apache.commons.csv.CSVRecord; +import org.opensearch.OpenSearchException; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.geospatial.annotation.VisibleForTesting; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.dao.GeoIpDataDao; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; + +@Log4j2 +public class DatasourceUpdateService { + private static final int SLEEP_TIME_IN_MILLIS = 5000; // 5 seconds + private static final int MAX_WAIT_TIME_FOR_REPLICATION_TO_COMPLETE_IN_MILLIS = 10 * 60 * 60 * 1000; // 10 hours + private final ClusterService clusterService; + private final ClusterSettings clusterSettings; + private final DatasourceDao datasourceDao; + private final GeoIpDataDao geoIpDataDao; + + public DatasourceUpdateService( + final ClusterService clusterService, + final DatasourceDao datasourceDao, + final GeoIpDataDao geoIpDataDao + ) { + this.clusterService = clusterService; + this.clusterSettings = clusterService.getClusterSettings(); + this.datasourceDao = datasourceDao; + this.geoIpDataDao = geoIpDataDao; + } + + /** + * Update GeoIp data + * + * The first column is ip range field regardless its header name. + * Therefore, we don't store the first column's header name. + * + * @param datasource the datasource + * @param renewLock runnable to renew lock + * + * @throws IOException + */ + public void updateOrCreateGeoIpData(final Datasource datasource, final Runnable renewLock) throws IOException { + URL url = new URL(datasource.getEndpoint()); + DatasourceManifest manifest = DatasourceManifest.Builder.build(url); + + if (shouldUpdate(datasource, manifest) == false) { + log.info("Skipping GeoIP database update. Update is not required for {}", datasource.getName()); + datasource.getUpdateStats().setLastSkippedAt(Instant.now()); + datasourceDao.updateDatasource(datasource); + return; + } + + Instant startTime = Instant.now(); + String indexName = setupIndex(datasource); + String[] header; + List fieldsToStore; + try (CSVParser reader = geoIpDataDao.getDatabaseReader(manifest)) { + CSVRecord headerLine = reader.iterator().next(); + header = validateHeader(headerLine).values(); + fieldsToStore = Arrays.asList(header).subList(1, header.length); + if (datasource.isCompatible(fieldsToStore) == false) { + throw new OpenSearchException( + "new fields [{}] does not contain all old fields [{}]", + fieldsToStore.toString(), + datasource.getDatabase().getFields().toString() + ); + } + geoIpDataDao.putGeoIpData(indexName, header, reader.iterator(), renewLock); + } + + waitUntilAllShardsStarted(indexName, MAX_WAIT_TIME_FOR_REPLICATION_TO_COMPLETE_IN_MILLIS); + Instant endTime = Instant.now(); + updateDatasourceAsSucceeded(indexName, datasource, manifest, fieldsToStore, startTime, endTime); + } + + /** + * We wait until all shards are ready to serve search requests before updating datasource metadata to + * point to a new index so that there won't be latency degradation during GeoIP data update + * + * @param indexName the indexName + */ + @VisibleForTesting + protected void waitUntilAllShardsStarted(final String indexName, final int timeout) { + Instant start = Instant.now(); + try { + while (Instant.now().toEpochMilli() - start.toEpochMilli() < timeout) { + if (clusterService.state().routingTable().allShards(indexName).stream().allMatch(shard -> shard.started())) { + return; + } + Thread.sleep(SLEEP_TIME_IN_MILLIS); + } + throw new OpenSearchException( + "index[{}] replication did not complete after {} millis", + MAX_WAIT_TIME_FOR_REPLICATION_TO_COMPLETE_IN_MILLIS + ); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + } + + /** + * Return header fields of geo data with given url of a manifest file + * + * The first column is ip range field regardless its header name. + * Therefore, we don't store the first column's header name. + * + * @param manifestUrl the url of a manifest file + * @return header fields of geo data + */ + public List getHeaderFields(String manifestUrl) throws IOException { + URL url = new URL(manifestUrl); + DatasourceManifest manifest = DatasourceManifest.Builder.build(url); + + try (CSVParser reader = geoIpDataDao.getDatabaseReader(manifest)) { + String[] fields = reader.iterator().next().values(); + return Arrays.asList(fields).subList(1, fields.length); + } + } + + /** + * Delete all indices except the one which are being used + * + * @param datasource + */ + public void deleteUnusedIndices(final Datasource datasource) { + try { + List indicesToDelete = datasource.getIndices() + .stream() + .filter(index -> index.equals(datasource.currentIndexName()) == false) + .collect(Collectors.toList()); + + List deletedIndices = deleteIndices(indicesToDelete); + + if (deletedIndices.isEmpty() == false) { + datasource.getIndices().removeAll(deletedIndices); + datasourceDao.updateDatasource(datasource); + } + } catch (Exception e) { + log.error("Failed to delete old indices for {}", datasource.getName(), e); + } + } + + /** + * Update datasource with given systemSchedule and task + * + * @param datasource datasource to update + * @param systemSchedule new system schedule value + * @param task new task value + */ + public void updateDatasource(final Datasource datasource, final IntervalSchedule systemSchedule, final DatasourceTask task) { + boolean updated = false; + if (datasource.getSystemSchedule().equals(systemSchedule) == false) { + datasource.setSystemSchedule(systemSchedule); + updated = true; + } + if (datasource.getTask().equals(task) == false) { + datasource.setTask(task); + updated = true; + } + + if (updated) { + datasourceDao.updateDatasource(datasource); + } + } + + private List deleteIndices(final List indicesToDelete) { + List deletedIndices = new ArrayList<>(indicesToDelete.size()); + for (String index : indicesToDelete) { + if (clusterService.state().metadata().hasIndex(index) == false) { + deletedIndices.add(index); + continue; + } + + try { + geoIpDataDao.deleteIp2GeoDataIndex(index); + deletedIndices.add(index); + } catch (Exception e) { + log.error("Failed to delete an index [{}]", index, e); + } + } + return deletedIndices; + } + + /** + * Validate header + * + * 1. header should not be null + * 2. the number of values in header should be more than one + * + * @param header the header + * @return CSVRecord the input header + */ + private CSVRecord validateHeader(CSVRecord header) { + if (header == null) { + throw new OpenSearchException("geoip database is empty"); + } + 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 + */ + private void updateDatasourceAsSucceeded( + final String newIndexName, + final Datasource datasource, + final DatasourceManifest manifest, + final List fields, + final Instant startTime, + final Instant endTime + ) { + datasource.setCurrentIndex(newIndexName); + datasource.setDatabase(manifest, fields); + datasource.getUpdateStats().setLastSucceededAt(endTime); + datasource.getUpdateStats().setLastProcessingTimeInMillis(endTime.toEpochMilli() - startTime.toEpochMilli()); + datasource.enable(); + datasource.setState(DatasourceState.AVAILABLE); + datasourceDao.updateDatasource(datasource); + log.info( + "GeoIP database creation succeeded for {} and took {} seconds", + datasource.getName(), + Duration.between(startTime, endTime) + ); + } + + /*** + * Setup index to add a new geoip data + * + * @param datasource the datasource + * @return new index name + */ + private String setupIndex(final Datasource datasource) { + String indexName = datasource.newIndexName(UUID.randomUUID().toString()); + datasource.getIndices().add(indexName); + datasourceDao.updateDatasource(datasource); + geoIpDataDao.createIndexIfNotExists(indexName); + return indexName; + } + + /** + * Determine if update is needed or not + * + * Update is needed when all following conditions are met + * 1. updatedAt value in datasource is equal or before updateAt value in manifest + * 2. SHA256 hash value in datasource is different with SHA256 hash value in manifest + * + * @param datasource + * @param manifest + * @return + */ + private boolean shouldUpdate(final Datasource datasource, final DatasourceManifest manifest) { + if (datasource.getDatabase().getUpdatedAt() != null + && datasource.getDatabase().getUpdatedAt().toEpochMilli() > manifest.getUpdatedAt()) { + return false; + } + + if (manifest.getSha256Hash().equals(datasource.getDatabase().getSha256Hash())) { + return false; + } + return true; + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/listener/Ip2GeoListener.java b/src/main/java/org/opensearch/geospatial/ip2geo/listener/Ip2GeoListener.java new file mode 100644 index 0000000000..e6fe98f9e0 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/listener/Ip2GeoListener.java @@ -0,0 +1,123 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.listener; + +import static org.opensearch.geospatial.ip2geo.jobscheduler.Datasource.IP2GEO_DATA_INDEX_NAME_PREFIX; + +import java.io.IOException; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.List; +import java.util.stream.Collectors; + +import lombok.AllArgsConstructor; +import lombok.extern.log4j.Log4j2; + +import org.opensearch.action.ActionListener; +import org.opensearch.action.bulk.BulkResponse; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterStateListener; +import org.opensearch.cluster.RestoreInProgress; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.component.AbstractLifecycleComponent; +import org.opensearch.common.inject.Inject; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.dao.GeoIpDataDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceTask; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; +import org.opensearch.threadpool.ThreadPool; + +@Log4j2 +@AllArgsConstructor(onConstructor = @__(@Inject)) +public class Ip2GeoListener extends AbstractLifecycleComponent implements ClusterStateListener { + private static final int SCHEDULE_IN_MIN = 15; + private static final int DELAY_IN_MILLIS = 10000; + private final ClusterService clusterService; + private final ThreadPool threadPool; + private final DatasourceDao datasourceDao; + private final GeoIpDataDao geoIpDataDao; + + @Override + public void clusterChanged(final ClusterChangedEvent event) { + if (event.localNodeClusterManager() == false) { + return; + } + + for (RestoreInProgress.Entry entry : event.state().custom(RestoreInProgress.TYPE, RestoreInProgress.EMPTY)) { + if (RestoreInProgress.State.SUCCESS.equals(entry.state()) == false) { + continue; + } + + if (entry.indices().stream().anyMatch(index -> DatasourceExtension.JOB_INDEX_NAME.equals(index))) { + threadPool.generic().submit(() -> forceUpdateGeoIpData()); + } + + List ip2GeoDataIndices = entry.indices() + .stream() + .filter(index -> index.startsWith(IP2GEO_DATA_INDEX_NAME_PREFIX)) + .collect(Collectors.toList()); + if (ip2GeoDataIndices.isEmpty() == false) { + threadPool.generic().submit(() -> geoIpDataDao.deleteIp2GeoDataIndex(ip2GeoDataIndices)); + } + } + } + + private void forceUpdateGeoIpData() { + datasourceDao.getAllDatasources(new ActionListener<>() { + @Override + public void onResponse(final List datasources) { + datasources.stream().forEach(Ip2GeoListener.this::scheduleForceUpdate); + datasourceDao.updateDatasource(datasources, new ActionListener<>() { + @Override + public void onResponse(final BulkResponse bulkItemResponses) { + log.info("Datasources are updated for cleanup"); + } + + @Override + public void onFailure(final Exception e) { + log.error("Failed to update datasource for cleanup after restoring", e); + } + }); + } + + @Override + public void onFailure(final Exception e) { + log.error("Failed to get datasource after restoring", e); + } + }); + } + + /** + * Give a delay so that job scheduler can schedule the job right after the delay. Otherwise, it schedules + * the job after specified update interval. + */ + private void scheduleForceUpdate(Datasource datasource) { + IntervalSchedule schedule = new IntervalSchedule(Instant.now(), SCHEDULE_IN_MIN, ChronoUnit.MINUTES, DELAY_IN_MILLIS); + datasource.resetDatabase(); + datasource.setSystemSchedule(schedule); + datasource.setTask(DatasourceTask.ALL); + } + + @Override + protected void doStart() { + if (DiscoveryNode.isClusterManagerNode(clusterService.getSettings())) { + clusterService.addListener(this); + } + } + + @Override + protected void doStop() { + clusterService.removeListener(this); + } + + @Override + protected void doClose() throws IOException { + + } +} diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessor.java b/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessor.java new file mode 100644 index 0000000000..56100c0b67 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessor.java @@ -0,0 +1,274 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ +package org.opensearch.geospatial.ip2geo.processor; + +import static org.opensearch.ingest.ConfigurationUtils.readBooleanProperty; +import static org.opensearch.ingest.ConfigurationUtils.readOptionalList; +import static org.opensearch.ingest.ConfigurationUtils.readStringProperty; + +import java.io.IOException; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.function.BiConsumer; +import java.util.stream.Collectors; + +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.extern.log4j.Log4j2; + +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.dao.GeoIpDataDao; +import org.opensearch.geospatial.ip2geo.dao.Ip2GeoCachedDao; +import org.opensearch.ingest.AbstractProcessor; +import org.opensearch.ingest.IngestDocument; +import org.opensearch.ingest.IngestService; +import org.opensearch.ingest.Processor; + +/** + * Ip2Geo processor + */ +@Log4j2 +public final class Ip2GeoProcessor extends AbstractProcessor { + private static final Map DATA_EXPIRED = Map.of("error", "ip2geo_data_expired"); + public static final String CONFIG_FIELD = "field"; + public static final String CONFIG_TARGET_FIELD = "target_field"; + public static final String CONFIG_DATASOURCE = "datasource"; + public static final String CONFIG_PROPERTIES = "properties"; + public static final String CONFIG_IGNORE_MISSING = "ignore_missing"; + + private final String field; + private final String targetField; + /** + * @return The datasource name + */ + @Getter + private final String datasourceName; + private final Set properties; + private final boolean ignoreMissing; + private final ClusterSettings clusterSettings; + private final DatasourceDao datasourceDao; + private final GeoIpDataDao geoIpDataDao; + private final Ip2GeoCachedDao ip2GeoCachedDao; + + /** + * Ip2Geo processor type + */ + public static final String TYPE = "ip2geo"; + + /** + * Construct an Ip2Geo processor. + * @param tag the processor tag + * @param description the processor description + * @param field the source field to geo-IP map + * @param targetField the target field + * @param datasourceName the datasourceName + * @param properties the properties + * @param ignoreMissing true if documents with a missing value for the field should be ignored + * @param clusterSettings the cluster settings + * @param datasourceDao the datasource facade + * @param geoIpDataDao the geoip data facade + * @param ip2GeoCachedDao the cache + */ + public Ip2GeoProcessor( + final String tag, + final String description, + final String field, + final String targetField, + final String datasourceName, + final Set properties, + final boolean ignoreMissing, + final ClusterSettings clusterSettings, + final DatasourceDao datasourceDao, + final GeoIpDataDao geoIpDataDao, + final Ip2GeoCachedDao ip2GeoCachedDao + ) { + super(tag, description); + this.field = field; + this.targetField = targetField; + this.datasourceName = datasourceName; + this.properties = properties; + this.ignoreMissing = ignoreMissing; + this.clusterSettings = clusterSettings; + this.datasourceDao = datasourceDao; + this.geoIpDataDao = geoIpDataDao; + this.ip2GeoCachedDao = ip2GeoCachedDao; + } + + /** + * Add geo data of a given ip address to ingestDocument in asynchronous way + * + * @param ingestDocument the document + * @param handler the handler + */ + @Override + public void execute(IngestDocument ingestDocument, BiConsumer handler) { + try { + Object ip = ingestDocument.getFieldValue(field, Object.class, ignoreMissing); + + if (ip == null) { + handler.accept(ingestDocument, null); + return; + } + + if (ip instanceof String) { + executeInternal(ingestDocument, handler, (String) ip); + } else if (ip instanceof List) { + executeInternal(ingestDocument, handler, ((List) ip)); + } else { + handler.accept( + null, + new IllegalArgumentException( + String.format(Locale.ROOT, "field [%s] should contain only string or array of strings", field) + ) + ); + } + } catch (Exception e) { + handler.accept(null, e); + } + } + + /** + * Use {@code execute(IngestDocument, BiConsumer)} instead + * + * @param ingestDocument the document + * @return none + */ + @Override + public IngestDocument execute(IngestDocument ingestDocument) { + throw new IllegalStateException("Not implemented"); + } + + private void executeInternal( + final IngestDocument ingestDocument, + final BiConsumer handler, + final String ip + ) { + validateDatasourceIsInAvailableState(datasourceName); + String indexName = ip2GeoCachedDao.getIndexName(datasourceName); + if (ip2GeoCachedDao.isExpired(datasourceName) || indexName == null) { + handleExpiredData(ingestDocument, handler); + return; + } + + Map geoData = ip2GeoCachedDao.getGeoData(indexName, ip); + if (geoData.isEmpty() == false) { + ingestDocument.setFieldValue(targetField, filteredGeoData(geoData)); + } + handler.accept(ingestDocument, null); + } + + private Map filteredGeoData(final Map geoData) { + if (properties == null) { + return geoData; + } + + return properties.stream().filter(p -> geoData.containsKey(p)).collect(Collectors.toMap(p -> p, p -> geoData.get(p))); + } + + private void validateDatasourceIsInAvailableState(final String datasourceName) { + if (ip2GeoCachedDao.has(datasourceName) == false) { + throw new IllegalStateException("datasource does not exist"); + } + + if (DatasourceState.AVAILABLE.equals(ip2GeoCachedDao.getState(datasourceName)) == false) { + throw new IllegalStateException("datasource is not in an available state"); + } + } + + private void handleExpiredData(final IngestDocument ingestDocument, final BiConsumer handler) { + ingestDocument.setFieldValue(targetField, DATA_EXPIRED); + handler.accept(ingestDocument, null); + } + + /** + * Handle multiple ips + * + * @param ingestDocument the document + * @param handler the handler + * @param ips the ip list + */ + private void executeInternal( + final IngestDocument ingestDocument, + final BiConsumer handler, + final List ips + ) { + for (Object ip : ips) { + if (ip instanceof String == false) { + throw new IllegalArgumentException("array in field [" + field + "] should only contain strings"); + } + } + + validateDatasourceIsInAvailableState(datasourceName); + String indexName = ip2GeoCachedDao.getIndexName(datasourceName); + if (ip2GeoCachedDao.isExpired(datasourceName) || indexName == null) { + handleExpiredData(ingestDocument, handler); + return; + } + + List> geoDataList = ips.stream() + .map(ip -> ip2GeoCachedDao.getGeoData(indexName, (String) ip)) + .filter(geoData -> geoData.isEmpty() == false) + .map(this::filteredGeoData) + .collect(Collectors.toList()); + + if (geoDataList.isEmpty() == false) { + ingestDocument.setFieldValue(targetField, geoDataList); + } + handler.accept(ingestDocument, null); + } + + @Override + public String getType() { + return TYPE; + } + + /** + * Ip2Geo processor factory + */ + @AllArgsConstructor + public static final class Factory implements Processor.Factory { + private final IngestService ingestService; + private final DatasourceDao datasourceDao; + private final GeoIpDataDao geoIpDataDao; + private final Ip2GeoCachedDao ip2GeoCachedDao; + + /** + * Within this method, blocking request cannot be called because this method is executed in a transport thread. + * This means, validation using data in an index won't work. + */ + @Override + public Ip2GeoProcessor create( + final Map registry, + final String processorTag, + final String description, + final Map config + ) throws IOException { + String ipField = readStringProperty(TYPE, processorTag, config, CONFIG_FIELD); + String targetField = readStringProperty(TYPE, processorTag, config, CONFIG_TARGET_FIELD, "ip2geo"); + String datasourceName = readStringProperty(TYPE, processorTag, config, CONFIG_DATASOURCE); + List propertyNames = readOptionalList(TYPE, processorTag, config, CONFIG_PROPERTIES); + boolean ignoreMissing = readBooleanProperty(TYPE, processorTag, config, CONFIG_IGNORE_MISSING, false); + + return new Ip2GeoProcessor( + processorTag, + description, + ipField, + targetField, + datasourceName, + propertyNames == null ? null : new HashSet<>(propertyNames), + ignoreMissing, + ingestService.getClusterService().getClusterSettings(), + datasourceDao, + geoIpDataDao, + ip2GeoCachedDao + ); + } + } +} diff --git a/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java b/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java index 085bbf9275..de69daa6ea 100644 --- a/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java +++ b/src/main/java/org/opensearch/geospatial/plugin/GeospatialPlugin.java @@ -5,11 +5,16 @@ package org.opensearch.geospatial.plugin; +import static org.opensearch.geospatial.ip2geo.jobscheduler.Datasource.IP2GEO_DATA_INDEX_NAME_PREFIX; + +import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; import java.util.function.Supplier; +import lombok.extern.log4j.Log4j2; + import org.opensearch.action.ActionRequest; import org.opensearch.action.ActionResponse; import org.opensearch.client.Client; @@ -17,8 +22,10 @@ import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.collect.MapBuilder; +import org.opensearch.common.component.LifecycleComponent; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.IndexScopedSettings; +import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; import org.opensearch.common.settings.SettingsFilter; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; @@ -32,6 +39,29 @@ import org.opensearch.geospatial.index.mapper.xyshape.XYShapeFieldMapper; import org.opensearch.geospatial.index.mapper.xyshape.XYShapeFieldTypeParser; import org.opensearch.geospatial.index.query.xyshape.XYShapeQueryBuilder; +import org.opensearch.geospatial.ip2geo.action.DeleteDatasourceAction; +import org.opensearch.geospatial.ip2geo.action.DeleteDatasourceTransportAction; +import org.opensearch.geospatial.ip2geo.action.GetDatasourceAction; +import org.opensearch.geospatial.ip2geo.action.GetDatasourceTransportAction; +import org.opensearch.geospatial.ip2geo.action.PutDatasourceAction; +import org.opensearch.geospatial.ip2geo.action.PutDatasourceTransportAction; +import org.opensearch.geospatial.ip2geo.action.RestDeleteDatasourceHandler; +import org.opensearch.geospatial.ip2geo.action.RestGetDatasourceHandler; +import org.opensearch.geospatial.ip2geo.action.RestPutDatasourceHandler; +import org.opensearch.geospatial.ip2geo.action.RestUpdateDatasourceHandler; +import org.opensearch.geospatial.ip2geo.action.UpdateDatasourceAction; +import org.opensearch.geospatial.ip2geo.action.UpdateDatasourceTransportAction; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.dao.GeoIpDataDao; +import org.opensearch.geospatial.ip2geo.dao.Ip2GeoCachedDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceRunner; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; +import org.opensearch.geospatial.ip2geo.listener.Ip2GeoListener; +import org.opensearch.geospatial.ip2geo.processor.Ip2GeoProcessor; import org.opensearch.geospatial.processor.FeatureProcessor; import org.opensearch.geospatial.rest.action.upload.geojson.RestUploadGeoJSONAction; import org.opensearch.geospatial.search.aggregations.bucket.geogrid.GeoHexGrid; @@ -40,17 +70,21 @@ import org.opensearch.geospatial.stats.upload.UploadStats; import org.opensearch.geospatial.stats.upload.UploadStatsAction; import org.opensearch.geospatial.stats.upload.UploadStatsTransportAction; +import org.opensearch.index.IndexModule; import org.opensearch.index.mapper.Mapper; +import org.opensearch.indices.SystemIndexDescriptor; import org.opensearch.ingest.Processor; import org.opensearch.plugins.ActionPlugin; import org.opensearch.plugins.IngestPlugin; import org.opensearch.plugins.MapperPlugin; import org.opensearch.plugins.Plugin; import org.opensearch.plugins.SearchPlugin; +import org.opensearch.plugins.SystemIndexPlugin; import org.opensearch.repositories.RepositoriesService; import org.opensearch.rest.RestController; import org.opensearch.rest.RestHandler; import org.opensearch.script.ScriptService; +import org.opensearch.threadpool.ExecutorBuilder; import org.opensearch.threadpool.ThreadPool; import org.opensearch.watcher.ResourceWatcherService; @@ -58,15 +92,53 @@ * Entry point for Geospatial features. It provides additional Processors, Actions * to interact with Cluster. */ -public class GeospatialPlugin extends Plugin implements IngestPlugin, ActionPlugin, MapperPlugin, SearchPlugin { +@Log4j2 +public class GeospatialPlugin extends Plugin implements IngestPlugin, ActionPlugin, MapperPlugin, SearchPlugin, SystemIndexPlugin { + private Ip2GeoCachedDao ip2GeoCachedDao; + private DatasourceDao datasourceDao; + private GeoIpDataDao geoIpDataDao; + + @Override + public Collection getSystemIndexDescriptors(Settings settings) { + return List.of(new SystemIndexDescriptor(IP2GEO_DATA_INDEX_NAME_PREFIX, "System index used for Ip2Geo data")); + } @Override public Map getProcessors(Processor.Parameters parameters) { + this.datasourceDao = new DatasourceDao(parameters.client, parameters.ingestService.getClusterService()); + this.geoIpDataDao = new GeoIpDataDao(parameters.ingestService.getClusterService(), parameters.client); + this.ip2GeoCachedDao = new Ip2GeoCachedDao(parameters.ingestService.getClusterService(), datasourceDao, geoIpDataDao); return MapBuilder.newMapBuilder() .put(FeatureProcessor.TYPE, new FeatureProcessor.Factory()) + .put(Ip2GeoProcessor.TYPE, new Ip2GeoProcessor.Factory(parameters.ingestService, datasourceDao, geoIpDataDao, ip2GeoCachedDao)) .immutableMap(); } + @Override + public void onIndexModule(IndexModule indexModule) { + if (DatasourceExtension.JOB_INDEX_NAME.equals(indexModule.getIndex().getName())) { + indexModule.addIndexOperationListener(ip2GeoCachedDao); + log.info("Ip2GeoListener started listening to operations on index {}", DatasourceExtension.JOB_INDEX_NAME); + } + } + + @Override + public Collection> getGuiceServiceClasses() { + return List.of(Ip2GeoListener.class); + } + + @Override + public List> getExecutorBuilders(Settings settings) { + List> executorBuilders = new ArrayList<>(); + executorBuilders.add(Ip2GeoExecutor.executorBuilder(settings)); + return executorBuilders; + } + + @Override + public List> getSettings() { + return Ip2GeoSettings.settings(); + } + @Override public Collection createComponents( Client client, @@ -81,7 +153,25 @@ public Collection createComponents( IndexNameExpressionResolver indexNameExpressionResolver, Supplier repositoriesServiceSupplier ) { - return List.of(UploadStats.getInstance()); + DatasourceUpdateService datasourceUpdateService = new DatasourceUpdateService(clusterService, datasourceDao, geoIpDataDao); + Ip2GeoExecutor ip2GeoExecutor = new Ip2GeoExecutor(threadPool); + Ip2GeoLockService ip2GeoLockService = new Ip2GeoLockService(clusterService, client); + /** + * 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, datasourceUpdateService, ip2GeoExecutor, datasourceDao, ip2GeoLockService); + + return List.of( + UploadStats.getInstance(), + datasourceUpdateService, + datasourceDao, + ip2GeoExecutor, + geoIpDataDao, + ip2GeoLockService, + ip2GeoCachedDao + ); } @Override @@ -94,17 +184,39 @@ public List getRestHandlers( IndexNameExpressionResolver indexNameExpressionResolver, Supplier nodesInCluster ) { - RestUploadGeoJSONAction uploadGeoJSONAction = new RestUploadGeoJSONAction(); - RestUploadStatsAction statsAction = new RestUploadStatsAction(); - return List.of(statsAction, uploadGeoJSONAction); + List geoJsonHandlers = List.of(new RestUploadStatsAction(), new RestUploadGeoJSONAction()); + + List ip2geoHandlers = List.of( + new RestPutDatasourceHandler(clusterSettings), + new RestGetDatasourceHandler(), + new RestUpdateDatasourceHandler(), + new RestDeleteDatasourceHandler() + ); + + List allHandlers = new ArrayList<>(); + allHandlers.addAll(geoJsonHandlers); + allHandlers.addAll(ip2geoHandlers); + return allHandlers; } @Override public List> getActions() { - return List.of( + List> geoJsonHandlers = List.of( new ActionHandler<>(UploadGeoJSONAction.INSTANCE, UploadGeoJSONTransportAction.class), new ActionHandler<>(UploadStatsAction.INSTANCE, UploadStatsTransportAction.class) ); + + List> ip2geoHandlers = List.of( + new ActionHandler<>(PutDatasourceAction.INSTANCE, PutDatasourceTransportAction.class), + new ActionHandler<>(GetDatasourceAction.INSTANCE, GetDatasourceTransportAction.class), + new ActionHandler<>(UpdateDatasourceAction.INSTANCE, UpdateDatasourceTransportAction.class), + new ActionHandler<>(DeleteDatasourceAction.INSTANCE, DeleteDatasourceTransportAction.class) + ); + + List> allHandlers = new ArrayList<>(); + allHandlers.addAll(geoJsonHandlers); + allHandlers.addAll(ip2geoHandlers); + return allHandlers; } @Override diff --git a/src/main/java/org/opensearch/geospatial/shared/Constants.java b/src/main/java/org/opensearch/geospatial/shared/Constants.java new file mode 100644 index 0000000000..7b6488a480 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/shared/Constants.java @@ -0,0 +1,15 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.shared; + +import java.util.Locale; + +import org.opensearch.Version; + +public class Constants { + public static final String USER_AGENT_KEY = "User-Agent"; + public static final String USER_AGENT_VALUE = String.format(Locale.ROOT, "OpenSearch/%s vanilla", Version.CURRENT.toString()); +} diff --git a/src/main/java/org/opensearch/geospatial/shared/StashedThreadContext.java b/src/main/java/org/opensearch/geospatial/shared/StashedThreadContext.java new file mode 100644 index 0000000000..1ee5929795 --- /dev/null +++ b/src/main/java/org/opensearch/geospatial/shared/StashedThreadContext.java @@ -0,0 +1,41 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.shared; + +import java.util.function.Supplier; + +import org.opensearch.client.Client; +import org.opensearch.common.util.concurrent.ThreadContext; + +/** + * Helper class to run code with stashed thread context + * + * Code need to be run with stashed thread context if it interacts with system index + * when security plugin is enabled. + */ +public class StashedThreadContext { + /** + * Set the thread context to default, this is needed to allow actions on model system index + * when security plugin is enabled + * @param function runnable that needs to be executed after thread context has been stashed, accepts and returns nothing + */ + public static void run(final Client client, final Runnable function) { + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + function.run(); + } + } + + /** + * Set the thread context to default, this is needed to allow actions on model system index + * when security plugin is enabled + * @param function supplier function that needs to be executed after thread context has been stashed, return object + */ + public static T run(final Client client, final Supplier function) { + try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { + return function.get(); + } + } +} diff --git a/src/main/plugin-metadata/plugin-security.policy b/src/main/plugin-metadata/plugin-security.policy new file mode 100644 index 0000000000..6e9e103065 --- /dev/null +++ b/src/main/plugin-metadata/plugin-security.policy @@ -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. + */ + +grant { + // needed by Ip2Geo datasource to get GeoIP database + permission java.net.SocketPermission "*", "connect,resolve"; +}; diff --git a/src/main/resources/META-INF/services/org.opensearch.jobscheduler.spi.JobSchedulerExtension b/src/main/resources/META-INF/services/org.opensearch.jobscheduler.spi.JobSchedulerExtension new file mode 100644 index 0000000000..e3d6fe6f1b --- /dev/null +++ b/src/main/resources/META-INF/services/org.opensearch.jobscheduler.spi.JobSchedulerExtension @@ -0,0 +1,11 @@ +# +# 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. +# + +# This file is needed to register DatasourceExtension in job scheduler framework +# See https://github.com/opensearch-project/job-scheduler/blob/main/README.md#getting-started +org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceExtension diff --git a/src/main/resources/mappings/ip2geo_datasource.json b/src/main/resources/mappings/ip2geo_datasource.json new file mode 100644 index 0000000000..567052d61e --- /dev/null +++ b/src/main/resources/mappings/ip2geo_datasource.json @@ -0,0 +1,132 @@ +{ + "properties": { + "database": { + "properties": { + "fields": { + "type": "text" + }, + "provider": { + "type": "text" + }, + "sha256_hash": { + "type": "text" + }, + "updated_at_in_epoch_millis": { + "type": "long" + }, + "valid_for_in_days": { + "type": "long" + } + } + }, + "enabled_time": { + "type": "long" + }, + "endpoint": { + "type": "text" + }, + "indices": { + "type": "text" + }, + "last_update_time": { + "type": "long" + }, + "name": { + "type": "text" + }, + "schedule": { + "properties": { + "interval": { + "properties": { + "period": { + "type": "long" + }, + "start_time": { + "type": "long" + }, + "unit": { + "type": "text" + } + } + } + } + }, + "state": { + "type": "text" + }, + "system_schedule": { + "properties": { + "interval": { + "properties": { + "period": { + "type": "long" + }, + "start_time": { + "type": "long" + }, + "unit": { + "type": "text", + "fields": { + "keyword": { + "type": "keyword", + "ignore_above": 256 + } + } + } + } + } + } + }, + "task": { + "type": "text", + "fields": { + "keyword": { + "type": "keyword", + "ignore_above": 256 + } + } + }, + "update_enabled": { + "type": "boolean" + }, + "update_stats": { + "properties": { + "last_failed_at_in_epoch_millis": { + "type": "long" + }, + "last_processing_time_in_millis": { + "type": "long" + }, + "last_skipped_at_in_epoch_millis": { + "type": "long" + }, + "last_succeeded_at_in_epoch_millis": { + "type": "long" + } + } + }, + "user_schedule": { + "properties": { + "interval": { + "properties": { + "period": { + "type": "long" + }, + "start_time": { + "type": "long" + }, + "unit": { + "type": "text", + "fields": { + "keyword": { + "type": "keyword", + "ignore_above": 256 + } + } + } + } + } + } + } + } +} \ No newline at end of file diff --git a/src/main/resources/mappings/ip2geo_geoip.json b/src/main/resources/mappings/ip2geo_geoip.json new file mode 100644 index 0000000000..3179ef0db6 --- /dev/null +++ b/src/main/resources/mappings/ip2geo_geoip.json @@ -0,0 +1,9 @@ +{ + "dynamic": false, + "properties": { + "_cidr": { + "type": "ip_range", + "doc_values": false + } + } +} diff --git a/src/test/java/org/opensearch/geospatial/GeospatialRestTestCase.java b/src/test/java/org/opensearch/geospatial/GeospatialRestTestCase.java index 1e5c1439ab..deae4225ab 100644 --- a/src/test/java/org/opensearch/geospatial/GeospatialRestTestCase.java +++ b/src/test/java/org/opensearch/geospatial/GeospatialRestTestCase.java @@ -17,9 +17,12 @@ import static org.opensearch.search.aggregations.Aggregations.AGGREGATIONS_FIELD; import java.io.IOException; +import java.time.Duration; +import java.time.Instant; import java.util.Collections; import java.util.HashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Optional; import java.util.stream.IntStream; @@ -47,12 +50,11 @@ import org.opensearch.geospatial.action.upload.geojson.UploadGeoJSONRequestContent; import org.opensearch.geospatial.index.mapper.xyshape.XYShapeFieldMapper; import org.opensearch.geospatial.index.query.xyshape.XYShapeQueryBuilder; -import org.opensearch.geospatial.processor.FeatureProcessor; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; import org.opensearch.geospatial.rest.action.upload.geojson.RestUploadGeoJSONAction; import org.opensearch.ingest.Pipeline; public abstract class GeospatialRestTestCase extends OpenSearchSecureRestTestCase { - public static final String SOURCE = "_source"; public static final String DOC = "_doc"; public static final String URL_DELIMITER = "/"; @@ -70,12 +72,24 @@ public abstract class GeospatialRestTestCase extends OpenSearchSecureRestTestCas public static final String SHAPE_ID_FIELD = "id"; public static final String SHAPE_INDEX_PATH_FIELD = "path"; public static final String QUERY_PARAM_TOKEN = "?"; + private static final String SETTINGS = "_settings"; + private static final String SIMULATE = "_simulate"; + private static final String DOCS = "docs"; + private static final String DATASOURCES = "datasources"; + private static final String STATE = "state"; + private static final String PUT = "PUT"; + private static final String GET = "GET"; + private static final String DELETE = "DELETE"; private static String buildPipelinePath(String name) { return String.join(URL_DELIMITER, "_ingest", "pipeline", name); } - protected static void createPipeline(String name, Optional description, List> processorConfigs) + private static String buildDatasourcePath(String name) { + return String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource", name); + } + + protected static Response createPipeline(String name, Optional description, List> processorConfigs) throws IOException { XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); if (description.isPresent()) { @@ -88,7 +102,7 @@ protected static void createPipeline(String name, Optional description, Request request = new Request("PUT", buildPipelinePath(name)); request.setJsonEntity(org.opensearch.common.Strings.toString(builder)); - client().performRequest(request); + return client().performRequest(request); } protected static void deletePipeline(String name) throws IOException { @@ -96,6 +110,89 @@ protected static void deletePipeline(String name) throws IOException { client().performRequest(request); } + protected Response createDatasource(final String name, Map properties) throws IOException { + XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); + for (Map.Entry config : properties.entrySet()) { + builder.field(config.getKey(), config.getValue()); + } + builder.endObject(); + + Request request = new Request(PUT, buildDatasourcePath(name)); + request.setJsonEntity(org.opensearch.common.Strings.toString(builder)); + return client().performRequest(request); + } + + protected void waitForDatasourceToBeAvailable(final String name, final Duration timeout) throws Exception { + Instant start = Instant.now(); + while (DatasourceState.AVAILABLE.equals(getDatasourceState(name)) == false) { + if (Duration.between(start, Instant.now()).compareTo(timeout) > 0) { + throw new RuntimeException( + String.format( + Locale.ROOT, + "Datasource state didn't change to %s after %d seconds", + DatasourceState.AVAILABLE.name(), + timeout.toSeconds() + ) + ); + } + Thread.sleep(1000); + } + } + + private DatasourceState getDatasourceState(final String name) throws Exception { + List> datasources = (List>) getDatasource(name).get(DATASOURCES); + return DatasourceState.valueOf((String) datasources.get(0).get(STATE)); + } + + protected Response deleteDatasource(final String name) throws IOException { + Request request = new Request(DELETE, buildDatasourcePath(name)); + return client().performRequest(request); + } + + protected Response deleteDatasource(final String name, final int retry) throws Exception { + for (int i = 0; i < retry; i++) { + try { + Request request = new Request(DELETE, buildDatasourcePath(name)); + return client().performRequest(request); + } catch (Exception e) { + if (i + 1 == retry) { + throw e; + } + Thread.sleep(1000); + } + } + throw new RuntimeException("should not reach here"); + } + + protected Map getDatasource(final String name) throws Exception { + Request request = new Request(GET, buildDatasourcePath(name)); + Response response = client().performRequest(request); + return createParser(XContentType.JSON.xContent(), EntityUtils.toString(response.getEntity())).map(); + } + + protected Response updateDatasource(final String name, Map properties) throws IOException { + XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); + for (Map.Entry config : properties.entrySet()) { + builder.field(config.getKey(), config.getValue()); + } + builder.endObject(); + + Request request = new Request(PUT, String.join(URL_DELIMITER, buildDatasourcePath(name), SETTINGS)); + request.setJsonEntity(org.opensearch.common.Strings.toString(builder)); + return client().performRequest(request); + } + + protected Map simulatePipeline(final String name, List docs) throws Exception { + XContentBuilder builder = XContentFactory.jsonBuilder().startObject(); + builder.field(DOCS, docs); + builder.endObject(); + + Request request = new Request(GET, String.join(URL_DELIMITER, buildPipelinePath(name), SIMULATE)); + request.setJsonEntity(org.opensearch.common.Strings.toString(builder)); + Response response = client().performRequest(request); + return createParser(XContentType.JSON.xContent(), EntityUtils.toString(response.getEntity())).map(); + } + protected static void createIndex(String name, Settings settings, Map fieldMap) throws IOException { XContentBuilder xContentBuilder = XContentFactory.jsonBuilder().startObject().startObject(MAPPING_PROPERTIES_KEY); for (Map.Entry entry : fieldMap.entrySet()) { @@ -137,9 +234,9 @@ public static String indexDocument(String indexName, String docID, String body, return docID; } - protected Map buildGeoJSONFeatureProcessorConfig(Map properties) { + protected Map buildProcessorConfig(final String processorType, final Map properties) { Map featureProcessor = new HashMap<>(); - featureProcessor.put(FeatureProcessor.TYPE, properties); + featureProcessor.put(processorType, properties); return featureProcessor; } diff --git a/src/test/java/org/opensearch/geospatial/exceptions/ConcurrentModificationExceptionTests.java b/src/test/java/org/opensearch/geospatial/exceptions/ConcurrentModificationExceptionTests.java new file mode 100644 index 0000000000..ef57b94ffc --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/exceptions/ConcurrentModificationExceptionTests.java @@ -0,0 +1,42 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.exceptions; + +import lombok.SneakyThrows; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.test.OpenSearchTestCase; + +public class ConcurrentModificationExceptionTests extends OpenSearchTestCase { + public void testConstructor_whenCreated_thenSucceed() { + ConcurrentModificationException exception = new ConcurrentModificationException("Resource is being modified by another processor"); + assertEquals(RestStatus.BAD_REQUEST, exception.status()); + } + + public void testConstructor_whenCreatedWithRootCause_thenSucceed() { + ConcurrentModificationException exception = new ConcurrentModificationException( + "Resource is being modified by another processor", + new RuntimeException() + ); + assertEquals(RestStatus.BAD_REQUEST, exception.status()); + } + + @SneakyThrows + public void testConstructor_whenCreatedWithStream_thenSucceed() { + ConcurrentModificationException exception = new ConcurrentModificationException( + "New datasource is not compatible with existing datasource" + ); + + BytesStreamOutput output = new BytesStreamOutput(); + exception.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + ConcurrentModificationException copiedException = new ConcurrentModificationException(input); + assertEquals(exception.getMessage(), copiedException.getMessage()); + assertEquals(exception.status(), copiedException.status()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceExceptionTests.java b/src/test/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceExceptionTests.java new file mode 100644 index 0000000000..69d026a700 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/exceptions/IncompatibleDatasourceExceptionTests.java @@ -0,0 +1,44 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.exceptions; + +import lombok.SneakyThrows; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.test.OpenSearchTestCase; + +public class IncompatibleDatasourceExceptionTests extends OpenSearchTestCase { + public void testConstructor_whenCreated_thenSucceed() { + IncompatibleDatasourceException exception = new IncompatibleDatasourceException( + "New datasource is not compatible with existing datasource" + ); + assertEquals(RestStatus.BAD_REQUEST, exception.status()); + } + + public void testConstructor_whenCreatedWithRootCause_thenSucceed() { + IncompatibleDatasourceException exception = new IncompatibleDatasourceException( + "New datasource is not compatible with existing datasource", + new RuntimeException() + ); + assertEquals(RestStatus.BAD_REQUEST, exception.status()); + } + + @SneakyThrows + public void testConstructor_whenCreatedWithStream_thenSucceed() { + IncompatibleDatasourceException exception = new IncompatibleDatasourceException( + "New datasource is not compatible with existing datasource" + ); + + BytesStreamOutput output = new BytesStreamOutput(); + exception.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + IncompatibleDatasourceException copiedException = new IncompatibleDatasourceException(input); + assertEquals(exception.getMessage(), copiedException.getMessage()); + assertEquals(exception.status(), copiedException.status()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/exceptions/ResourceInUseExceptionTests.java b/src/test/java/org/opensearch/geospatial/exceptions/ResourceInUseExceptionTests.java new file mode 100644 index 0000000000..f63cdbc044 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/exceptions/ResourceInUseExceptionTests.java @@ -0,0 +1,37 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.exceptions; + +import lombok.SneakyThrows; + +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.test.OpenSearchTestCase; + +public class ResourceInUseExceptionTests extends OpenSearchTestCase { + public void testConstructor_whenCreated_thenSucceed() { + ResourceInUseException exception = new ResourceInUseException("Resource is in use"); + assertEquals(RestStatus.BAD_REQUEST, exception.status()); + } + + public void testConstructor_whenCreatedWithRootCause_thenSucceed() { + ResourceInUseException exception = new ResourceInUseException("Resource is in use", new RuntimeException()); + assertEquals(RestStatus.BAD_REQUEST, exception.status()); + } + + @SneakyThrows + public void testConstructor_whenCreatedWithStream_thenSucceed() { + ResourceInUseException exception = new ResourceInUseException("New datasource is not compatible with existing datasource"); + + BytesStreamOutput output = new BytesStreamOutput(); + exception.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + ResourceInUseException copiedException = new ResourceInUseException(input); + assertEquals(exception.getMessage(), copiedException.getMessage()); + assertEquals(exception.status(), copiedException.status()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoDataServer.java b/src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoDataServer.java new file mode 100644 index 0000000000..ba1e49098a --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoDataServer.java @@ -0,0 +1,127 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.nio.file.Files; +import java.nio.file.Paths; + +import lombok.SneakyThrows; +import lombok.extern.log4j.Log4j2; + +import org.opensearch.common.SuppressForbidden; + +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +/** + * Simple http server to serve static files under test/java/resources/ip2geo/server for integration testing + */ +@Log4j2 +@SuppressForbidden(reason = "used only for testing") +public class Ip2GeoDataServer { + private static final String SYS_PROPERTY_KEY_CLUSTER_ENDPOINT = "tests.rest.cluster"; + private static final String LOCAL_CLUSTER_ENDPOINT = "127.0.0.1"; + private static final String ROOT = "ip2geo/server"; + private static final int PORT = 8001; + private static final String EXTERNAL_ENDPOINT_PREFIX = + "https://github.com/opensearch-project/geospatial/blob/main/src/test/resources/ip2geo/server"; + + private static HttpServer server; + private static volatile int counter = 0; + private static String endpointPrefix = "http://localhost:" + PORT; + private static String cityFilePath = endpointPrefix + "/city/manifest_local.json"; + private static String countryFilePath = endpointPrefix + "/country/manifest_local.json"; + + /** + * Return an endpoint to a manifest file for a sample city data + * The sample data should contain three lines as follows + * + * cidr,city,country + * 10.0.0.0/8,Seattle,USA + * 127.0.0.0/12,Vancouver,Canada + * fd12:2345:6789:1::/64,Bengaluru,India + * + */ + public static String getEndpointCity() { + return cityFilePath; + } + + /** + * Return an endpoint to a manifest file for a sample country data + * The sample data should contain three lines as follows + * + * cidr,country + * 10.0.0.0/8,USA + * 127.0.0.0/12,Canada + * fd12:2345:6789:1::/64,India + * + */ + public static String getEndpointCountry() { + return countryFilePath; + } + + @SneakyThrows + synchronized public static void start() { + log.info("Start server is called"); + // If it is remote cluster test, use external endpoint and do not launch local server + if (System.getProperty(SYS_PROPERTY_KEY_CLUSTER_ENDPOINT).contains(LOCAL_CLUSTER_ENDPOINT) == false) { + log.info("Remote cluster[{}] testing. Skip launching local server", System.getProperty(SYS_PROPERTY_KEY_CLUSTER_ENDPOINT)); + cityFilePath = EXTERNAL_ENDPOINT_PREFIX + "/city/manifest.json"; + countryFilePath = EXTERNAL_ENDPOINT_PREFIX + "/country/manifest.json"; + return; + } + + counter++; + if (server != null) { + log.info("Server has started already"); + return; + } + server = HttpServer.create(new InetSocketAddress("localhost", PORT), 0); + server.createContext("/", new Ip2GeoHttpHandler()); + server.start(); + log.info("Local file server started on port {}", PORT); + } + + synchronized public static void stop() { + log.info("Stop server is called"); + if (server == null) { + log.info("Server has stopped already"); + return; + } + counter--; + if (counter > 0) { + log.info("[{}] processors are still using the server", counter); + return; + } + + server.stop(0); + server = null; + log.info("Server stopped"); + } + + @SuppressForbidden(reason = "used only for testing") + private static class Ip2GeoHttpHandler implements HttpHandler { + @Override + public void handle(final HttpExchange exchange) throws IOException { + try { + byte[] data = Files.readAllBytes( + Paths.get(this.getClass().getClassLoader().getResource(ROOT + exchange.getRequestURI().getPath()).toURI()) + ); + exchange.sendResponseHeaders(200, data.length); + OutputStream outputStream = exchange.getResponseBody(); + outputStream.write(data); + outputStream.flush(); + outputStream.close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + } +} 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 0000000000..7bf2961fbc --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/Ip2GeoTestCase.java @@ -0,0 +1,355 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Locale; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; +import java.util.stream.Collectors; + +import lombok.SneakyThrows; + +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.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.dao.GeoIpDataDao; +import org.opensearch.geospatial.ip2geo.dao.Ip2GeoCachedDao; +import org.opensearch.geospatial.ip2geo.dao.Ip2GeoProcessorDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceTask; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; +import org.opensearch.geospatial.ip2geo.processor.Ip2GeoProcessor; +import org.opensearch.ingest.IngestMetadata; +import org.opensearch.ingest.IngestService; +import org.opensearch.jobscheduler.spi.LockModel; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; +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 DatasourceDao datasourceDao; + @Mock + protected Ip2GeoExecutor ip2GeoExecutor; + @Mock + protected GeoIpDataDao geoIpDataDao; + @Mock + protected Ip2GeoCachedDao ip2GeoCachedDao; + @Mock + protected ClusterState clusterState; + @Mock + protected Metadata metadata; + @Mock + protected IngestService ingestService; + @Mock + protected ActionFilters actionFilters; + @Mock + protected ThreadPool threadPool; + @Mock + protected TransportService transportService; + @Mock + protected Ip2GeoLockService ip2GeoLockService; + @Mock + protected Ip2GeoProcessorDao ip2GeoProcessorDao; + @Mock + protected RoutingTable routingTable; + protected IngestMetadata ingestMetadata; + 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); + ingestMetadata = new IngestMetadata(Collections.emptyMap()); + when(metadata.custom(IngestMetadata.TYPE)).thenReturn(ingestMetadata); + when(clusterService.getSettings()).thenReturn(Settings.EMPTY); + when(clusterService.getClusterSettings()).thenReturn(clusterSettings); + when(clusterService.state()).thenReturn(clusterState); + when(clusterState.metadata()).thenReturn(metadata); + when(clusterState.getMetadata()).thenReturn(metadata); + when(clusterState.routingTable()).thenReturn(routingTable); + when(ip2GeoExecutor.forDatasourceUpdate()).thenReturn(OpenSearchExecutors.newDirectExecutorService()); + when(ingestService.getClusterService()).thenReturn(clusterService); + when(threadPool.generic()).thenReturn(OpenSearchExecutors.newDirectExecutorService()); + } + + @After + public void clean() throws Exception { + openMocks.close(); + client.close(); + verifyingClient.close(); + } + + protected 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)); + } + + protected DatasourceState randomState() { + return Arrays.stream(DatasourceState.values()) + .sequential() + .collect(Collectors.toList()) + .get(Randomness.createSecure().nextInt(DatasourceState.values().length - 1)); + } + + protected DatasourceTask randomTask() { + return Arrays.stream(DatasourceTask.values()) + .sequential() + .collect(Collectors.toList()) + .get(Randomness.createSecure().nextInt(DatasourceTask.values().length - 1)); + } + + protected 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) + ); + } + + @SneakyThrows + @SuppressForbidden(reason = "unit test") + protected String sampleManifestUrl() { + return Paths.get(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").toURI()).toUri().toURL().toExternalForm(); + } + + @SuppressForbidden(reason = "unit test") + protected String sampleManifestUrlWithInvalidUrl() throws Exception { + return Paths.get(this.getClass().getClassLoader().getResource("ip2geo/manifest_invalid_url.json").toURI()) + .toUri() + .toURL() + .toExternalForm(); + } + + @SuppressForbidden(reason = "unit test") + protected File sampleIp2GeoFile() { + return new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.csv").getFile()); + } + + protected long randomPositiveLong() { + long value = Randomness.get().nextLong(); + return value < 0 ? -value : value; + } + + /** + * Update interval should be > 0 and < validForInDays. + * For an update test to work, there should be at least one eligible value other than current update interval. + * Therefore, the smallest value for validForInDays is 2. + * Update interval is random value from 1 to validForInDays - 2. + * The new update value will be validForInDays - 1. + */ + protected Datasource randomDatasource(final Instant updateStartTime) { + int validForInDays = 3 + Randomness.get().nextInt(30); + Instant now = Instant.now().truncatedTo(ChronoUnit.MILLIS); + Datasource datasource = new Datasource(); + datasource.setName(GeospatialTestHelper.randomLowerCaseString()); + datasource.setUserSchedule( + new IntervalSchedule( + updateStartTime.truncatedTo(ChronoUnit.MILLIS), + 1 + Randomness.get().nextInt(validForInDays - 2), + ChronoUnit.DAYS + ) + ); + datasource.setSystemSchedule(datasource.getUserSchedule()); + datasource.setTask(randomTask()); + datasource.setState(randomState()); + datasource.setCurrentIndex(datasource.newIndexName(UUID.randomUUID().toString())); + datasource.setIndices(Arrays.asList(GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString())); + datasource.setEndpoint(String.format(Locale.ROOT, "https://%s.com/manifest.json", GeospatialTestHelper.randomLowerCaseString())); + datasource.getDatabase() + .setFields(Arrays.asList(GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString())); + datasource.getDatabase().setProvider(GeospatialTestHelper.randomLowerCaseString()); + datasource.getDatabase().setUpdatedAt(now); + datasource.getDatabase().setSha256Hash(GeospatialTestHelper.randomLowerCaseString()); + datasource.getDatabase().setValidForInDays((long) validForInDays); + datasource.getUpdateStats().setLastSkippedAt(now); + datasource.getUpdateStats().setLastSucceededAt(now); + datasource.getUpdateStats().setLastFailedAt(now); + datasource.getUpdateStats().setLastProcessingTimeInMillis(randomPositiveLong()); + datasource.setLastUpdateTime(now); + if (Randomness.get().nextInt() % 2 == 0) { + datasource.enable(); + } else { + datasource.disable(); + } + return datasource; + } + + protected Datasource randomDatasource() { + return randomDatasource(Instant.now()); + } + + protected LockModel randomLockModel() { + LockModel lockModel = new LockModel( + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString(), + Instant.now(), + randomPositiveLong(), + false + ); + return lockModel; + } + + protected Ip2GeoProcessor randomIp2GeoProcessor(String datasourceName) { + String tag = GeospatialTestHelper.randomLowerCaseString(); + String description = GeospatialTestHelper.randomLowerCaseString(); + String field = GeospatialTestHelper.randomLowerCaseString(); + String targetField = GeospatialTestHelper.randomLowerCaseString(); + Set properties = Set.of(GeospatialTestHelper.randomLowerCaseString()); + Ip2GeoProcessor ip2GeoProcessor = new Ip2GeoProcessor( + tag, + description, + field, + targetField, + datasourceName, + properties, + true, + clusterSettings, + datasourceDao, + geoIpDataDao, + ip2GeoCachedDao + ); + return ip2GeoProcessor; + } + + /** + * 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 + ) { + try { + listener.onResponse((Response) executeVerifier.get().apply(action, request)); + } catch (Exception e) { + listener.onFailure(e); + } + } + + /** + * 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/DeleteDatasourceRequestTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceRequestTests.java new file mode 100644 index 0000000000..a3bc17fa1b --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceRequestTests.java @@ -0,0 +1,53 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import lombok.SneakyThrows; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; + +public class DeleteDatasourceRequestTests extends Ip2GeoTestCase { + @SneakyThrows + public void testStreamInOut_whenValidInput_thenSucceed() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + DeleteDatasourceRequest request = new DeleteDatasourceRequest(datasourceName); + + // Run + BytesStreamOutput output = new BytesStreamOutput(); + request.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + DeleteDatasourceRequest copiedRequest = new DeleteDatasourceRequest(input); + + // Verify + assertEquals(request.getName(), copiedRequest.getName()); + } + + public void testValidate_whenNull_thenError() { + DeleteDatasourceRequest request = new DeleteDatasourceRequest((String) null); + + // Run + ActionRequestValidationException error = request.validate(); + + // Verify + assertNotNull(error.validationErrors()); + assertFalse(error.validationErrors().isEmpty()); + } + + public void testValidate_whenBlank_thenError() { + DeleteDatasourceRequest request = new DeleteDatasourceRequest(" "); + + // Run + ActionRequestValidationException error = request.validate(); + + // Verify + assertNotNull(error.validationErrors()); + assertFalse(error.validationErrors().isEmpty()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceTransportActionTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceTransportActionTests.java new file mode 100644 index 0000000000..3abf3c9da4 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/DeleteDatasourceTransportActionTests.java @@ -0,0 +1,165 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +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 java.io.IOException; +import java.time.Instant; +import java.util.Arrays; +import java.util.Collections; + +import lombok.SneakyThrows; + +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.mockito.InOrder; +import org.mockito.Mockito; +import org.opensearch.OpenSearchException; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.support.master.AcknowledgedResponse; +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.jobscheduler.spi.LockModel; +import org.opensearch.tasks.Task; + +public class DeleteDatasourceTransportActionTests extends Ip2GeoTestCase { + private DeleteDatasourceTransportAction action; + + @Before + public void init() { + action = new DeleteDatasourceTransportAction( + transportService, + actionFilters, + ip2GeoLockService, + ingestService, + datasourceDao, + geoIpDataDao, + ip2GeoProcessorDao, + threadPool + ); + } + + @SneakyThrows + public void testDoExecute_whenFailedToAcquireLock_thenError() { + validateDoExecute(null, null); + } + + @SneakyThrows + public void testDoExecute_whenValidInput_thenSucceed() { + String jobIndexName = GeospatialTestHelper.randomLowerCaseString(); + String jobId = GeospatialTestHelper.randomLowerCaseString(); + LockModel lockModel = new LockModel(jobIndexName, jobId, Instant.now(), randomPositiveLong(), false); + validateDoExecute(lockModel, null); + } + + @SneakyThrows + public void testDoExecute_whenException_thenError() { + validateDoExecute(null, new RuntimeException()); + } + + private void validateDoExecute(final LockModel lockModel, final Exception exception) throws IOException { + Task task = mock(Task.class); + Datasource datasource = randomDatasource(); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + DeleteDatasourceRequest request = new DeleteDatasourceRequest(datasource.getName()); + ActionListener listener = mock(ActionListener.class); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + if (exception == null) { + // Run + captor.getValue().onResponse(lockModel); + + // Verify + if (lockModel == null) { + verify(listener).onFailure(any(OpenSearchException.class)); + } else { + verify(listener).onResponse(new AcknowledgedResponse(true)); + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + } + } else { + // Run + captor.getValue().onFailure(exception); + // Verify + verify(listener).onFailure(exception); + } + } + + @SneakyThrows + public void testDeleteDatasource_whenNull_thenThrowException() { + Datasource datasource = randomDatasource(); + expectThrows(ResourceNotFoundException.class, () -> action.deleteDatasource(datasource.getName())); + } + + @SneakyThrows + public void testDeleteDatasource_whenSafeToDelete_thenDelete() { + Datasource datasource = randomDatasource(); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + when(ip2GeoProcessorDao.getProcessors(datasource.getName())).thenReturn(Collections.emptyList()); + + // Run + action.deleteDatasource(datasource.getName()); + + // Verify + assertEquals(DatasourceState.DELETING, datasource.getState()); + verify(datasourceDao).updateDatasource(datasource); + InOrder inOrder = Mockito.inOrder(geoIpDataDao, datasourceDao); + inOrder.verify(geoIpDataDao).deleteIp2GeoDataIndex(datasource.getIndices()); + inOrder.verify(datasourceDao).deleteDatasource(datasource); + } + + @SneakyThrows + public void testDeleteDatasource_whenProcessorIsUsingDatasource_thenThrowException() { + Datasource datasource = randomDatasource(); + datasource.setState(DatasourceState.AVAILABLE); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + when(ip2GeoProcessorDao.getProcessors(datasource.getName())).thenReturn(Arrays.asList(randomIp2GeoProcessor(datasource.getName()))); + + // Run + expectThrows(OpenSearchException.class, () -> action.deleteDatasource(datasource.getName())); + + // Verify + assertEquals(DatasourceState.AVAILABLE, datasource.getState()); + verify(datasourceDao, never()).updateDatasource(datasource); + verify(geoIpDataDao, never()).deleteIp2GeoDataIndex(datasource.getIndices()); + verify(datasourceDao, never()).deleteDatasource(datasource); + } + + @SneakyThrows + public void testDeleteDatasource_whenProcessorIsCreatedDuringDeletion_thenThrowException() { + Datasource datasource = randomDatasource(); + datasource.setState(DatasourceState.AVAILABLE); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + when(ip2GeoProcessorDao.getProcessors(datasource.getName())).thenReturn( + Collections.emptyList(), + Arrays.asList(randomIp2GeoProcessor(datasource.getName())) + ); + + // Run + expectThrows(OpenSearchException.class, () -> action.deleteDatasource(datasource.getName())); + + // Verify + verify(datasourceDao, times(2)).updateDatasource(datasource); + verify(geoIpDataDao, never()).deleteIp2GeoDataIndex(datasource.getIndices()); + verify(datasourceDao, never()).deleteDatasource(datasource); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceRequestTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceRequestTests.java new file mode 100644 index 0000000000..7ee19c6362 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceRequestTests.java @@ -0,0 +1,55 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; + +public class GetDatasourceRequestTests extends Ip2GeoTestCase { + public void testStreamInOut_whenEmptyNames_thenSucceed() throws Exception { + String[] names = new String[0]; + GetDatasourceRequest request = new GetDatasourceRequest(names); + assertNull(request.validate()); + + // Run + BytesStreamOutput output = new BytesStreamOutput(); + request.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + GetDatasourceRequest copiedRequest = new GetDatasourceRequest(input); + + // Verify + assertArrayEquals(request.getNames(), copiedRequest.getNames()); + } + + public void testStreamInOut_whenNames_thenSucceed() throws Exception { + String[] names = { GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString() }; + GetDatasourceRequest request = new GetDatasourceRequest(names); + assertNull(request.validate()); + + // Run + BytesStreamOutput output = new BytesStreamOutput(); + request.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + GetDatasourceRequest copiedRequest = new GetDatasourceRequest(input); + + // Verify + assertArrayEquals(request.getNames(), copiedRequest.getNames()); + } + + public void testValidate_whenNull_thenError() { + GetDatasourceRequest request = new GetDatasourceRequest((String[]) null); + + // Run + ActionRequestValidationException error = request.validate(); + + // Verify + assertNotNull(error.validationErrors()); + assertFalse(error.validationErrors().isEmpty()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceResponseTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceResponseTests.java new file mode 100644 index 0000000000..f1efad9375 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceResponseTests.java @@ -0,0 +1,98 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.util.Arrays; +import java.util.List; +import java.util.Locale; + +import org.opensearch.common.Strings; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; + +public class GetDatasourceResponseTests extends Ip2GeoTestCase { + + public void testStreamInOut_whenValidInput_thenSucceed() throws Exception { + List datasourceList = Arrays.asList(randomDatasource(), randomDatasource()); + GetDatasourceResponse response = new GetDatasourceResponse(datasourceList); + + // Run + BytesStreamOutput output = new BytesStreamOutput(); + response.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + GetDatasourceResponse copiedResponse = new GetDatasourceResponse(input); + + // Verify + assertArrayEquals(response.getDatasources().toArray(), copiedResponse.getDatasources().toArray()); + } + + public void testToXContent_whenValidInput_thenSucceed() throws Exception { + List datasourceList = Arrays.asList(randomDatasource(), randomDatasource()); + GetDatasourceResponse response = new GetDatasourceResponse(datasourceList); + String json = Strings.toString(response.toXContent(JsonXContent.contentBuilder(), null)); + for (Datasource datasource : datasourceList) { + assertTrue(json.contains(String.format(Locale.ROOT, "\"name\":\"%s\"", datasource.getName()))); + assertTrue(json.contains(String.format(Locale.ROOT, "\"state\":\"%s\"", datasource.getState()))); + assertTrue(json.contains(String.format(Locale.ROOT, "\"endpoint\":\"%s\"", datasource.getEndpoint()))); + assertTrue( + json.contains(String.format(Locale.ROOT, "\"update_interval_in_days\":%d", datasource.getUserSchedule().getInterval())) + ); + assertTrue(json.contains(String.format(Locale.ROOT, "\"next_update_at_in_epoch_millis\""))); + assertTrue(json.contains(String.format(Locale.ROOT, "\"provider\":\"%s\"", datasource.getDatabase().getProvider()))); + assertTrue(json.contains(String.format(Locale.ROOT, "\"sha256_hash\":\"%s\"", datasource.getDatabase().getSha256Hash()))); + assertTrue( + json.contains( + String.format(Locale.ROOT, "\"updated_at_in_epoch_millis\":%d", datasource.getDatabase().getUpdatedAt().toEpochMilli()) + ) + ); + assertTrue(json.contains(String.format(Locale.ROOT, "\"valid_for_in_days\":%d", datasource.getDatabase().getValidForInDays()))); + for (String field : datasource.getDatabase().getFields()) { + assertTrue(json.contains(field)); + } + assertTrue( + json.contains( + String.format( + Locale.ROOT, + "\"last_succeeded_at_in_epoch_millis\":%d", + datasource.getUpdateStats().getLastSucceededAt().toEpochMilli() + ) + ) + ); + assertTrue( + json.contains( + String.format( + Locale.ROOT, + "\"last_processing_time_in_millis\":%d", + datasource.getUpdateStats().getLastProcessingTimeInMillis() + ) + ) + ); + assertTrue( + json.contains( + String.format( + Locale.ROOT, + "\"last_failed_at_in_epoch_millis\":%d", + datasource.getUpdateStats().getLastFailedAt().toEpochMilli() + ) + ) + ); + assertTrue( + json.contains( + String.format( + Locale.ROOT, + "\"last_skipped_at_in_epoch_millis\":%d", + datasource.getUpdateStats().getLastSkippedAt().toEpochMilli() + ) + ) + ); + + } + } + +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceTransportActionTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceTransportActionTests.java new file mode 100644 index 0000000000..581fd3b846 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/GetDatasourceTransportActionTests.java @@ -0,0 +1,102 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import org.junit.Before; +import org.opensearch.OpenSearchException; +import org.opensearch.action.ActionListener; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.index.IndexNotFoundException; +import org.opensearch.tasks.Task; + +public class GetDatasourceTransportActionTests extends Ip2GeoTestCase { + private GetDatasourceTransportAction action; + + @Before + public void init() { + action = new GetDatasourceTransportAction(transportService, actionFilters, datasourceDao); + } + + public void testDoExecute_whenAll_thenSucceed() { + Task task = mock(Task.class); + GetDatasourceRequest request = new GetDatasourceRequest(new String[] { "_all" }); + ActionListener listener = mock(ActionListener.class); + + // Run + action.doExecute(task, request, listener); + + // Verify + verify(datasourceDao).getAllDatasources(any(ActionListener.class)); + } + + public void testDoExecute_whenNames_thenSucceed() { + Task task = mock(Task.class); + List datasources = Arrays.asList(randomDatasource(), randomDatasource()); + String[] datasourceNames = datasources.stream().map(Datasource::getName).toArray(String[]::new); + + GetDatasourceRequest request = new GetDatasourceRequest(datasourceNames); + ActionListener listener = mock(ActionListener.class); + + // Run + action.doExecute(task, request, listener); + + // Verify + verify(datasourceDao).getDatasources(eq(datasourceNames), any(ActionListener.class)); + } + + public void testDoExecute_whenNull_thenException() { + Task task = mock(Task.class); + GetDatasourceRequest request = new GetDatasourceRequest((String[]) null); + ActionListener listener = mock(ActionListener.class); + + // Run + Exception exception = expectThrows(OpenSearchException.class, () -> action.doExecute(task, request, listener)); + + // Verify + assertTrue(exception.getMessage().contains("should not be null")); + } + + public void testNewActionListener_whenOnResponse_thenSucceed() { + List datasources = Arrays.asList(randomDatasource(), randomDatasource()); + ActionListener actionListener = mock(ActionListener.class); + + // Run + action.newActionListener(actionListener).onResponse(datasources); + + // Verify + verify(actionListener).onResponse(new GetDatasourceResponse(datasources)); + } + + public void testNewActionListener_whenOnFailureWithNoSuchIndexException_thenEmptyDatasource() { + ActionListener actionListener = mock(ActionListener.class); + + // Run + action.newActionListener(actionListener).onFailure(new IndexNotFoundException("no index")); + + // Verify + verify(actionListener).onResponse(new GetDatasourceResponse(Collections.emptyList())); + } + + public void testNewActionListener_whenOnFailure_thenFails() { + ActionListener actionListener = mock(ActionListener.class); + + // Run + action.newActionListener(actionListener).onFailure(new RuntimeException()); + + // Verify + verify(actionListener).onFailure(any(RuntimeException.class)); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequestTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequestTests.java new file mode 100644 index 0000000000..b182b3c1b2 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceRequestTests.java @@ -0,0 +1,171 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.util.Arrays; +import java.util.Locale; +import java.util.Map; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.Randomness; +import org.opensearch.common.Strings; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; + +public class PutDatasourceRequestTests extends Ip2GeoTestCase { + + public void testValidate_whenInvalidUrl_thenFails() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint("invalidUrl"); + 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 testValidate_whenInvalidManifestFile_thenFails() { + 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)); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertTrue(exception.validationErrors().get(0).contains("Error occurred while reading a file")); + } + + public void testValidate_whenValidInput_thenSucceed() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(datasourceName); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(1)); + assertNull(request.validate()); + } + + public void testValidate_whenZeroUpdateInterval_thenFails() 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 testValidate_whenLargeUpdateInterval_thenFail() 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 testValidate_whenInvalidUrlInsideManifest_thenFail() 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 testValidate_whenInvalidDatasourceNames_thenFails() throws Exception { + String validDatasourceName = GeospatialTestHelper.randomLowerCaseString(); + String domain = GeospatialTestHelper.randomLowerCaseString(); + PutDatasourceRequest request = new PutDatasourceRequest(validDatasourceName); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(Randomness.get().nextInt(10) + 1)); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertNull(exception); + + String fileNameChar = validDatasourceName + Strings.INVALID_FILENAME_CHARS.stream() + .skip(Randomness.get().nextInt(Strings.INVALID_FILENAME_CHARS.size() - 1)) + .findFirst(); + String startsWith = Arrays.asList("_", "-", "+").get(Randomness.get().nextInt(3)) + validDatasourceName; + String empty = ""; + String hash = validDatasourceName + "#"; + String colon = validDatasourceName + ":"; + StringBuilder longName = new StringBuilder(); + while (longName.length() < 256) { + longName.append(GeospatialTestHelper.randomLowerCaseString()); + } + String point = Arrays.asList(".", "..").get(Randomness.get().nextInt(2)); + Map nameToError = Map.of( + fileNameChar, + "not contain the following characters", + empty, + "must not be empty", + hash, + "must not contain '#'", + colon, + "must not contain ':'", + startsWith, + "must not start with", + longName.toString(), + "name is too long", + point, + "must not be '.' or '..'" + ); + + for (Map.Entry entry : nameToError.entrySet()) { + request.setName(entry.getKey()); + + // Run + exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertTrue(exception.validationErrors().get(0).contains(entry.getValue())); + } + } + + public void testStreamInOut_whenValidInput_thenSucceed() 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(29) + 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 0000000000..ef426cc568 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/PutDatasourceTransportActionTests.java @@ -0,0 +1,195 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +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 java.io.IOException; + +import lombok.SneakyThrows; + +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.opensearch.ResourceAlreadyExistsException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.StepListener; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.exceptions.ConcurrentModificationException; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.index.engine.VersionConflictEngineException; +import org.opensearch.jobscheduler.spi.LockModel; +import org.opensearch.tasks.Task; + +public class PutDatasourceTransportActionTests extends Ip2GeoTestCase { + private PutDatasourceTransportAction action; + + @Before + public void init() { + action = new PutDatasourceTransportAction( + transportService, + actionFilters, + threadPool, + datasourceDao, + datasourceUpdateService, + ip2GeoLockService + ); + } + + @SneakyThrows + public void testDoExecute_whenFailedToAcquireLock_thenError() { + validateDoExecute(null, null, null); + } + + @SneakyThrows + public void testDoExecute_whenAcquiredLock_thenSucceed() { + validateDoExecute(randomLockModel(), null, null); + } + + @SneakyThrows + public void testDoExecute_whenExceptionBeforeAcquiringLock_thenError() { + validateDoExecute(randomLockModel(), new RuntimeException(), null); + } + + @SneakyThrows + public void testDoExecute_whenExceptionAfterAcquiringLock_thenError() { + validateDoExecute(randomLockModel(), null, new RuntimeException()); + } + + private void validateDoExecute(final LockModel lockModel, final Exception before, final Exception after) throws IOException { + Task task = mock(Task.class); + Datasource datasource = randomDatasource(); + PutDatasourceRequest request = new PutDatasourceRequest(datasource.getName()); + ActionListener listener = mock(ActionListener.class); + if (after != null) { + doThrow(after).when(datasourceDao).createIndexIfNotExists(any(StepListener.class)); + } + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + if (before == null) { + // Run + captor.getValue().onResponse(lockModel); + + // Verify + if (lockModel == null) { + verify(listener).onFailure(any(ConcurrentModificationException.class)); + } + if (after != null) { + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + verify(listener).onFailure(after); + } else { + verify(ip2GeoLockService, never()).releaseLock(eq(lockModel)); + } + } else { + // Run + captor.getValue().onFailure(before); + // Verify + verify(listener).onFailure(before); + } + } + + @SneakyThrows + public void testInternalDoExecute_whenValidInput_thenSucceed() { + PutDatasourceRequest request = new PutDatasourceRequest(GeospatialTestHelper.randomLowerCaseString()); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(1)); + ActionListener listener = mock(ActionListener.class); + + // Run + action.internalDoExecute(request, randomLockModel(), listener); + + // Verify + ArgumentCaptor captor = ArgumentCaptor.forClass(StepListener.class); + verify(datasourceDao).createIndexIfNotExists(captor.capture()); + + // Run + captor.getValue().onResponse(null); + // Verify + ArgumentCaptor datasourceCaptor = ArgumentCaptor.forClass(Datasource.class); + ArgumentCaptor actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class); + verify(datasourceDao).putDatasource(datasourceCaptor.capture(), actionListenerCaptor.capture()); + assertEquals(request.getName(), datasourceCaptor.getValue().getName()); + assertEquals(request.getEndpoint(), datasourceCaptor.getValue().getEndpoint()); + assertEquals(request.getUpdateInterval().days(), datasourceCaptor.getValue().getUserSchedule().getInterval()); + + // Run next listener.onResponse + actionListenerCaptor.getValue().onResponse(null); + // Verify + verify(listener).onResponse(new AcknowledgedResponse(true)); + } + + public void testGetIndexResponseListener_whenVersionConflict_thenFailure() { + Datasource datasource = new Datasource(); + ActionListener listener = mock(ActionListener.class); + action.getIndexResponseListener(datasource, randomLockModel(), listener) + .onFailure( + new VersionConflictEngineException( + null, + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString() + ) + ); + verify(listener).onFailure(any(ResourceAlreadyExistsException.class)); + } + + @SneakyThrows + public void testCreateDatasource_whenInvalidState_thenUpdateStateAsFailed() { + Datasource datasource = new Datasource(); + datasource.setState(randomStateExcept(DatasourceState.CREATING)); + datasource.getUpdateStats().setLastFailedAt(null); + + // Run + action.createDatasource(datasource, mock(Runnable.class)); + + // Verify + assertEquals(DatasourceState.CREATE_FAILED, datasource.getState()); + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceDao).updateDatasource(datasource); + verify(datasourceUpdateService, never()).updateOrCreateGeoIpData(any(Datasource.class), any(Runnable.class)); + } + + @SneakyThrows + public void testCreateDatasource_whenExceptionHappens_thenUpdateStateAsFailed() { + Datasource datasource = new Datasource(); + doThrow(new RuntimeException()).when(datasourceUpdateService).updateOrCreateGeoIpData(any(Datasource.class), any(Runnable.class)); + + // Run + action.createDatasource(datasource, mock(Runnable.class)); + + // Verify + assertEquals(DatasourceState.CREATE_FAILED, datasource.getState()); + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceDao).updateDatasource(datasource); + } + + @SneakyThrows + public void testCreateDatasource_whenValidInput_thenUpdateStateAsCreating() { + Datasource datasource = new Datasource(); + + Runnable renewLock = mock(Runnable.class); + // Run + action.createDatasource(datasource, renewLock); + + // Verify + verify(datasourceUpdateService).updateOrCreateGeoIpData(datasource, renewLock); + assertEquals(DatasourceState.CREATING, datasource.getState()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/RestDeleteDatasourceHandlerTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestDeleteDatasourceHandlerTests.java new file mode 100644 index 0000000000..937c5532ed --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestDeleteDatasourceHandlerTests.java @@ -0,0 +1,49 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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.Locale; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.junit.Before; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.rest.RestRequest; +import org.opensearch.test.rest.FakeRestRequest; +import org.opensearch.test.rest.RestActionTestCase; + +public class RestDeleteDatasourceHandlerTests extends RestActionTestCase { + private String path; + private RestDeleteDatasourceHandler action; + + @Before + public void setupAction() { + action = new RestDeleteDatasourceHandler(); + controller().registerHandler(action); + path = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/%s"); + } + + public void testPrepareRequest_whenValidInput_thenSucceed() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.DELETE) + .withPath(String.format(Locale.ROOT, path, datasourceName)) + .build(); + AtomicBoolean isExecuted = new AtomicBoolean(false); + + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof DeleteDatasourceRequest); + DeleteDatasourceRequest deleteDatasourceRequest = (DeleteDatasourceRequest) actionRequest; + assertEquals(datasourceName, deleteDatasourceRequest.getName()); + isExecuted.set(true); + return null; + }); + + dispatchRequest(request); + assertTrue(isExecuted.get()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/RestGetDatasourceHandlerTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestGetDatasourceHandlerTests.java new file mode 100644 index 0000000000..e1177da6a5 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestGetDatasourceHandlerTests.java @@ -0,0 +1,78 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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.Locale; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.apache.commons.lang3.StringUtils; +import org.junit.Before; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.rest.RestRequest; +import org.opensearch.test.rest.FakeRestRequest; +import org.opensearch.test.rest.RestActionTestCase; + +public class RestGetDatasourceHandlerTests extends RestActionTestCase { + private String PATH_FOR_ALL = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource"); + private String path; + private RestGetDatasourceHandler action; + + @Before + public void setupAction() { + action = new RestGetDatasourceHandler(); + controller().registerHandler(action); + path = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/%s"); + } + + public void testPrepareRequest_whenNames_thenSucceed() { + String dsName1 = GeospatialTestHelper.randomLowerCaseString(); + String dsName2 = GeospatialTestHelper.randomLowerCaseString(); + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.GET) + .withPath(String.format(Locale.ROOT, path, StringUtils.joinWith(",", dsName1, dsName2))) + .build(); + + AtomicBoolean isExecuted = new AtomicBoolean(false); + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { + // Verifying + assertTrue(actionRequest instanceof GetDatasourceRequest); + GetDatasourceRequest getDatasourceRequest = (GetDatasourceRequest) actionRequest; + assertArrayEquals(new String[] { dsName1, dsName2 }, getDatasourceRequest.getNames()); + isExecuted.set(true); + return null; + }); + + // Run + dispatchRequest(request); + + // Verify + assertTrue(isExecuted.get()); + } + + public void testPrepareRequest_whenAll_thenSucceed() { + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.GET) + .withPath(PATH_FOR_ALL) + .build(); + + AtomicBoolean isExecuted = new AtomicBoolean(false); + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { + // Verifying + assertTrue(actionRequest instanceof GetDatasourceRequest); + GetDatasourceRequest getDatasourceRequest = (GetDatasourceRequest) actionRequest; + assertArrayEquals(new String[] {}, getDatasourceRequest.getNames()); + isExecuted.set(true); + return null; + }); + + // Run + dispatchRequest(request); + + // Verify + assertTrue(isExecuted.get()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandlerTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandlerTests.java new file mode 100644 index 0000000000..ecbba86539 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestPutDatasourceHandlerTests.java @@ -0,0 +1,83 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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.Locale; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.junit.Before; +import org.opensearch.common.SuppressForbidden; +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.core.common.bytes.BytesArray; +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(new ClusterSettings(Settings.EMPTY, new HashSet(Ip2GeoSettings.settings()))); + controller().registerHandler(action); + path = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/%s"); + } + + public void testPrepareRequest() { + 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(Locale.ROOT, path, datasourceName)) + .withContent(new BytesArray(content), XContentType.JSON) + .build(); + AtomicBoolean isExecuted = new AtomicBoolean(false); + + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof PutDatasourceRequest); + PutDatasourceRequest putDatasourceRequest = (PutDatasourceRequest) actionRequest; + assertEquals("https://test.com", putDatasourceRequest.getEndpoint()); + assertEquals(TimeValue.timeValueDays(1), putDatasourceRequest.getUpdateInterval()); + assertEquals(datasourceName, putDatasourceRequest.getName()); + isExecuted.set(true); + return null; + }); + + dispatchRequest(request); + assertTrue(isExecuted.get()); + } + + public void testPrepareRequestDefaultValue() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) + .withPath(String.format(Locale.ROOT, path, datasourceName)) + .withContent(new BytesArray("{}"), XContentType.JSON) + .build(); + AtomicBoolean isExecuted = new AtomicBoolean(false); + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof PutDatasourceRequest); + PutDatasourceRequest putDatasourceRequest = (PutDatasourceRequest) actionRequest; + assertEquals("https://geoip.maps.opensearch.org/v1/geolite2-city/manifest.json", putDatasourceRequest.getEndpoint()); + assertEquals(TimeValue.timeValueDays(3), putDatasourceRequest.getUpdateInterval()); + assertEquals(datasourceName, putDatasourceRequest.getName()); + isExecuted.set(true); + return null; + }); + + dispatchRequest(request); + assertTrue(isExecuted.get()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandlerTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandlerTests.java new file mode 100644 index 0000000000..ef15d03039 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/RestUpdateDatasourceHandlerTests.java @@ -0,0 +1,79 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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.Locale; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.junit.Before; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.core.common.bytes.BytesArray; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.rest.RestRequest; +import org.opensearch.test.rest.FakeRestRequest; +import org.opensearch.test.rest.RestActionTestCase; + +public class RestUpdateDatasourceHandlerTests extends RestActionTestCase { + private String path; + private RestUpdateDatasourceHandler handler; + + @Before + public void setupAction() { + handler = new RestUpdateDatasourceHandler(); + controller().registerHandler(handler); + path = String.join(URL_DELIMITER, getPluginURLPrefix(), "ip2geo/datasource/%s/_settings"); + } + + public void testPrepareRequest_whenValidInput_thenSucceed() { + 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(Locale.ROOT, path, datasourceName)) + .withContent(new BytesArray(content), XContentType.JSON) + .build(); + AtomicBoolean isExecuted = new AtomicBoolean(false); + + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof UpdateDatasourceRequest); + UpdateDatasourceRequest updateDatasourceRequest = (UpdateDatasourceRequest) actionRequest; + assertEquals("https://test.com", updateDatasourceRequest.getEndpoint()); + assertEquals(TimeValue.timeValueDays(1), updateDatasourceRequest.getUpdateInterval()); + assertEquals(datasourceName, updateDatasourceRequest.getName()); + isExecuted.set(true); + return null; + }); + + dispatchRequest(request); + assertTrue(isExecuted.get()); + } + + public void testPrepareRequest_whenNullInput_thenSucceed() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String content = "{}"; + RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) + .withPath(String.format(Locale.ROOT, path, datasourceName)) + .withContent(new BytesArray(content), XContentType.JSON) + .build(); + AtomicBoolean isExecuted = new AtomicBoolean(false); + + verifyingClient.setExecuteLocallyVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof UpdateDatasourceRequest); + UpdateDatasourceRequest updateDatasourceRequest = (UpdateDatasourceRequest) actionRequest; + assertNull(updateDatasourceRequest.getEndpoint()); + assertNull(updateDatasourceRequest.getUpdateInterval()); + assertEquals(datasourceName, updateDatasourceRequest.getName()); + isExecuted.set(true); + return null; + }); + + dispatchRequest(request); + assertTrue(isExecuted.get()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceIT.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceIT.java new file mode 100644 index 0000000000..6c7baa2f7a --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceIT.java @@ -0,0 +1,104 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.io.IOException; +import java.time.Duration; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +import lombok.SneakyThrows; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.opensearch.client.ResponseException; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.geospatial.GeospatialRestTestCase; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoDataServer; + +public class UpdateDatasourceIT extends GeospatialRestTestCase { + // Use this value in resource name to avoid name conflict among tests + private static final String PREFIX = UpdateDatasourceIT.class.getSimpleName().toLowerCase(Locale.ROOT); + + @BeforeClass + public static void start() { + Ip2GeoDataServer.start(); + } + + @AfterClass + public static void stop() { + Ip2GeoDataServer.stop(); + } + + @SneakyThrows + public void testUpdateDatasource_whenValidInput_thenUpdated() { + boolean isDatasourceCreated = false; + String datasourceName = PREFIX + GeospatialTestHelper.randomLowerCaseString(); + try { + Map datasourceProperties = Map.of( + PutDatasourceRequest.ENDPOINT_FIELD.getPreferredName(), + Ip2GeoDataServer.getEndpointCountry() + ); + + // Create datasource and wait for it to be available + createDatasource(datasourceName, datasourceProperties); + isDatasourceCreated = true; + waitForDatasourceToBeAvailable(datasourceName, Duration.ofSeconds(10)); + + int updateIntervalInDays = 1; + updateDatasourceEndpoint(datasourceName, Ip2GeoDataServer.getEndpointCity(), updateIntervalInDays); + List> datasources = (List>) getDatasource(datasourceName).get("datasources"); + + assertEquals(Ip2GeoDataServer.getEndpointCity(), datasources.get(0).get("endpoint")); + assertEquals(updateIntervalInDays, datasources.get(0).get("update_interval_in_days")); + } finally { + if (isDatasourceCreated) { + deleteDatasource(datasourceName, 3); + } + } + } + + @SneakyThrows + public void testUpdateDatasource_whenIncompatibleFields_thenFails() { + boolean isDatasourceCreated = false; + String datasourceName = PREFIX + GeospatialTestHelper.randomLowerCaseString(); + try { + Map datasourceProperties = Map.of( + PutDatasourceRequest.ENDPOINT_FIELD.getPreferredName(), + Ip2GeoDataServer.getEndpointCity() + ); + + // Create datasource and wait for it to be available + createDatasource(datasourceName, datasourceProperties); + isDatasourceCreated = true; + waitForDatasourceToBeAvailable(datasourceName, Duration.ofSeconds(10)); + + // Update should fail as country data does not have every fields that city data has + int updateIntervalInDays = 1; + ResponseException exception = expectThrows( + ResponseException.class, + () -> updateDatasourceEndpoint(datasourceName, Ip2GeoDataServer.getEndpointCountry(), updateIntervalInDays) + ); + assertEquals(RestStatus.BAD_REQUEST.getStatus(), exception.getResponse().getStatusLine().getStatusCode()); + } finally { + if (isDatasourceCreated) { + deleteDatasource(datasourceName, 3); + } + } + } + + private void updateDatasourceEndpoint(final String datasourceName, final String endpoint, final int updateInterval) throws IOException { + Map properties = Map.of( + UpdateDatasourceRequest.ENDPOINT_FIELD.getPreferredName(), + endpoint, + UpdateDatasourceRequest.UPDATE_INTERVAL_IN_DAYS_FIELD.getPreferredName(), + updateInterval + ); + updateDatasource(datasourceName, properties); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceRequestTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceRequestTests.java new file mode 100644 index 0000000000..c37b6830c5 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceRequestTests.java @@ -0,0 +1,121 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import java.util.Locale; + +import lombok.SneakyThrows; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.common.Randomness; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.common.io.stream.BytesStreamInput; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; + +public class UpdateDatasourceRequestTests extends Ip2GeoTestCase { + + public void testValidate_whenNullValues_thenFails() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasourceName); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertEquals("no values to update", exception.validationErrors().get(0)); + } + + public void testValidate_whenInvalidUrl_thenFails() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasourceName); + request.setEndpoint("invalidUrl"); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertEquals("Invalid URL format is provided", exception.validationErrors().get(0)); + } + + public void testValidate_whenInvalidManifestFile_thenFails() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String domain = GeospatialTestHelper.randomLowerCaseString(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasourceName); + request.setEndpoint(String.format(Locale.ROOT, "https://%s.com", domain)); + + // Run + ActionRequestValidationException exception = request.validate(); + + // Verify + assertEquals(1, exception.validationErrors().size()); + assertTrue(exception.validationErrors().get(0).contains("Error occurred while reading a file")); + } + + @SneakyThrows + public void testValidate_whenValidInput_thenSucceed() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasourceName); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(1)); + + // Run and verify + assertNull(request.validate()); + } + + @SneakyThrows + public void testValidate_whenZeroUpdateInterval_thenFails() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasourceName); + 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) + ); + } + + @SneakyThrows + public void testValidate_whenInvalidUrlInsideManifest_thenFail() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(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")); + } + + @SneakyThrows + public void testStreamInOut_whenValidInput_thenSucceed() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String domain = GeospatialTestHelper.randomLowerCaseString(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasourceName); + request.setEndpoint(String.format(Locale.ROOT, "https://%s.com", domain)); + request.setUpdateInterval(TimeValue.timeValueDays(Randomness.get().nextInt(29) + 1)); + + // Run + BytesStreamOutput output = new BytesStreamOutput(); + request.writeTo(output); + BytesStreamInput input = new BytesStreamInput(output.bytes().toBytesRef().bytes); + UpdateDatasourceRequest copiedRequest = new UpdateDatasourceRequest(input); + + // Verify + assertEquals(request, copiedRequest); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceTransportActionTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceTransportActionTests.java new file mode 100644 index 0000000000..e0a94a75ea --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/action/UpdateDatasourceTransportActionTests.java @@ -0,0 +1,270 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.action; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +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 java.security.InvalidParameterException; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.List; + +import lombok.SneakyThrows; + +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.opensearch.OpenSearchException; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.geospatial.exceptions.IncompatibleDatasourceException; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceTask; +import org.opensearch.jobscheduler.spi.LockModel; +import org.opensearch.tasks.Task; + +public class UpdateDatasourceTransportActionTests extends Ip2GeoTestCase { + private UpdateDatasourceTransportAction action; + + @Before + public void init() { + action = new UpdateDatasourceTransportAction( + transportService, + actionFilters, + ip2GeoLockService, + datasourceDao, + datasourceUpdateService, + threadPool + ); + } + + public void testDoExecute_whenFailedToAcquireLock_thenError() { + validateDoExecuteWithLockError(null); + } + + public void testDoExecute_whenExceptionToAcquireLock_thenError() { + validateDoExecuteWithLockError(new RuntimeException()); + } + + private void validateDoExecuteWithLockError(final Exception exception) { + Task task = mock(Task.class); + Datasource datasource = randomDatasource(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasource.getName()); + ActionListener listener = mock(ActionListener.class); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + if (exception == null) { + // Run + captor.getValue().onResponse(null); + // Verify + verify(listener).onFailure(any(OpenSearchException.class)); + } else { + // Run + captor.getValue().onFailure(exception); + // Verify + verify(listener).onFailure(exception); + } + } + + @SneakyThrows + public void testDoExecute_whenValidInput_thenUpdate() { + Datasource datasource = randomDatasource(Instant.now().minusSeconds(60)); + datasource.setTask(DatasourceTask.DELETE_UNUSED_INDICES); + Instant originalStartTime = datasource.getSchedule().getStartTime(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasource.getName()); + request.setEndpoint(sampleManifestUrl()); + request.setUpdateInterval(TimeValue.timeValueDays(datasource.getSchedule().getInterval())); + + Task task = mock(Task.class); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + when(datasourceUpdateService.getHeaderFields(request.getEndpoint())).thenReturn(datasource.getDatabase().getFields()); + ActionListener listener = mock(ActionListener.class); + LockModel lockModel = randomLockModel(); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + // Run + captor.getValue().onResponse(lockModel); + + // Verify + verify(datasourceDao).getDatasource(datasource.getName()); + verify(datasourceDao).updateDatasource(datasource); + verify(datasourceUpdateService).getHeaderFields(request.getEndpoint()); + assertEquals(request.getEndpoint(), datasource.getEndpoint()); + assertEquals(request.getUpdateInterval().days(), datasource.getUserSchedule().getInterval()); + verify(listener).onResponse(new AcknowledgedResponse(true)); + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + assertTrue(originalStartTime.isBefore(datasource.getSchedule().getStartTime())); + assertEquals(DatasourceTask.ALL, datasource.getTask()); + } + + @SneakyThrows + public void testDoExecute_whenNoChangesInValues_thenNoUpdate() { + Datasource datasource = randomDatasource(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasource.getName()); + request.setEndpoint(datasource.getEndpoint()); + + Task task = mock(Task.class); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + ActionListener listener = mock(ActionListener.class); + LockModel lockModel = randomLockModel(); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + // Run + captor.getValue().onResponse(lockModel); + + // Verify + verify(datasourceDao).getDatasource(datasource.getName()); + verify(datasourceUpdateService, never()).getHeaderFields(anyString()); + verify(datasourceDao, never()).updateDatasource(datasource); + verify(listener).onResponse(new AcknowledgedResponse(true)); + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + } + + @SneakyThrows + public void testDoExecute_whenNoDatasource_thenError() { + Datasource datasource = randomDatasource(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasource.getName()); + + Task task = mock(Task.class); + ActionListener listener = mock(ActionListener.class); + LockModel lockModel = randomLockModel(); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + // Run + captor.getValue().onResponse(lockModel); + + // Verify + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener).onFailure(exceptionCaptor.capture()); + assertEquals(ResourceNotFoundException.class, exceptionCaptor.getValue().getClass()); + exceptionCaptor.getValue().getMessage().contains("no such datasource exist"); + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + } + + @SneakyThrows + public void testDoExecute_whenIncompatibleFields_thenError() { + Datasource datasource = randomDatasource(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasource.getName()); + request.setEndpoint(sampleManifestUrl()); + + Task task = mock(Task.class); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + List newFields = datasource.getDatabase().getFields().subList(0, 0); + when(datasourceUpdateService.getHeaderFields(request.getEndpoint())).thenReturn(newFields); + ActionListener listener = mock(ActionListener.class); + LockModel lockModel = randomLockModel(); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + // Run + captor.getValue().onResponse(lockModel); + + // Verify + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener).onFailure(exceptionCaptor.capture()); + assertEquals(IncompatibleDatasourceException.class, exceptionCaptor.getValue().getClass()); + exceptionCaptor.getValue().getMessage().contains("does not contain"); + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + } + + @SneakyThrows + public void testDoExecute_whenLargeUpdateInterval_thenError() { + Datasource datasource = randomDatasource(); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasource.getName()); + request.setUpdateInterval(TimeValue.timeValueDays(datasource.getDatabase().getValidForInDays())); + + Task task = mock(Task.class); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + ActionListener listener = mock(ActionListener.class); + LockModel lockModel = randomLockModel(); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + // Run + captor.getValue().onResponse(lockModel); + + // Verify + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener).onFailure(exceptionCaptor.capture()); + assertEquals(InvalidParameterException.class, exceptionCaptor.getValue().getClass()); + exceptionCaptor.getValue().getMessage().contains("should be smaller"); + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + } + + @SneakyThrows + public void testDoExecute_whenExpireWithNewUpdateInterval_thenError() { + Datasource datasource = randomDatasource(); + datasource.getUpdateStats().setLastSkippedAt(null); + datasource.getUpdateStats().setLastSucceededAt(Instant.now().minus(datasource.getDatabase().getValidForInDays(), ChronoUnit.DAYS)); + UpdateDatasourceRequest request = new UpdateDatasourceRequest(datasource.getName()); + request.setUpdateInterval(TimeValue.timeValueDays(1)); + + Task task = mock(Task.class); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + ActionListener listener = mock(ActionListener.class); + LockModel lockModel = randomLockModel(); + + // Run + action.doExecute(task, request, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(ip2GeoLockService).acquireLock(eq(datasource.getName()), anyLong(), captor.capture()); + + // Run + captor.getValue().onResponse(lockModel); + + // Verify + ArgumentCaptor exceptionCaptor = ArgumentCaptor.forClass(Exception.class); + verify(listener).onFailure(exceptionCaptor.capture()); + assertEquals(IllegalArgumentException.class, exceptionCaptor.getValue().getClass()); + exceptionCaptor.getValue().getMessage().contains("will expire"); + verify(ip2GeoLockService).releaseLock(eq(lockModel)); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceManifestTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceManifestTests.java new file mode 100644 index 0000000000..f7b689e10f --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/common/DatasourceManifestTests.java @@ -0,0 +1,38 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.File; +import java.io.FileInputStream; +import java.net.URLConnection; + +import lombok.SneakyThrows; + +import org.opensearch.common.SuppressForbidden; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.shared.Constants; + +@SuppressForbidden(reason = "unit test") +public class DatasourceManifestTests extends Ip2GeoTestCase { + + @SneakyThrows + public void testInternalBuild_whenCalled_thenCorrectUserAgentValueIsSet() { + URLConnection connection = mock(URLConnection.class); + File manifestFile = new File(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").getFile()); + when(connection.getInputStream()).thenReturn(new FileInputStream(manifestFile)); + + // Run + DatasourceManifest manifest = DatasourceManifest.Builder.internalBuild(connection); + + // Verify + verify(connection).addRequestProperty(Constants.USER_AGENT_KEY, Constants.USER_AGENT_VALUE); + assertEquals("https://test.com/db.zip", manifest.getUrl()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoLockServiceTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoLockServiceTests.java new file mode 100644 index 0000000000..971bec0d5e --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoLockServiceTests.java @@ -0,0 +1,116 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +import static org.mockito.Mockito.mock; +import static org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService.LOCK_DURATION_IN_SECONDS; +import static org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService.RENEW_AFTER_IN_SECONDS; + +import java.time.Instant; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.Before; +import org.opensearch.action.ActionListener; +import org.opensearch.action.DocWriteResponse; +import org.opensearch.action.update.UpdateRequest; +import org.opensearch.action.update.UpdateResponse; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.jobscheduler.spi.LockModel; + +public class Ip2GeoLockServiceTests extends Ip2GeoTestCase { + private Ip2GeoLockService ip2GeoLockService; + private Ip2GeoLockService noOpsLockService; + + @Before + public void init() { + ip2GeoLockService = new Ip2GeoLockService(clusterService, verifyingClient); + noOpsLockService = new Ip2GeoLockService(clusterService, client); + } + + public void testAcquireLock_whenValidInput_thenSucceed() { + // Cannot test because LockService is final class + // Simply calling method to increase coverage + noOpsLockService.acquireLock(GeospatialTestHelper.randomLowerCaseString(), randomPositiveLong(), mock(ActionListener.class)); + } + + public void testAcquireLock_whenCalled_thenNotBlocked() { + long expectedDurationInMillis = 1000; + Instant before = Instant.now(); + assertTrue(ip2GeoLockService.acquireLock(null, null).isEmpty()); + Instant after = Instant.now(); + assertTrue(after.toEpochMilli() - before.toEpochMilli() < expectedDurationInMillis); + } + + public void testReleaseLock_whenValidInput_thenSucceed() { + // Cannot test because LockService is final class + // Simply calling method to increase coverage + noOpsLockService.releaseLock(null); + } + + public void testRenewLock_whenCalled_thenNotBlocked() { + long expectedDurationInMillis = 1000; + Instant before = Instant.now(); + assertNull(ip2GeoLockService.renewLock(null)); + Instant after = Instant.now(); + assertTrue(after.toEpochMilli() - before.toEpochMilli() < expectedDurationInMillis); + } + + public void testGetRenewLockRunnable_whenLockIsFresh_thenDoNotRenew() { + LockModel lockModel = new LockModel( + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString(), + Instant.now(), + LOCK_DURATION_IN_SECONDS, + false + ); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + // Verifying + assertTrue(actionRequest instanceof UpdateRequest); + return new UpdateResponse( + mock(ShardId.class), + GeospatialTestHelper.randomLowerCaseString(), + randomPositiveLong(), + randomPositiveLong(), + randomPositiveLong(), + DocWriteResponse.Result.UPDATED + ); + }); + + AtomicReference reference = new AtomicReference<>(lockModel); + ip2GeoLockService.getRenewLockRunnable(reference).run(); + assertEquals(lockModel, reference.get()); + } + + public void testGetRenewLockRunnable_whenLockIsStale_thenRenew() { + LockModel lockModel = new LockModel( + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString(), + Instant.now().minusSeconds(RENEW_AFTER_IN_SECONDS), + LOCK_DURATION_IN_SECONDS, + false + ); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + // Verifying + assertTrue(actionRequest instanceof UpdateRequest); + return new UpdateResponse( + mock(ShardId.class), + GeospatialTestHelper.randomLowerCaseString(), + randomPositiveLong(), + randomPositiveLong(), + randomPositiveLong(), + DocWriteResponse.Result.UPDATED + ); + }); + + AtomicReference reference = new AtomicReference<>(lockModel); + ip2GeoLockService.getRenewLockRunnable(reference).run(); + assertNotEquals(lockModel, reference.get()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettingsTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettingsTests.java new file mode 100644 index 0000000000..ffd40f5f58 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/common/Ip2GeoSettingsTests.java @@ -0,0 +1,21 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.common; + +import org.opensearch.test.OpenSearchTestCase; + +public class Ip2GeoSettingsTests extends OpenSearchTestCase { + public void testValidateInvalidUrl() { + Ip2GeoSettings.DatasourceEndpointValidator validator = new Ip2GeoSettings.DatasourceEndpointValidator(); + Exception e = expectThrows(IllegalArgumentException.class, () -> validator.validate("InvalidUrl")); + assertEquals("Invalid URL format is provided", e.getMessage()); + } + + public void testValidateValidUrl() { + Ip2GeoSettings.DatasourceEndpointValidator validator = new Ip2GeoSettings.DatasourceEndpointValidator(); + validator.validate("https://test.com"); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/dao/DatasourceDaoTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/dao/DatasourceDaoTests.java new file mode 100644 index 0000000000..88329bb399 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/dao/DatasourceDaoTests.java @@ -0,0 +1,394 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +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.Arrays; +import java.util.List; + +import lombok.SneakyThrows; + +import org.apache.lucene.search.TotalHits; +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.opensearch.ResourceAlreadyExistsException; +import org.opensearch.ResourceNotFoundException; +import org.opensearch.action.ActionListener; +import org.opensearch.action.DocWriteRequest; +import org.opensearch.action.StepListener; +import org.opensearch.action.admin.indices.create.CreateIndexRequest; +import org.opensearch.action.bulk.BulkRequest; +import org.opensearch.action.delete.DeleteRequest; +import org.opensearch.action.delete.DeleteResponse; +import org.opensearch.action.get.GetRequest; +import org.opensearch.action.get.GetResponse; +import org.opensearch.action.get.MultiGetItemResponse; +import org.opensearch.action.get.MultiGetRequest; +import org.opensearch.action.get.MultiGetResponse; +import org.opensearch.action.index.IndexRequest; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.WriteRequest; +import org.opensearch.cluster.routing.Preference; +import org.opensearch.common.Randomness; +import org.opensearch.common.xcontent.json.JsonXContent; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.rest.RestStatus; +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.index.query.QueryBuilders; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; + +public class DatasourceDaoTests extends Ip2GeoTestCase { + private DatasourceDao datasourceDao; + + @Before + public void init() { + datasourceDao = new DatasourceDao(verifyingClient, clusterService); + } + + public void testCreateIndexIfNotExists_whenIndexExist_thenCreateRequestIsNotCalled() { + when(metadata.hasIndex(DatasourceExtension.JOB_INDEX_NAME)).thenReturn(true); + + // Verify + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { throw new RuntimeException("Shouldn't get called"); }); + + // Run + StepListener stepListener = new StepListener<>(); + datasourceDao.createIndexIfNotExists(stepListener); + + // Verify stepListener is called + stepListener.result(); + } + + public void testCreateIndexIfNotExists_whenIndexExist_thenCreateRequestIsCalled() { + when(metadata.hasIndex(DatasourceExtension.JOB_INDEX_NAME)).thenReturn(false); + + // Verify + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof CreateIndexRequest); + CreateIndexRequest request = (CreateIndexRequest) actionRequest; + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); + assertEquals("1", request.settings().get("index.number_of_shards")); + assertEquals("0-all", request.settings().get("index.auto_expand_replicas")); + assertEquals("true", request.settings().get("index.hidden")); + assertNotNull(request.mappings()); + return null; + }); + + // Run + StepListener stepListener = new StepListener<>(); + datasourceDao.createIndexIfNotExists(stepListener); + + // Verify stepListener is called + stepListener.result(); + } + + public void testCreateIndexIfNotExists_whenIndexCreatedAlready_thenExceptionIsIgnored() { + when(metadata.hasIndex(DatasourceExtension.JOB_INDEX_NAME)).thenReturn(false); + verifyingClient.setExecuteVerifier( + (actionResponse, actionRequest) -> { throw new ResourceAlreadyExistsException(DatasourceExtension.JOB_INDEX_NAME); } + ); + + // Run + StepListener stepListener = new StepListener<>(); + datasourceDao.createIndexIfNotExists(stepListener); + + // Verify stepListener is called + stepListener.result(); + } + + public void testCreateIndexIfNotExists_whenExceptionIsThrown_thenExceptionIsThrown() { + when(metadata.hasIndex(DatasourceExtension.JOB_INDEX_NAME)).thenReturn(false); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { throw new RuntimeException(); }); + + // Run + StepListener stepListener = new StepListener<>(); + datasourceDao.createIndexIfNotExists(stepListener); + + // Verify stepListener is called + expectThrows(RuntimeException.class, () -> stepListener.result()); + } + + public void testUpdateDatasource_whenValidInput_thenSucceed() 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.getName(), request.id()); + assertEquals(DocWriteRequest.OpType.INDEX, request.opType()); + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); + assertEquals(WriteRequest.RefreshPolicy.IMMEDIATE, request.getRefreshPolicy()); + return null; + }); + + datasourceDao.updateDatasource(datasource); + assertTrue(previousTime.isBefore(datasource.getLastUpdateTime())); + } + + @SneakyThrows + public void testPutDatasource_whenValidInpu_thenSucceed() { + Datasource datasource = randomDatasource(); + Instant previousTime = Instant.now().minusMillis(1); + datasource.setLastUpdateTime(previousTime); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof IndexRequest); + IndexRequest indexRequest = (IndexRequest) actionRequest; + assertEquals(DatasourceExtension.JOB_INDEX_NAME, indexRequest.index()); + assertEquals(datasource.getName(), indexRequest.id()); + assertEquals(WriteRequest.RefreshPolicy.IMMEDIATE, indexRequest.getRefreshPolicy()); + assertEquals(DocWriteRequest.OpType.CREATE, indexRequest.opType()); + return null; + }); + + datasourceDao.putDatasource(datasource, mock(ActionListener.class)); + assertTrue(previousTime.isBefore(datasource.getLastUpdateTime())); + } + + public void testGetDatasource_whenException_thenNull() throws Exception { + Datasource datasource = setupClientForGetRequest(true, new IndexNotFoundException(DatasourceExtension.JOB_INDEX_NAME)); + assertNull(datasourceDao.getDatasource(datasource.getName())); + } + + public void testGetDatasource_whenExist_thenReturnDatasource() throws Exception { + Datasource datasource = setupClientForGetRequest(true, null); + assertEquals(datasource, datasourceDao.getDatasource(datasource.getName())); + } + + public void testGetDatasource_whenNotExist_thenNull() throws Exception { + Datasource datasource = setupClientForGetRequest(false, null); + assertNull(datasourceDao.getDatasource(datasource.getName())); + } + + public void testGetDatasource_whenExistWithListener_thenListenerIsCalledWithDatasource() { + Datasource datasource = setupClientForGetRequest(true, null); + ActionListener listener = mock(ActionListener.class); + datasourceDao.getDatasource(datasource.getName(), listener); + verify(listener).onResponse(eq(datasource)); + } + + public void testGetDatasource_whenNotExistWithListener_thenListenerIsCalledWithNull() { + Datasource datasource = setupClientForGetRequest(false, null); + ActionListener listener = mock(ActionListener.class); + datasourceDao.getDatasource(datasource.getName(), listener); + verify(listener).onResponse(null); + } + + private Datasource setupClientForGetRequest(final boolean isExist, final RuntimeException exception) { + Datasource datasource = randomDatasource(); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assertTrue(actionRequest instanceof GetRequest); + GetRequest request = (GetRequest) actionRequest; + assertEquals(datasource.getName(), request.id()); + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); + assertEquals(Preference.PRIMARY.type(), request.preference()); + GetResponse response = getMockedGetResponse(isExist ? datasource : null); + if (exception != null) { + throw exception; + } + return response; + }); + return datasource; + } + + public void testDeleteDatasource_whenValidInput_thenSucceed() { + Datasource datasource = randomDatasource(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + // Verify + assertTrue(actionRequest instanceof DeleteRequest); + DeleteRequest request = (DeleteRequest) actionRequest; + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); + assertEquals(DocWriteRequest.OpType.DELETE, request.opType()); + assertEquals(datasource.getName(), request.id()); + assertEquals(WriteRequest.RefreshPolicy.IMMEDIATE, request.getRefreshPolicy()); + + DeleteResponse response = mock(DeleteResponse.class); + when(response.status()).thenReturn(RestStatus.OK); + return response; + }); + + // Run + datasourceDao.deleteDatasource(datasource); + } + + public void testDeleteDatasource_whenIndexNotFound_thenThrowException() { + Datasource datasource = randomDatasource(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + DeleteResponse response = mock(DeleteResponse.class); + when(response.status()).thenReturn(RestStatus.NOT_FOUND); + return response; + }); + + // Run + expectThrows(ResourceNotFoundException.class, () -> datasourceDao.deleteDatasource(datasource)); + } + + public void testGetDatasources_whenValidInput_thenSucceed() { + List datasources = Arrays.asList(randomDatasource(), randomDatasource()); + String[] names = datasources.stream().map(Datasource::getName).toArray(String[]::new); + ActionListener> listener = mock(ActionListener.class); + MultiGetItemResponse[] multiGetItemResponses = datasources.stream().map(datasource -> { + GetResponse getResponse = getMockedGetResponse(datasource); + MultiGetItemResponse multiGetItemResponse = mock(MultiGetItemResponse.class); + when(multiGetItemResponse.getResponse()).thenReturn(getResponse); + return multiGetItemResponse; + }).toArray(MultiGetItemResponse[]::new); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + // Verify + assertTrue(actionRequest instanceof MultiGetRequest); + MultiGetRequest request = (MultiGetRequest) actionRequest; + assertEquals(2, request.getItems().size()); + assertEquals(Preference.PRIMARY.type(), request.preference()); + for (MultiGetRequest.Item item : request.getItems()) { + assertEquals(DatasourceExtension.JOB_INDEX_NAME, item.index()); + assertTrue(datasources.stream().filter(datasource -> datasource.getName().equals(item.id())).findAny().isPresent()); + } + + MultiGetResponse response = mock(MultiGetResponse.class); + when(response.getResponses()).thenReturn(multiGetItemResponses); + return response; + }); + + // Run + datasourceDao.getDatasources(names, listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(List.class); + verify(listener).onResponse(captor.capture()); + assertEquals(datasources, captor.getValue()); + + } + + public void testGetAllDatasources_whenAsynchronous_thenSucceed() { + List datasources = Arrays.asList(randomDatasource(), randomDatasource()); + ActionListener> listener = mock(ActionListener.class); + SearchHits searchHits = getMockedSearchHits(datasources); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + // Verify + assertTrue(actionRequest instanceof SearchRequest); + SearchRequest request = (SearchRequest) actionRequest; + assertEquals(1, request.indices().length); + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.indices()[0]); + assertEquals(QueryBuilders.matchAllQuery(), request.source().query()); + assertEquals(1000, request.source().size()); + assertEquals(Preference.PRIMARY.type(), request.preference()); + + SearchResponse response = mock(SearchResponse.class); + when(response.getHits()).thenReturn(searchHits); + return response; + }); + + // Run + datasourceDao.getAllDatasources(listener); + + // Verify + ArgumentCaptor> captor = ArgumentCaptor.forClass(List.class); + verify(listener).onResponse(captor.capture()); + assertEquals(datasources, captor.getValue()); + } + + public void testGetAllDatasources_whenSynchronous_thenSucceed() { + List datasources = Arrays.asList(randomDatasource(), randomDatasource()); + SearchHits searchHits = getMockedSearchHits(datasources); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + // Verify + assertTrue(actionRequest instanceof SearchRequest); + SearchRequest request = (SearchRequest) actionRequest; + assertEquals(1, request.indices().length); + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.indices()[0]); + assertEquals(QueryBuilders.matchAllQuery(), request.source().query()); + assertEquals(1000, request.source().size()); + assertEquals(Preference.PRIMARY.type(), request.preference()); + + SearchResponse response = mock(SearchResponse.class); + when(response.getHits()).thenReturn(searchHits); + return response; + }); + + // Run + datasourceDao.getAllDatasources(); + + // Verify + assertEquals(datasources, datasourceDao.getAllDatasources()); + } + + public void testUpdateDatasource_whenValidInput_thenUpdate() { + List datasources = Arrays.asList(randomDatasource(), randomDatasource()); + + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + // Verify + assertTrue(actionRequest instanceof BulkRequest); + BulkRequest bulkRequest = (BulkRequest) actionRequest; + assertEquals(2, bulkRequest.requests().size()); + for (int i = 0; i < bulkRequest.requests().size(); i++) { + IndexRequest request = (IndexRequest) bulkRequest.requests().get(i); + assertEquals(DatasourceExtension.JOB_INDEX_NAME, request.index()); + assertEquals(datasources.get(i).getName(), request.id()); + assertEquals(DocWriteRequest.OpType.INDEX, request.opType()); + assertTrue(request.source().utf8ToString().contains(datasources.get(i).getEndpoint())); + } + return null; + }); + + datasourceDao.updateDatasource(datasources, mock(ActionListener.class)); + } + + private SearchHits getMockedSearchHits(List datasources) { + SearchHit[] searchHitArray = datasources.stream().map(this::toBytesReference).map(this::toSearchHit).toArray(SearchHit[]::new); + + return new SearchHits(searchHitArray, new TotalHits(1l, TotalHits.Relation.EQUAL_TO), 1); + } + + private GetResponse getMockedGetResponse(Datasource datasource) { + GetResponse response = mock(GetResponse.class); + when(response.isExists()).thenReturn(datasource != null); + when(response.getSourceAsBytesRef()).thenReturn(toBytesReference(datasource)); + return response; + } + + private BytesReference toBytesReference(Datasource datasource) { + if (datasource == null) { + return null; + } + + try { + return BytesReference.bytes(datasource.toXContent(JsonXContent.contentBuilder(), null)); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + private SearchHit toSearchHit(BytesReference bytesReference) { + SearchHit searchHit = new SearchHit(Randomness.get().nextInt()); + searchHit.sourceRef(bytesReference); + return searchHit; + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/dao/GeoIpDataDaoTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/dao/GeoIpDataDaoTests.java new file mode 100644 index 0000000000..8007d4bf11 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/dao/GeoIpDataDaoTests.java @@ -0,0 +1,291 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +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.io.FileInputStream; +import java.net.URLConnection; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.Iterator; +import java.util.Locale; +import java.util.Map; + +import lombok.SneakyThrows; + +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.opensearch.OpenSearchException; +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.admin.indices.settings.put.UpdateSettingsRequest; +import org.opensearch.action.bulk.BulkRequest; +import org.opensearch.action.bulk.BulkResponse; +import org.opensearch.action.search.SearchRequest; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.cluster.routing.Preference; +import org.opensearch.common.Strings; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.common.DatasourceManifest; +import org.opensearch.geospatial.shared.Constants; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.search.SearchHit; +import org.opensearch.search.SearchHits; + +@SuppressForbidden(reason = "unit test") +public class GeoIpDataDaoTests extends Ip2GeoTestCase { + private static final String IP_RANGE_FIELD_NAME = "_cidr"; + private static final String DATA_FIELD_NAME = "_data"; + private GeoIpDataDao noOpsGeoIpDataDao; + private GeoIpDataDao verifyingGeoIpDataDao; + + @Before + public void init() { + noOpsGeoIpDataDao = new GeoIpDataDao(clusterService, client); + verifyingGeoIpDataDao = new GeoIpDataDao(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"); }); + verifyingGeoIpDataDao.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, (int) request.settings().getAsInt("index.number_of_shards", 0)); + assertNull(request.settings().get("index.auto_expand_replicas")); + assertEquals(0, (int) request.settings().getAsInt("index.number_of_replicas", 1)); + assertEquals(-1, (int) request.settings().getAsInt("index.refresh_interval", 0)); + assertEquals(true, request.settings().getAsBoolean("index.hidden", false)); + + assertEquals( + "{\"dynamic\": false,\"properties\": {\"_cidr\": {\"type\": \"ip_range\",\"doc_values\": false}}}", + request.mappings() + ); + return null; + }); + verifyingGeoIpDataDao.createIndexIfNotExists(index); + } + + @SneakyThrows + public void testCreateDocument_whenBlankValue_thenDoNotAdd() { + String[] names = { "ip", "country", "location", "city" }; + String[] values = { "1.0.0.0/25", "USA", " ", "Seattle" }; + assertEquals( + "{\"_cidr\":\"1.0.0.0/25\",\"_data\":{\"country\":\"USA\",\"city\":\"Seattle\"}}", + Strings.toString(noOpsGeoIpDataDao.createDocument(names, values)) + ); + } + + @SneakyThrows + public void testCreateDocument_whenFieldsAndValuesLengthDoesNotMatch_thenThrowException() { + String[] names = { "ip", "country", "location", "city" }; + String[] values = { "1.0.0.0/25", "USA", " " }; + + // Run + Exception e = expectThrows(OpenSearchException.class, () -> noOpsGeoIpDataDao.createDocument(names, values)); + + // Verify + assertTrue(e.getMessage().contains("does not match")); + } + + 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_sha256", + 1l, + Instant.now().toEpochMilli(), + "tester" + ); + CSVParser parser = noOpsGeoIpDataDao.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_sha256", + 1l, + Instant.now().toEpochMilli(), + "tester" + ); + OpenSearchException exception = expectThrows(OpenSearchException.class, () -> noOpsGeoIpDataDao.getDatabaseReader(manifest)); + assertTrue(exception.getMessage().contains("does not exist")); + } + + @SneakyThrows + public void testInternalGetDatabaseReader_whenCalled_thenSetUserAgent() { + 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_sha256", + 1l, + Instant.now().toEpochMilli(), + "tester" + ); + + URLConnection connection = mock(URLConnection.class); + when(connection.getInputStream()).thenReturn(new FileInputStream(zipFile)); + + // Run + noOpsGeoIpDataDao.internalGetDatabaseReader(manifest, connection); + + // Verify + verify(connection).addRequestProperty(Constants.USER_AGENT_KEY, Constants.USER_AGENT_VALUE); + } + + public void testDeleteIp2GeoDataIndex_whenCalled_thenDeleteIndex() { + 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]); + return new AcknowledgedResponse(true); + }); + verifyingGeoIpDataDao.deleteIp2GeoDataIndex(index); + } + + public void testDeleteIp2GeoDataIndexWithNonIp2GeoDataIndex() { + String index = GeospatialTestHelper.randomLowerCaseString(); + Exception e = expectThrows(OpenSearchException.class, () -> verifyingGeoIpDataDao.deleteIp2GeoDataIndex(index)); + assertTrue(e.getMessage().contains("not ip2geo data index")); + verify(verifyingClient, never()).index(any()); + } + + @SneakyThrows + public void testPutGeoIpData_whenValidInput_thenSucceed() { + String index = GeospatialTestHelper.randomLowerCaseString(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + if (actionRequest instanceof BulkRequest) { + BulkRequest request = (BulkRequest) actionRequest; + assertEquals(2, request.numberOfActions()); + 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 if (actionRequest instanceof UpdateSettingsRequest) { + UpdateSettingsRequest request = (UpdateSettingsRequest) actionRequest; + assertEquals(1, request.indices().length); + assertEquals(index, request.indices()[0]); + assertEquals(true, request.settings().getAsBoolean("index.blocks.write", false)); + assertNull(request.settings().get("index.num_of_replica")); + assertEquals("0-all", request.settings().get("index.auto_expand_replicas")); + return null; + } else { + throw new RuntimeException("invalid request is called"); + } + }); + Runnable renewLock = mock(Runnable.class); + try (CSVParser csvParser = CSVParser.parse(sampleIp2GeoFile(), StandardCharsets.UTF_8, CSVFormat.RFC4180)) { + Iterator iterator = csvParser.iterator(); + String[] fields = iterator.next().values(); + verifyingGeoIpDataDao.putGeoIpData(index, fields, iterator, renewLock); + verify(renewLock, times(2)).run(); + } + } + + public void testGetGeoIpData_whenDataExist_thenReturnTheData() { + String indexName = GeospatialTestHelper.randomLowerCaseString(); + String ip = randomIpAddress(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assert actionRequest instanceof SearchRequest; + SearchRequest request = (SearchRequest) actionRequest; + assertEquals(Preference.LOCAL.type(), 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; + }); + + // Run + Map geoData = verifyingGeoIpDataDao.getGeoIpData(indexName, ip); + + // Verify + assertEquals("seattle", geoData.get("city")); + } + + public void testGetGeoIpData_whenNoData_thenReturnEmpty() { + String indexName = GeospatialTestHelper.randomLowerCaseString(); + String ip = randomIpAddress(); + verifyingClient.setExecuteVerifier((actionResponse, actionRequest) -> { + assert actionRequest instanceof SearchRequest; + SearchRequest request = (SearchRequest) actionRequest; + assertEquals(Preference.LOCAL.type(), request.preference()); + assertEquals(1, request.source().size()); + assertEquals(QueryBuilders.termQuery(IP_RANGE_FIELD_NAME, ip), request.source().query()); + + SearchHit[] searchHitArray = {}; + SearchHits searchHits = new SearchHits(searchHitArray, new TotalHits(0l, TotalHits.Relation.EQUAL_TO), 0); + + SearchResponse response = mock(SearchResponse.class); + when(response.getHits()).thenReturn(searchHits); + return response; + }); + + // Run + Map geoData = verifyingGeoIpDataDao.getGeoIpData(indexName, ip); + + // Verify + assertTrue(geoData.isEmpty()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoCachedDaoTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoCachedDaoTests.java new file mode 100644 index 0000000000..9ce2f792d7 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoCachedDaoTests.java @@ -0,0 +1,265 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import lombok.SneakyThrows; + +import org.junit.Before; +import org.opensearch.common.network.NetworkAddress; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.index.shard.ShardId; +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.index.IndexNotFoundException; +import org.opensearch.index.engine.Engine; + +public class Ip2GeoCachedDaoTests extends Ip2GeoTestCase { + private Ip2GeoCachedDao ip2GeoCachedDao; + + @Before + public void init() { + ip2GeoCachedDao = new Ip2GeoCachedDao(clusterService, datasourceDao, geoIpDataDao); + } + + public void testGetIndexName_whenCalled_thenReturnIndexName() { + Datasource datasource = randomDatasource(); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + // Run + String indexName = ip2GeoCachedDao.getIndexName(datasource.getName()); + + // Verify + assertEquals(datasource.currentIndexName(), indexName); + } + + public void testGetIndexName_whenIndexNotFound_thenReturnNull() { + when(datasourceDao.getAllDatasources()).thenThrow(new IndexNotFoundException("not found")); + + // Run + String indexName = ip2GeoCachedDao.getIndexName(GeospatialTestHelper.randomLowerCaseString()); + + // Verify + assertNull(indexName); + } + + public void testIsExpired_whenExpired_thenReturnTrue() { + Datasource datasource = randomDatasource(); + datasource.getUpdateStats().setLastSucceededAt(Instant.MIN); + datasource.getUpdateStats().setLastSkippedAt(null); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + // Run + boolean isExpired = ip2GeoCachedDao.isExpired(datasource.getName()); + + // Verify + assertTrue(isExpired); + } + + public void testIsExpired_whenNotExpired_thenReturnFalse() { + Datasource datasource = randomDatasource(); + datasource.getUpdateStats().setLastSucceededAt(Instant.now()); + datasource.getUpdateStats().setLastSkippedAt(null); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + // Run + boolean isExpired = ip2GeoCachedDao.isExpired(datasource.getName()); + + // Verify + assertFalse(isExpired); + } + + public void testHas_whenHasDatasource_thenReturnTrue() { + Datasource datasource = randomDatasource(); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + // Run + boolean hasDatasource = ip2GeoCachedDao.has(datasource.getName()); + + // Verify + assertTrue(hasDatasource); + } + + public void testHas_whenNoDatasource_thenReturnFalse() { + Datasource datasource = randomDatasource(); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + // Run + boolean hasDatasource = ip2GeoCachedDao.has(datasourceName); + + // Verify + assertFalse(hasDatasource); + } + + public void testGetState_whenCalled_thenReturnState() { + Datasource datasource = randomDatasource(); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + // Run + DatasourceState state = ip2GeoCachedDao.getState(datasource.getName()); + + // Verify + assertEquals(datasource.getState(), state); + } + + public void testGetGeoData_whenCalled_thenReturnGeoData() { + Datasource datasource = randomDatasource(); + String ip = NetworkAddress.format(randomIp(false)); + Map expectedGeoData = Map.of("city", "Seattle"); + when(geoIpDataDao.getGeoIpData(datasource.currentIndexName(), ip)).thenReturn(expectedGeoData); + + // Run + Map geoData = ip2GeoCachedDao.getGeoData(datasource.currentIndexName(), ip); + + // Verify + assertEquals(expectedGeoData, geoData); + } + + @SneakyThrows + public void testPostIndex_whenFailed_thenNoUpdate() { + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList()); + Datasource datasource = randomDatasource(); + + ShardId shardId = mock(ShardId.class); + Engine.Index index = mock(Engine.Index.class); + BytesReference bytesReference = BytesReference.bytes(datasource.toXContent(XContentFactory.jsonBuilder(), null)); + when(index.source()).thenReturn(bytesReference); + Engine.IndexResult result = mock(Engine.IndexResult.class); + when(result.getResultType()).thenReturn(Engine.Result.Type.FAILURE); + + // Run + ip2GeoCachedDao.postIndex(shardId, index, result); + + // Verify + assertFalse(ip2GeoCachedDao.has(datasource.getName())); + assertTrue(ip2GeoCachedDao.isExpired(datasource.getName())); + assertNull(ip2GeoCachedDao.getIndexName(datasource.getName())); + assertNull(ip2GeoCachedDao.getState(datasource.getName())); + } + + @SneakyThrows + public void testPostIndex_whenSucceed_thenUpdate() { + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList()); + Datasource datasource = randomDatasource(); + + ShardId shardId = mock(ShardId.class); + Engine.Index index = mock(Engine.Index.class); + BytesReference bytesReference = BytesReference.bytes(datasource.toXContent(XContentFactory.jsonBuilder(), null)); + when(index.source()).thenReturn(bytesReference); + Engine.IndexResult result = mock(Engine.IndexResult.class); + when(result.getResultType()).thenReturn(Engine.Result.Type.SUCCESS); + + // Run + ip2GeoCachedDao.postIndex(shardId, index, result); + + // Verify + assertTrue(ip2GeoCachedDao.has(datasource.getName())); + assertFalse(ip2GeoCachedDao.isExpired(datasource.getName())); + assertEquals(datasource.currentIndexName(), ip2GeoCachedDao.getIndexName(datasource.getName())); + assertEquals(datasource.getState(), ip2GeoCachedDao.getState(datasource.getName())); + } + + public void testPostDelete_whenFailed_thenNoUpdate() { + Datasource datasource = randomDatasource(); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + ShardId shardId = mock(ShardId.class); + Engine.Delete index = mock(Engine.Delete.class); + Engine.DeleteResult result = mock(Engine.DeleteResult.class); + when(result.getResultType()).thenReturn(Engine.Result.Type.FAILURE); + + // Run + ip2GeoCachedDao.postDelete(shardId, index, result); + + // Verify + assertTrue(ip2GeoCachedDao.has(datasource.getName())); + } + + public void testPostDelete_whenSucceed_thenUpdate() { + Datasource datasource = randomDatasource(); + when(datasourceDao.getAllDatasources()).thenReturn(Arrays.asList(datasource)); + + ShardId shardId = mock(ShardId.class); + Engine.Delete index = mock(Engine.Delete.class); + when(index.id()).thenReturn(datasource.getName()); + Engine.DeleteResult result = mock(Engine.DeleteResult.class); + when(result.getResultType()).thenReturn(Engine.Result.Type.SUCCESS); + + // Run + ip2GeoCachedDao.postDelete(shardId, index, result); + + // Verify + assertFalse(ip2GeoCachedDao.has(datasource.getName())); + } + + @SneakyThrows + public void testUpdateMaxSize_whenBiggerSize_thenContainsAllData() { + int cacheSize = 10; + String datasource = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoCachedDao.GeoDataCache geoDataCache = new Ip2GeoCachedDao.GeoDataCache(cacheSize); + List ips = new ArrayList<>(cacheSize); + for (int i = 0; i < cacheSize; i++) { + String ip = NetworkAddress.format(randomIp(false)); + ips.add(ip); + geoDataCache.putIfAbsent(datasource, ip, addr -> Collections.emptyMap()); + } + + // Verify all data exist in the cache + assertTrue(ips.stream().allMatch(ip -> geoDataCache.get(datasource, ip) != null)); + + // Update cache size + int newCacheSize = 15; + geoDataCache.updateMaxSize(newCacheSize); + + // Verify all data exist in the cache + assertTrue(ips.stream().allMatch(ip -> geoDataCache.get(datasource, ip) != null)); + + // Add (newCacheSize - cacheSize + 1) data and the first data should not be available in the cache + for (int i = 0; i < newCacheSize - cacheSize + 1; i++) { + geoDataCache.putIfAbsent(datasource, NetworkAddress.format(randomIp(false)), addr -> Collections.emptyMap()); + } + assertNull(geoDataCache.get(datasource, ips.get(0))); + } + + @SneakyThrows + public void testUpdateMaxSize_whenSmallerSize_thenContainsPartialData() { + int cacheSize = 10; + String datasource = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoCachedDao.GeoDataCache geoDataCache = new Ip2GeoCachedDao.GeoDataCache(cacheSize); + List ips = new ArrayList<>(cacheSize); + for (int i = 0; i < cacheSize; i++) { + String ip = NetworkAddress.format(randomIp(false)); + ips.add(ip); + geoDataCache.putIfAbsent(datasource, ip, addr -> Collections.emptyMap()); + } + + // Verify all data exist in the cache + assertTrue(ips.stream().allMatch(ip -> geoDataCache.get(datasource, ip) != null)); + + // Update cache size + int newCacheSize = 5; + geoDataCache.updateMaxSize(newCacheSize); + + // Verify the last (cacheSize - newCacheSize) data is available in the cache + List deleted = ips.subList(0, ips.size() - newCacheSize); + List retained = ips.subList(ips.size() - newCacheSize, ips.size()); + assertTrue(deleted.stream().allMatch(ip -> geoDataCache.get(datasource, ip) == null)); + assertTrue(retained.stream().allMatch(ip -> geoDataCache.get(datasource, ip) != null)); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoProcessorDaoTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoProcessorDaoTests.java new file mode 100644 index 0000000000..9088b0defb --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/dao/Ip2GeoProcessorDaoTests.java @@ -0,0 +1,78 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.dao; + +import static org.mockito.Mockito.when; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import org.junit.Before; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.processor.Ip2GeoProcessor; +import org.opensearch.ingest.IngestMetadata; +import org.opensearch.ingest.PipelineConfiguration; + +public class Ip2GeoProcessorDaoTests extends Ip2GeoTestCase { + private Ip2GeoProcessorDao ip2GeoProcessorDao; + + @Before + public void init() { + ip2GeoProcessorDao = new Ip2GeoProcessorDao(ingestService); + } + + public void testGetProcessors_whenNullMetadata_thenReturnEmpty() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + when(metadata.custom(IngestMetadata.TYPE)).thenReturn(null); + + List ip2GeoProcessorList = ip2GeoProcessorDao.getProcessors(datasourceName); + assertTrue(ip2GeoProcessorList.isEmpty()); + } + + public void testGetProcessors_whenNoProcessorForGivenDatasource_thenReturnEmpty() { + String datasourceBeingUsed = GeospatialTestHelper.randomLowerCaseString(); + String datasourceNotBeingUsed = GeospatialTestHelper.randomLowerCaseString(); + String pipelineId = GeospatialTestHelper.randomLowerCaseString(); + Map pipelines = new HashMap<>(); + pipelines.put(pipelineId, createPipelineConfiguration()); + IngestMetadata ingestMetadata = new IngestMetadata(pipelines); + when(metadata.custom(IngestMetadata.TYPE)).thenReturn(ingestMetadata); + Ip2GeoProcessor ip2GeoProcessor = randomIp2GeoProcessor(datasourceBeingUsed); + when(ingestService.getProcessorsInPipeline(pipelineId, Ip2GeoProcessor.class)).thenReturn(Arrays.asList(ip2GeoProcessor)); + + List ip2GeoProcessorList = ip2GeoProcessorDao.getProcessors(datasourceNotBeingUsed); + assertTrue(ip2GeoProcessorList.isEmpty()); + } + + public void testGetProcessors_whenProcessorsForGivenDatasource_thenReturnProcessors() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + String pipelineId = GeospatialTestHelper.randomLowerCaseString(); + Map pipelines = new HashMap<>(); + pipelines.put(pipelineId, createPipelineConfiguration()); + IngestMetadata ingestMetadata = new IngestMetadata(pipelines); + when(metadata.custom(IngestMetadata.TYPE)).thenReturn(ingestMetadata); + Ip2GeoProcessor ip2GeoProcessor = randomIp2GeoProcessor(datasourceName); + when(ingestService.getProcessorsInPipeline(pipelineId, Ip2GeoProcessor.class)).thenReturn(Arrays.asList(ip2GeoProcessor)); + + List ip2GeoProcessorList = ip2GeoProcessorDao.getProcessors(datasourceName); + assertEquals(1, ip2GeoProcessorList.size()); + assertEquals(ip2GeoProcessor.getDatasourceName(), ip2GeoProcessorList.get(0).getDatasourceName()); + } + + private PipelineConfiguration createPipelineConfiguration() { + String id = GeospatialTestHelper.randomLowerCaseString(); + ByteBuffer byteBuffer = ByteBuffer.wrap(GeospatialTestHelper.randomLowerCaseString().getBytes(StandardCharsets.US_ASCII)); + BytesReference config = BytesReference.fromByteBuffer(byteBuffer); + return new PipelineConfiguration(id, config, XContentType.JSON); + } +} 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 0000000000..0ea221170e --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceExtensionTests.java @@ -0,0 +1,43 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +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.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() { + 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(randomPositiveLong(), randomPositiveLong(), randomPositiveLong()) + ); + + 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 0000000000..d4c460c4e2 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceRunnerTests.java @@ -0,0 +1,230 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +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.Duration; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.Optional; + +import lombok.SneakyThrows; + +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.geospatial.ip2geo.common.DatasourceState; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.jobscheduler.spi.JobDocVersion; +import org.opensearch.jobscheduler.spi.JobExecutionContext; +import org.opensearch.jobscheduler.spi.LockModel; +import org.opensearch.jobscheduler.spi.ScheduledJobParameter; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; + +public class DatasourceRunnerTests extends Ip2GeoTestCase { + @Before + public void init() { + DatasourceRunner.getJobRunnerInstance() + .initialize(clusterService, datasourceUpdateService, ip2GeoExecutor, datasourceDao, ip2GeoLockService); + } + + public void testGetJobRunnerInstance_whenCalledAgain_thenReturnSameInstance() { + assertTrue(DatasourceRunner.getJobRunnerInstance() == DatasourceRunner.getJobRunnerInstance()); + } + + public void testRunJob_whenInvalidClass_thenThrowException() { + JobDocVersion jobDocVersion = new JobDocVersion(randomInt(), randomInt(), randomInt()); + String jobIndexName = randomLowerCaseString(); + String jobId = randomLowerCaseString(); + JobExecutionContext jobExecutionContext = new JobExecutionContext(Instant.now(), jobDocVersion, lockService, jobIndexName, jobId); + ScheduledJobParameter jobParameter = mock(ScheduledJobParameter.class); + + // Run + expectThrows(IllegalStateException.class, () -> DatasourceRunner.getJobRunnerInstance().runJob(jobParameter, jobExecutionContext)); + } + + @SneakyThrows + public void testRunJob_whenValidInput_thenSucceed() { + 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 = randomDatasource(); + + LockModel lockModel = randomLockModel(); + when(ip2GeoLockService.acquireLock(datasource.getName(), Ip2GeoLockService.LOCK_DURATION_IN_SECONDS)).thenReturn( + Optional.of(lockModel) + ); + + // Run + DatasourceRunner.getJobRunnerInstance().runJob(datasource, jobExecutionContext); + + // Verify + verify(ip2GeoLockService).acquireLock(datasource.getName(), Ip2GeoLockService.LOCK_DURATION_IN_SECONDS); + verify(datasourceDao).getDatasource(datasource.getName()); + verify(ip2GeoLockService).releaseLock(lockModel); + } + + @SneakyThrows + public void testUpdateDatasourceRunner_whenExceptionBeforeAcquiringLock_thenNoReleaseLock() { + ScheduledJobParameter jobParameter = mock(ScheduledJobParameter.class); + when(jobParameter.getName()).thenReturn(GeospatialTestHelper.randomLowerCaseString()); + when(ip2GeoLockService.acquireLock(jobParameter.getName(), Ip2GeoLockService.LOCK_DURATION_IN_SECONDS)).thenThrow( + new RuntimeException() + ); + + // Run + expectThrows(Exception.class, () -> DatasourceRunner.getJobRunnerInstance().updateDatasourceRunner(jobParameter).run()); + + // Verify + verify(ip2GeoLockService, never()).releaseLock(any()); + } + + @SneakyThrows + public void testUpdateDatasourceRunner_whenExceptionAfterAcquiringLock_thenReleaseLock() { + ScheduledJobParameter jobParameter = mock(ScheduledJobParameter.class); + when(jobParameter.getName()).thenReturn(GeospatialTestHelper.randomLowerCaseString()); + LockModel lockModel = randomLockModel(); + when(ip2GeoLockService.acquireLock(jobParameter.getName(), Ip2GeoLockService.LOCK_DURATION_IN_SECONDS)).thenReturn( + Optional.of(lockModel) + ); + when(datasourceDao.getDatasource(jobParameter.getName())).thenThrow(new RuntimeException()); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasourceRunner(jobParameter).run(); + + // Verify + verify(ip2GeoLockService).releaseLock(any()); + } + + @SneakyThrows + public void testUpdateDatasource_whenDatasourceDoesNotExist_thenDoNothing() { + Datasource datasource = new Datasource(); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource, mock(Runnable.class)); + + // Verify + verify(datasourceUpdateService, never()).deleteUnusedIndices(any()); + } + + @SneakyThrows + public void testUpdateDatasource_whenInvalidState_thenUpdateLastFailedAt() { + Datasource datasource = new Datasource(); + datasource.enable(); + datasource.getUpdateStats().setLastFailedAt(null); + datasource.setState(randomStateExcept(DatasourceState.AVAILABLE)); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource, mock(Runnable.class)); + + // Verify + assertFalse(datasource.isEnabled()); + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceDao).updateDatasource(datasource); + } + + @SneakyThrows + public void testUpdateDatasource_whenValidInput_thenSucceed() { + Datasource datasource = randomDatasource(); + datasource.setState(DatasourceState.AVAILABLE); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + Runnable renewLock = mock(Runnable.class); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource, renewLock); + + // Verify + verify(datasourceUpdateService, times(2)).deleteUnusedIndices(datasource); + verify(datasourceUpdateService).updateOrCreateGeoIpData(datasource, renewLock); + verify(datasourceUpdateService).updateDatasource(datasource, datasource.getUserSchedule(), DatasourceTask.ALL); + } + + @SneakyThrows + public void testUpdateDatasource_whenDeleteTask_thenDeleteOnly() { + Datasource datasource = randomDatasource(); + datasource.setState(DatasourceState.AVAILABLE); + datasource.setTask(DatasourceTask.DELETE_UNUSED_INDICES); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + Runnable renewLock = mock(Runnable.class); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource, renewLock); + + // Verify + verify(datasourceUpdateService, times(2)).deleteUnusedIndices(datasource); + verify(datasourceUpdateService, never()).updateOrCreateGeoIpData(datasource, renewLock); + verify(datasourceUpdateService).updateDatasource(datasource, datasource.getUserSchedule(), DatasourceTask.ALL); + } + + @SneakyThrows + public void testUpdateDatasource_whenExpired_thenDeleteIndicesAgain() { + Datasource datasource = randomDatasource(); + datasource.getUpdateStats().setLastSkippedAt(null); + datasource.getUpdateStats() + .setLastSucceededAt(Instant.now().minus(datasource.getDatabase().getValidForInDays() + 1, ChronoUnit.DAYS)); + datasource.setState(DatasourceState.AVAILABLE); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + Runnable renewLock = mock(Runnable.class); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource, renewLock); + + // Verify + verify(datasourceUpdateService, times(3)).deleteUnusedIndices(datasource); + verify(datasourceUpdateService).updateOrCreateGeoIpData(datasource, renewLock); + verify(datasourceUpdateService).updateDatasource(datasource, datasource.getUserSchedule(), DatasourceTask.ALL); + } + + @SneakyThrows + public void testUpdateDatasource_whenWillExpire_thenScheduleDeleteTask() { + Datasource datasource = randomDatasource(); + datasource.getUpdateStats().setLastSkippedAt(null); + datasource.getUpdateStats() + .setLastSucceededAt(Instant.now().minus(datasource.getDatabase().getValidForInDays(), ChronoUnit.DAYS).plusSeconds(60)); + datasource.setState(DatasourceState.AVAILABLE); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + Runnable renewLock = mock(Runnable.class); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource, renewLock); + + // Verify + verify(datasourceUpdateService, times(2)).deleteUnusedIndices(datasource); + verify(datasourceUpdateService).updateOrCreateGeoIpData(datasource, renewLock); + + ArgumentCaptor captor = ArgumentCaptor.forClass(IntervalSchedule.class); + verify(datasourceUpdateService).updateDatasource(eq(datasource), captor.capture(), eq(DatasourceTask.DELETE_UNUSED_INDICES)); + assertTrue(Duration.between(datasource.expirationDay(), captor.getValue().getNextExecutionTime(Instant.now())).getSeconds() < 30); + } + + @SneakyThrows + public void testUpdateDatasourceExceptionHandling() { + Datasource datasource = new Datasource(); + datasource.setName(randomLowerCaseString()); + datasource.getUpdateStats().setLastFailedAt(null); + when(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + doThrow(new RuntimeException("test failure")).when(datasourceUpdateService).deleteUnusedIndices(any()); + + // Run + DatasourceRunner.getJobRunnerInstance().updateDatasource(datasource, mock(Runnable.class)); + + // Verify + assertNotNull(datasource.getUpdateStats().getLastFailedAt()); + verify(datasourceDao).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 new file mode 100644 index 0000000000..aaa0d29bd1 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceTests.java @@ -0,0 +1,190 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import static org.opensearch.geospatial.ip2geo.jobscheduler.Datasource.IP2GEO_DATA_INDEX_NAME_PREFIX; + +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Locale; + +import lombok.SneakyThrows; + +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoTestCase; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; + +public class DatasourceTests extends Ip2GeoTestCase { + + @SneakyThrows + public void testParser_whenAllValueIsFilled_thenSucceed() { + 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.setCurrentIndex(GeospatialTestHelper.randomLowerCaseString()); + datasource.getDatabase().setFields(Arrays.asList("field1", "field2")); + datasource.getDatabase().setProvider("test_provider"); + datasource.getDatabase().setUpdatedAt(Instant.now().truncatedTo(ChronoUnit.MILLIS)); + datasource.getDatabase().setSha256Hash(GeospatialTestHelper.randomLowerCaseString()); + datasource.getDatabase().setValidForInDays(1l); + datasource.getUpdateStats().setLastProcessingTimeInMillis(randomPositiveLong()); + 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)); + } + + @SneakyThrows + public void testParser_whenNullForOptionalFields_thenSucceed() { + 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.PARSER.parse( + createParser(datasource.toXContent(XContentFactory.jsonBuilder(), null)), + null + ); + assertTrue(datasource.equals(anotherDatasource)); + } + + public void testCurrentIndexName_whenNotExpired_thenReturnName() { + String id = GeospatialTestHelper.randomLowerCaseString(); + Instant now = Instant.now(); + Datasource datasource = new Datasource(); + datasource.setName(id); + datasource.setCurrentIndex(datasource.newIndexName(GeospatialTestHelper.randomLowerCaseString())); + datasource.getDatabase().setProvider("provider"); + datasource.getDatabase().setSha256Hash("sha256Hash"); + datasource.getDatabase().setUpdatedAt(now); + datasource.getDatabase().setFields(new ArrayList<>()); + + assertNotNull(datasource.currentIndexName()); + } + + public void testCurrentIndexName_whenExpired_thenReturnNull() { + String id = GeospatialTestHelper.randomLowerCaseString(); + Instant now = Instant.now(); + Datasource datasource = new Datasource(); + datasource.setName(id); + datasource.setCurrentIndex(datasource.newIndexName(GeospatialTestHelper.randomLowerCaseString())); + datasource.getDatabase().setProvider("provider"); + datasource.getDatabase().setSha256Hash("sha256Hash"); + datasource.getDatabase().setUpdatedAt(now); + datasource.getDatabase().setValidForInDays(1l); + datasource.getUpdateStats().setLastSucceededAt(Instant.now().minus(25, ChronoUnit.HOURS)); + datasource.getDatabase().setFields(new ArrayList<>()); + + assertTrue(datasource.isExpired()); + assertNull(datasource.currentIndexName()); + } + + @SneakyThrows + public void testNewIndexName_whenCalled_thenReturnedExpectedValue() { + String name = GeospatialTestHelper.randomLowerCaseString(); + String suffix = GeospatialTestHelper.randomLowerCaseString(); + Datasource datasource = new Datasource(); + datasource.setName(name); + assertEquals(String.format(Locale.ROOT, "%s.%s.%s", IP2GEO_DATA_INDEX_NAME_PREFIX, name, suffix), datasource.newIndexName(suffix)); + } + + public void testResetDatabase_whenCalled_thenNullifySomeFields() { + Datasource datasource = randomDatasource(); + assertNotNull(datasource.getDatabase().getSha256Hash()); + assertNotNull(datasource.getDatabase().getUpdatedAt()); + + // Run + datasource.resetDatabase(); + + // Verify + assertNull(datasource.getDatabase().getSha256Hash()); + assertNull(datasource.getDatabase().getUpdatedAt()); + } + + public void testIsExpired_whenCalled_thenExpectedValue() { + 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 testWillExpired_whenCalled_thenExpectedValue() { + Datasource datasource = new Datasource(); + // never expire if validForInDays is null + assertFalse(datasource.willExpire(Instant.MAX)); + + long validForInDays = 1; + datasource.getDatabase().setValidForInDays(validForInDays); + + // if last skipped date is null, use only last succeeded date to determine + datasource.getUpdateStats().setLastSucceededAt(Instant.now().minus(1, ChronoUnit.DAYS)); + assertTrue( + datasource.willExpire(datasource.getUpdateStats().getLastSucceededAt().plus(validForInDays, ChronoUnit.DAYS).plusSeconds(1)) + ); + assertFalse(datasource.willExpire(datasource.getUpdateStats().getLastSucceededAt().plus(validForInDays, ChronoUnit.DAYS))); + + // use the latest date between last skipped date and last succeeded date to determine + datasource.getUpdateStats().setLastSkippedAt(Instant.now()); + assertTrue( + datasource.willExpire(datasource.getUpdateStats().getLastSkippedAt().plus(validForInDays, ChronoUnit.DAYS).plusSeconds(1)) + ); + assertFalse(datasource.willExpire(datasource.getUpdateStats().getLastSkippedAt().plus(validForInDays, ChronoUnit.DAYS))); + + datasource.getUpdateStats().setLastSkippedAt(Instant.now().minus(1, ChronoUnit.HOURS)); + datasource.getUpdateStats().setLastSucceededAt(Instant.now()); + assertTrue( + datasource.willExpire(datasource.getUpdateStats().getLastSucceededAt().plus(validForInDays, ChronoUnit.DAYS).plusSeconds(1)) + ); + assertFalse(datasource.willExpire(datasource.getUpdateStats().getLastSucceededAt().plus(validForInDays, ChronoUnit.DAYS))); + } + + public void testExpirationDay_whenCalled_thenExpectedValue() { + Datasource datasource = new Datasource(); + datasource.getDatabase().setValidForInDays(null); + assertEquals(Instant.MAX, datasource.expirationDay()); + + long validForInDays = 1; + datasource.getDatabase().setValidForInDays(validForInDays); + + // if last skipped date is null, use only last succeeded date to determine + datasource.getUpdateStats().setLastSucceededAt(Instant.now().minus(1, ChronoUnit.DAYS)); + assertEquals(datasource.getUpdateStats().getLastSucceededAt().plus(validForInDays, ChronoUnit.DAYS), datasource.expirationDay()); + + // use the latest date between last skipped date and last succeeded date to determine + datasource.getUpdateStats().setLastSkippedAt(Instant.now()); + assertEquals(datasource.getUpdateStats().getLastSkippedAt().plus(validForInDays, ChronoUnit.DAYS), datasource.expirationDay()); + + datasource.getUpdateStats().setLastSkippedAt(Instant.now().minus(1, ChronoUnit.HOURS)); + datasource.getUpdateStats().setLastSucceededAt(Instant.now()); + assertEquals(datasource.getUpdateStats().getLastSucceededAt().plus(validForInDays, ChronoUnit.DAYS), datasource.expirationDay()); + } + + 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 0000000000..5e08400b27 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/jobscheduler/DatasourceUpdateServiceTests.java @@ -0,0 +1,272 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.jobscheduler; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isA; +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 java.io.File; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.time.temporal.ChronoUnit; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; + +import lombok.SneakyThrows; + +import org.apache.commons.csv.CSVFormat; +import org.apache.commons.csv.CSVParser; +import org.junit.Before; +import org.opensearch.OpenSearchException; +import org.opensearch.cluster.routing.ShardRouting; +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; +import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule; + +@SuppressForbidden(reason = "unit test") +public class DatasourceUpdateServiceTests extends Ip2GeoTestCase { + private DatasourceUpdateService datasourceUpdateService; + + @Before + public void init() { + datasourceUpdateService = new DatasourceUpdateService(clusterService, datasourceDao, geoIpDataDao); + } + + @SneakyThrows + public void testUpdateOrCreateGeoIpData_whenHashValueIsSame_thenSkipUpdate() { + 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().setSha256Hash(manifest.getSha256Hash()); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + + // Run + datasourceUpdateService.updateOrCreateGeoIpData(datasource, mock(Runnable.class)); + + // Verify + assertNotNull(datasource.getUpdateStats().getLastSkippedAt()); + verify(datasourceDao).updateDatasource(datasource); + } + + @SneakyThrows + public void testUpdateOrCreateGeoIpData_whenExpired_thenUpdate() { + 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(geoIpDataDao.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())); + datasource.getDatabase().setSha256Hash(manifest.getSha256Hash()); + datasource.getDatabase().setValidForInDays(1l); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + datasource.resetDatabase(); + + // Run + datasourceUpdateService.updateOrCreateGeoIpData(datasource, mock(Runnable.class)); + + // Verify + verify(geoIpDataDao).putGeoIpData(eq(datasource.currentIndexName()), isA(String[].class), any(Iterator.class), any(Runnable.class)); + } + + @SneakyThrows + public void testUpdateOrCreateGeoIpData_whenInvalidData_thenThrowException() { + 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(geoIpDataDao.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().setSha256Hash(manifest.getSha256Hash().substring(1)); + datasource.getDatabase().setFields(Arrays.asList("country_name")); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + + // Run + expectThrows(OpenSearchException.class, () -> datasourceUpdateService.updateOrCreateGeoIpData(datasource, mock(Runnable.class))); + } + + @SneakyThrows + public void testUpdateOrCreateGeoIpData_whenIncompatibleFields_thenThrowException() { + 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(geoIpDataDao.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().setSha256Hash(manifest.getSha256Hash().substring(1)); + datasource.getDatabase().setFields(Arrays.asList("country_name", "additional_field")); + datasource.setEndpoint(manifestFile.toURI().toURL().toExternalForm()); + + // Run + expectThrows(OpenSearchException.class, () -> datasourceUpdateService.updateOrCreateGeoIpData(datasource, mock(Runnable.class))); + } + + @SneakyThrows + public void testUpdateOrCreateGeoIpData_whenValidInput_thenSucceed() { + 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(geoIpDataDao.getDatabaseReader(any())).thenReturn(CSVParser.parse(sampleFile, StandardCharsets.UTF_8, CSVFormat.RFC4180)); + ShardRouting shardRouting = mock(ShardRouting.class); + when(shardRouting.started()).thenReturn(true); + when(routingTable.allShards(anyString())).thenReturn(Arrays.asList(shardRouting)); + + Datasource datasource = new Datasource(); + datasource.setState(DatasourceState.AVAILABLE); + datasource.getDatabase().setUpdatedAt(Instant.ofEpochMilli(manifest.getUpdatedAt() - 1)); + datasource.getDatabase().setSha256Hash(manifest.getSha256Hash().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, mock(Runnable.class)); + + // Verify + assertEquals(manifest.getProvider(), datasource.getDatabase().getProvider()); + assertEquals(manifest.getSha256Hash(), datasource.getDatabase().getSha256Hash()); + assertEquals(Instant.ofEpochMilli(manifest.getUpdatedAt()), datasource.getDatabase().getUpdatedAt()); + assertEquals(manifest.getValidForInDays(), datasource.getDatabase().getValidForInDays()); + assertNotNull(datasource.getUpdateStats().getLastSucceededAt()); + assertNotNull(datasource.getUpdateStats().getLastProcessingTimeInMillis()); + verify(datasourceDao, times(2)).updateDatasource(datasource); + verify(geoIpDataDao).putGeoIpData(eq(datasource.currentIndexName()), isA(String[].class), any(Iterator.class), any(Runnable.class)); + } + + public void testWaitUntilAllShardsStarted_whenTimedOut_thenThrowException() { + String indexName = GeospatialTestHelper.randomLowerCaseString(); + ShardRouting shardRouting = mock(ShardRouting.class); + when(shardRouting.started()).thenReturn(false); + when(routingTable.allShards(indexName)).thenReturn(Arrays.asList(shardRouting)); + + // Run + Exception e = expectThrows(OpenSearchException.class, () -> datasourceUpdateService.waitUntilAllShardsStarted(indexName, 10)); + + // Verify + assertTrue(e.getMessage().contains("did not complete")); + } + + @SneakyThrows + public void testWaitUntilAllShardsStarted_whenInterrupted_thenThrowException() { + String indexName = GeospatialTestHelper.randomLowerCaseString(); + ShardRouting shardRouting = mock(ShardRouting.class); + when(shardRouting.started()).thenReturn(false); + when(routingTable.allShards(indexName)).thenReturn(Arrays.asList(shardRouting)); + + // Run + Thread.currentThread().interrupt(); + Exception e = expectThrows(RuntimeException.class, () -> datasourceUpdateService.waitUntilAllShardsStarted(indexName, 10)); + + // Verify + assertEquals(InterruptedException.class, e.getCause().getClass()); + } + + @SneakyThrows + public void testGetHeaderFields_whenValidInput_thenReturnCorrectValue() { + File manifestFile = new File(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").getFile()); + + File sampleFile = new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.csv").getFile()); + when(geoIpDataDao.getDatabaseReader(any())).thenReturn(CSVParser.parse(sampleFile, StandardCharsets.UTF_8, CSVFormat.RFC4180)); + + // Run + assertEquals(Arrays.asList("country_name"), datasourceUpdateService.getHeaderFields(manifestFile.toURI().toURL().toExternalForm())); + } + + @SneakyThrows + public void testDeleteUnusedIndices_whenValidInput_thenSucceed() { + 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.setName(datasourceName); + datasource.setCurrentIndex(currentIndex); + 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); + + datasourceUpdateService.deleteUnusedIndices(datasource); + + assertEquals(1, datasource.getIndices().size()); + assertEquals(currentIndex, datasource.getIndices().get(0)); + verify(datasourceDao).updateDatasource(datasource); + verify(geoIpDataDao).deleteIp2GeoDataIndex(oldIndex); + } + + public void testUpdateDatasource_whenNoChange_thenNoUpdate() { + Datasource datasource = randomDatasource(); + + // Run + datasourceUpdateService.updateDatasource(datasource, datasource.getSystemSchedule(), datasource.getTask()); + + // Verify + verify(datasourceDao, never()).updateDatasource(any()); + } + + public void testUpdateDatasource_whenChange_thenUpdate() { + Datasource datasource = randomDatasource(); + datasource.setTask(DatasourceTask.ALL); + + // Run + datasourceUpdateService.updateDatasource( + datasource, + new IntervalSchedule(Instant.now(), datasource.getSystemSchedule().getInterval() + 1, ChronoUnit.DAYS), + datasource.getTask() + ); + datasourceUpdateService.updateDatasource(datasource, datasource.getSystemSchedule(), DatasourceTask.DELETE_UNUSED_INDICES); + + // Verify + verify(datasourceDao, times(2)).updateDatasource(any()); + } + + @SneakyThrows + public void testGetHeaderFields_whenValidInput_thenSucceed() { + File manifestFile = new File(this.getClass().getClassLoader().getResource("ip2geo/manifest.json").getFile()); + File sampleFile = new File(this.getClass().getClassLoader().getResource("ip2geo/sample_valid.csv").getFile()); + when(geoIpDataDao.getDatabaseReader(any())).thenReturn(CSVParser.parse(sampleFile, StandardCharsets.UTF_8, CSVFormat.RFC4180)); + + // Run + List fields = datasourceUpdateService.getHeaderFields(manifestFile.toURI().toURL().toExternalForm()); + + // Verify + List expectedFields = Arrays.asList("country_name"); + assertEquals(expectedFields, fields); + } +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/listener/Ip2GeoListenerTests.java b/src/test/java/org/opensearch/geospatial/ip2geo/listener/Ip2GeoListenerTests.java new file mode 100644 index 0000000000..d31f38bcc9 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/listener/Ip2GeoListenerTests.java @@ -0,0 +1,199 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.listener; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +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 java.time.Instant; +import java.util.Arrays; +import java.util.List; + +import org.junit.Before; +import org.mockito.ArgumentCaptor; +import org.opensearch.action.ActionListener; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.RestoreInProgress; +import org.opensearch.common.settings.Settings; +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.geospatial.ip2geo.jobscheduler.DatasourceTask; +import org.opensearch.snapshots.Snapshot; +import org.opensearch.snapshots.SnapshotId; + +public class Ip2GeoListenerTests extends Ip2GeoTestCase { + private Ip2GeoListener ip2GeoListener; + + @Before + public void init() { + ip2GeoListener = new Ip2GeoListener(clusterService, threadPool, datasourceDao, geoIpDataDao); + } + + public void testDoStart_whenClusterManagerNode_thenAddListener() { + Settings settings = Settings.builder().put("node.roles", "cluster_manager").build(); + when(clusterService.getSettings()).thenReturn(settings); + + // Run + ip2GeoListener.doStart(); + + // Verify + verify(clusterService).addListener(ip2GeoListener); + } + + public void testDoStart_whenNotClusterManagerNode_thenDoNotAddListener() { + Settings settings = Settings.builder().put("node.roles", "data").build(); + when(clusterService.getSettings()).thenReturn(settings); + + // Run + ip2GeoListener.doStart(); + + // Verify + verify(clusterService, never()).addListener(ip2GeoListener); + } + + public void testDoStop_whenCalled_thenRemoveListener() { + // Run + ip2GeoListener.doStop(); + + // Verify + verify(clusterService).removeListener(ip2GeoListener); + } + + public void testClusterChanged_whenNotClusterManagerNode_thenDoNothing() { + ClusterChangedEvent event = mock(ClusterChangedEvent.class); + when(event.localNodeClusterManager()).thenReturn(false); + + // Run + ip2GeoListener.clusterChanged(event); + + // Verify + verify(threadPool, never()).generic(); + } + + public void testClusterChanged_whenNotComplete_thenDoNothing() { + SnapshotId snapshotId = new SnapshotId(GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString()); + Snapshot snapshot = new Snapshot(GeospatialTestHelper.randomLowerCaseString(), snapshotId); + RestoreInProgress.Entry entry = new RestoreInProgress.Entry( + GeospatialTestHelper.randomLowerCaseString(), + snapshot, + RestoreInProgress.State.STARTED, + Arrays.asList(DatasourceExtension.JOB_INDEX_NAME), + null + ); + RestoreInProgress restoreInProgress = new RestoreInProgress.Builder().add(entry).build(); + ClusterState clusterState = mock(ClusterState.class); + when(clusterState.custom(RestoreInProgress.TYPE, RestoreInProgress.EMPTY)).thenReturn(restoreInProgress); + ClusterChangedEvent event = mock(ClusterChangedEvent.class); + when(event.localNodeClusterManager()).thenReturn(true); + when(event.state()).thenReturn(clusterState); + + // Run + ip2GeoListener.clusterChanged(event); + + // Verify + verify(threadPool, never()).generic(); + } + + public void testClusterChanged_whenNotDatasourceIndex_thenDoNothing() { + SnapshotId snapshotId = new SnapshotId(GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString()); + Snapshot snapshot = new Snapshot(GeospatialTestHelper.randomLowerCaseString(), snapshotId); + RestoreInProgress.Entry entry = new RestoreInProgress.Entry( + GeospatialTestHelper.randomLowerCaseString(), + snapshot, + RestoreInProgress.State.FAILURE, + Arrays.asList(GeospatialTestHelper.randomLowerCaseString()), + null + ); + RestoreInProgress restoreInProgress = new RestoreInProgress.Builder().add(entry).build(); + ClusterState clusterState = mock(ClusterState.class); + when(clusterState.custom(RestoreInProgress.TYPE, RestoreInProgress.EMPTY)).thenReturn(restoreInProgress); + ClusterChangedEvent event = mock(ClusterChangedEvent.class); + when(event.localNodeClusterManager()).thenReturn(true); + when(event.state()).thenReturn(clusterState); + + // Run + ip2GeoListener.clusterChanged(event); + + // Verify + verify(threadPool, never()).generic(); + } + + public void testClusterChanged_whenDatasourceIndexIsRestored_thenUpdate() { + SnapshotId snapshotId = new SnapshotId(GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString()); + Snapshot snapshot = new Snapshot(GeospatialTestHelper.randomLowerCaseString(), snapshotId); + RestoreInProgress.Entry entry = new RestoreInProgress.Entry( + GeospatialTestHelper.randomLowerCaseString(), + snapshot, + RestoreInProgress.State.SUCCESS, + Arrays.asList(DatasourceExtension.JOB_INDEX_NAME), + null + ); + RestoreInProgress restoreInProgress = new RestoreInProgress.Builder().add(entry).build(); + ClusterState clusterState = mock(ClusterState.class); + when(clusterState.custom(RestoreInProgress.TYPE, RestoreInProgress.EMPTY)).thenReturn(restoreInProgress); + ClusterChangedEvent event = mock(ClusterChangedEvent.class); + when(event.localNodeClusterManager()).thenReturn(true); + when(event.state()).thenReturn(clusterState); + + // Run + ip2GeoListener.clusterChanged(event); + + // Verify + verify(threadPool).generic(); + ArgumentCaptor>> captor = ArgumentCaptor.forClass(ActionListener.class); + verify(datasourceDao).getAllDatasources(captor.capture()); + + // Run + List datasources = Arrays.asList(randomDatasource(), randomDatasource()); + datasources.stream().forEach(datasource -> { datasource.setTask(DatasourceTask.DELETE_UNUSED_INDICES); }); + + captor.getValue().onResponse(datasources); + + // Verify + datasources.stream().forEach(datasource -> { + assertEquals(DatasourceTask.ALL, datasource.getTask()); + assertNull(datasource.getDatabase().getUpdatedAt()); + assertNull(datasource.getDatabase().getSha256Hash()); + assertTrue(datasource.getSystemSchedule().getNextExecutionTime(Instant.now()).isAfter(Instant.now())); + assertTrue(datasource.getSystemSchedule().getNextExecutionTime(Instant.now()).isBefore(Instant.now().plusSeconds(60))); + }); + verify(datasourceDao).updateDatasource(eq(datasources), any()); + } + + public void testClusterChanged_whenGeoIpDataIsRestored_thenDelete() { + Datasource datasource = randomDatasource(); + SnapshotId snapshotId = new SnapshotId(GeospatialTestHelper.randomLowerCaseString(), GeospatialTestHelper.randomLowerCaseString()); + Snapshot snapshot = new Snapshot(GeospatialTestHelper.randomLowerCaseString(), snapshotId); + RestoreInProgress.Entry entry = new RestoreInProgress.Entry( + GeospatialTestHelper.randomLowerCaseString(), + snapshot, + RestoreInProgress.State.SUCCESS, + Arrays.asList(datasource.currentIndexName()), + null + ); + RestoreInProgress restoreInProgress = new RestoreInProgress.Builder().add(entry).build(); + ClusterState clusterState = mock(ClusterState.class); + when(clusterState.custom(RestoreInProgress.TYPE, RestoreInProgress.EMPTY)).thenReturn(restoreInProgress); + ClusterChangedEvent event = mock(ClusterChangedEvent.class); + when(event.localNodeClusterManager()).thenReturn(true); + when(event.state()).thenReturn(clusterState); + + // Run + ip2GeoListener.clusterChanged(event); + + // Verify + verify(threadPool).generic(); + verify(geoIpDataDao).deleteIp2GeoDataIndex(Arrays.asList(datasource.currentIndexName())); + } + +} diff --git a/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorIT.java b/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorIT.java new file mode 100644 index 0000000000..d73893f396 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorIT.java @@ -0,0 +1,225 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.processor; + +import java.io.IOException; +import java.time.Duration; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Collectors; + +import lombok.SneakyThrows; + +import org.opensearch.client.Response; +import org.opensearch.client.ResponseException; +import org.opensearch.common.Randomness; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.geospatial.GeospatialRestTestCase; +import org.opensearch.geospatial.GeospatialTestHelper; +import org.opensearch.geospatial.ip2geo.Ip2GeoDataServer; +import org.opensearch.geospatial.ip2geo.action.PutDatasourceRequest; + +public class Ip2GeoProcessorIT extends GeospatialRestTestCase { + // Use this value in resource name to avoid name conflict among tests + private static final String PREFIX = Ip2GeoProcessorIT.class.getSimpleName().toLowerCase(Locale.ROOT); + private static final String CITY = "city"; + private static final String COUNTRY = "country"; + private static final String IP = "ip"; + private static final String SOURCE = "_source"; + + @SneakyThrows + public void testCreateIp2GeoProcessor_whenValidInput_thenAddData() { + Ip2GeoDataServer.start(); + boolean isDatasourceCreated = false; + boolean isProcessorCreated = false; + String pipelineName = PREFIX + GeospatialTestHelper.randomLowerCaseString(); + String datasourceName = PREFIX + GeospatialTestHelper.randomLowerCaseString(); + try { + String targetField = GeospatialTestHelper.randomLowerCaseString(); + String field = GeospatialTestHelper.randomLowerCaseString(); + + Map datasourceProperties = Map.of( + PutDatasourceRequest.ENDPOINT_FIELD.getPreferredName(), + Ip2GeoDataServer.getEndpointCity() + ); + + // Create datasource and wait for it to be available + createDatasource(datasourceName, datasourceProperties); + isDatasourceCreated = true; + // Creation of datasource with same name should fail + ResponseException createException = expectThrows( + ResponseException.class, + () -> createDatasource(datasourceName, datasourceProperties) + ); + // Verify + assertEquals(RestStatus.BAD_REQUEST.getStatus(), createException.getResponse().getStatusLine().getStatusCode()); + waitForDatasourceToBeAvailable(datasourceName, Duration.ofSeconds(10)); + + Map processorProperties = Map.of( + Ip2GeoProcessor.CONFIG_FIELD, + field, + Ip2GeoProcessor.CONFIG_DATASOURCE, + datasourceName, + Ip2GeoProcessor.CONFIG_TARGET_FIELD, + targetField, + Ip2GeoProcessor.CONFIG_PROPERTIES, + Arrays.asList(CITY) + ); + + // Create ip2geo processor + createIp2GeoProcessorPipeline(pipelineName, processorProperties); + isProcessorCreated = true; + + Map> sampleData = getSampleData(); + List docs = sampleData.entrySet() + .stream() + .map(entry -> createDocument(field, entry.getKey())) + .collect(Collectors.toList()); + + // Simulate processor + Map response = simulatePipeline(pipelineName, docs); + + // Verify data added to document + List> sources = convertToListOfSources(response, targetField); + sources.stream().allMatch(source -> source.size() == 1); + List cities = sources.stream().map(value -> value.get(CITY)).collect(Collectors.toList()); + List expectedCities = sampleData.values().stream().map(value -> value.get(CITY)).collect(Collectors.toList()); + assertEquals(expectedCities, cities); + + // Delete datasource fails when there is a process using it + ResponseException deleteException = expectThrows(ResponseException.class, () -> deleteDatasource(datasourceName)); + // Verify + assertEquals(RestStatus.BAD_REQUEST.getStatus(), deleteException.getResponse().getStatusLine().getStatusCode()); + } finally { + Exception exception = null; + try { + if (isProcessorCreated) { + deletePipeline(pipelineName); + } + if (isDatasourceCreated) { + deleteDatasource(datasourceName, 3); + } + } catch (Exception e) { + exception = e; + } + Ip2GeoDataServer.stop(); + if (exception != null) { + throw exception; + } + } + } + + private Response createIp2GeoProcessorPipeline(final String pipelineName, final Map properties) throws IOException { + String field = GeospatialTestHelper.randomLowerCaseString(); + String datasourceName = PREFIX + GeospatialTestHelper.randomLowerCaseString(); + Map defaultProperties = Map.of( + Ip2GeoProcessor.CONFIG_FIELD, + field, + Ip2GeoProcessor.CONFIG_DATASOURCE, + datasourceName + ); + Map baseProperties = new HashMap<>(); + baseProperties.putAll(defaultProperties); + baseProperties.putAll(properties); + Map processorConfig = buildProcessorConfig(Ip2GeoProcessor.TYPE, baseProperties); + + return createPipeline(pipelineName, Optional.empty(), Arrays.asList(processorConfig)); + } + + private Map> getSampleData() { + Map> sampleData = new HashMap<>(); + sampleData.put( + String.format( + Locale.ROOT, + "10.%d.%d.%d", + Randomness.get().nextInt(255), + Randomness.get().nextInt(255), + Randomness.get().nextInt(255) + ), + Map.of(CITY, "Seattle", COUNTRY, "USA") + ); + sampleData.put( + String.format( + Locale.ROOT, + "127.%d.%d.%d", + Randomness.get().nextInt(15), + Randomness.get().nextInt(255), + Randomness.get().nextInt(255) + ), + Map.of(CITY, "Vancouver", COUNTRY, "Canada") + ); + sampleData.put( + String.format( + Locale.ROOT, + "fd12:2345:6789:1:%x:%x:%x:%x", + Randomness.get().nextInt(65535), + Randomness.get().nextInt(65535), + Randomness.get().nextInt(65535), + Randomness.get().nextInt(65535) + ), + Map.of(CITY, "Bengaluru", COUNTRY, "India") + ); + return sampleData; + } + + private Map> createDocument(String... args) { + if (args.length % 2 == 1) { + throw new RuntimeException("Number of arguments should be even"); + } + Map source = new HashMap<>(); + for (int i = 0; i < args.length; i += 2) { + source.put(args[0], args[1]); + } + return Map.of(SOURCE, source); + } + + /** + * This method convert returned value of simulatePipeline method to a list of sources + * + * For example, + * Input: + * { + * "docs" : [ + * { + * "doc" : { + * "_index" : "_index", + * "_id" : "_id", + * "_source" : { + * "ip2geo" : { + * "ip" : "127.0.0.1", + * "city" : "Seattle" + * }, + * "_ip" : "127.0.0.1" + * }, + * "_ingest" : { + * "timestamp" : "2023-05-12T17:41:42.939703Z" + * } + * } + * } + * ] + * } + * + * Output: + * [ + * { + * "ip" : "127.0.0.1", + * "city" : "Seattle" + * } + * ] + * + */ + private List> convertToListOfSources(final Map data, final String targetField) { + List>> docs = (List>>) data.get("docs"); + return docs.stream() + .map(doc -> (Map>) doc.get("doc").get(SOURCE)) + .map(source -> source.get(targetField)) + .collect(Collectors.toList()); + } +} 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 0000000000..82233c6638 --- /dev/null +++ b/src/test/java/org/opensearch/geospatial/ip2geo/processor/Ip2GeoProcessorTests.java @@ -0,0 +1,301 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.geospatial.ip2geo.processor; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.BiConsumer; + +import lombok.SneakyThrows; + +import org.junit.Before; +import org.mockito.ArgumentCaptor; +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 List SUPPORTED_FIELDS = Arrays.asList("city", "country"); + private Ip2GeoProcessor.Factory factory; + + @Before + public void init() { + factory = new Ip2GeoProcessor.Factory(ingestService, datasourceDao, geoIpDataDao, ip2GeoCachedDao); + } + + 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 testExecute_whenNoIp_thenException() 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 = mock(BiConsumer.class); + + // Run + processor.execute(document, handler); + + // Verify + verify(handler).accept(isNull(), any(IllegalArgumentException.class)); + } + + public void testExecute_whenNonStringValue_thenException() 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 = mock(BiConsumer.class); + + // Run + processor.execute(document, handler); + + // Verify + verify(handler).accept(isNull(), any(IllegalArgumentException.class)); + } + + @SneakyThrows + public void testExecute_whenNoDatasource_thenNotExistError() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + + Map source = new HashMap<>(); + String ip = randomIpAddress(); + source.put("ip", ip); + IngestDocument document = new IngestDocument(source, new HashMap<>()); + + when(ip2GeoCachedDao.has(datasourceName)).thenReturn(false); + BiConsumer handler = mock(BiConsumer.class); + + // Run + processor.execute(document, handler); + + // Verify + ArgumentCaptor captor = ArgumentCaptor.forClass(Exception.class); + verify(handler).accept(isNull(), captor.capture()); + captor.getValue().getMessage().contains("not exist"); + } + + @SneakyThrows + public void testExecute_whenExpired_thenExpiredMsg() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + BiConsumer handler = mock(BiConsumer.class); + + String indexName = GeospatialTestHelper.randomLowerCaseString(); + when(ip2GeoCachedDao.getIndexName(datasourceName)).thenReturn(indexName); + when(ip2GeoCachedDao.has(datasourceName)).thenReturn(true); + when(ip2GeoCachedDao.getState(datasourceName)).thenReturn(DatasourceState.AVAILABLE); + when(ip2GeoCachedDao.isExpired(datasourceName)).thenReturn(true); + Map geoData = Map.of("city", "Seattle", "country", "USA"); + when(ip2GeoCachedDao.getGeoData(eq(indexName), any())).thenReturn(geoData); + + // Run for single ip + String ip = randomIpAddress(); + IngestDocument documentWithIp = createDocument(ip); + processor.execute(documentWithIp, handler); + + // Verify + verify(handler).accept(documentWithIp, null); + assertEquals("ip2geo_data_expired", documentWithIp.getFieldValue(DEFAULT_TARGET_FIELD + ".error", String.class)); + + // Run for multi ips + List ips = Arrays.asList(randomIpAddress(), randomIpAddress()); + IngestDocument documentWithIps = createDocument(ips); + processor.execute(documentWithIps, handler); + + // Verify + verify(handler).accept(documentWithIps, null); + assertEquals("ip2geo_data_expired", documentWithIp.getFieldValue(DEFAULT_TARGET_FIELD + ".error", String.class)); + } + + @SneakyThrows + public void testExecute_whenNotAvailable_thenException() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + BiConsumer handler = mock(BiConsumer.class); + + String indexName = GeospatialTestHelper.randomLowerCaseString(); + when(ip2GeoCachedDao.getIndexName(datasourceName)).thenReturn(indexName); + when(ip2GeoCachedDao.has(datasourceName)).thenReturn(true); + when(ip2GeoCachedDao.getState(datasourceName)).thenReturn(DatasourceState.CREATE_FAILED); + when(ip2GeoCachedDao.isExpired(datasourceName)).thenReturn(false); + Map geoData = Map.of("city", "Seattle", "country", "USA"); + when(ip2GeoCachedDao.getGeoData(eq(indexName), any())).thenReturn(geoData); + + // Run for single ip + String ip = randomIpAddress(); + IngestDocument documentWithIp = createDocument(ip); + processor.execute(documentWithIp, handler); + + // Run for multi ips + List ips = Arrays.asList(randomIpAddress(), randomIpAddress()); + IngestDocument documentWithIps = createDocument(ips); + processor.execute(documentWithIps, handler); + + // Verify + ArgumentCaptor captor = ArgumentCaptor.forClass(IllegalStateException.class); + verify(handler, times(2)).accept(isNull(), captor.capture()); + assertTrue(captor.getAllValues().stream().allMatch(e -> e.getMessage().contains("not in an available state"))); + } + + @SneakyThrows + public void testExecute_whenCalled_thenGeoIpDataIsAdded() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + BiConsumer handler = mock(BiConsumer.class); + + String indexName = GeospatialTestHelper.randomLowerCaseString(); + when(ip2GeoCachedDao.getIndexName(datasourceName)).thenReturn(indexName); + when(ip2GeoCachedDao.has(datasourceName)).thenReturn(true); + when(ip2GeoCachedDao.getState(datasourceName)).thenReturn(DatasourceState.AVAILABLE); + when(ip2GeoCachedDao.isExpired(datasourceName)).thenReturn(false); + Map geoData = Map.of("city", "Seattle", "country", "USA"); + when(ip2GeoCachedDao.getGeoData(eq(indexName), any())).thenReturn(geoData); + + // Run for single ip + String ip = randomIpAddress(); + IngestDocument documentWithIp = createDocument(ip); + processor.execute(documentWithIp, handler); + + // Verify + assertEquals(geoData.get("city"), documentWithIp.getFieldValue("ip2geo.city", String.class)); + assertEquals(geoData.get("country"), documentWithIp.getFieldValue("ip2geo.country", String.class)); + + // Run for multi ips + List ips = Arrays.asList(randomIpAddress(), randomIpAddress()); + IngestDocument documentWithIps = createDocument(ips); + processor.execute(documentWithIps, handler); + + // Verify + assertEquals(2, documentWithIps.getFieldValue("ip2geo", List.class).size()); + Map addedValue = (Map) documentWithIps.getFieldValue("ip2geo", List.class).get(0); + assertEquals(geoData.get("city"), addedValue.get("city")); + assertEquals(geoData.get("country"), addedValue.get("country")); + } + + @SneakyThrows + public void testExecute_whenPropertiesSet_thenFilteredGeoIpDataIsAdded() { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Map.of(Ip2GeoProcessor.CONFIG_PROPERTIES, Arrays.asList("country"))); + BiConsumer handler = mock(BiConsumer.class); + + String indexName = GeospatialTestHelper.randomLowerCaseString(); + when(ip2GeoCachedDao.getIndexName(datasourceName)).thenReturn(indexName); + when(ip2GeoCachedDao.has(datasourceName)).thenReturn(true); + when(ip2GeoCachedDao.getState(datasourceName)).thenReturn(DatasourceState.AVAILABLE); + when(ip2GeoCachedDao.isExpired(datasourceName)).thenReturn(false); + Map geoData = Map.of("city", "Seattle", "country", "USA"); + when(ip2GeoCachedDao.getGeoData(eq(indexName), any())).thenReturn(geoData); + + // Run for single ip + String ip = randomIpAddress(); + IngestDocument documentWithIp = createDocument(ip); + processor.execute(documentWithIp, handler); + + // Verify + assertFalse(documentWithIp.hasField("ip2geo.city")); + assertEquals(geoData.get("country"), documentWithIp.getFieldValue("ip2geo.country", String.class)); + + // Run for multi ips + List ips = Arrays.asList(randomIpAddress(), randomIpAddress()); + IngestDocument documentWithIps = createDocument(ips); + processor.execute(documentWithIps, handler); + + // Verify + assertEquals(2, documentWithIps.getFieldValue("ip2geo", List.class).size()); + Map addedValue = (Map) documentWithIps.getFieldValue("ip2geo", List.class).get(0); + assertFalse(addedValue.containsKey("city")); + assertEquals(geoData.get("country"), addedValue.get("country")); + } + + public void testExecute_whenNoHandler_thenException() 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 testExecute_whenContainsNonString_thenException() throws Exception { + String datasourceName = GeospatialTestHelper.randomLowerCaseString(); + Ip2GeoProcessor processor = createProcessor(datasourceName, Collections.emptyMap()); + List ips = Arrays.asList(randomIpAddress(), 1); + Map source = new HashMap<>(); + source.put("ip", ips); + IngestDocument document = new IngestDocument(source, new HashMap<>()); + BiConsumer handler = mock(BiConsumer.class); + + // Run + processor.execute(document, handler); + + // Verify + ArgumentCaptor captor = ArgumentCaptor.forClass(IllegalArgumentException.class); + verify(handler).accept(isNull(), captor.capture()); + assertTrue(captor.getValue().getMessage().contains("should only contain strings")); + } + + private Ip2GeoProcessor createProcessor(final String datasourceName, final Map config) throws Exception { + Datasource datasource = new Datasource(); + datasource.setName(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(datasourceDao.getDatasource(datasource.getName())).thenReturn(datasource); + Map baseConfig = new HashMap<>(); + baseConfig.put(Ip2GeoProcessor.CONFIG_FIELD, "ip"); + baseConfig.put(Ip2GeoProcessor.CONFIG_DATASOURCE, datasource.getName()); + baseConfig.putAll(config); + + return factory.create( + Collections.emptyMap(), + GeospatialTestHelper.randomLowerCaseString(), + GeospatialTestHelper.randomLowerCaseString(), + baseConfig + ); + } + + private IngestDocument createDocument(String ip) { + Map source = new HashMap<>(); + source.put("ip", ip); + return new IngestDocument(source, new HashMap<>()); + } + + private IngestDocument createDocument(List ips) { + Map source = new HashMap<>(); + source.put("ip", ips); + return new IngestDocument(source, new HashMap<>()); + } +} diff --git a/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java b/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java index 1a77595583..bbe51267f4 100644 --- a/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java +++ b/src/test/java/org/opensearch/geospatial/plugin/GeospatialPluginTests.java @@ -5,46 +5,203 @@ package org.opensearch.geospatial.plugin; +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.component.LifecycleComponent; +import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +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.RestDeleteDatasourceHandler; +import org.opensearch.geospatial.ip2geo.action.RestGetDatasourceHandler; +import org.opensearch.geospatial.ip2geo.action.RestPutDatasourceHandler; +import org.opensearch.geospatial.ip2geo.action.RestUpdateDatasourceHandler; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoExecutor; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoLockService; +import org.opensearch.geospatial.ip2geo.common.Ip2GeoSettings; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; +import org.opensearch.geospatial.ip2geo.dao.GeoIpDataDao; +import org.opensearch.geospatial.ip2geo.dao.Ip2GeoCachedDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.DatasourceUpdateService; +import org.opensearch.geospatial.ip2geo.listener.Ip2GeoListener; 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(clusterSettings), + new RestGetDatasourceHandler(), + new RestUpdateDatasourceHandler(), + new RestDeleteDatasourceHandler() + ); + + 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, + DatasourceDao.class, + Ip2GeoExecutor.class, + GeoIpDataDao.class, + Ip2GeoLockService.class, + Ip2GeoCachedDao.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; + private GeospatialPlugin plugin; + + @Before + public void init() { + openMocks = MockitoAnnotations.openMocks(this); + settings = Settings.EMPTY; + when(client.settings()).thenReturn(settings); + when(clusterService.getClusterSettings()).thenReturn(clusterSettings); + when(clusterService.getSettings()).thenReturn(settings); + when(ingestService.getClusterService()).thenReturn(clusterService); + nodeEnvironment = null; + plugin = new GeospatialPlugin(); + // Need to call getProcessors to initialize few instances in plugin class + plugin.getProcessors(getProcessorParameter()); + } + + @After + public void close() throws Exception { + openMocks.close(); + } + + public void testSystemIndexDescriptors() { + Set registeredSystemIndexPatterns = new HashSet<>(); + for (SystemIndexDescriptor descriptor : plugin.getSystemIndexDescriptors(Settings.EMPTY)) { + registeredSystemIndexPatterns.add(descriptor.getIndexPattern()); + } + assertEquals(SUPPORTED_SYSTEM_INDEX_PATTERN, registeredSystemIndexPatterns); - private final List SUPPORTED_REST_HANDLERS = List.of(new RestUploadGeoJSONAction(), new RestUploadStatsAction()); + } + + public void testExecutorBuilders() { + assertEquals(1, plugin.getExecutorBuilders(Settings.EMPTY).size()); + } + + public void testCreateComponents() { + 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 testGetGuiceServiceClasses() { + Collection> classes = List.of(Ip2GeoListener.class); + assertEquals(classes, plugin.getGuiceServiceClasses()); + } public void testIsAnIngestPlugin() { - GeospatialPlugin plugin = new GeospatialPlugin(); assertTrue(plugin instanceof IngestPlugin); } public void testFeatureProcessorIsAdded() { - GeospatialPlugin plugin = new GeospatialPlugin(); - Map processors = plugin.getProcessors(null); + Map processors = plugin.getProcessors(getProcessorParameter()); assertTrue(processors.containsKey(FeatureProcessor.TYPE)); assertTrue(processors.get(FeatureProcessor.TYPE) instanceof FeatureProcessor.Factory); } public void testTotalRestHandlers() { - GeospatialPlugin plugin = new GeospatialPlugin(); - assertEquals(SUPPORTED_REST_HANDLERS.size(), plugin.getRestHandlers(Settings.EMPTY, null, null, null, null, null, null).size()); + assertEquals( + SUPPORTED_REST_HANDLERS.size(), + plugin.getRestHandlers(Settings.EMPTY, null, clusterSettings, null, null, null, null).size() + ); } public void testUploadGeoJSONTransportIsAdded() { - GeospatialPlugin plugin = new GeospatialPlugin(); final List> actions = plugin.getActions(); assertEquals(1, actions.stream().filter(actionHandler -> actionHandler.getAction() instanceof UploadGeoJSONAction).count()); } + + private Processor.Parameters getProcessorParameter() { + return new Processor.Parameters( + mock(Environment.class), + mock(ScriptService.class), + null, + null, + null, + null, + ingestService, + client, + null + ); + } } diff --git a/src/test/java/org/opensearch/geospatial/processor/FeatureProcessorIT.java b/src/test/java/org/opensearch/geospatial/processor/FeatureProcessorIT.java index 585f8628d2..5b75292c0c 100644 --- a/src/test/java/org/opensearch/geospatial/processor/FeatureProcessorIT.java +++ b/src/test/java/org/opensearch/geospatial/processor/FeatureProcessorIT.java @@ -47,9 +47,9 @@ public void testIndexGeoJSONSuccess() throws Exception { Map geoFields = new HashMap<>(); geoFields.put(geoShapeField, "geo_shape"); - Map processorProperties = new HashMap<>(); + Map processorProperties = new HashMap<>(); processorProperties.put(FeatureProcessor.FIELD_KEY, geoShapeField); - Map geoJSONProcessorConfig = buildGeoJSONFeatureProcessorConfig(processorProperties); + Map geoJSONProcessorConfig = buildProcessorConfig(FeatureProcessor.TYPE, processorProperties); List> configs = new ArrayList<>(); configs.add(geoJSONProcessorConfig); diff --git a/src/test/resources/ip2geo/manifest.json b/src/test/resources/ip2geo/manifest.json new file mode 100644 index 0000000000..86a76e4723 --- /dev/null +++ b/src/test/resources/ip2geo/manifest.json @@ -0,0 +1,8 @@ +{ + "url": "https://test.com/db.zip", + "db_name": "sample_valid.csv", + "sha256_hash": "safasdfaskkkesadfasdf", + "valid_for_in_days": 30, + "updated_at_in_epoch_milli": 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 0000000000..c9f1723e3a --- /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", + "sha256_hash": "safasdfaskkkesadfasdf", + "valid_for_in_days": 30, + "updated_at_in_epoch_milli": 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 0000000000..39665b747e --- /dev/null +++ b/src/test/resources/ip2geo/manifest_template.json @@ -0,0 +1,8 @@ +{ + "url": "URL", + "db_name": "sample_valid.csv", + "sha256_hash": "safasdfaskkkesadfasdf", + "valid_for_in_days": 30, + "updated_at_in_epoch_milli": 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 0000000000..08670061c8 --- /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 0000000000..a6d0893579 --- /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 0000000000..0bdeeadbf1 Binary files /dev/null and b/src/test/resources/ip2geo/sample_valid.zip differ diff --git a/src/test/resources/ip2geo/server/city/city.zip b/src/test/resources/ip2geo/server/city/city.zip new file mode 100644 index 0000000000..12fbd7198d Binary files /dev/null and b/src/test/resources/ip2geo/server/city/city.zip differ diff --git a/src/test/resources/ip2geo/server/city/manifest.json b/src/test/resources/ip2geo/server/city/manifest.json new file mode 100644 index 0000000000..de1e3f3b5e --- /dev/null +++ b/src/test/resources/ip2geo/server/city/manifest.json @@ -0,0 +1,8 @@ +{ + "url": "https://github.com/opensearch-project/geospatial/blob/main/src/test/resources/ip2geo/server/city/city.zip", + "db_name": "data.csv", + "sha256_hash": "oDPgEv+9+kNov7bdQQiLrhr8jQeEPdLnuJ22Hz5npvk=", + "valid_for_in_days": 30, + "updated_at_in_epoch_milli": 1683590400000, + "provider": "opensearch" +} diff --git a/src/test/resources/ip2geo/server/city/manifest_local.json b/src/test/resources/ip2geo/server/city/manifest_local.json new file mode 100644 index 0000000000..a69ccbefde --- /dev/null +++ b/src/test/resources/ip2geo/server/city/manifest_local.json @@ -0,0 +1,8 @@ +{ + "url": "http://localhost:8001/city/city.zip", + "db_name": "data.csv", + "sha256_hash": "oDPgEv+9+kNov7bdQQiLrhr8jQeEPdLnuJ22Hz5npvk=", + "valid_for_in_days": 30, + "updated_at_in_epoch_milli": 1683590400000, + "provider": "opensearch" +} diff --git a/src/test/resources/ip2geo/server/country/country.zip b/src/test/resources/ip2geo/server/country/country.zip new file mode 100644 index 0000000000..1c930b1a74 Binary files /dev/null and b/src/test/resources/ip2geo/server/country/country.zip differ diff --git a/src/test/resources/ip2geo/server/country/manifest.json b/src/test/resources/ip2geo/server/country/manifest.json new file mode 100644 index 0000000000..25460e5b0d --- /dev/null +++ b/src/test/resources/ip2geo/server/country/manifest.json @@ -0,0 +1,8 @@ +{ + "url": "https://github.com/opensearch-project/geospatial/blob/main/src/test/resources/ip2geo/server/country/country.zip", + "db_name": "data.csv", + "sha256_hash": "oDPgEv+4+kNov7bdQQiLrhr8jQeEPdLnuJ11Hz5npvk=", + "valid_for_in_days": 30, + "updated_at_in_epoch_milli": 1683590400000, + "provider": "opensearch" +} diff --git a/src/test/resources/ip2geo/server/country/manifest_local.json b/src/test/resources/ip2geo/server/country/manifest_local.json new file mode 100644 index 0000000000..4c63840b7b --- /dev/null +++ b/src/test/resources/ip2geo/server/country/manifest_local.json @@ -0,0 +1,8 @@ +{ + "url": "http://localhost:8001/country/country.zip", + "db_name": "data.csv", + "sha256_hash": "oDPgEv+4+kNov7bdQQiLrhr8jQeEPdLnuJ11Hz5npvk=", + "valid_for_in_days": 30, + "updated_at_in_epoch_milli": 1683590400000, + "provider": "opensearch" +} diff --git a/src/yamlRestTest/resources/rest-api-spec/test/10_basic.yml b/src/yamlRestTest/resources/rest-api-spec/test/10_basic.yml index 21c13a60c8..264ec58564 100644 --- a/src/yamlRestTest/resources/rest-api-spec/test/10_basic.yml +++ b/src/yamlRestTest/resources/rest-api-spec/test/10_basic.yml @@ -1,8 +1,8 @@ -"Test that geospatial plugin is loaded in OpenSearch": +"Test that geospatial and job scheduler plugins are loaded in OpenSearch": - do: cat.plugins: local: true h: component - match: - $body: /^opensearch-geospatial\n$/ + $body: /^opensearch-geospatial\nopensearch-job-scheduler\n$/