
io.kestra.plugin.gcp.firestore.Query Maven / Gradle / Ivy
package io.kestra.plugin.gcp.firestore;
import com.google.cloud.firestore.CollectionReference;
import com.google.cloud.firestore.Query.Direction;
import com.google.cloud.firestore.QueryDocumentSnapshot;
import io.kestra.core.exceptions.IllegalVariableEvaluationException;
import io.kestra.core.models.annotations.Example;
import io.kestra.core.models.annotations.Plugin;
import io.kestra.core.models.annotations.PluginProperty;
import io.kestra.core.models.executions.metrics.Counter;
import io.kestra.core.models.tasks.RunnableTask;
import io.kestra.core.models.tasks.common.FetchOutput;
import io.kestra.core.models.tasks.common.FetchType;
import io.kestra.core.runners.RunContext;
import io.kestra.core.serializers.FileSerde;
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.*;
import lombok.experimental.SuperBuilder;
import org.apache.commons.lang3.tuple.Pair;
import java.io.*;
import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import jakarta.validation.constraints.NotNull;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import static io.kestra.core.utils.Rethrow.throwConsumer;
@SuperBuilder
@ToString
@EqualsAndHashCode
@Getter
@NoArgsConstructor
@Schema(
title = "Query documents of a collection."
)
@Plugin(
examples = {
@Example(
full = true,
code = """
id: gcp_firestore_query
namespace: company.team
tasks:
- id: query
type: io.kestra.plugin.gcp.firestore.Query
collection: "persons"
filters:
- field: "lastname"
value: "Doe"
"""
)
}
)
public class Query extends AbstractFirestore implements RunnableTask {
@Schema(
title = "The way you want to store the data",
description = "FETCH_ONE output the first row, "
+ "FETCH output all the rows, "
+ "STORE store all rows in a file, "
+ "NONE do nothing."
)
@Builder.Default
@PluginProperty
private FetchType fetchType = FetchType.STORE;
@Schema(
title = "List of query filters that will be added as a where clause."
)
@PluginProperty
private List filters;
@Schema(
title = "Field name for the order by clause."
)
@PluginProperty
private String orderBy;
@Schema(
title = "Field name for the order by clause."
)
@PluginProperty
@Builder.Default
private Direction orderDirection = Direction.ASCENDING;
@Schema(
title = "Start offset for pagination of the query results."
)
@PluginProperty
private Integer offset;
@Schema(
title = "Maximum numbers of returned results."
)
@PluginProperty
private Integer limit;
@Override
public FetchOutput run(RunContext runContext) throws Exception {
try (var firestore = this.connection(runContext)) {
var collectionRef = this.collection(runContext, firestore);
var query = getQuery(runContext, collectionRef, this.filters);
if (this.orderBy != null) {
query.orderBy(this.orderBy, this.orderDirection);
}
if (this.offset != null) {
query.offset(this.offset);
}
if (this.limit != null) {
query.limit(this.limit);
}
var queryDocumentSnapshots = query.get().get().getDocuments();
var outputBuilder = FetchOutput.builder();
switch (fetchType) {
case FETCH:
Pair, Long> fetch = this.fetch(queryDocumentSnapshots);
outputBuilder
.rows(fetch.getLeft())
.size(fetch.getRight());
break;
case FETCH_ONE:
var o = this.fetchOne(queryDocumentSnapshots);
outputBuilder
.row(o)
.size(o != null ? 1L : 0L);
break;
case STORE:
Pair store = this.store(runContext, queryDocumentSnapshots);
outputBuilder
.uri(store.getLeft())
.size(store.getRight());
break;
}
var output = outputBuilder.build();
runContext.metric(Counter.of(
"records", output.getSize(),
"collection", collectionRef.getId()
));
return output;
}
}
private com.google.cloud.firestore.Query getQuery(RunContext runContext, CollectionReference collectionRef, List filters)
throws IllegalVariableEvaluationException {
// this is a no-op but allow to create an empty query
var query = collectionRef.offset(0);
if (this.filters == null || this.filters.isEmpty()) {
return query;
}
for(Filter option : filters) {
query = appendQueryPart(runContext, query, option);
}
return query;
}
private com.google.cloud.firestore.Query appendQueryPart(RunContext runContext, com.google.cloud.firestore.Query query,
Filter filter)
throws IllegalVariableEvaluationException {
switch (filter.getOperator()) {
case EQUAL_TO: {
return query.whereEqualTo(filter.getField(), runContext.render(filter.getValue()));
}
case NOT_EQUAL_TO: {
return query.whereNotEqualTo(filter.getField(), runContext.render(filter.getValue()));
}
case LESS_THAN: {
return query.whereLessThan(filter.getField(), runContext.render(filter.getValue()));
}
case LESS_THAN_OR_EQUAL_TO: {
return query.whereLessThanOrEqualTo(filter.getField(), runContext.render(filter.getValue()));
}
case GREATER_THAN: {
return query.whereGreaterThan(filter.getField(), runContext.render(filter.getValue()));
}
case GREATER_THAN_OR_EQUAL_TO: {
return query.whereGreaterThanOrEqualTo(filter.getField(), runContext.render(filter.getValue()));
}
// more where clause could be supported, but we need to validate that the value is a List
}
// Should never occur
throw new IllegalArgumentException("Unknown QueryOperator: " + filter.getOperator());
}
private Pair store(RunContext runContext, List documents) throws IOException {
File tempFile = runContext.workingDir().createTempFile(".ion").toFile();
try (var output = new BufferedWriter(new FileWriter(tempFile), FileSerde.BUFFER_SIZE)) {
Flux
© 2015 - 2025 Weber Informatics LLC | Privacy Policy