API Signature

Signed-off-by: Andy Kwok <[email protected]>
Bit-Quill · Nov 21, 2024 · 68d192f · 68d192f
1 parent fba033b
commit 68d192f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentActionClient.java b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentActionClient.java
@@ -12,6 +12,7 @@
 import org.opensearch.core.action.ActionResponse;
 
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 
 /**
@@ -23,13 +24,25 @@ public class IpEnrichmentActionClient {
 
     NodeClient nodeClient;
 
+
+    /**
+     * IpEnrichment with default datasource.
+     * @param ipString Ip String to resolve.
+     * @return A map instance which contain GeoLocation data for the given Ip address.
+     */
+    public Map<String, Object> getGeoLocationData (String ipString) {
+        return getGeoLocationData(ipString, "defaultDataSource");
+    }
+
     /**
      * Client facing method, which read an IP in String form and return a map instance which contain the associated GeoLocation data.
      * @param ipString IP v4 || v6 address in String form.
+     * @param datasourceName datasourceName in String form.
      * @return A map instance which contain GeoLocation data for the given Ip address.
      */
-    public Map<String, Object> getGeoLocationData (String ipString) {
-        ActionFuture<ActionResponse> responseActionFuture = nodeClient.execute(IpEnrichmentAction.INSTANCE, new IpEnrichmentRequest(ipString));
+    public Map<String, Object> getGeoLocationData (String ipString, String datasourceName) {
+        ActionFuture<ActionResponse> responseActionFuture = nodeClient.execute(
+                IpEnrichmentAction.INSTANCE, new IpEnrichmentRequest(ipString, datasourceName));
         try {
             ActionResponse genericActionResponse = responseActionFuture.get();
             IpEnrichmentResponse enrichmentResponse = IpEnrichmentResponse.fromActionResponse(genericActionResponse);

diff --git a/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java b/client/src/main/java/org/opensearch/geospatial/action/IpEnrichmentRequest.java
@@ -19,6 +19,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.Optional;
 
 /**
  * Wrapper for the IP 2 GeoLocation action request.
@@ -30,6 +31,8 @@ public class IpEnrichmentRequest extends ActionRequest {
 
     private String ipString;
 
+    private String datasourceName;
+
     /**
      * Constructor for TransportAction.
      * @param streamInput the streamInput.

diff --git a/src/main/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportAction.java b/src/main/java/org/opensearch/geospatial/ip2geo/action/IpEnrichmentTransportAction.java
@@ -57,10 +57,11 @@ public IpEnrichmentTransportAction(
     protected void doExecute(Task task, ActionRequest request, ActionListener<ActionResponse> listener) {
         IpEnrichmentRequest enrichmentRequest = IpEnrichmentRequest.fromActionRequest(request);
         String ipString = enrichmentRequest.getIpString();
-        Map<String, Object> testResult = ip2GeoCachedDao.getGeoData(".geospatial-ip2geo-data.my-datasource.ef3486f8-401b-4d77-b89b-3a4cd19eda04", ipString);
-        System.out.println(testResult);
-        log.debug("GeoSpatial IP lookup on IP: [{}], and result [{}]", ipString, testResult);
-        listener.onResponse(new IpEnrichmentResponse(testResult));
+        String indexName = ip2GeoCachedDao.getIndexName(enrichmentRequest.getDatasourceName());
+        Map<String, Object> 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));
     }
 
 }