Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 75 additions & 86 deletions backend/src/main/java/com/bakdata/conquery/apiv1/QueryProcessor.java
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,13 @@
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import jakarta.inject.Inject;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.validation.Validator;
import jakarta.ws.rs.BadRequestException;
import jakarta.ws.rs.WebApplicationException;
import jakarta.ws.rs.core.Response;
import jakarta.ws.rs.core.UriBuilder;

import com.bakdata.conquery.apiv1.execution.ExecutionStatus;
import com.bakdata.conquery.apiv1.execution.FullExecutionStatus;
Expand All @@ -27,7 +34,6 @@
import com.bakdata.conquery.models.auth.AuthorizationHelper;
import com.bakdata.conquery.models.auth.entities.Group;
import com.bakdata.conquery.models.auth.entities.Subject;
import com.bakdata.conquery.models.auth.entities.User;
import com.bakdata.conquery.models.auth.permissions.Ability;
import com.bakdata.conquery.models.auth.permissions.ConqueryPermission;
import com.bakdata.conquery.models.common.Range;
Expand All @@ -40,7 +46,10 @@
import com.bakdata.conquery.models.execution.ManagedExecution;
import com.bakdata.conquery.models.i18n.I18n;
import com.bakdata.conquery.models.identifiable.ids.Id;
import com.bakdata.conquery.models.identifiable.ids.specific.*;
import com.bakdata.conquery.models.identifiable.ids.specific.ConnectorId;
import com.bakdata.conquery.models.identifiable.ids.specific.DatasetId;
import com.bakdata.conquery.models.identifiable.ids.specific.GroupId;
import com.bakdata.conquery.models.identifiable.ids.specific.ManagedExecutionId;
import com.bakdata.conquery.models.identifiable.mapping.IdPrinter;
import com.bakdata.conquery.models.query.*;
import com.bakdata.conquery.models.query.preview.EntityPreviewExecution;
Expand All @@ -57,12 +66,6 @@
import com.bakdata.conquery.util.QueryUtils;
import com.bakdata.conquery.util.QueryUtils.NamespacedIdentifiableCollector;
import com.bakdata.conquery.util.io.IdColumnUtil;
import jakarta.inject.Inject;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.validation.Validator;
import jakarta.ws.rs.BadRequestException;
import jakarta.ws.rs.core.Response;
import jakarta.ws.rs.core.UriBuilder;
import lombok.AllArgsConstructor;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
Expand Down Expand Up @@ -129,6 +132,31 @@ public static List<ResultAsset> getResultAssets(

}

private static List<QueryVisitor> visitQuery(QueryDescription queryContent, ExecutionManager executionManager, String primaryGroupName) {
final List<QueryVisitor> visitors = new ArrayList<>();

// This maps works as long as we have query visitors that are not configured in anyway.
// So adding a visitor twice would replace the previous one but both would have yielded the same result.
// For the future a better data structure might be desired that also regards similar QueryVisitors of different configuration
queryContent.addVisitors(visitors);

// Initialize checks that need to traverse the query tree
visitors.add(new QueryUtils.OnlyReusingChecker(executionManager));
visitors.add(new NamespacedIdentifiableCollector());
visitors.add(new ExecutionMetrics.QueryMetricsReporter(primaryGroupName));


// Chain all Consumers
Consumer<Visitable> consumerChain = QueryUtils.getNoOpEntryPoint();
for (QueryVisitor visitor : visitors) {
consumerChain = consumerChain.andThen(visitor);
}

// Apply consumers to the query tree
queryContent.visit(consumerChain);
return visitors;
}

public List<? extends ExecutionStatus> getAllQueries(DatasetId dataset, HttpServletRequest req, Subject subject, boolean allProviders) {
try (Stream<ManagedExecutionId> allQueries = storage.getAllExecutionIds()) {
return getQueriesFiltered(dataset, RequestAwareUriBuilder.fromRequest(req), subject, allQueries, allProviders).toList();
Expand Down Expand Up @@ -263,7 +291,6 @@ public void deleteQuery(Subject subject, ManagedExecutionId executionId) {
storage.removeExecution(executionId);
}


public FullExecutionStatus getQueryFullStatus(ManagedExecutionId queryId, Subject subject, UriBuilder url, Boolean allProviders, boolean await) {
final Namespace namespace = datasetRegistry.get(queryId.getDataset());

Expand Down Expand Up @@ -312,7 +339,7 @@ public ExternalUploadResult uploadEntities(Subject subject, DatasetId dataset, E
execution =
((ManagedQuery) namespace
.getExecutionManager()
.createExecution(query, subject.getId(), namespace, false));
.createExecution(query, subject.getId(), namespace, false, UUID.randomUUID()));

if (upload.getLabel() != null) {
execution.setLabel(upload.getLabel());
Expand Down Expand Up @@ -345,7 +372,8 @@ public FullExecutionStatus getSingleEntityExport(

// TODO make sure that subqueries are also system
// TODO do not persist system queries
final EntityPreviewExecution execution = (EntityPreviewExecution) postQuery(dataset, form, subject, true);
final EntityPreviewExecution execution = (EntityPreviewExecution) createExecution(dataset, form, subject, true, Optional.empty());
runExecution(execution);

final ExecutionManager executionManager = namespace.getExecutionManager();
if (executionManager.awaitDone(execution.getId(), 10, TimeUnit.SECONDS) == ExecutionState.RUNNING) {
Expand All @@ -370,117 +398,77 @@ public FullExecutionStatus getSingleEntityExport(
* Creates a query for all datasets, then submits it for execution on the
* intended dataset.
*/
public ManagedExecution postQuery(DatasetId dataset, QueryDescription queryContent, Subject subject, boolean system) {
public ManagedExecution createExecution(DatasetId dataset, QueryDescription queryContent, Subject subject, boolean system, Optional<UUID> maybeQueryId) {

log.info("Query posted on Dataset[{}] by User[{{}].", dataset, subject.getId());

// This maps works as long as we have query visitors that are not configured in anyway.
// So adding a visitor twice would replace the previous one but both would have yielded the same result.
// For the future a better data structure might be desired that also regards similar QueryVisitors of different configuration
final List<QueryVisitor> visitors = new ArrayList<>();
queryContent.addVisitors(visitors);

// Initialize checks that need to traverse the query tree
final QueryUtils.OnlyReusingChecker onlyReusingChecker = new QueryUtils.OnlyReusingChecker();
visitors.add(onlyReusingChecker);
final NamespacedIdentifiableCollector namespacedIdentifiableCollector = new NamespacedIdentifiableCollector();
visitors.add(namespacedIdentifiableCollector);
final Namespace namespace = datasetRegistry.get(dataset);
final ExecutionManager executionManager = namespace.getExecutionManager();

final String primaryGroupName = AuthorizationHelper.getPrimaryGroup(subject, storage).map(Group::getName).orElse("none");
final ExecutionMetrics.QueryMetricsReporter queryMetricsReporter = new ExecutionMetrics.QueryMetricsReporter(primaryGroupName);
visitors.add(queryMetricsReporter);


// Chain all Consumers
Consumer<Visitable> consumerChain = QueryUtils.getNoOpEntryPoint();
for (QueryVisitor visitor : visitors) {
consumerChain = consumerChain.andThen(visitor);
}

// Apply consumers to the query tree
queryContent.visit(consumerChain);
final List<QueryVisitor> visitors = visitQuery(queryContent, executionManager, primaryGroupName);

final NamespacedIdentifiableCollector namespacedIdentifiableCollector = QueryUtils.getVisitor(visitors, NamespacedIdentifiableCollector.class);
final QueryUtils.OnlyReusingChecker onlyReusingChecker = QueryUtils.getVisitor(visitors, QueryUtils.OnlyReusingChecker.class);

queryContent.authorize(subject, dataset, visitors, storage);
// After all authorization checks we can now use the actual subject to invoke the query and do not to bubble down the Userish in methods

// After all authorization checks we can now use the actual subject to invoke the query and do not to bubble down the User-ish in method

if (maybeQueryId.filter(id -> storage.getExecution(new ManagedExecutionId(dataset, id)) != null).isPresent()) {
throw new WebApplicationException("Query[%s] already exists.".formatted(maybeQueryId.get()), Response.Status.CONFLICT);
}

ExecutionMetrics.reportNamespacedIds(namespacedIdentifiableCollector.getIdentifiables(), primaryGroupName);

ExecutionMetrics.reportQueryClassUsage(queryContent.getClass(), primaryGroupName);

final Namespace namespace = datasetRegistry.get(dataset);
final ExecutionManager executionManager = namespace.getExecutionManager();


// If this is only a re-executing query, try to execute the underlying query instead.
{
final Optional<ManagedExecutionId> executionId = onlyReusingChecker.getOnlyReused();

final Optional<ManagedExecution>
execution =
executionId.map(id -> tryReuse(queryContent, id, namespace, executionManager, subject.getUser()));

if (execution.isPresent()) {
return execution.get();
}
}

// Execute the query
return executionManager.runQuery(namespace, queryContent, subject.getId(), system);
}

/**
* Determine if the submitted query does reuse ONLY another query and restart that instead of creating another one.
*/
private ManagedExecution tryReuse(QueryDescription query, ManagedExecutionId executionId, Namespace namespace, ExecutionManager executionManager, User user) {
final UUID queryId = maybeQueryId.orElseGet(UUID::randomUUID);

ManagedExecution execution = storage.getExecution(executionId);
final Optional<ManagedQuery> maybeReused = onlyReusingChecker.getOnlyReused(queryContent);

if (execution == null) {
return null;
}

// Direct reuse only works if the queries are of the same type (As reuse reconstructs the Query for different types)
if (!query.getClass().equals(execution.getSubmitted().getClass())) {
return null;
// Never reuse if a query-Id is provided
// If this is only a re-executing query, try to execute the underlying query instead.
if (maybeQueryId.isPresent() || maybeReused.isEmpty()) {
return executionManager.createExecution(queryContent, subject.getId(), namespace, system, queryId);
}

// If SecondaryIds differ from selected and prior, we cannot reuse them.
if (query instanceof SecondaryIdQuery secondaryIdQuery) {
final SecondaryIdDescriptionId selectedSecondaryId = secondaryIdQuery.getSecondaryId();
final SecondaryIdDescriptionId reusedSecondaryId = ((SecondaryIdQuery) execution.getSubmitted()).getSecondaryId();

if (!selectedSecondaryId.equals(reusedSecondaryId)) {
return null;
}
}
ManagedExecution execution = maybeReused.get();

// If the user is not the owner of the execution, we definitely create a new Execution, so the owner can cancel it
if (!user.isOwner(execution)) {
if (!subject.isOwner(execution)) {
final ManagedExecution
newExecution =
executionManager.createExecution(execution.getSubmitted(), user.getId(), namespace, false);
executionManager.createExecution(execution.getSubmitted(), subject.getId(), namespace, false, queryId);

newExecution.setLabel(execution.getLabel());
newExecution.setTags(execution.getTags().clone());

storage.updateExecution(newExecution);
execution = newExecution;
return newExecution;
}

final ExecutionState state = execution.getState();
if (state.equals(ExecutionState.RUNNING)) {
log.trace("The Execution[{}] was already started and its state is: {}", execution.getId(), state);
return execution;
}

log.trace("Re-executing Query {}", execution);

executionManager.execute(execution);
if (!state.equals(ExecutionState.RUNNING)) {
log.trace("Reusing Query {}", execution.getId());
}
Comment thread
thoniTUB marked this conversation as resolved.

return execution;

}

public void runExecution(ManagedExecution execution) {
if (execution.getState().equals(ExecutionState.RUNNING)) {
log.trace("The Execution[{}] is already RUNNING.", execution.getId());
return;
}
execution.getNamespace().getExecutionManager().execute(execution);
}

/**
* Execute a basic query on a single concept and return only the included entities Id's.
*/
Expand Down Expand Up @@ -511,7 +499,8 @@ public Stream<Map<String, String>> resolveEntities(Subject subject, List<FilterV

final QueryDescription query = new ConceptQuery(new CQOr(queries, Optional.of(false), DateAggregationAction.BLOCK));

final ManagedExecution execution = postQuery(dataset, query, subject, true);
final ManagedExecution execution = createExecution(dataset, query, subject, true, Optional.empty());
runExecution(execution);

if (namespace.getExecutionManager().awaitDone(execution.getId(), 10, TimeUnit.SECONDS) == ExecutionState.RUNNING) {
log.warn("Still waiting for {} after 10 Seconds.", execution.getId());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -142,17 +142,9 @@ public void addState(ManagedExecutionId id, ExecutionInfo result) {
executionInfosL2.put(id, result);
}

public final ManagedExecution runQuery(Namespace namespace, QueryDescription query, UserId user, boolean system) {
final ManagedExecution execution = createExecution(query, user, namespace, system);

execute(execution);

return execution;
}

// Visible for testing
public final ManagedExecution createExecution(QueryDescription query, UserId user, Namespace namespace, boolean system) {
return createExecution(query, UUID.randomUUID(), user, namespace, system);
public final ManagedExecution createExecution(QueryDescription query, UserId user, Namespace namespace, boolean system, UUID queryId) {
return createExecution(query, queryId, user, namespace, system);
}

public final void execute(ManagedExecution execution) {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,23 +1,16 @@
package com.bakdata.conquery.resources.api;

import static com.bakdata.conquery.resources.ResourceConstants.DATASET;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.stream.Stream;
import jakarta.inject.Inject;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.validation.Valid;
import jakarta.validation.constraints.NotEmpty;
import jakarta.validation.constraints.NotNull;
import jakarta.ws.rs.Consumes;
import jakarta.ws.rs.GET;
import jakarta.ws.rs.POST;
import jakarta.ws.rs.Path;
import jakarta.ws.rs.PathParam;
import jakarta.ws.rs.Produces;
import jakarta.ws.rs.QueryParam;
import jakarta.ws.rs.*;
import jakarta.ws.rs.core.Context;
import jakarta.ws.rs.core.Response;
import jakarta.ws.rs.core.UriBuilder;
Expand All @@ -35,10 +28,13 @@
import com.bakdata.conquery.models.auth.permissions.Ability;
import com.bakdata.conquery.models.execution.ManagedExecution;
import com.bakdata.conquery.models.identifiable.ids.specific.DatasetId;
import com.bakdata.conquery.util.validation.ValidUUID4;
import io.dropwizard.auth.Auth;
import lombok.Data;
import lombok.RequiredArgsConstructor;

import static com.bakdata.conquery.resources.ResourceConstants.DATASET;

@Path("datasets/{" + DATASET + "}/queries")
@Data
@Consumes(AdditionalMediaTypes.JSON)
Expand Down Expand Up @@ -75,7 +71,6 @@ public Stream<Map<String, String>> resolveEntities(@Auth Subject subject, @Valid
}



@POST
@Path("/upload")
public ExternalUploadResult upload(@Auth Subject subject, @Valid ExternalUpload upload) {
Expand All @@ -94,20 +89,43 @@ public List<? extends ExecutionStatus> getAllQueries(@Auth Subject subject, @Que
return processor.getAllQueries(dataset, servletRequest, subject, allProviders.orElse(false));
}

/***
* Create and run the submitted query
*/
@POST
public Response postQuery(@Auth Subject subject, @QueryParam("all-providers") Optional<Boolean> allProviders, @NotNull @Valid QueryDescription query) {
public Response postQuery(@Auth Subject subject, @QueryParam("all-providers") Optional<Boolean> allProviders, @QueryParam("queryId") Optional<@ValidUUID4 UUID> queryId, @NotNull @Valid QueryDescription query) {
subject.authorize(dataset, Ability.READ);

final ManagedExecution execution = processor.createExecution(dataset, query, subject, false, queryId);
processor.runExecution(execution);

return Response.ok(processor.getQueryFullStatus(execution.getId(),
subject,
RequestAwareUriBuilder.fromRequest(servletRequest),
allProviders.orElse(false),
false
))
.status(Response.Status.CREATED)
.build();
}

/**
* Only create the query.
*/
@PUT
public Response putQuery(@Auth Subject subject, @QueryParam("all-providers") Optional<Boolean> allProviders, @QueryParam("queryId") Optional<@ValidUUID4 UUID> queryId, @NotNull @Valid QueryDescription query) {
subject.authorize(dataset, Ability.READ);

final ManagedExecution execution = processor.postQuery(dataset, query, subject, false);
final ManagedExecution execution = processor.createExecution(dataset, query, subject, false, queryId);

return Response.ok(processor.getQueryFullStatus(execution.getId(),
subject,
RequestAwareUriBuilder.fromRequest(servletRequest),
allProviders.orElse(false),
false
))
.status(Response.Status.CREATED)
.build();
subject,
RequestAwareUriBuilder.fromRequest(servletRequest),
allProviders.orElse(false),
false
))
.status(Response.Status.CREATED)
.build();
}

}
Loading
Loading