Skip to content

Commit

Permalink
Function - Geoip
Browse files Browse the repository at this point in the history
Signed-off-by: Andy Kwok <[email protected]>
  • Loading branch information
andy-k-improving committed Jan 13, 2025
1 parent a135f1d commit 6abfd6d
Show file tree
Hide file tree
Showing 23 changed files with 669 additions and 24 deletions.
4 changes: 4 additions & 0 deletions core/src/main/java/org/opensearch/sql/expression/DSL.java
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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")),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -83,6 +84,7 @@ public static synchronized BuiltinFunctionRepository getInstance() {
SystemFunctions.register(instance);
OpenSearchFunctions.register(instance);
IPFunctions.register(instance);
GeoIPFunctions.register(instance);
}
return instance;
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
};
}
}
Original file line number Diff line number Diff line change
@@ -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;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
/*
* 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.assertNull;
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());
}

/** To make sure no logic being evaluated when no environment being passed. */
@Test
public void testDefaultValueOf() {
var geoip = DSL.geoip(DSL.literal("HARDCODED_DATASOURCE_NAME"), DSL.ref("ip_address", STRING));
assertNull(geoip.valueOf());
}
}
2 changes: 2 additions & 0 deletions docs/user/ppl/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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>`_
Expand Down
51 changes: 28 additions & 23 deletions integ-test/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,12 @@ Map dummyLoginConfig = [
password: "admin"
]

Map dynamicLoginConfig = [
https: System.getProperty("https"),
user: System.getProperty("user"),
password: System.getProperty("password")
]

repositories {
mavenCentral()
maven { url 'https://jitpack.io' }
Expand Down Expand Up @@ -261,6 +267,10 @@ static def getAllTransportSocketURI(cluster) {
.collect(Collectors.joining(","))
}

static def getConcatClusterName(clusters) {
return clusters.stream().map(cluster -> cluster.getName()).collect(Collectors.joining(","))
}

def getPluginDownloadLink(pluginName) {
var repo = "https://aws.oss.sonatype.org/content/repositories/snapshots/org/opensearch/plugin/" +
pluginName + "/$opensearch_build_snapshot/"
Expand Down Expand Up @@ -337,6 +347,7 @@ task stopPrometheus(type: KillProcessTask) {
stopPrometheus.mustRunAfter startPrometheus

task integJdbcTest(type: RestIntegTestTask) {

testClusters.findAll {c -> c.clusterName == "integJdbcTest"}.first().
plugin ":opensearch-sql-plugin"

Expand All @@ -360,9 +371,7 @@ task integJdbcTest(type: RestIntegTestTask) {
systemProperty 'tests.security.manager', 'false'
systemProperty('project.root', project.projectDir.absolutePath)

systemProperty "https", System.getProperty("https")
systemProperty "user", System.getProperty("user")
systemProperty "password", System.getProperty("password")
systemProperties dynamicLoginConfig

// Set default query size limit
systemProperty 'defaultQuerySizeLimit', '10000'
Expand All @@ -381,11 +390,11 @@ task integJdbcTest(type: RestIntegTestTask) {
}

task integTestWithSecurity(type: RestIntegTestTask) {

useCluster testClusters.integTestWithSecurity
useCluster testClusters.remoteIntegTestWithSecurity

systemProperty "cluster.names",
getClusters().stream().map(cluster -> cluster.getName()).collect(Collectors.joining(","))
systemProperty "cluster.names", getConcatClusterName(getClusters())

getClusters().forEach { cluster ->
configureSecurityPlugin(cluster)
Expand Down Expand Up @@ -424,7 +433,7 @@ task integTestWithSecurity(type: RestIntegTestTask) {
// NOTE: this IT config discovers only junit5 (jupiter) tests.
// https://github.com/opensearch-project/sql/issues/1974
filter {
includeTestsMatching 'org.opensearch.sql.security.CrossClusterSearchIT'
includeTestsMatching 'org.opensearch.sql.geo.PplIpEnrichmentIT'
}
}

Expand Down Expand Up @@ -477,15 +486,8 @@ integTest {
// Set properties for connection to clusters and between clusters
doFirst {
getClusters().forEach { cluster ->
String allTransportSocketURI = cluster.nodes.stream().flatMap { node ->
node.getAllTransportPortURI().stream()
}.collect(Collectors.joining(","))
String allHttpSocketURI = cluster.nodes.stream().flatMap { node ->
node.getAllHttpSocketURI().stream()
}.collect(Collectors.joining(","))

systemProperty "tests.rest.${cluster.name}.http_hosts", "${-> allHttpSocketURI}"
systemProperty "tests.rest.${cluster.name}.transport_hosts", "${-> allTransportSocketURI}"
systemProperty "tests.rest.${cluster.name}.http_hosts", "${-> getAllHttpSocketURI(cluster)}"
systemProperty "tests.rest.${cluster.name}.transport_hosts", "${-> getAllTransportSocketURI(cluster)}"
}
}

Expand All @@ -501,9 +503,7 @@ integTest {
systemProperty 'tests.security.manager', 'false'
systemProperty('project.root', project.projectDir.absolutePath)

systemProperty "https", System.getProperty("https")
systemProperty "user", System.getProperty("user")
systemProperty "password", System.getProperty("password")
systemProperties dynamicLoginConfig

// Set default query size limit
systemProperty 'defaultQuerySizeLimit', '10000'
Expand All @@ -517,7 +517,6 @@ integTest {
systemProperty 'cluster.debug', getDebug()
}


if (System.getProperty("test.debug") != null) {
jvmArgs '-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=*:5006'
}
Expand Down Expand Up @@ -554,6 +553,9 @@ integTest {

// Exclude this IT, because they executed in another task (:integTestWithSecurity)
exclude 'org/opensearch/sql/security/**'

// Exclude this IT, because they executed in another task (:integTestWithGeo)
exclude 'org/opensearch/sql/geo/**'
}


Expand Down Expand Up @@ -759,10 +761,7 @@ task integTestRemote(type: RestIntegTestTask) {
systemProperty 'tests.security.manager', 'false'
systemProperty('project.root', project.projectDir.absolutePath)
systemProperty 'java.io.tmpdir', opensearch_tmp_dir.absolutePath

systemProperty "https", System.getProperty("https")
systemProperty "user", System.getProperty("user")
systemProperty "password", System.getProperty("password")
systemProperties dynamicLoginConfig

// Set default query size limit
systemProperty 'defaultQuerySizeLimit', '10000'
Expand All @@ -783,3 +782,9 @@ task integTestRemote(type: RestIntegTestTask) {
exclude 'org/opensearch/sql/legacy/OrderIT.class'
exclude 'org/opensearch/sql/jdbc/**'
}






Loading

0 comments on commit 6abfd6d

Please sign in to comment.