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
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.
132 changes: 112 additions & 20 deletions ...nfig/config-persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java
View file
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
24 changes: 2 additions & 22 deletions ...ig-persistence/src/main/java/io/airbyte/config/persistence/DatabaseConfigPersistence.java
View file
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
30 changes: 28 additions & 2 deletions ...te-config/config-persistence/src/main/java/io/airbyte/config/persistence/DbConverter.java
View file
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.