PPL command expression implementation for geoip

.github/workflows/integ-tests-with-geo.yml

@@ -0,0 +1,89 @@
+name: GeoSpatial Plugin IT
+
+on:
+  pull_request:
+  push:
+    branches-ignore:
+      - 'dependabot/**'
+    paths:
+      - 'integ-test/**'
+      - '.github/workflows/integ-tests-with-geo.yml'
+
+jobs:
+  Get-CI-Image-Tag:
+    uses: opensearch-project/opensearch-build/.github/workflows/get-ci-image-tag.yml@main
+    with:
+      product: opensearch
+
+  security-it-linux:
+    needs: Get-CI-Image-Tag
+    strategy:
+      fail-fast: false
+      matrix:
+        java: [21]
+    runs-on: ubuntu-latest
+    container:
+      # using the same image which is used by opensearch-build team to build the OpenSearch Distribution
+      # this image tag is subject to change as more dependencies and updates will arrive over time
+      image: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-version-linux }}
+      options: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-options }}
+
+    steps:
+    - name: Run start commands
+      run: ${{ needs.Get-CI-Image-Tag.outputs.ci-image-start-command }}
+
+    - uses: actions/checkout@v4
+
+    - name: Set up JDK ${{ matrix.java }}
+      uses: actions/setup-java@v4
+      with:
+        distribution: 'temurin'
+        java-version: ${{ matrix.java }}
+
+    - name: Build with Gradle
+      run: |
+        chown -R 1000:1000 `pwd`
+        su `id -un 1000` -c "./gradlew integTestWithGeo"
+
+    - name: Upload test reports
+      if: ${{ always() }}
+      uses: actions/upload-artifact@v4
+      continue-on-error: true
+      with:
+        name: test-reports-${{ matrix.os }}-${{ matrix.java }}
+        path: |
+          integ-test/build/reports/**
+          integ-test/build/testclusters/*/logs/*
+          integ-test/build/testclusters/*/config/*
+
+  security-it-windows-macos:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ windows-latest, macos-13 ]
+        java: [21]
+
+    runs-on: ${{ matrix.os }}
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up JDK ${{ matrix.java }}
+      uses: actions/setup-java@v4
+      with:
+        distribution: 'temurin'
+        java-version: ${{ matrix.java }}
+
+    - name: Build with Gradle
+      run: ./gradlew integTestWithGeo
+
+    - name: Upload test reports
+      if: ${{ always() }}
+      uses: actions/upload-artifact@v4
+      continue-on-error: true
+      with:
+        name: test-reports-${{ matrix.os }}-${{ matrix.java }}
+        path: |
+          integ-test/build/reports/**
+          integ-test/build/testclusters/*/logs/*
+          integ-test/build/testclusters/*/config/*

core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java

@@ -21,6 +21,7 @@
 import org.opensearch.sql.analysis.symbol.Namespace;
 import org.opensearch.sql.analysis.symbol.Symbol;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.Node;
 import org.opensearch.sql.ast.expression.AggregateFunction;
 import org.opensearch.sql.ast.expression.AllFields;
 import org.opensearch.sql.ast.expression.And;
@@ -101,6 +102,11 @@ public Expression visitUnresolvedAttribute(UnresolvedAttribute node, AnalysisCon
     return visitIdentifier(node.getAttr(), context);
   }
 
+  @Override
+  public Expression visitGeoIp(Node node, AnalysisContext context) {
+    return super.visitGeoIp(node, context);
+  }
+
   @Override
   public Expression visitEqualTo(EqualTo node, AnalysisContext context) {
     Expression left = node.getLeft().accept(this, context);

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -131,6 +131,10 @@ public T visitEqualTo(EqualTo node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitGeoIp(Node node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitLiteral(Literal node, C context) {
     return visitChildren(node, context);
   }

core/src/main/java/org/opensearch/sql/expression/DSL.java

@@ -969,6 +969,10 @@ public static FunctionExpression utc_timestamp(
     return compile(functionProperties, BuiltinFunctionName.UTC_TIMESTAMP, args);
   }
 
+  public static FunctionExpression geoip(Expression... args) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.GEOIP, args);
+  }
+
   @SuppressWarnings("unchecked")
   private static <T extends FunctionImplementation> T compile(
       FunctionProperties functionProperties, BuiltinFunctionName bfn, Expression... args) {

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -204,6 +204,9 @@ public enum BuiltinFunctionName {
   TRIM(FunctionName.of("trim")),
   UPPER(FunctionName.of("upper")),
 
+  /** GEOSPATIAL Functions. */
+  GEOIP(FunctionName.of("geoip")),
+
   /** NULL Test. */
   IS_NULL(FunctionName.of("is null")),
   IS_NOT_NULL(FunctionName.of("is not null")),

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionRepository.java

@@ -27,6 +27,7 @@
 import org.opensearch.sql.expression.aggregation.AggregatorFunctions;
 import org.opensearch.sql.expression.datetime.DateTimeFunctions;
 import org.opensearch.sql.expression.datetime.IntervalClause;
+import org.opensearch.sql.expression.ip.GeoIPFunctions;
 import org.opensearch.sql.expression.ip.IPFunctions;
 import org.opensearch.sql.expression.operator.arthmetic.ArithmeticFunctions;
 import org.opensearch.sql.expression.operator.arthmetic.MathematicalFunctions;
@@ -83,6 +84,7 @@ public static synchronized BuiltinFunctionRepository getInstance() {
       SystemFunctions.register(instance);
       OpenSearchFunctions.register(instance);
       IPFunctions.register(instance);
+      GeoIPFunctions.register(instance);
     }
     return instance;
   }

core/src/main/java/org/opensearch/sql/expression/ip/GeoIPFunctions.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.ip;
+
+import static org.opensearch.sql.data.type.ExprCoreType.BOOLEAN;
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+import static org.opensearch.sql.expression.function.FunctionDSL.define;
+
+import java.util.Arrays;
+import java.util.List;
+import lombok.experimental.UtilityClass;
+import org.apache.commons.lang3.tuple.Pair;
+import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
+import org.opensearch.sql.expression.function.DefaultFunctionResolver;
+import org.opensearch.sql.expression.function.FunctionBuilder;
+import org.opensearch.sql.expression.function.FunctionName;
+import org.opensearch.sql.expression.function.FunctionSignature;
+import org.opensearch.sql.expression.function.SerializableFunction;
+
+/**
+ * Utility class to register the method signature for geoip( ) expression, concreted reallocated to
+ * `opensearch` module, as this Ip location require GeoSpatial Plugin runtime support.
+ */
+@UtilityClass
+public class GeoIPFunctions {
+
+  public void register(BuiltinFunctionRepository repository) {
+    repository.register(geoIp());
+  }
+
+  /**
+   * To register all method signatures related to geoip( ) expression under eval.
+   *
+   * @return Resolver for geoip( ) expression.
+   */
+  private DefaultFunctionResolver geoIp() {
+    return define(
+        BuiltinFunctionName.GEOIP.getName(),
+        openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING)),
+        openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING, STRING)));
+  }
+
+  /**
+   * Util method to generate probe implementation with given list of argument types, with marker
+   * class `OpenSearchFunctionExpression` to annotate this is an OpenSearch specific expression.
+   *
+   * @param returnType return type.
+   * @return Binary Function Implementation.
+   */
+  public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
+      openSearchImpl(ExprType returnType, List<ExprType> args) {
+    return functionName -> {
+      FunctionSignature functionSignature = new FunctionSignature(functionName, args);
+      FunctionBuilder functionBuilder =
+          (functionProperties, arguments) ->
+              new OpenSearchFunctionExpression(functionName, arguments, returnType);
+      return Pair.of(functionSignature, functionBuilder);
+    };
+  }
+}

core/src/main/java/org/opensearch/sql/expression/ip/OpenSearchFunctionExpression.java

@@ -0,0 +1,47 @@
+/*
+ *
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.opensearch.sql.expression.ip;
+
+import java.util.List;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.FunctionExpression;
+import org.opensearch.sql.expression.env.Environment;
+import org.opensearch.sql.expression.function.FunctionName;
+
+/**
+ * Marker class to identify functions only compatible with OpenSearch storage engine. Any attempt to
+ * invoke the method different from OpenSearch will result in UnsupportedOperationException.
+ */
+public class OpenSearchFunctionExpression extends FunctionExpression {
+
+  private final ExprType returnType;
+
+  public OpenSearchFunctionExpression(
+      FunctionName functionName, List<Expression> arguments, ExprType returnType) {
+    super(functionName, arguments);
+    this.returnType = returnType;
+  }
+
+  @Override
+  public ExprValue valueOf() {
+    return null;
+  }
+
+  @Override
+  public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+    throw new UnsupportedOperationException(
+        "OpenSearch runtime specific function, no default implementation available");
+  }
+
+  @Override
+  public ExprType type() {
+    return returnType;
+  }
+}

core/src/test/java/org/opensearch/sql/expression/ip/GeoIPFunctionTest.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.ip;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.opensearch.sql.data.type.ExprCoreType.BOOLEAN;
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.expression.DSL;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.env.Environment;
+
+@ExtendWith(MockitoExtension.class)
+public class GeoIPFunctionTest {
+
+  // Mock value environment for testing.
+  @Mock private Environment<Expression, ExprValue> env;
+
+  @Test
+  public void geoIpDefaultImplementation() {
+    UnsupportedOperationException exception =
+        assertThrows(
+            UnsupportedOperationException.class,
+            () ->
+                DSL.geoip(DSL.literal("HARDCODED_DATASOURCE_NAME"), DSL.ref("ip_address", STRING))
+                    .valueOf(env));
+    assertTrue(exception.getMessage().matches(".*no default implementation available"));
+  }
+
+  @Test
+  public void testGeoipFnctionSignature() {
+    var geoip = DSL.geoip(DSL.literal("HARDCODED_DATASOURCE_NAME"), DSL.ref("ip_address", STRING));
+    assertEquals(BOOLEAN, geoip.type());
+  }
+}

docs/user/ppl/functions/geoip.rst

@@ -0,0 +1,38 @@
+====================
+Geo IP Address Functions
+====================
+
+.. rubric:: Table of contents
+
+.. contents::
+   :local:
+   :depth: 1
+
+GEOIP
+---------
+
+Description
+>>>>>>>>>>>
+
+Usage: `geoip(dataSourceName, ipAddress, options)` to lookup location information from given IP addresses via OpenSearch GeoSpatial plugin API.
+
+Argument type: STRING, STRING, STRING
+
+Return type: Tuple
+
+Example:
+
+    os> source=weblogs | eval LookupResult = geoip("dataSourceName", "50.68.18.229", "country_iso_code,city_name")
+    fetched rows / total rows = 1/1
+    +-------------------------------------------------------------+
+    | LookupResult                                                        |
+    |-------------------------------------------------------------|
+    | {'city_name': 'Vancouver', 'country_iso_code': 'CA'}        |
+    +-------------------------------------------------------------+
+
+
+Note:
+ - `dataSourceName` must be an established dataSource on OpenSearch GeoSpatial plugin, detail of configuration can be found: https://opensearch.org/docs/latest/ingest-pipelines/processors/ip2geo/
+ - `ip` can be an IPv4 or an IPv6 address
+ - `options` is a comma separated String for user to specify fields to output, the selection of fields subject to dataSourceProvider's schema, the list of geolite2-city dataset provide fields: "country_iso_code", "country_name", "continent_name", "region_iso_code", "region_name", "city_name", "time_zone", "location"
+

docs/user/ppl/index.rst

@@ -106,6 +106,8 @@ The query start with search command and then flowing a set of command delimited
 
   - `IP Address Functions <functions/ip.rst>`_
 
+  - `Geo IP Address Functions <functions/geoip.rst>`_
+
 * **Optimization**
 
   - `Optimization <../../user/optimization/optimization.rst>`_