Skip to content

Commit

Permalink
Efficient queries for connection list (#17360)
Browse files Browse the repository at this point in the history
* query once for all needed models, instead of querying within connections loop

* cleanup and fix failing tests

* pmd fix

* fix query and add test

* return empty if input list is empty

* undo aggressive autoformatting

* don't query for connection operations in a loop, instead query once and group-by connectionID in memory

* try handling operationIds in a single query instead of two

* remove optional

* fix operationIds query

* very annoying, test was failing because operationIds can be listed in a different order. verify operationIds separately from rest of object

* combined queries/functions instead of separate queries for actor and definition

* remove leftover lines that aren't doing anything

* format

* add javadoc

* format

* use leftjoin so that connections that lack operations aren't left out

* clean up comments and format
  • Loading branch information
pmossman authored Oct 10, 2022
1 parent fb523af commit 39a14b7
Show file tree
Hide file tree
Showing 12 changed files with 524 additions and 78 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,12 @@
import static io.airbyte.db.instance.configs.jooq.generated.Tables.OPERATION;
import static io.airbyte.db.instance.configs.jooq.generated.Tables.WORKSPACE;
import static org.jooq.impl.DSL.asterisk;
import static org.jooq.impl.DSL.groupConcat;
import static org.jooq.impl.DSL.noCondition;
import static org.jooq.impl.DSL.select;

import com.fasterxml.jackson.databind.JsonNode;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Charsets;
import com.google.common.collect.Sets;
import com.google.common.hash.HashFunction;
Expand Down Expand Up @@ -54,6 +57,8 @@
import java.io.IOException;
import java.time.OffsetDateTime;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -82,6 +87,8 @@
public class ConfigRepository {

private static final Logger LOGGER = LoggerFactory.getLogger(ConfigRepository.class);
private static final String OPERATION_IDS_AGG_FIELD = "operation_ids_agg";
private static final String OPERATION_IDS_AGG_DELIMITER = ",";

private final ConfigPersistence persistence;
private final ExceptionWrappingDatabase database;
Expand Down Expand Up @@ -613,43 +620,73 @@ public void writeStandardSync(final StandardSync standardSync) throws JsonValida
persistence.writeConfig(ConfigSchema.STANDARD_SYNC, standardSync.getConnectionId().toString(), standardSync);
}

public List<StandardSync> listStandardSyncs() throws ConfigNotFoundException, IOException, JsonValidationException {
public List<StandardSync> listStandardSyncs() throws IOException, JsonValidationException {
return persistence.listConfigs(ConfigSchema.STANDARD_SYNC, StandardSync.class);
}

public List<StandardSync> listStandardSyncsUsingOperation(final UUID operationId)
throws IOException {
final Result<Record> result = database.query(ctx -> ctx.select(CONNECTION.asterisk())

final Result<Record> connectionAndOperationIdsResult = database.query(ctx -> ctx
// SELECT connection.* plus the connection's associated operationIds as a concatenated list
.select(
CONNECTION.asterisk(),
groupConcat(CONNECTION_OPERATION.OPERATION_ID).separator(OPERATION_IDS_AGG_DELIMITER).as(OPERATION_IDS_AGG_FIELD))
.from(CONNECTION)
.join(CONNECTION_OPERATION)
.on(CONNECTION_OPERATION.CONNECTION_ID.eq(CONNECTION.ID))
.where(CONNECTION_OPERATION.OPERATION_ID.eq(operationId))).fetch();
return getStandardSyncsFromResult(result);

// inner join with all connection_operation rows that match the connection's id
.join(CONNECTION_OPERATION).on(CONNECTION_OPERATION.CONNECTION_ID.eq(CONNECTION.ID))

// only keep rows for connections that have an operationId that matches the input.
// needs to be a sub query because we want to keep all operationIds for matching connections
// in the main query
.where(CONNECTION.ID.in(
select(CONNECTION.ID).from(CONNECTION).join(CONNECTION_OPERATION).on(CONNECTION_OPERATION.CONNECTION_ID.eq(CONNECTION.ID))
.where(CONNECTION_OPERATION.OPERATION_ID.eq(operationId))))

// group by connection.id so that the groupConcat above works
.groupBy(CONNECTION.ID)).fetch();

return getStandardSyncsFromResult(connectionAndOperationIdsResult);
}

public List<StandardSync> listWorkspaceStandardSyncs(final UUID workspaceId, final boolean includeDeleted) throws IOException {
final Result<Record> result = database.query(ctx -> ctx.select(CONNECTION.asterisk())
final Result<Record> connectionAndOperationIdsResult = database.query(ctx -> ctx
// SELECT connection.* plus the connection's associated operationIds as a concatenated list
.select(
CONNECTION.asterisk(),
groupConcat(CONNECTION_OPERATION.OPERATION_ID).separator(OPERATION_IDS_AGG_DELIMITER).as(OPERATION_IDS_AGG_FIELD))
.from(CONNECTION)

// left join with all connection_operation rows that match the connection's id.
// left join includes connections that don't have any connection_operations
.leftJoin(CONNECTION_OPERATION).on(CONNECTION_OPERATION.CONNECTION_ID.eq(CONNECTION.ID))

// join with source actors so that we can filter by workspaceId
.join(ACTOR).on(CONNECTION.SOURCE_ID.eq(ACTOR.ID))
.where(ACTOR.WORKSPACE_ID.eq(workspaceId)
.and(includeDeleted ? noCondition() : CONNECTION.STATUS.notEqual(StatusType.deprecated))))
.fetch();
return getStandardSyncsFromResult(result);
.and(includeDeleted ? noCondition() : CONNECTION.STATUS.notEqual(StatusType.deprecated)))

// group by connection.id so that the groupConcat above works
.groupBy(CONNECTION.ID)).fetch();

return getStandardSyncsFromResult(connectionAndOperationIdsResult);
}

private List<StandardSync> getStandardSyncsFromResult(final Result<Record> result) throws IOException {
private List<StandardSync> getStandardSyncsFromResult(final Result<Record> connectionAndOperationIdsResult) {
final List<StandardSync> standardSyncs = new ArrayList<>();
for (final Record record : result) {
final UUID connectionId = record.get(CONNECTION.ID);
final Result<Record> connectionOperationRecords = database.query(ctx -> ctx.select(asterisk())
.from(CONNECTION_OPERATION)
.where(CONNECTION_OPERATION.CONNECTION_ID.eq(connectionId))
.fetch());

final List<UUID> connectionOperationIds =
connectionOperationRecords.stream().map(r -> r.get(CONNECTION_OPERATION.OPERATION_ID)).collect(Collectors.toList());
standardSyncs.add(DbConverter.buildStandardSync(record, connectionOperationIds));
for (final Record record : connectionAndOperationIdsResult) {
final String operationIdsFromRecord = record.get(OPERATION_IDS_AGG_FIELD, String.class);

// can be null when connection has no connectionOperations
final List<UUID> operationIds = operationIdsFromRecord == null
? Collections.emptyList()
: Arrays.stream(operationIdsFromRecord.split(OPERATION_IDS_AGG_DELIMITER)).map(UUID::fromString).toList();

standardSyncs.add(DbConverter.buildStandardSync(record, operationIds));
}

return standardSyncs;
}

Expand Down Expand Up @@ -798,6 +835,61 @@ private Map<UUID, AirbyteCatalog> findCatalogByHash(final String catalogHash, fi
return result;
}

// Data-carrier records to hold combined result of query for a Source or Destination and its
// corresponding Definition. This enables the API layer to
// process combined information about a Source/Destination/Definition pair without requiring two
// separate queries and in-memory join operation,
// because the config models are grouped immediately in the repository layer.
@VisibleForTesting
public record SourceAndDefinition(SourceConnection source, StandardSourceDefinition definition) {

}

@VisibleForTesting
public record DestinationAndDefinition(DestinationConnection destination, StandardDestinationDefinition definition) {

}

public List<SourceAndDefinition> getSourceAndDefinitionsFromSourceIds(final List<UUID> sourceIds) throws IOException {
final Result<Record> records = database.query(ctx -> ctx
.select(ACTOR.asterisk(), ACTOR_DEFINITION.asterisk())
.from(ACTOR)
.join(ACTOR_DEFINITION)
.on(ACTOR.ACTOR_DEFINITION_ID.eq(ACTOR_DEFINITION.ID))
.where(ACTOR.ACTOR_TYPE.eq(ActorType.source), ACTOR.ID.in(sourceIds))
.fetch());

final List<SourceAndDefinition> sourceAndDefinitions = new ArrayList<>();

for (final Record record : records) {
final SourceConnection source = DbConverter.buildSourceConnection(record);
final StandardSourceDefinition definition = DbConverter.buildStandardSourceDefinition(record);
sourceAndDefinitions.add(new SourceAndDefinition(source, definition));
}

return sourceAndDefinitions;
}

public List<DestinationAndDefinition> getDestinationAndDefinitionsFromDestinationIds(final List<UUID> destinationIds) throws IOException {
final Result<Record> records = database.query(ctx -> ctx
.select(ACTOR.asterisk(), ACTOR_DEFINITION.asterisk())
.from(ACTOR)
.join(ACTOR_DEFINITION)
.on(ACTOR.ACTOR_DEFINITION_ID.eq(ACTOR_DEFINITION.ID))
.where(ACTOR.ACTOR_TYPE.eq(ActorType.destination), ACTOR.ID.in(destinationIds))
.fetch());

final List<DestinationAndDefinition> destinationAndDefinitions = new ArrayList<>();

for (final Record record : records) {
final DestinationConnection destination = DbConverter.buildDestinationConnection(record);
final StandardDestinationDefinition definition = DbConverter.buildStandardDestinationDefinition(record);
destinationAndDefinitions.add(new DestinationAndDefinition(destination, definition));
}

return destinationAndDefinitions;
}

public ActorCatalog getActorCatalogById(final UUID actorCatalogId)
throws IOException, ConfigNotFoundException {
final Result<Record> result = database.query(ctx -> ctx.select(ACTOR_CATALOG.asterisk())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -456,7 +456,7 @@ private List<ConfigWithMetadata<SourceConnection>> listSourceConnectionWithMetad

final List<ConfigWithMetadata<SourceConnection>> sourceConnections = new ArrayList<>();
for (final Record record : result) {
final SourceConnection sourceConnection = buildSourceConnection(record);
final SourceConnection sourceConnection = DbConverter.buildSourceConnection(record);
sourceConnections.add(new ConfigWithMetadata<>(
record.get(ACTOR.ID).toString(),
ConfigSchema.SOURCE_CONNECTION.name(),
Expand All @@ -467,16 +467,6 @@ private List<ConfigWithMetadata<SourceConnection>> listSourceConnectionWithMetad
return sourceConnections;
}

private SourceConnection buildSourceConnection(final Record record) {
return new SourceConnection()
.withSourceId(record.get(ACTOR.ID))
.withConfiguration(Jsons.deserialize(record.get(ACTOR.CONFIGURATION).data()))
.withWorkspaceId(record.get(ACTOR.WORKSPACE_ID))
.withSourceDefinitionId(record.get(ACTOR.ACTOR_DEFINITION_ID))
.withTombstone(record.get(ACTOR.TOMBSTONE))
.withName(record.get(ACTOR.NAME));
}

private List<ConfigWithMetadata<DestinationConnection>> listDestinationConnectionWithMetadata() throws IOException {
return listDestinationConnectionWithMetadata(Optional.empty());
}
Expand All @@ -492,7 +482,7 @@ private List<ConfigWithMetadata<DestinationConnection>> listDestinationConnectio

final List<ConfigWithMetadata<DestinationConnection>> destinationConnections = new ArrayList<>();
for (final Record record : result) {
final DestinationConnection destinationConnection = buildDestinationConnection(record);
final DestinationConnection destinationConnection = DbConverter.buildDestinationConnection(record);
destinationConnections.add(new ConfigWithMetadata<>(
record.get(ACTOR.ID).toString(),
ConfigSchema.DESTINATION_CONNECTION.name(),
Expand All @@ -503,16 +493,6 @@ private List<ConfigWithMetadata<DestinationConnection>> listDestinationConnectio
return destinationConnections;
}

private DestinationConnection buildDestinationConnection(final Record record) {
return new DestinationConnection()
.withDestinationId(record.get(ACTOR.ID))
.withConfiguration(Jsons.deserialize(record.get(ACTOR.CONFIGURATION).data()))
.withWorkspaceId(record.get(ACTOR.WORKSPACE_ID))
.withDestinationDefinitionId(record.get(ACTOR.ACTOR_DEFINITION_ID))
.withTombstone(record.get(ACTOR.TOMBSTONE))
.withName(record.get(ACTOR.NAME));
}

private List<ConfigWithMetadata<SourceOAuthParameter>> listSourceOauthParamWithMetadata() throws IOException {
return listSourceOauthParamWithMetadata(Optional.empty());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

package io.airbyte.config.persistence;

import static io.airbyte.db.instance.configs.jooq.generated.Tables.ACTOR;
import static io.airbyte.db.instance.configs.jooq.generated.Tables.ACTOR_CATALOG;
import static io.airbyte.db.instance.configs.jooq.generated.Tables.ACTOR_DEFINITION;
import static io.airbyte.db.instance.configs.jooq.generated.Tables.ACTOR_OAUTH_PARAMETER;
Expand All @@ -15,12 +16,14 @@
import io.airbyte.commons.json.Jsons;
import io.airbyte.config.ActorCatalog;
import io.airbyte.config.ActorDefinitionResourceRequirements;
import io.airbyte.config.DestinationConnection;
import io.airbyte.config.DestinationOAuthParameter;
import io.airbyte.config.JobSyncConfig.NamespaceDefinitionType;
import io.airbyte.config.Notification;
import io.airbyte.config.ResourceRequirements;
import io.airbyte.config.Schedule;
import io.airbyte.config.ScheduleData;
import io.airbyte.config.SourceConnection;
import io.airbyte.config.SourceOAuthParameter;
import io.airbyte.config.StandardDestinationDefinition;
import io.airbyte.config.StandardSourceDefinition;
Expand All @@ -32,15 +35,18 @@
import io.airbyte.config.WorkspaceServiceAccount;
import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
import io.airbyte.protocol.models.ConnectorSpecification;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import org.jooq.Record;

/**
* Provides static methods for converting from repository layer results (often in the form of a jooq
* {@link Record}) to config models.
*/
public class DbConverter {

public static StandardSync buildStandardSync(final Record record, final List<UUID> connectionOperationId) throws IOException {
public static StandardSync buildStandardSync(final Record record, final List<UUID> connectionOperationId) {
return new StandardSync()
.withConnectionId(record.get(CONNECTION.ID))
.withNamespaceDefinition(
Expand Down Expand Up @@ -89,6 +95,26 @@ public static StandardWorkspace buildStandardWorkspace(final Record record) {
.withFeedbackDone(record.get(WORKSPACE.FEEDBACK_COMPLETE));
}

public static SourceConnection buildSourceConnection(final Record record) {
return new SourceConnection()
.withSourceId(record.get(ACTOR.ID))
.withConfiguration(Jsons.deserialize(record.get(ACTOR.CONFIGURATION).data()))
.withWorkspaceId(record.get(ACTOR.WORKSPACE_ID))
.withSourceDefinitionId(record.get(ACTOR.ACTOR_DEFINITION_ID))
.withTombstone(record.get(ACTOR.TOMBSTONE))
.withName(record.get(ACTOR.NAME));
}

public static DestinationConnection buildDestinationConnection(final Record record) {
return new DestinationConnection()
.withDestinationId(record.get(ACTOR.ID))
.withConfiguration(Jsons.deserialize(record.get(ACTOR.CONFIGURATION).data()))
.withWorkspaceId(record.get(ACTOR.WORKSPACE_ID))
.withDestinationDefinitionId(record.get(ACTOR.ACTOR_DEFINITION_ID))
.withTombstone(record.get(ACTOR.TOMBSTONE))
.withName(record.get(ACTOR.NAME));
}

public static StandardSourceDefinition buildStandardSourceDefinition(final Record record) {
return new StandardSourceDefinition()
.withSourceDefinitionId(record.get(ACTOR_DEFINITION.ID))
Expand Down
Loading

0 comments on commit 39a14b7

Please sign in to comment.