Skip to content

Commit b8c3820

Browse files
andrii0lomakinandrii0lomakin
authored
Merge pull request #550 from JetBrains/ytdb-406-has-label-improvements
YTDB-406 Improve hasLabel implementation to work properly with polymorphic queries
2 parents d0eb4fb + ee0ef6b commit b8c3820

File tree

17 files changed

+1134
-647
lines changed

17 files changed

+1134
-647
lines changed

‎core/src/main/java/com/jetbrains/youtrackdb/api/gremlin/YTDBGraphTraversalSourceDSL.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package com.jetbrains.youtrackdb.api.gremlin;
22

3+
import com.jetbrains.youtrackdb.api.gremlin.service.YTDBCommandService;
34
import com.jetbrains.youtrackdb.api.gremlin.tokens.YTDBQueryConfigParam;
45

56
import com.jetbrains.youtrackdb.internal.core.gremlin.YTDBTransaction;
7+
import java.util.Map;
68
import javax.annotation.Nonnull;
79
import org.apache.commons.lang3.function.FailableConsumer;
810
import org.apache.commons.lang3.function.FailableFunction;
@@ -76,4 +78,26 @@ public <X extends Exception, R> R computeInTx(
7678
var tx = tx();
7779
return YTDBTransaction.computeInTx(code, (YTDBTransaction) tx);
7880
}
81+
82+
/// Execute a generic YouTrackDB command. The result of the execution is ignored, so it only makes
83+
/// sense to use this method for running non-idempotent commands.
84+
///
85+
/// @param command The command to execute.
86+
public void command(@Nonnull String command) {
87+
command(command, Map.of());
88+
}
89+
90+
/// Execute a generic parameterized YouTrackDB command. The result of the execution is ignored, so
91+
/// it only makes sense to use this method for running non-idempotent commands.
92+
///
93+
/// @param command The command to execute.
94+
/// @param arguments The arguments to pass to the command.
95+
public void command(@Nonnull String command, @Nonnull Map<?, ?> arguments) {
96+
call(
97+
YTDBCommandService.NAME, Map.of(
98+
YTDBCommandService.COMMAND, command,
99+
YTDBCommandService.ARGUMENTS, arguments
100+
)
101+
).iterate();
102+
}
79103
}

‎core/src/main/java/com/jetbrains/youtrackdb/api/gremlin/service/YTDBCommandService.java‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
package com.jetbrains.youtrackdb.api.gremlin.service;
2+
3+
import com.jetbrains.youtrackdb.api.gremlin.embedded.YTDBElement;
4+
import com.jetbrains.youtrackdb.internal.core.gremlin.YTDBGraphInternal;
5+
import java.util.Map;
6+
import java.util.Set;
7+
import org.apache.tinkerpop.gremlin.process.traversal.Traversal.Admin;
8+
import org.apache.tinkerpop.gremlin.process.traversal.traverser.TraverserRequirement;
9+
import org.apache.tinkerpop.gremlin.structure.service.Service;
10+
import org.apache.tinkerpop.gremlin.structure.util.CloseableIterator;
11+
12+
/// TinkerPop service that allows running any YouTrackDB non-idempotent command via GraphTraversal.
13+
///
14+
/// This is a `Start` service, meaning that it is not allowed to be used mid-traversal. It always
15+
/// produces an empty result.
16+
public class YTDBCommandService<E extends YTDBElement> implements Service<E, E> {
17+
18+
public static final String NAME = "ytdbCommand";
19+
public static final String COMMAND = "command";
20+
public static final String ARGUMENTS = "args";
21+
22+
private final String command;
23+
private final Map<?, ?> commandParams;
24+
25+
public YTDBCommandService(String command, Map<?, ?> commandParams) {
26+
this.command = command;
27+
this.commandParams = commandParams;
28+
}
29+
30+
public static class Factory<E extends YTDBElement> implements ServiceFactory<E, E> {
31+
32+
@Override
33+
public String getName() {
34+
return NAME;
35+
}
36+
37+
@Override
38+
public Set<Type> getSupportedTypes() {
39+
return Set.of(Type.Start);
40+
}
41+
42+
@Override
43+
public Service<E, E> createService(boolean isStart, Map params) {
44+
if (!isStart) {
45+
throw new UnsupportedOperationException(Exceptions.cannotUseMidTraversal);
46+
}
47+
48+
final String command;
49+
final Map<?, ?> commandParams;
50+
51+
if (params.get(COMMAND) instanceof String c) {
52+
command = c;
53+
} else {
54+
throw new IllegalArgumentException(params.get(COMMAND) + " is not a String");
55+
}
56+
57+
if (params.get(ARGUMENTS) instanceof Map<?, ?> m) {
58+
commandParams = m;
59+
} else if (params.get(ARGUMENTS) == null) {
60+
commandParams = Map.of();
61+
} else {
62+
throw new IllegalArgumentException(params.get(ARGUMENTS) + " is not a Map");
63+
}
64+
65+
return new YTDBCommandService<>(command, commandParams);
66+
}
67+
}
68+
69+
@Override
70+
public Type getType() {
71+
return Type.Start;
72+
}
73+
74+
@Override
75+
public Set<TraverserRequirement> getRequirements() {
76+
return Set.of();
77+
}
78+
79+
@Override
80+
public CloseableIterator<E> execute(ServiceCallContext ctx, Map params) {
81+
82+
final var graph = (((Admin<?, ?>) ctx.getTraversal()))
83+
.getGraph()
84+
.orElseThrow(() -> new IllegalStateException("Graph is not available"));
85+
86+
(((YTDBGraphInternal) graph)).executeCommand(command, commandParams);
87+
88+
return CloseableIterator.empty();
89+
}
90+
}

‎core/src/main/java/com/jetbrains/youtrackdb/api/gremlin/service/YTDBServices.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ private YTDBServices() {
1414

1515
{
1616
registerService(new YTDBRemovePropertyService.Factory<>());
17+
registerService(new YTDBCommandService.Factory<>());
1718
frozen = true;
1819
}
1920

‎core/src/main/java/com/jetbrains/youtrackdb/internal/core/gremlin/YTDBGraphImplAbstract.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.jetbrains.youtrackdb.internal.core.gremlin.traversal.strategy.optimization.YTDBGraphStepStrategy;
1919
import com.jetbrains.youtrackdb.internal.core.id.RecordIdInternal;
2020
import java.util.Iterator;
21+
import java.util.Map;
2122
import java.util.Objects;
2223
import java.util.function.Consumer;
2324
import java.util.function.Function;
@@ -211,6 +212,13 @@ public <R> R computeSchemaCode(Function<DatabaseSessionEmbedded, R> code) {
211212
}
212213
}
213214

215+
@Override
216+
public void executeCommand(String command, Map<?, ?> params) {
217+
try (var session = acquireSession()) {
218+
session.command(command, params);
219+
}
220+
}
221+
214222
@Override
215223
public Variables variables() {
216224
throw new NotImplementedException();

‎core/src/main/java/com/jetbrains/youtrackdb/internal/core/gremlin/YTDBGraphInternal.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.jetbrains.youtrackdb.api.gremlin.YTDBGraph;
44
import com.jetbrains.youtrackdb.internal.core.db.DatabaseSessionEmbedded;
5+
import java.util.Map;
56
import java.util.function.Consumer;
67
import java.util.function.Function;
78
import javax.annotation.Nullable;
@@ -16,4 +17,6 @@ public interface YTDBGraphInternal extends YTDBGraph {
1617

1718
@Nullable
1819
<R> R computeSchemaCode(Function<DatabaseSessionEmbedded, R> code);
20+
21+
void executeCommand(String command, Map<?, ?> params);
1922
}

‎core/src/main/java/com/jetbrains/youtrackdb/internal/core/gremlin/traversal/step/filter/YTDBHasLabelStep.java‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
package com.jetbrains.youtrackdb.internal.core.gremlin.traversal.step.filter;
2+
3+
import com.jetbrains.youtrackdb.api.gremlin.embedded.YTDBElement;
4+
import com.jetbrains.youtrackdb.internal.core.gremlin.YTDBElementImpl;
5+
import java.util.EnumSet;
6+
import java.util.List;
7+
import java.util.Set;
8+
import org.apache.tinkerpop.gremlin.process.traversal.P;
9+
import org.apache.tinkerpop.gremlin.process.traversal.Traversal;
10+
import org.apache.tinkerpop.gremlin.process.traversal.Traverser.Admin;
11+
import org.apache.tinkerpop.gremlin.process.traversal.step.filter.FilterStep;
12+
import org.apache.tinkerpop.gremlin.process.traversal.step.util.AbstractStep;
13+
import org.apache.tinkerpop.gremlin.process.traversal.traverser.TraverserRequirement;
14+
import org.apache.tinkerpop.gremlin.structure.util.StringFactory;
15+
16+
/// A filtering step that replaces the standard TinkerPop "hasLabel" step. It tests "~label"
17+
/// predicates natively (via YouTrackDB API) and respects the "polymorphicQuery" traversal setting.
18+
public class YTDBHasLabelStep<S extends YTDBElement> extends FilterStep<S> {
19+
20+
private List<P<? super String>> predicates;
21+
private boolean polymorphic;
22+
23+
public YTDBHasLabelStep(
24+
final Traversal.Admin<?, ?> traversal,
25+
List<P<? super String>> predicates,
26+
boolean polymorphic
27+
) {
28+
super(traversal);
29+
this.predicates = predicates;
30+
this.polymorphic = polymorphic;
31+
}
32+
33+
@Override
34+
public Set<TraverserRequirement> getRequirements() {
35+
return EnumSet.of(TraverserRequirement.OBJECT);
36+
}
37+
38+
private boolean test(String className) {
39+
return predicates.stream().anyMatch(p -> p.test(className));
40+
}
41+
42+
@Override
43+
protected boolean filter(Admin<S> traverser) {
44+
45+
// will it make sense to add a step-level cache for storing the names of the classes
46+
// that fit or don't fit the predicate? should it be thread-safe?
47+
48+
if (traverser.get() instanceof YTDBElementImpl ytdbElement) {
49+
50+
final var entity = ytdbElement.getRawEntity();
51+
final var schemaClass = entity.getSchemaClass();
52+
if (schemaClass == null) {
53+
// this shouldn't ever happen, I suppose
54+
return false;
55+
}
56+
if (test(schemaClass.getName())) {
57+
return true;
58+
}
59+
if (!polymorphic) {
60+
return false;
61+
}
62+
63+
for (var c : schemaClass.getAllSuperClasses()) {
64+
if (test(c.getName())) {
65+
return true;
66+
}
67+
}
68+
69+
return false;
70+
} else {
71+
return test(traverser.get().label());
72+
}
73+
}
74+
75+
@Override
76+
public String toString() {
77+
return StringFactory.stepString(this, predicates, polymorphic);
78+
}
79+
80+
@Override
81+
public boolean equals(Object o) {
82+
return super.equals(o);
83+
}
84+
85+
@Override
86+
public int hashCode() {
87+
return super.hashCode() ^
88+
predicates.hashCode() ^
89+
Boolean.hashCode(polymorphic);
90+
}
91+
92+
@Override
93+
public AbstractStep<S, S> clone() {
94+
final var clone = (YTDBHasLabelStep<S>) super.clone();
95+
clone.polymorphic = this.polymorphic;
96+
// predicates appear to be mutable (setValue method)
97+
clone.predicates = this.predicates.stream()
98+
.<P<? super String>>map(P::clone)
99+
.toList();
100+
return clone;
101+
}
102+
}

‎core/src/main/java/com/jetbrains/youtrackdb/internal/core/gremlin/traversal/step/map/YTDBClassCountStep.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,20 @@ public class YTDBClassCountStep<S> extends AbstractStep<S, Long> {
1515

1616
private boolean vertexStep;
1717
private List<String> klasses;
18+
private boolean polymorphic;
1819

1920
protected boolean done = false;
2021

21-
public YTDBClassCountStep(Traversal.Admin traversal, List<String> klasses, boolean vertexStep) {
22+
public YTDBClassCountStep(
23+
Traversal.Admin traversal,
24+
List<String> klasses,
25+
boolean vertexStep,
26+
boolean polymorphic
27+
) {
2228
super(traversal);
2329
this.klasses = klasses;
2430
this.vertexStep = vertexStep;
31+
this.polymorphic = polymorphic;
2532
}
2633

2734
@Override
@@ -32,7 +39,7 @@ protected Traverser.Admin<Long> processNextStart() throws NoSuchElementException
3239
Long v =
3340
klasses.stream()
3441
.filter(this::filterClass)
35-
.mapToLong(session::countClass)
42+
.mapToLong(cl -> session.countClass(cl, polymorphic))
3643
.reduce(0, Long::sum);
3744
//noinspection unchecked
3845
return this.traversal.getTraverserGenerator().generate(v, (Step<Long, ?>) this, 1L);
@@ -74,6 +81,7 @@ public YTDBClassCountStep<S> clone() {
7481
newCount.klasses = this.klasses;
7582
newCount.vertexStep = this.vertexStep;
7683
newCount.done = false;
84+
newCount.polymorphic = this.polymorphic;
7785
return newCount;
7886
}
7987
}

‎core/src/main/java/com/jetbrains/youtrackdb/internal/core/gremlin/traversal/strategy/optimization/YTDBGraphCountStrategy.java‎

Lines changed: 24 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -47,36 +47,36 @@ public void apply(Traversal.Admin<?, ?> traversal) {
4747
return;
4848
}
4949

50-
var steps = traversal.getSteps();
51-
if (steps.size() < 2) {
50+
final var polymorphicSetting = YTDBStrategyUtil.isPolymorphic(traversal);
51+
if (polymorphicSetting == null) {
52+
// means we couldn't access the graph from the traversal
5253
return;
5354
}
5455

55-
var startStep = traversal.getStartStep();
56-
var endStep = traversal.getEndStep();
56+
final var steps = traversal.getSteps();
57+
5758
if (steps.size() == 2
58-
&& startStep instanceof YTDBGraphStep<?, ?> step
59-
&& endStep instanceof CountGlobalStep) {
60-
61-
if (step.getHasContainers().size() == 1) {
62-
var hasContainers = step.getHasContainers();
63-
var classes =
64-
hasContainers.stream()
65-
.filter(YTDBGraphCountStrategy::isLabelFilter)
66-
.map(YTDBGraphCountStrategy::extractLabels)
67-
.flatMap(Collection::stream)
68-
.collect(Collectors.toList());
69-
if (!classes.isEmpty()) {
70-
TraversalHelper.removeAllSteps(traversal);
71-
traversal.addStep(new YTDBClassCountStep<>(traversal, classes, step.isVertexStep()));
72-
}
73-
} else if (step.getHasContainers().isEmpty() && step.getIds().length == 0) {
59+
&& steps.getFirst() instanceof YTDBGraphStep<?, ?> step
60+
&& steps.getLast() instanceof CountGlobalStep) {
61+
62+
final var hasContainers = step.getHasContainers();
63+
List<String> classes = List.of();
64+
boolean polymorphic = polymorphicSetting;
65+
66+
if (hasContainers.size() == 1 && isLabelFilter(hasContainers.getFirst())) {
67+
// g.V().hasLabel('Foo').count()
68+
classes = extractLabels(hasContainers.getFirst());
69+
} else if (hasContainers.isEmpty() && step.getIds().length == 0) {
70+
// g.V().count()
71+
classes = List.of(
72+
step.isVertexStep() ? SchemaClass.VERTEX_CLASS_NAME : SchemaClass.EDGE_CLASS_NAME);
73+
polymorphic = true; // should be polymorphic, because we want to see all vertices or edges
74+
}
75+
76+
if (!classes.isEmpty()) {
7477
TraversalHelper.removeAllSteps(traversal);
75-
var baseClass =
76-
step.isVertexStep() ? SchemaClass.VERTEX_CLASS_NAME : SchemaClass.EDGE_CLASS_NAME;
7778
traversal.addStep(
78-
new YTDBClassCountStep<>(
79-
traversal, Collections.singletonList(baseClass), step.isVertexStep()));
79+
new YTDBClassCountStep<>(traversal, classes, step.isVertexStep(), polymorphic));
8080
}
8181
}
8282
}

0 commit comments

Comments
 (0)