Fetch default data source

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

Hide file tree

Showing 3 changed files with 28 additions and 7 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
@@ -31,7 +31,7 @@ public class IpEnrichmentActionClient {
      * @return A map instance which contain GeoLocation data for the given Ip address.
      */
     public Map<String, Object> 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
@@ -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
@@ -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<ActionRe
 
     private Ip2GeoCachedDao ip2GeoCachedDao;
 
+    private DatasourceDao datasourceDao;
+
+    private String defaultDataSourceName;
+
     /**
      * Constructor
      * @param transportService the transport service
@@ -40,9 +48,13 @@ public class IpEnrichmentTransportAction extends HandledTransportAction<ActionRe
     public IpEnrichmentTransportAction(
             TransportService transportService,
             ActionFilters actionFilters,
-            Ip2GeoCachedDao cachedDao) {
+            Ip2GeoCachedDao cachedDao,
+            DatasourceDao datasourceDao) {
         super(IpEnrichmentAction.NAME, transportService, actionFilters, IpEnrichmentRequest::new);
         this.ip2GeoCachedDao = cachedDao;
+        this.datasourceDao = datasourceDao;
+        List<Datasource> 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<ActionResponse> listener) {
         IpEnrichmentRequest enrichmentRequest = IpEnrichmentRequest.fromActionRequest(request);
         String ipString = enrichmentRequest.getIpString();
-        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));
+        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<String, Object> geoLocationData = ip2GeoCachedDao.getGeoData(indexName, ipString);
+            log.debug("GeoSpatial IP lookup on IP: [{}], and result [{}]", ipString, geoLocationData);
+            listener.onResponse(new IpEnrichmentResponse(geoLocationData));
+        }
     }
 
 }