diff --git a/docs/user/ppl/functions/ip.rst b/docs/user/ppl/functions/ip.rst
index 897d6ccfad0..ec853c27093 100644
--- a/docs/user/ppl/functions/ip.rst
+++ b/docs/user/ppl/functions/ip.rst
@@ -45,7 +45,7 @@ Description
Usage: `geoip(dataSourceName, ipAddress[, options])` to lookup location information from given IP addresses via OpenSearch GeoSpatial plugin API.
-Argument type: STRING, STRING, STRING
+Argument type: STRING, STRING/IP, STRING
Return type: OBJECT
diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java
index 9de93a515a2..78233d8dd52 100644
--- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java
+++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java
@@ -5,12 +5,52 @@
package org.opensearch.sql.calcite.remote;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_WEBLOGS;
+import static org.opensearch.sql.util.MatcherUtils.rows;
+import static org.opensearch.sql.util.MatcherUtils.schema;
+import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
+import static org.opensearch.sql.util.MatcherUtils.verifySchema;
+
+import java.io.IOException;
+import java.util.Map;
+import org.json.JSONObject;
+import org.junit.jupiter.api.Test;
+import org.opensearch.client.Request;
import org.opensearch.sql.ppl.GeoIpFunctionsIT;
public class CalciteGeoIpFunctionsIT extends GeoIpFunctionsIT {
@Override
public void init() throws Exception {
super.init();
+ loadIndex(Index.WEBLOG);
enableCalcite();
+
+ // Only limited IPs are loaded into geospatial data sources. Therefore, we insert IPs that match
+ // those known ones for test purpose
+ Request bulkRequest = new Request("POST", "/_bulk?refresh=true");
+ bulkRequest.setJsonEntity(
+ String.format(
+ "{\"index\":{\"_index\":\"%s\",\"_id\":6}}\n"
+ + "{\"host\":\"10.0.0.1\",\"method\":\"POST\"}\n"
+ + "{\"index\":{\"_index\":\"%s\",\"_id\":7}}\n"
+ + "{\"host\":\"fd12:2345:6789:1:a1b2:c3d4:e5f6:789a\",\"method\":\"POST\"}\n",
+ TEST_INDEX_WEBLOGS, TEST_INDEX_WEBLOGS));
+ client().performRequest(bulkRequest);
+ }
+
+ // In v2 it supports only string as IP inputs
+ @Test
+ public void testGeoIpEnrichmentWithIpFieldAsInput() throws IOException {
+ JSONObject result =
+ executeQuery(
+ String.format(
+ "source=%s | where method='POST' | eval ip_to_country = geoip('%s', host,"
+ + " 'country') | fields host, ip_to_country",
+ TEST_INDEX_WEBLOGS, DATASOURCE_NAME));
+ verifySchema(result, schema("host", "ip"), schema("ip_to_country", "struct"));
+ verifyDataRows(
+ result,
+ rows("10.0.0.1", Map.of("country", "USA")),
+ rows("fd12:2345:6789:1:a1b2:c3d4:e5f6:789a", Map.of("country", "India")));
}
}
diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java
index fe747fcdc03..14c6e2fb8f8 100644
--- a/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java
+++ b/integ-test/src/test/java/org/opensearch/sql/ppl/GeoIpFunctionsIT.java
@@ -40,7 +40,7 @@ public class GeoIpFunctionsIT extends PPLIntegTestCase {
"endpoint",
"https://raw.githubusercontent.com/opensearch-project/geospatial/main/src/test/resources/ip2geo/server/city/manifest.json");
- private static String DATASOURCE_NAME = "dummycityindex";
+ protected static String DATASOURCE_NAME = "dummycityindex";
private static String PLUGIN_NAME = "opensearch-geospatial";
@@ -83,7 +83,7 @@ public void testGeoIpEnrichment() {
String.format(
"search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s) | fields name, ip,"
+ " enrichmentResult",
- TEST_INDEX_GEOIP, "dummycityindex", "ip"));
+ TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip"));
verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
@@ -101,7 +101,7 @@ public void testGeoIpEnrichmentWithSingleOption() {
String.format(
"search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\") |"
+ " fields name, ip, enrichmentResult",
- TEST_INDEX_GEOIP, "dummycityindex", "ip", "city"));
+ TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip", "city"));
verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
@@ -119,7 +119,7 @@ public void testGeoIpEnrichmentWithSpaceSeparatedMultipleOptions() {
String.format(
"search source=%s | eval enrichmentResult = geoip(\\\"%s\\\",%s,\\\"%s\\\") |"
+ " fields name, ip, enrichmentResult",
- TEST_INDEX_GEOIP, "dummycityindex", "ip", "city , country"));
+ TEST_INDEX_GEOIP, DATASOURCE_NAME, "ip", "city , country"));
verifyColumn(resultGeoIp, columnName("name"), columnName("ip"), columnName("enrichmentResult"));
verifyDataRows(
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java
index 9b3c4b0a1f5..001ed064d1d 100644
--- a/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java
+++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/functions/GeoIpFunction.java
@@ -16,16 +16,15 @@
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rex.RexCall;
-import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
-import org.apache.calcite.sql.type.OperandTypes;
import org.apache.calcite.sql.type.SqlReturnTypeInference;
-import org.apache.calcite.sql.type.SqlTypeFamily;
import org.apache.calcite.sql.type.SqlTypeName;
import org.opensearch.geospatial.action.IpEnrichmentActionClient;
import org.opensearch.sql.common.utils.StringUtils;
+import org.opensearch.sql.data.model.ExprIpValue;
import org.opensearch.sql.data.model.ExprStringValue;
import org.opensearch.sql.data.model.ExprTupleValue;
import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprCoreType;
import org.opensearch.sql.expression.function.ImplementorUDF;
import org.opensearch.sql.expression.function.UDFOperandMetadata;
import org.opensearch.transport.client.node.NodeClient;
@@ -38,8 +37,8 @@
* Signatures:
*
*
- * - (STRING, STRING) -> MAP
- *
- (STRING, STRING, STRING) -> MAP
+ *
- (STRING, IP) -> MAP
+ *
- (STRING, IP, STRING) -> MAP
*
*/
public class GeoIpFunction extends ImplementorUDF {
@@ -59,11 +58,10 @@ public SqlReturnTypeInference getReturnTypeInference() {
@Override
public UDFOperandMetadata getOperandMetadata() {
- return UDFOperandMetadata.wrap(
- (CompositeOperandTypeChecker)
- OperandTypes.CHARACTER_CHARACTER.or(
- OperandTypes.family(
- SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER)));
+ return UDFOperandMetadata.wrapUDT(
+ List.of(
+ List.of(ExprCoreType.STRING, ExprCoreType.IP),
+ List.of(ExprCoreType.STRING, ExprCoreType.IP, ExprCoreType.STRING)));
}
public static class GeoIPImplementor implements NotNullImplementor {
@@ -87,16 +85,20 @@ public Expression implement(
}
public static Map fetchIpEnrichment(
- String dataSource, String ipAddress, NodeClient nodeClient) {
- return fetchIpEnrichment(dataSource, ipAddress, Collections.emptySet(), nodeClient);
+ String dataSource, ExprIpValue ipAddress, NodeClient nodeClient) {
+ return fetchIpEnrichment(
+ dataSource, ipAddress.toString(), Collections.emptySet(), nodeClient);
}
public static Map fetchIpEnrichment(
- String dataSource, String ipAddress, String commaSeparatedOptions, NodeClient nodeClient) {
+ String dataSource,
+ ExprIpValue ipAddress,
+ String commaSeparatedOptions,
+ NodeClient nodeClient) {
String unquotedOptions = StringUtils.unquoteText(commaSeparatedOptions);
final Set options =
Arrays.stream(unquotedOptions.split(",")).map(String::trim).collect(Collectors.toSet());
- return fetchIpEnrichment(dataSource, ipAddress, options, nodeClient);
+ return fetchIpEnrichment(dataSource, ipAddress.toString(), options, nodeClient);
}
private static Map fetchIpEnrichment(