Skip to content

Commit

Permalink
Add OpenSearch query builders and tests
Browse files Browse the repository at this point in the history
  • Loading branch information
harshach committed Sep 13, 2024
1 parent a81f891 commit 9a5ce15
Show file tree
Hide file tree
Showing 9 changed files with 449 additions and 63 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import org.openmetadata.service.search.queries.OMQueryBuilder;

public class ElasticQueryBuilder implements OMQueryBuilder {

private boolean isMustNot = false;
private QueryBuilder query;

public ElasticQueryBuilder() {
Expand Down Expand Up @@ -85,11 +85,6 @@ public OMQueryBuilder should(OMQueryBuilder query) {
return should(List.of(query));
}

@Override
public OMQueryBuilder mustNot(OMQueryBuilder query) {
return mustNot(List.of(query));
}

@Override
public boolean hasClauses() {
if (query instanceof BoolQueryBuilder) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,151 @@
package org.openmetadata.service.search.opensearch.queries;

import java.util.List;
import org.openmetadata.service.search.queries.OMQueryBuilder;
import os.org.opensearch.index.query.BoolQueryBuilder;
import os.org.opensearch.index.query.MatchAllQueryBuilder;
import os.org.opensearch.index.query.QueryBuilder;
import os.org.opensearch.index.query.QueryBuilders;

public class OpenSearchQueryBuilder implements OMQueryBuilder {
private boolean isMustNot = false;
private QueryBuilder query;

public OpenSearchQueryBuilder() {
// Default constructor
}

public OpenSearchQueryBuilder(QueryBuilder query) {
this.query = query;
}

@Override
public boolean isEmpty() {
return query == null;
}

@Override
public boolean isMatchNone() {
// Check if the query is a bool query with must_not match_all
if (query instanceof BoolQueryBuilder) {
BoolQueryBuilder boolQuery = (BoolQueryBuilder) query;
return boolQuery.must().isEmpty()
&& boolQuery.should().isEmpty()
&& boolQuery.mustNot().size() == 1
&& boolQuery.mustNot().get(0) instanceof MatchAllQueryBuilder;
}
return false;
}

@Override
public boolean isMatchAll() {
return query instanceof MatchAllQueryBuilder;
}

@Override
public OMQueryBuilder must(List<OMQueryBuilder> queries) {
BoolQueryBuilder boolQuery = getOrCreateBoolQuery();
for (OMQueryBuilder q : queries) {
OpenSearchQueryBuilder eqb = (OpenSearchQueryBuilder) q;
boolQuery.must(eqb.build());
}
this.query = boolQuery;
return this;
}

@Override
public OMQueryBuilder should(List<OMQueryBuilder> queries) {
BoolQueryBuilder boolQuery = getOrCreateBoolQuery();
for (OMQueryBuilder q : queries) {
OpenSearchQueryBuilder eqb = (OpenSearchQueryBuilder) q;
boolQuery.should(eqb.build());
}
this.query = boolQuery;
return this;
}

@Override
public OMQueryBuilder mustNot(List<OMQueryBuilder> queries) {
BoolQueryBuilder boolQuery = getOrCreateBoolQuery();
for (OMQueryBuilder q : queries) {
OpenSearchQueryBuilder eqb = (OpenSearchQueryBuilder) q;
boolQuery.mustNot(eqb.build());
}
this.query = boolQuery;
return this;
}

@Override
public OMQueryBuilder must(OMQueryBuilder query) {
return must(List.of(query));
}

@Override
public OMQueryBuilder should(OMQueryBuilder query) {
return should(List.of(query));
}

@Override
public boolean hasClauses() {
if (query instanceof BoolQueryBuilder) {
BoolQueryBuilder boolQuery = (BoolQueryBuilder) query;
return !boolQuery.must().isEmpty()
|| !boolQuery.should().isEmpty()
|| !boolQuery.mustNot().isEmpty();
}
return query != null;
}

public QueryBuilder build() {
return query;
}

public OpenSearchQueryBuilder setQuery(QueryBuilder query) {
this.query = query;
return this;
}

// Helper methods

public OpenSearchQueryBuilder matchNoneQuery() {
this.query = QueryBuilders.boolQuery().mustNot(QueryBuilders.matchAllQuery());
return this;
}

public OpenSearchQueryBuilder matchAllQuery() {
this.query = QueryBuilders.matchAllQuery();
return this;
}

public OpenSearchQueryBuilder boolQuery() {
this.query = QueryBuilders.boolQuery();
return this;
}

public OpenSearchQueryBuilder termQuery(String field, String value) {
this.query = QueryBuilders.termQuery(field, value);
return this;
}

public OpenSearchQueryBuilder termsQuery(String field, List<String> values) {
this.query = QueryBuilders.termsQuery(field, values);
return this;
}

public OpenSearchQueryBuilder existsQuery(String field) {
this.query = QueryBuilders.existsQuery(field);
return this;
}

private BoolQueryBuilder getOrCreateBoolQuery() {
if (query instanceof BoolQueryBuilder) {
return (BoolQueryBuilder) query;
} else {
BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
if (query != null) {
boolQuery.must(query);
}
return boolQuery;
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
package org.openmetadata.service.search.opensearch.queries;

import java.util.List;
import org.openmetadata.service.search.queries.OMQueryBuilder;
import org.openmetadata.service.search.queries.QueryBuilderFactory;

public class OpenSearchQueryBuilderFactory implements QueryBuilderFactory {

@Override
public OMQueryBuilder matchNoneQuery() {
return new OpenSearchQueryBuilder().matchNoneQuery();
}

@Override
public OMQueryBuilder matchAllQuery() {
return new OpenSearchQueryBuilder().matchAllQuery();
}

@Override
public OMQueryBuilder boolQuery() {
return new OpenSearchQueryBuilder().boolQuery();
}

@Override
public OMQueryBuilder termQuery(String field, String value) {
return new OpenSearchQueryBuilder().termQuery(field, value);
}

@Override
public OMQueryBuilder termsQuery(String field, List<String> values) {
return new OpenSearchQueryBuilder().termsQuery(field, values);
}

@Override
public OMQueryBuilder existsQuery(String field) {
return new OpenSearchQueryBuilder().existsQuery(field);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,5 @@ public interface OMQueryBuilder {

OMQueryBuilder should(OMQueryBuilder query);

OMQueryBuilder mustNot(OMQueryBuilder query);

boolean hasClauses();
}
Original file line number Diff line number Diff line change
Expand Up @@ -46,32 +46,6 @@ public void addMustNot(OMQueryBuilder query) {
}
}

public void mergeMust(ConditionCollector other) {
if (other.isMatchNothing()) {
this.setMatchNothing(true);
} else {
this.mustQueries.addAll(other.mustQueries);
this.mustNotQueries.addAll(other.mustNotQueries);
this.shouldQueries.addAll(other.shouldQueries);
}
}

public void mergeShould(ConditionCollector other) {
if (!other.isMatchNothing()) {
this.shouldQueries.addAll(other.mustQueries);
this.shouldQueries.addAll(other.shouldQueries);
this.mustNotQueries.addAll(other.mustNotQueries);
}
}

public void mergeMustNot(ConditionCollector other) {
if (!other.isMatchNothing()) {
this.mustNotQueries.addAll(other.mustQueries);
this.mustNotQueries.addAll(other.shouldQueries);
this.mustQueries.addAll(other.mustNotQueries);
}
}

public boolean isMatchAllQuery() {
if (mustQueries.size() == 1 && shouldQueries.isEmpty() && mustNotQueries.isEmpty()) {
OMQueryBuilder query = mustQueries.get(0);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,13 @@ private void preprocessExpression(SpelNode node, ConditionCollector collector) {
preprocessExpression(node.getChild(0), subCollector);

if (subCollector.isMatchAllQuery()) {
collector.setMatchNothing(true); // NOT TRUE == FALSE
collector.setMatchNothing(true);
} else if (subCollector.isMatchNothing()) {
collector.addMust(queryBuilderFactory.matchAllQuery());
} else {
OMQueryBuilder subQuery = subCollector.buildFinalQuery();
if (subQuery != null && !subQuery.isEmpty()) {
collector.addMustNot(subQuery); // Add to must_not for negation
collector.addMustNot(subQuery); // Add must_not without extra nesting
}
}
} else if (node instanceof MethodReference) {
Expand Down Expand Up @@ -214,13 +214,13 @@ public void hasAnyRole(List<String> roles, ConditionCollector collector) {
public void hasDomain(ConditionCollector collector) {
User user = (User) spelContext.lookupVariable("user");
if (user.getDomain() == null) {
collector.setMatchNothing(true);
return;
OMQueryBuilder existsQuery = queryBuilderFactory.existsQuery("domain.id");
collector.addMustNot(existsQuery);
} else {
String userDomainId = user.getDomain().getId().toString();
OMQueryBuilder domainQuery = queryBuilderFactory.termQuery("domain.id", userDomainId);
collector.addMust(domainQuery);
}
String userDomainId = user.getDomain().getId().toString();

OMQueryBuilder domainQuery = queryBuilderFactory.termQuery("domain.id", userDomainId);
collector.addMust(domainQuery);
}

public void inAnyTeam(List<String> teamNames, ConditionCollector collector) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,16 @@
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;
import static org.openmetadata.service.util.TestUtils.assertFieldDoesNotExist;
import static org.openmetadata.service.util.TestUtils.assertFieldExists;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jayway.jsonpath.DocumentContext;
import com.jayway.jsonpath.JsonPath;
import com.jayway.jsonpath.PathNotFoundException;
import es.org.elasticsearch.index.query.QueryBuilder;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.jupiter.api.BeforeEach;
Expand All @@ -28,7 +28,7 @@
import org.openmetadata.service.security.policyevaluator.CompiledRule;
import org.openmetadata.service.security.policyevaluator.SubjectContext;

class RBACConditionEvaluatorTest {
class ElasticSearchRBACConditionEvaluatorTest {

private RBACConditionEvaluator evaluator;
private User mockUser;
Expand Down Expand Up @@ -333,12 +333,17 @@ void testConditionUserLacksDomain() throws IOException {
OMQueryBuilder finalQuery = evaluator.evaluateConditions(mockSubjectContext);
QueryBuilder elasticQuery = ((ElasticQueryBuilder) finalQuery).build();
String generatedQuery = elasticQuery.toString();
DocumentContext jsonContext = JsonPath.parse(generatedQuery);
assertFieldExists(
jsonContext, "$.bool.must_not[?(@.exists.field=='domain.id')]", "must_not for domain.id");

assertTrue(
generatedQuery.contains("\"must_not\""), "The query should contain 'must_not' clause.");
assertTrue(
generatedQuery.contains("\"match_all\""),
"The must_not clause should contain 'match_all' query.");
// Check for owner ID and Public tag in the query
assertFieldExists(
jsonContext,
"$.bool.should[?(@.term['owner.id'].value=='" + mockUser.getId().toString() + "')]",
"owner.id");
assertFieldExists(
jsonContext, "$.bool.should[?(@.term['tags.tagFQN'].value=='Public')]", "Public tag");
}

@Test
Expand Down Expand Up @@ -690,19 +695,26 @@ void testComplexConditionWithMultipleNegations() throws IOException {
assertFieldExists(jsonContext, "$.bool.should[?(@.term['owner.id'])]", "owner.id");
}

private void assertFieldExists(DocumentContext jsonContext, String jsonPath, String fieldName) {
List<Map<String, Object>> result = jsonContext.read(jsonPath, List.class);
assertTrue(result.size() > 0, "The query should contain '" + fieldName + "' term.");
}
@Test
void testNotHasDomainWhenUserHasNoDomain() throws IOException {
setupMockPolicies("!hasDomain() && isOwner()", "ALLOW");
when(mockUser.getDomain()).thenReturn(null);
when(mockUser.getId()).thenReturn(UUID.randomUUID());

private void assertFieldDoesNotExist(
DocumentContext jsonContext, String jsonPath, String fieldName) {
try {
List<Map<String, Object>> result = jsonContext.read(jsonPath, List.class);
assertTrue(result.isEmpty(), "The query should not contain '" + fieldName + "' term.");
} catch (PathNotFoundException e) {
// If the path doesn't exist, this is expected behavior, so the test should pass.
assertTrue(true, "The path does not exist as expected: " + jsonPath);
}
OMQueryBuilder finalQuery = evaluator.evaluateConditions(mockSubjectContext);
QueryBuilder elasticQuery = ((ElasticQueryBuilder) finalQuery).build();
String generatedQuery = elasticQuery.toString();

DocumentContext jsonContext = JsonPath.parse(generatedQuery);

assertFieldExists(
jsonContext,
"$.bool.must_not[0].bool.must_not[?(@.exists.field=='domain.id')]",
"must_not for hasDomain");
assertFieldExists(
jsonContext,
"$.bool.should[?(@.term['owner.id'].value=='" + mockUser.getId().toString() + "')]",
"owner.id");
assertFieldDoesNotExist(jsonContext, "$.bool[?(@.match_none)]", "match_none should not exist");
}
}
Loading

0 comments on commit 9a5ce15

Please sign in to comment.