Skip to content

Commit

Permalink
handle tracking for sources / destinations not yet created (airbytehq…
Browse files Browse the repository at this point in the history 
…#1677)
  • Loading branch information
@cgardens
cgardens authored Jan 18, 2021
1 parent 0a38379 commit 196afe5
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 51 deletions.
11 changes: 11 additions & 0 deletions airbyte-scheduler/src/main/java/io/airbyte/scheduler/JobSubmitter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import org.apache.logging.log4j.util.Strings;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
Expand Down Expand Up @@ -137,6 +138,11 @@ private void assertSameIds(long expectedAttemptId, long actualAttemptId) {
@VisibleForTesting
void trackSubmission(Job job) {
try {
// if there is no scope, do not track. this is the case where we are running check for sources /
// destinations that don't exist.
if (Strings.isEmpty(job.getScope())) {
return;
}
final Builder<String, Object> metadataBuilder = generateMetadata(job);
metadataBuilder.put("attempt_stage", "STARTED");
track(metadataBuilder.build());
Expand All @@ -148,6 +154,11 @@ void trackSubmission(Job job) {
@VisibleForTesting
void trackCompletion(Job job, io.airbyte.workers.JobStatus status) {
try {
// if there is no scope, do not track. this is the case where we are running check for sources /
// destinations that don't exist.
if (Strings.isEmpty(job.getScope())) {
return;
}
final Builder<String, Object> metadataBuilder = generateMetadata(job);
metadataBuilder.put("attempt_stage", "ENDED");
metadataBuilder.put("attempt_completion_status", status);
Expand Down
9 changes: 6 additions & 3 deletions airbyte-scheduler/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,8 @@ public long createSourceCheckConnectionJob(SourceConnection source, String docke
.withConfigType(ConfigType.CHECK_CONNECTION_SOURCE)
.withCheckConnection(jobCheckConnectionConfig);

return jobPersistence.enqueueJob(source.getSourceId().toString(), jobConfig).orElseThrow();
final String sourceId = source.getSourceId() != null ? source.getSourceId().toString() : "";
return jobPersistence.enqueueJob(sourceId, jobConfig).orElseThrow();
}

@Override
Expand All @@ -71,7 +72,8 @@ public long createDestinationCheckConnectionJob(DestinationConnection destinatio
.withConfigType(ConfigType.CHECK_CONNECTION_DESTINATION)
.withCheckConnection(jobCheckConnectionConfig);

return jobPersistence.enqueueJob(destination.getDestinationId().toString(), jobConfig).orElseThrow();
final String destinationId = destination.getDestinationId() != null ? destination.getDestinationId().toString() : "";
return jobPersistence.enqueueJob(destinationId, jobConfig).orElseThrow();
}

@Override
Expand All @@ -84,7 +86,8 @@ public long createDiscoverSchemaJob(SourceConnection source, String dockerImageN
.withConfigType(ConfigType.DISCOVER_SCHEMA)
.withDiscoverCatalog(jobDiscoverCatalogConfig);

return jobPersistence.enqueueJob(source.getSourceId().toString(), jobConfig).orElseThrow();
final String sourceId = source.getSourceId() != null ? source.getSourceId().toString() : "";
return jobPersistence.enqueueJob(sourceId, jobConfig).orElseThrow();
}

@Override
Expand Down
35 changes: 7 additions & 28 deletions airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@

package io.airbyte.server.handlers;

import com.google.common.annotations.VisibleForTesting;
import io.airbyte.api.model.CheckConnectionRead;
import io.airbyte.api.model.ConnectionIdRequestBody;
import io.airbyte.api.model.DestinationCoreConfig;
Expand Down Expand Up @@ -61,23 +60,15 @@
import io.airbyte.validation.json.JsonValidationException;
import java.io.IOException;
import java.util.UUID;
import java.util.function.Supplier;

public class SchedulerHandler {

private final ConfigRepository configRepository;
private final SchedulerJobClient schedulerJobClient;
private final Supplier<UUID> uuidSupplier;

@VisibleForTesting
SchedulerHandler(final ConfigRepository configRepository, SchedulerJobClient schedulerJobClient, Supplier<UUID> uuidSupplier) {
public SchedulerHandler(final ConfigRepository configRepository, SchedulerJobClient schedulerJobClient) {
this.configRepository = configRepository;
this.schedulerJobClient = schedulerJobClient;
this.uuidSupplier = uuidSupplier;
}

public SchedulerHandler(final ConfigRepository configRepository, SchedulerJobClient schedulerJobClient) {
this(configRepository, schedulerJobClient, UUID::randomUUID);
}

public CheckConnectionRead checkSourceConnectionFromSourceId(SourceIdRequestBody sourceIdRequestBody)
Expand All @@ -93,14 +84,10 @@ public CheckConnectionRead checkSourceConnectionFromSourceCreate(SourceCoreConfi
throws ConfigNotFoundException, IOException, JsonValidationException {
final StandardSourceDefinition sourceDef = configRepository.getStandardSourceDefinition(sourceCreate.getSourceDefinitionId());
final String imageName = DockerUtils.getTaggedImageName(sourceDef.getDockerRepository(), sourceDef.getDockerImageTag());
final UUID sourceUuid = uuidSupplier.get();
// todo (cgardens) - narrow the struct passed to the client. we are not setting fields that are
// technically declared as required.
final SourceConnection source = new SourceConnection()
.withName("source:" + sourceUuid) // todo (cgardens) - narrow the struct passed to the client.
.withSourceDefinitionId(sourceCreate.getSourceDefinitionId())
.withWorkspaceId(sourceUuid) // todo (cgardens) - narrow the struct passed to the client.
.withTombstone(false)
// todo (cgardens) - used to create the scope so we want a random value.
.withSourceId(sourceUuid)
.withConfiguration(sourceCreate.getConnectionConfiguration());

return reportConnectionStatus(schedulerJobClient.createSourceCheckConnectionJob(source, imageName));
Expand All @@ -118,14 +105,10 @@ public CheckConnectionRead checkDestinationConnectionFromDestinationCreate(Desti
throws ConfigNotFoundException, IOException, JsonValidationException {
final StandardDestinationDefinition destDef = configRepository.getStandardDestinationDefinition(destinationCreate.getDestinationDefinitionId());
final String imageName = DockerUtils.getTaggedImageName(destDef.getDockerRepository(), destDef.getDockerImageTag());
final UUID destinationUuid = uuidSupplier.get();
// todo (cgardens) - narrow the struct passed to the client. we are not setting fields that are
// technically declared as required.
final DestinationConnection destination = new DestinationConnection()
.withName("destination:" + destinationUuid) // todo (cgardens) - narrow the struct passed to the client.
.withDestinationDefinitionId(destinationCreate.getDestinationDefinitionId())
.withWorkspaceId(destinationUuid) // todo (cgardens) - narrow the struct passed to the client.
.withTombstone(false)
// todo (cgardens) - used to create the scope so we want a random value.
.withDestinationId(destinationUuid)
.withConfiguration(destinationCreate.getConnectionConfiguration());
return reportConnectionStatus(schedulerJobClient.createDestinationCheckConnectionJob(destination, imageName));
}
Expand All @@ -143,14 +126,10 @@ public SourceDiscoverSchemaRead discoverSchemaForSourceFromSourceCreate(SourceCo
throws ConfigNotFoundException, IOException, JsonValidationException {
final StandardSourceDefinition sourceDef = configRepository.getStandardSourceDefinition(sourceCreate.getSourceDefinitionId());
final String imageName = DockerUtils.getTaggedImageName(sourceDef.getDockerRepository(), sourceDef.getDockerImageTag());
final UUID sourceUuid = uuidSupplier.get();
// todo (cgardens) - narrow the struct passed to the client. we are not setting fields that are
// technically declared as required.
final SourceConnection source = new SourceConnection()
.withName("source:" + sourceUuid) // todo (cgardens) - narrow the struct passed to the client.
.withSourceDefinitionId(sourceCreate.getSourceDefinitionId())
.withWorkspaceId(sourceUuid) // todo (cgardens) - narrow the struct passed to the client.
.withTombstone(false)
// todo (cgardens) - used to create the scope so we want a random value.
.withSourceId(sourceUuid)
.withConfiguration(sourceCreate.getConnectionConfiguration());
final Job job = schedulerJobClient.createDiscoverSchemaJob(source, imageName);
return discoverJobToOutput(job);
Expand Down
31 changes: 11 additions & 20 deletions airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,6 @@
import java.util.HashMap;
import java.util.Optional;
import java.util.UUID;
import java.util.function.Supplier;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

Expand Down Expand Up @@ -109,21 +108,18 @@ class SchedulerHandlerTest {
private ConfigRepository configRepository;
private Job completedJob;
private SchedulerJobClient schedulerJobClient;
private Supplier<UUID> uuidSupplier;

@SuppressWarnings("unchecked")
@BeforeEach
void setup() {
completedJob = mock(Job.class, RETURNS_DEEP_STUBS);
when(completedJob.getStatus()).thenReturn(JobStatus.SUCCEEDED);
when(completedJob.getConfig().getConfigType()).thenReturn(ConfigType.SYNC);
when(completedJob.getScope()).thenReturn("sync:123");

uuidSupplier = mock(Supplier.class);
schedulerJobClient = spy(SchedulerJobClient.class);
configRepository = mock(ConfigRepository.class);

schedulerHandler = new SchedulerHandler(configRepository, schedulerJobClient, uuidSupplier);
schedulerHandler = new SchedulerHandler(configRepository, schedulerJobClient);
}

@Test
Expand All @@ -147,11 +143,9 @@ void testCheckSourceConnectionFromSourceId() throws JsonValidationException, IOE

@Test
void testCheckSourceConnectionFromSourceCreate() throws JsonValidationException, IOException, ConfigNotFoundException {
final SourceConnection source = Jsons.clone(SOURCE);
source.setName("source:" + source.getSourceId());
source.setWorkspaceId(source.getSourceId());

when(uuidSupplier.get()).thenReturn(source.getSourceId());
final SourceConnection source = new SourceConnection()
.withSourceDefinitionId(SOURCE.getSourceDefinitionId())
.withConfiguration(SOURCE.getConfiguration());

final SourceCoreConfig sourceCoreConfig = new SourceCoreConfig()
.sourceDefinitionId(source.getSourceDefinitionId())
Expand Down Expand Up @@ -261,12 +255,10 @@ void testCheckDestinationConnectionFromDestinationId() throws IOException, JsonV
}

@Test
void testCheckSourceConnectionFromDestinationCreate() throws JsonValidationException, IOException, ConfigNotFoundException {
final DestinationConnection destination = Jsons.clone(DESTINATION);
destination.setName("destination:" + destination.getDestinationId());
destination.setWorkspaceId(destination.getDestinationId());

when(uuidSupplier.get()).thenReturn(destination.getDestinationId());
void testCheckDestinationConnectionFromDestinationCreate() throws JsonValidationException, IOException, ConfigNotFoundException {
final DestinationConnection destination = new DestinationConnection()
.withDestinationDefinitionId(DESTINATION.getDestinationDefinitionId())
.withConfiguration(DESTINATION.getConfiguration());

final DestinationCoreConfig destinationCoreConfig = new DestinationCoreConfig()
.destinationDefinitionId(destination.getDestinationDefinitionId())
Expand Down Expand Up @@ -309,11 +301,10 @@ void testDiscoverSchemaForSourceFromSourceId() throws IOException, JsonValidatio

@Test
void testDiscoverSchemaForSourceFromSourceCreate() throws JsonValidationException, IOException, ConfigNotFoundException {
final SourceConnection source = Jsons.clone(SOURCE);
source.setName("source:" + source.getSourceId());
source.setWorkspaceId(source.getSourceId());
final SourceConnection source = new SourceConnection()
.withSourceDefinitionId(SOURCE.getSourceDefinitionId())
.withConfiguration(SOURCE.getConfiguration());

when(uuidSupplier.get()).thenReturn(source.getSourceId());
final JobOutput jobOutput = mock(JobOutput.class);
when(completedJob.getSuccessOutput()).thenReturn(Optional.of(jobOutput));
when(jobOutput.getDiscoverCatalog()).thenReturn(new StandardDiscoverCatalogOutput().withCatalog(new AirbyteCatalog()));
Expand Down

0 comments on commit 196afe5

Please sign in to comment.