diff --git a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentActionClient.java b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentActionClient.java index 83c7ec7f84..7a92f57359 100644 --- a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentActionClient.java +++ b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentActionClient.java @@ -31,7 +31,7 @@ public class IpEnrichmentActionClient { * @return A map instance which contain GeoLocation data for the given Ip address. */ public Map getGeoLocationData (String ipString) { - return getGeoLocationData(ipString, "defaultDataSource"); + return getGeoLocationData(ipString, null); } /** diff --git a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java index 20c2a035bf..487c582389 100644 --- a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java +++ b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java @@ -40,6 +40,7 @@ public class IpEnrichmentRequest extends ActionRequest { public IpEnrichmentRequest(StreamInput streamInput) throws IOException { super(streamInput); ipString = streamInput.readString(); + datasourceName= streamInput.readOptionalString(); } /** @@ -65,6 +66,7 @@ public ActionRequestValidationException validate() { public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeString(ipString); + out.writeOptionalString(datasourceName); } /** diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportAction.java index bf38065698..dde58f8fbe 100644 --- a/src/main/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportAction.java +++ b/src/main/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportAction.java @@ -15,11 +15,15 @@ import org.opensearch.geospatial.action.IpEnrichmentAction; import org.opensearch.geospatial.action.IpEnrichmentRequest; import org.opensearch.geospatial.action.IpEnrichmentResponse; +import org.opensearch.geospatial.ip2geo.dao.DatasourceDao; import org.opensearch.geospatial.ip2geo.dao.Ip2GeoCachedDao; +import org.opensearch.geospatial.ip2geo.jobscheduler.Datasource; import org.opensearch.tasks.Task; import org.opensearch.transport.TransportService; +import java.util.List; import java.util.Map; +import java.util.Optional; /** * Transport action to convert provided IP address String into GeoLocation data. @@ -30,6 +34,10 @@ public class IpEnrichmentTransportAction extends HandledTransportAction allDatasources = datasourceDao.getAllDatasources(); + this.defaultDataSourceName = (!allDatasources.isEmpty()) ? allDatasources.get(0).getName() : null; } @@ -57,11 +69,18 @@ public IpEnrichmentTransportAction( protected void doExecute(Task task, ActionRequest request, ActionListener listener) { IpEnrichmentRequest enrichmentRequest = IpEnrichmentRequest.fromActionRequest(request); String ipString = enrichmentRequest.getIpString(); - String indexName = ip2GeoCachedDao.getIndexName(enrichmentRequest.getDatasourceName()); - Map geoLocationData = ip2GeoCachedDao.getGeoData(indexName, ipString); - System.out.println(geoLocationData); - log.debug("GeoSpatial IP lookup on IP: [{}], and result [{}]", ipString, geoLocationData); - listener.onResponse(new IpEnrichmentResponse(geoLocationData)); + if (enrichmentRequest.getDatasourceName() == null && + defaultDataSourceName == null) { + log.error("No data source available, IpEnrichmentTransportAction aborted."); + listener.onFailure(new IllegalArgumentException()); + } else { + String dataSourceName = Optional.ofNullable(enrichmentRequest.getDatasourceName()) + .orElse(defaultDataSourceName); + String indexName = ip2GeoCachedDao.getIndexName(dataSourceName); + Map geoLocationData = ip2GeoCachedDao.getGeoData(indexName, ipString); + log.debug("GeoSpatial IP lookup on IP: [{}], and result [{}]", ipString, geoLocationData); + listener.onResponse(new IpEnrichmentResponse(geoLocationData)); + } } }