From 76c9fd6587026f1ded2483d7f27ab4941da3cdd2 Mon Sep 17 00:00:00 2001 From: Jonathan Pearlin Date: Wed, 14 Sep 2022 13:10:23 -0400 Subject: [PATCH] Convert airbyte-workers to Micronaut (#16434) * WIP Convert airbyte-workers to Micronaut framework * Rebase cleanup * Fix broken tests * Simplify code * Support control vs data plane configuration * make WORFKLOW_PROXY_CACHE non-static to avoid cacheing mocks in unit tests * Formatting * Pairing on Worker Micronaut (#16364) * add RouteToSyncTaskQueue activity * use new route activity in connection manager workflow * format * call router service for task queue * Revert temporal proxy changes * Formatting * Fix default value * register new route activity in test * fix SyncWorkflowTest now that it isn't doing any routing * Update dependencies * More dependency updates * Update dependencies * Improve conditional bean check * Match existing Optional functionality * Add notEquals check * Add missing env var to Helm chart * Fix typo * Mark LogConfigs as Singleton * Env vars for log/state storage type * Remove use of Optional in bean declarations * Fix typo in config property name * Support Temporal Cloud namespace * Change to @Value * Use correct value for conditional check * Upgrade Micronaut * Fix merge conflict * Formatting * Add missing env var * Use sync task queue environment variable * Handle sync task queue as set * format and force http * Handle case where sync task queue is empty * Add correct path to config property * Remove unused import * Remove unused parameter * Formatting * Use pattern for condition process factory beans * Cleanup * PR feedback * Revert hack for testing Co-authored-by: pmossman --- .env | 4 +- .env.dev | 6 + airbyte-config/config-models/build.gradle | 3 + .../java/io/airbyte/config/EnvConfigs.java | 2 +- .../io/airbyte/config/helpers/LogConfigs.java | 10 +- .../config/helpers/CloudLogsClientTest.java | 13 +- .../io/airbyte/config/helpers/S3LogsTest.java | 5 +- .../ContainerOrchestratorApp.java | 16 +- .../WorkerHeartbeatServer.java | 2 +- .../db/jdbc/TestStreamingJdbcDatabase.java | 2 +- airbyte-server/build.gradle | 13 +- .../java/io/airbyte/server/ServerApp.java | 22 +- .../server/handlers/ArchiveHandlerTest.java | 5 +- .../utils/AirbyteAcceptanceTestHarness.java | 6 +- airbyte-workers/build.gradle | 58 +- .../java/io/airbyte/workers/Application.java | 15 + .../workers/ApplicationInitializer.java | 357 +++++++++++ .../workers/ContainerOrchestratorConfig.java | 18 + .../workers/WorkerApiClientFactory.java | 13 - .../workers/WorkerApiClientFactoryImpl.java | 101 --- .../java/io/airbyte/workers/WorkerApp.java | 597 ------------------ .../workers/config/ActivityBeanFactory.java | 198 ++++++ .../workers/config/ApiClientBeanFactory.java | 121 ++++ .../config/ApplicationBeanFactory.java | 175 +++++ .../config/CloudStorageBeanFactory.java | 92 +++ ...ontainerOrchestratorConfigBeanFactory.java | 61 ++ .../workers/config/DatabaseBeanFactory.java | 166 +++++ .../config/JobErrorReportingBeanFactory.java | 79 +++ .../config/ProcessFactoryBeanFactory.java | 241 +++++++ .../config/SecretPersistenceBeanFactory.java | 88 +++ .../workers/config/TemporalBeanFactory.java | 92 +++ .../WorkerConfigurationBeanFactory.java | 392 ++++++++++++ .../controller/HeartbeatController.java | 51 ++ .../workers/helper/ConnectionHelper.java | 12 +- .../process/AsyncOrchestratorPodProcess.java | 8 +- .../workers/temporal/CancellationHandler.java | 2 +- .../temporal/ConnectionManagerUtils.java | 59 +- .../temporal/StreamResetRecordsHelper.java | 68 ++ .../workers/temporal/TemporalClient.java | 63 +- .../workers/temporal/TemporalUtils.java | 172 ++--- .../temporal/TemporalWorkflowUtils.java | 91 +++ .../annotations/TemporalActivityStub.java | 30 + .../CheckConnectionActivityImpl.java | 53 +- .../CheckConnectionWorkflowImpl.java | 11 +- .../catalog/DiscoverCatalogActivityImpl.java | 57 +- .../catalog/DiscoverCatalogWorkflowImpl.java | 14 +- .../ConnectionManagerWorkflowImpl.java | 187 +++--- .../AutoDisableConnectionActivityImpl.java | 56 +- .../activities/ConfigFetchActivityImpl.java | 47 +- .../ConnectionDeletionActivityImpl.java | 11 +- .../activities/GenerateInputActivityImpl.java | 13 +- ...obCreationAndStatusUpdateActivityImpl.java | 41 +- .../activities/RecordMetricActivity.java | 51 ++ .../activities/RecordMetricActivityImpl.java | 65 ++ .../RouteToSyncTaskQueueActivity.java | 38 ++ .../RouteToSyncTaskQueueActivityImpl.java | 24 + .../activities/StreamResetActivityImpl.java | 45 +- .../activities/WorkflowConfigActivity.java | 26 + .../WorkflowConfigActivityImpl.java | 38 ++ .../shared/ActivityConfiguration.java | 58 -- .../temporal/spec/SpecActivityImpl.java | 48 +- .../temporal/spec/SpecWorkflowImpl.java | 14 +- ...TemporalActivityStubGeneratorFunction.java | 21 + .../TemporalActivityStubInterceptor.java | 171 +++++ .../temporal/support/TemporalProxyHelper.java | 136 ++++ .../temporal/sync/DbtLauncherWorker.java | 13 +- .../sync/DbtTransformationActivityImpl.java | 92 +-- .../workers/temporal/sync/LauncherWorker.java | 22 +- .../temporal/sync/NormalizationActivity.java | 5 + .../sync/NormalizationActivityImpl.java | 104 +-- .../sync/NormalizationLauncherWorker.java | 13 +- .../sync/PersistStateActivityImpl.java | 11 +- .../sync/ReplicationActivityImpl.java | 101 ++- .../sync/ReplicationLauncherWorker.java | 13 +- .../sync/RouteToTaskQueueActivityImpl.java | 22 - .../workers/temporal/sync/RouterService.java | 32 +- .../temporal/sync/SyncWorkflowImpl.java | 72 +-- .../main/resources/application-control.yml | 40 ++ .../src/main/resources/application.yml | 219 +++++++ .../src/main/resources/micronaut-banner.txt | 8 + ...OrchestratorPodProcessIntegrationTest.java | 24 +- .../KubePodProcessIntegrationTest.java | 56 +- .../resources/application-test.yml | 11 + .../temporal/CancellationHandlerTest.java | 3 +- .../StreamResetRecordsHelperTest.java | 72 +++ .../workers/temporal/TemporalClientTest.java | 41 +- .../workers/temporal/TemporalUtilsTest.java | 23 +- .../CheckConnectionWorkflowTest.java | 32 +- .../ConnectionManagerWorkflowTest.java | 188 +++--- .../scheduling/WorkflowReplayingTest.java | 36 +- .../AutoDisableConnectionActivityTest.java | 29 +- .../activities/ConfigFetchActivityTest.java | 51 +- .../RecordMetricActivityImplTest.java | 86 +++ .../activities/StreamResetActivityTest.java | 49 +- .../WorkflowConfigActivityImplTest.java | 23 + .../temporal/stubs/ErrorTestWorkflowImpl.java | 24 + .../{ => stubs}/HeartbeatWorkflow.java | 5 +- .../stubs/InvalidTestWorkflowImpl.java | 27 + .../workers/temporal/stubs/TestActivity.java} | 7 +- .../workers/temporal/stubs/TestWorkflow.java | 24 + .../temporal/stubs/ValidTestWorkflowImpl.java | 41 ++ .../TemporalActivityStubInterceptorTest.java | 142 +++++ .../support/TemporalProxyHelperTest.java | 57 ++ .../sync/PersistStateActivityTest.java | 2 +- .../temporal/sync/RouterServiceTest.java | 60 ++ .../temporal/{ => sync}/SyncWorkflowTest.java | 99 +-- .../airbyte-worker/templates/deployment.yaml | 16 + charts/airbyte/templates/env-configmap.yaml | 3 + charts/airbyte/values.yaml | 9 + deps.toml | 25 +- docker-compose.yaml | 3 + 111 files changed, 4995 insertions(+), 1794 deletions(-) rename {airbyte-workers/src/main/java/io/airbyte/workers/process => airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator}/WorkerHeartbeatServer.java (98%) create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/Application.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/ApplicationInitializer.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/ContainerOrchestratorConfig.java delete mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactory.java delete mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactoryImpl.java delete mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/ActivityBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/ApiClientBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/ApplicationBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/CloudStorageBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/ContainerOrchestratorConfigBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/DatabaseBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/JobErrorReportingBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/ProcessFactoryBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/SecretPersistenceBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/TemporalBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/controller/HeartbeatController.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/StreamResetRecordsHelper.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalWorkflowUtils.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/annotations/TemporalActivityStub.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivity.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImpl.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivity.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivityImpl.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivity.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImpl.java delete mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/shared/ActivityConfiguration.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubGeneratorFunction.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptor.java create mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalProxyHelper.java delete mode 100644 airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivityImpl.java create mode 100644 airbyte-workers/src/main/resources/application-control.yml create mode 100644 airbyte-workers/src/main/resources/application.yml create mode 100644 airbyte-workers/src/main/resources/micronaut-banner.txt create mode 100644 airbyte-workers/src/test-integration/resources/application-test.yml create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/StreamResetRecordsHelperTest.java create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImplTest.java create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImplTest.java create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ErrorTestWorkflowImpl.java rename airbyte-workers/src/test/java/io/airbyte/workers/temporal/{ => stubs}/HeartbeatWorkflow.java (90%) create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/InvalidTestWorkflowImpl.java rename airbyte-workers/src/{main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivity.java => test/java/io/airbyte/workers/temporal/stubs/TestActivity.java} (55%) create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/TestWorkflow.java create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ValidTestWorkflowImpl.java create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptorTest.java create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalProxyHelperTest.java create mode 100644 airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/RouterServiceTest.java rename airbyte-workers/src/test/java/io/airbyte/workers/temporal/{ => sync}/SyncWorkflowTest.java (77%) diff --git a/.env b/.env index 700919c424b39..a569d10e163d6 100644 --- a/.env +++ b/.env @@ -81,6 +81,7 @@ LOG_LEVEL=INFO ### APPLICATIONS ### # Worker # +WORKERS_MICRONAUT_ENVIRONMENTS=control # Relevant to scaling. MAX_SYNC_WORKERS=5 MAX_SPEC_WORKERS=5 @@ -92,7 +93,6 @@ ACTIVITY_INITIAL_DELAY_BETWEEN_ATTEMPTS_SECONDS= ACTIVITY_MAX_DELAY_BETWEEN_ATTEMPTS_SECONDS= WORKFLOW_FAILURE_RESTART_DELAY_SECONDS= - ### FEATURE FLAGS ### AUTO_DISABLE_FAILING_CONNECTIONS=false EXPOSE_SECRETS_IN_EXPORT=false @@ -104,4 +104,4 @@ METRIC_CLIENT= # Useful only when metric client is set to be otel. Must start with http:// or https://. OTEL_COLLECTOR_ENDPOINT="http://host.docker.internal:4317" -USE_STREAM_CAPABLE_STATE=true +USE_STREAM_CAPABLE_STATE=true \ No newline at end of file diff --git a/.env.dev b/.env.dev index 2a6dc7eb8129b..b1780d7ee8c12 100644 --- a/.env.dev +++ b/.env.dev @@ -4,6 +4,7 @@ VERSION=dev DATABASE_USER=docker DATABASE_PASSWORD=docker DATABASE_DB=airbyte +DATABASE_URL=jdbc:postgresql://db:5432/airbyte CONFIG_ROOT=/data WORKSPACE_ROOT=/tmp/workspace DATA_DOCKER_MOUNT=airbyte_data_dev @@ -24,6 +25,11 @@ API_URL=/api/v1/ INTERNAL_API_HOST=airbyte-server:8001 SYNC_JOB_MAX_ATTEMPTS=3 SYNC_JOB_MAX_TIMEOUT_DAYS=3 +WORKERS_MICRONAUT_ENVIRONMENTS=control # Sentry SENTRY_DSN="" + +# Migration Configuration +CONFIGS_DATABASE_MINIMUM_FLYWAY_MIGRATION_VERSION=0.35.15.001 +JOBS_DATABASE_MINIMUM_FLYWAY_MIGRATION_VERSION=0.29.15.001 diff --git a/airbyte-config/config-models/build.gradle b/airbyte-config/config-models/build.gradle index fb1840fe9f488..b03a9e39d5b59 100644 --- a/airbyte-config/config-models/build.gradle +++ b/airbyte-config/config-models/build.gradle @@ -5,6 +5,9 @@ plugins { } dependencies { + annotationProcessor libs.bundles.micronaut.annotation.processor + implementation libs.bundles.micronaut.annotation + implementation project(':airbyte-json-validation') implementation project(':airbyte-protocol:protocol-models') implementation project(':airbyte-commons') diff --git a/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java b/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java index c92920f641e3f..10e4e751e0ab3 100644 --- a/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java +++ b/airbyte-config/config-models/src/main/java/io/airbyte/config/EnvConfigs.java @@ -227,7 +227,7 @@ public EnvConfigs() { public EnvConfigs(final Map envMap) { this.getEnv = envMap::get; this.getAllEnvKeys = envMap::keySet; - this.logConfigs = new LogConfigs(getLogConfiguration().orElse(null)); + this.logConfigs = new LogConfigs(getLogConfiguration()); this.stateStorageCloudConfigs = getStateStorageConfiguration().orElse(null); validateSyncWorkflowConfigs(); diff --git a/airbyte-config/config-models/src/main/java/io/airbyte/config/helpers/LogConfigs.java b/airbyte-config/config-models/src/main/java/io/airbyte/config/helpers/LogConfigs.java index b2fe82c5dfd1d..7495001227f3a 100644 --- a/airbyte-config/config-models/src/main/java/io/airbyte/config/helpers/LogConfigs.java +++ b/airbyte-config/config-models/src/main/java/io/airbyte/config/helpers/LogConfigs.java @@ -5,19 +5,23 @@ package io.airbyte.config.helpers; import io.airbyte.config.storage.CloudStorageConfigs; +import java.util.Optional; +import javax.inject.Named; +import javax.inject.Singleton; /** * Describes logging configuration. For now it just contains configuration around storage medium, * but in the future will have other configuration options (e.g. json logging, etc). */ +@Singleton public class LogConfigs { - public final static LogConfigs EMPTY = new LogConfigs(null); + public final static LogConfigs EMPTY = new LogConfigs(Optional.empty()); private final CloudStorageConfigs storageConfigs; - public LogConfigs(final CloudStorageConfigs storageConfigs) { - this.storageConfigs = storageConfigs; + public LogConfigs(@Named("logStorageConfigs") final Optional storageConfigs) { + this.storageConfigs = storageConfigs.orElse(null); } public CloudStorageConfigs getStorageConfigs() { diff --git a/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/CloudLogsClientTest.java b/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/CloudLogsClientTest.java index 2216c02465b6e..f4693ab7cd9a8 100644 --- a/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/CloudLogsClientTest.java +++ b/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/CloudLogsClientTest.java @@ -10,6 +10,7 @@ import io.airbyte.config.storage.CloudStorageConfigs.GcsConfig; import io.airbyte.config.storage.CloudStorageConfigs.MinioConfig; import io.airbyte.config.storage.CloudStorageConfigs.S3Config; +import java.util.Optional; import org.junit.jupiter.api.Test; @SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert") @@ -17,31 +18,31 @@ class CloudLogsClientTest { @Test void createCloudLogClientTestMinio() { - final var configs = new LogConfigs(CloudStorageConfigs.minio(new MinioConfig( + final var configs = new LogConfigs(Optional.of(CloudStorageConfigs.minio(new MinioConfig( "test-bucket", "access-key", "access-key-secret", - "minio-endpoint"))); + "minio-endpoint")))); assertEquals(S3Logs.class, CloudLogs.createCloudLogClient(configs).getClass()); } @Test void createCloudLogClientTestAws() { - final var configs = new LogConfigs(CloudStorageConfigs.s3(new S3Config( + final var configs = new LogConfigs(Optional.of(CloudStorageConfigs.s3(new S3Config( "test-bucket", "access-key", "access-key-secret", - "us-east-1"))); + "us-east-1")))); assertEquals(S3Logs.class, CloudLogs.createCloudLogClient(configs).getClass()); } @Test void createCloudLogClientTestGcs() { - final var configs = new LogConfigs(CloudStorageConfigs.gcs(new GcsConfig( + final var configs = new LogConfigs(Optional.of(CloudStorageConfigs.gcs(new GcsConfig( "storage-bucket", - "path/to/google/secret"))); + "path/to/google/secret")))); assertEquals(GcsLogs.class, CloudLogs.createCloudLogClient(configs).getClass()); } diff --git a/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/S3LogsTest.java b/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/S3LogsTest.java index b8c2412eaec54..6cff4287f225d 100644 --- a/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/S3LogsTest.java +++ b/airbyte-config/config-models/src/test/java/io/airbyte/config/helpers/S3LogsTest.java @@ -12,6 +12,7 @@ import java.nio.file.Files; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; @@ -27,11 +28,11 @@ class S3LogsTest { private static final Region REGION = Region.of(REGION_STRING); private static final String BUCKET_NAME = "airbyte-kube-integration-logging-test"; - private static final LogConfigs LOG_CONFIGS = new LogConfigs(CloudStorageConfigs.s3(new CloudStorageConfigs.S3Config( + private static final LogConfigs LOG_CONFIGS = new LogConfigs(Optional.of(CloudStorageConfigs.s3(new CloudStorageConfigs.S3Config( System.getenv(LogClientSingleton.S3_LOG_BUCKET), System.getenv(LogClientSingleton.AWS_ACCESS_KEY_ID), System.getenv(LogClientSingleton.AWS_SECRET_ACCESS_KEY), - System.getenv(LogClientSingleton.S3_LOG_BUCKET_REGION)))); + System.getenv(LogClientSingleton.S3_LOG_BUCKET_REGION))))); private S3Client s3Client; diff --git a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ContainerOrchestratorApp.java b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ContainerOrchestratorApp.java index e38e72b0929d9..7de2427651363 100644 --- a/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ContainerOrchestratorApp.java +++ b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/ContainerOrchestratorApp.java @@ -12,7 +12,6 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.helpers.LogClientSingleton; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.workers.WorkerApp; import io.airbyte.workers.WorkerConfigs; import io.airbyte.workers.WorkerUtils; import io.airbyte.workers.process.AsyncKubePodStatus; @@ -23,7 +22,6 @@ import io.airbyte.workers.process.KubePortManagerSingleton; import io.airbyte.workers.process.KubeProcessFactory; import io.airbyte.workers.process.ProcessFactory; -import io.airbyte.workers.process.WorkerHeartbeatServer; import io.airbyte.workers.storage.StateClients; import io.airbyte.workers.temporal.sync.DbtLauncherWorker; import io.airbyte.workers.temporal.sync.NormalizationLauncherWorker; @@ -58,6 +56,13 @@ public class ContainerOrchestratorApp { public static final int MAX_SECONDS_TO_WAIT_FOR_FILE_COPY = 60; + // TODO Move the following to configuration once converted to a Micronaut service + + // IMPORTANT: Changing the storage location will orphan already existing kube pods when the new + // version is deployed! + private static final Path STATE_STORAGE_PREFIX = Path.of("/state"); + private static final Integer KUBE_HEARTBEAT_PORT = 9000; + private final String application; private final Map envMap; private final JobRunConfig jobRunConfig; @@ -114,7 +119,7 @@ private void runInternal(final DefaultAsyncStateManager asyncStateManager) { throw new IllegalStateException("Could not find job orchestrator for application: " + application); } - final var heartbeatServer = new WorkerHeartbeatServer(WorkerApp.KUBE_HEARTBEAT_PORT); + final var heartbeatServer = new WorkerHeartbeatServer(KUBE_HEARTBEAT_PORT); heartbeatServer.startBackground(); asyncStateManager.write(kubePodInfo, AsyncKubePodStatus.RUNNING); @@ -146,7 +151,7 @@ public void run() { // IMPORTANT: Changing the storage location will orphan already existing kube pods when the new // version is deployed! - final var documentStoreClient = StateClients.create(configs.getStateStorageCloudConfigs(), WorkerApp.STATE_STORAGE_PREFIX); + final var documentStoreClient = StateClients.create(configs.getStateStorageCloudConfigs(), STATE_STORAGE_PREFIX); final var asyncStateManager = new DefaultAsyncStateManager(documentStoreClient); runInternal(asyncStateManager); @@ -212,7 +217,8 @@ private static ProcessFactory getProcessBuilderFactory(final Configs configs, fi if (configs.getWorkerEnvironment() == Configs.WorkerEnvironment.KUBERNETES) { final KubernetesClient fabricClient = new DefaultKubernetesClient(); final String localIp = InetAddress.getLocalHost().getHostAddress(); - final String kubeHeartbeatUrl = localIp + ":" + WorkerApp.KUBE_HEARTBEAT_PORT; + // TODO move port to configuration + final String kubeHeartbeatUrl = localIp + ":" + KUBE_HEARTBEAT_PORT; log.info("Using Kubernetes namespace: {}", configs.getJobKubeNamespace()); // this needs to have two ports for the source and two ports for the destination (all four must be diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/process/WorkerHeartbeatServer.java b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/WorkerHeartbeatServer.java similarity index 98% rename from airbyte-workers/src/main/java/io/airbyte/workers/process/WorkerHeartbeatServer.java rename to airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/WorkerHeartbeatServer.java index ba1abaf8b380f..1684abb963cea 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/process/WorkerHeartbeatServer.java +++ b/airbyte-container-orchestrator/src/main/java/io/airbyte/container_orchestrator/WorkerHeartbeatServer.java @@ -2,7 +2,7 @@ * Copyright (c) 2022 Airbyte, Inc., all rights reserved. */ -package io.airbyte.workers.process; +package io.airbyte.container_orchestrator; import com.google.common.collect.ImmutableMap; import com.google.common.net.HttpHeaders; diff --git a/airbyte-db/db-lib/src/test/java/io/airbyte/db/jdbc/TestStreamingJdbcDatabase.java b/airbyte-db/db-lib/src/test/java/io/airbyte/db/jdbc/TestStreamingJdbcDatabase.java index a2a36c63dc0de..bc5918cd7086d 100644 --- a/airbyte-db/db-lib/src/test/java/io/airbyte/db/jdbc/TestStreamingJdbcDatabase.java +++ b/airbyte-db/db-lib/src/test/java/io/airbyte/db/jdbc/TestStreamingJdbcDatabase.java @@ -24,10 +24,10 @@ import java.sql.SQLException; import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Set; import java.util.concurrent.atomic.AtomicReference; import javax.sql.DataSource; -import org.elasticsearch.common.collect.Map; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; diff --git a/airbyte-server/build.gradle b/airbyte-server/build.gradle index b6b18a01c557f..39653de5ddf0f 100644 --- a/airbyte-server/build.gradle +++ b/airbyte-server/build.gradle @@ -15,10 +15,19 @@ dependencies { implementation project(':airbyte-notification') implementation project(':airbyte-oauth') implementation project(':airbyte-protocol:protocol-models') - implementation project(':airbyte-scheduler:client') + implementation(project(':airbyte-scheduler:client')) { + exclude module: 'airbyte-workers' + } implementation project(':airbyte-scheduler:scheduler-models') implementation project(':airbyte-scheduler:scheduler-persistence') - implementation project(':airbyte-workers') + implementation(project(':airbyte-workers')) { + // Temporary hack to avoid dependency conflicts + exclude group: 'io.micronaut' + exclude group: 'io.micronaut.flyway' + exclude group: 'io.micronaut.jaxrs' + exclude group: 'io.micronaut.security' + exclude group: 'io.micronaut.sql' + } implementation libs.flyway.core implementation 'com.github.slugify:slugify:2.4' diff --git a/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java b/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java index dde9592ab6d56..39009fd39eb63 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java @@ -56,8 +56,11 @@ import io.airbyte.server.handlers.DbMigrationHandler; import io.airbyte.validation.json.JsonValidationException; import io.airbyte.workers.normalization.NormalizationRunnerFactory; +import io.airbyte.workers.temporal.ConnectionManagerUtils; +import io.airbyte.workers.temporal.StreamResetRecordsHelper; import io.airbyte.workers.temporal.TemporalClient; import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.TemporalWorkflowUtils; import io.temporal.serviceclient.WorkflowServiceStubs; import java.io.IOException; import java.net.http.HttpClient; @@ -224,14 +227,27 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, webUrlHelper, jobErrorReportingClient); + final TemporalUtils temporalUtils = new TemporalUtils( + configs.getTemporalCloudClientCert(), + configs.getTemporalCloudClientKey(), + configs.temporalCloudEnabled(), + configs.getTemporalCloudHost(), + configs.getTemporalCloudNamespace(), + configs.getTemporalHost(), + configs.getTemporalRetentionInDays()); + final StreamResetPersistence streamResetPersistence = new StreamResetPersistence(configsDatabase); - final WorkflowServiceStubs temporalService = TemporalUtils.createTemporalService(); + final WorkflowServiceStubs temporalService = temporalUtils.createTemporalService(); + final ConnectionManagerUtils connectionManagerUtils = new ConnectionManagerUtils(); + final StreamResetRecordsHelper streamResetRecordsHelper = new StreamResetRecordsHelper(jobPersistence, streamResetPersistence); final TemporalClient temporalClient = new TemporalClient( - TemporalUtils.createWorkflowClient(temporalService, TemporalUtils.getNamespace()), configs.getWorkspaceRoot(), + TemporalWorkflowUtils.createWorkflowClient(temporalService, temporalUtils.getNamespace()), temporalService, - streamResetPersistence); + streamResetPersistence, + connectionManagerUtils, + streamResetRecordsHelper); final OAuthConfigSupplier oAuthConfigSupplier = new OAuthConfigSupplier(configRepository, trackingClient); final DefaultSynchronousSchedulerClient syncSchedulerClient = diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/ArchiveHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/ArchiveHandlerTest.java index e4069c23ac69c..a6909856f7bbd 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/ArchiveHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/ArchiveHandlerTest.java @@ -98,6 +98,7 @@ class ArchiveHandlerTest { private JsonSecretsProcessor jsonSecretsProcessor; private ConfigRepository configRepository; private ArchiveHandler archiveHandler; + private WorkspaceHelper workspaceHelper; private static class NoOpFileTtlManager extends FileTtlManager { @@ -146,6 +147,8 @@ void setup() throws Exception { jobPersistence.setVersion(VERSION.serialize()); + workspaceHelper = new WorkspaceHelper(configRepository, jobPersistence); + archiveHandler = new ArchiveHandler( VERSION, configRepository, @@ -153,7 +156,7 @@ void setup() throws Exception { secretsRepositoryWriter, jobPersistence, seedPersistence, - new WorkspaceHelper(configRepository, jobPersistence), + workspaceHelper, new NoOpFileTtlManager(), true); } diff --git a/airbyte-test-utils/src/main/java/io/airbyte/test/utils/AirbyteAcceptanceTestHarness.java b/airbyte-test-utils/src/main/java/io/airbyte/test/utils/AirbyteAcceptanceTestHarness.java index 46e5106b7bf98..344401bc8a8a0 100644 --- a/airbyte-test-utils/src/main/java/io/airbyte/test/utils/AirbyteAcceptanceTestHarness.java +++ b/airbyte-test-utils/src/main/java/io/airbyte/test/utils/AirbyteAcceptanceTestHarness.java @@ -63,6 +63,7 @@ import io.airbyte.db.jdbc.JdbcUtils; import io.airbyte.test.airbyte_test_container.AirbyteTestContainer; import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.TemporalWorkflowUtils; import io.airbyte.workers.temporal.scheduling.ConnectionManagerWorkflow; import io.airbyte.workers.temporal.scheduling.state.WorkflowState; import io.fabric8.kubernetes.client.DefaultKubernetesClient; @@ -315,8 +316,9 @@ private void assignEnvVars() { } private WorkflowClient getWorkflowClient() { - final WorkflowServiceStubs temporalService = TemporalUtils.createTemporalService( - TemporalUtils.getAirbyteTemporalOptions("localhost:7233"), + final TemporalUtils temporalUtils = new TemporalUtils(null, null, null, null, null, null, null); + final WorkflowServiceStubs temporalService = temporalUtils.createTemporalService( + TemporalWorkflowUtils.getAirbyteTemporalOptions("localhost:7233"), TemporalUtils.DEFAULT_NAMESPACE); return WorkflowClient.newInstance(temporalService); } diff --git a/airbyte-workers/build.gradle b/airbyte-workers/build.gradle index 100e393bae91b..1c55283800dbd 100644 --- a/airbyte-workers/build.gradle +++ b/airbyte-workers/build.gradle @@ -11,6 +11,21 @@ configurations { } dependencies { + annotationProcessor platform(libs.micronaut.bom) + annotationProcessor libs.bundles.micronaut.annotation.processor + + implementation platform(libs.micronaut.bom) + implementation libs.bundles.micronaut + + // Ensure that the versions defined in deps.toml are used + // instead of versions from transitive dependencies + implementation (libs.flyway.core) { + force = true + } + implementation (libs.jooq) { + force = true + } + implementation 'com.google.auth:google-auth-library-oauth2-http:1.4.0' implementation 'com.auth0:java-jwt:3.19.2' implementation 'io.fabric8:kubernetes-client:5.12.2' @@ -18,11 +33,11 @@ dependencies { implementation 'org.apache.ant:ant:1.10.10' implementation 'org.apache.commons:commons-lang3:3.11' implementation 'org.apache.commons:commons-text:1.9' - implementation 'org.eclipse.jetty:jetty-server:9.4.31.v20200723' - implementation 'org.eclipse.jetty:jetty-servlet:9.4.31.v20200723' implementation 'org.quartz-scheduler:quartz:2.3.2' - implementation libs.flyway.core implementation libs.micrometer.statsd + implementation 'io.sentry:sentry:6.3.1' + implementation 'net.bytebuddy:byte-buddy:1.12.14' + implementation 'org.springframework:spring-core:5.3.22' implementation project(':airbyte-analytics') implementation project(':airbyte-api') @@ -34,14 +49,27 @@ dependencies { implementation project(':airbyte-metrics:metrics-lib') implementation project(':airbyte-json-validation') implementation project(':airbyte-protocol:protocol-models') - implementation project(':airbyte-scheduler:scheduler-persistence') + implementation (project(':airbyte-scheduler:scheduler-persistence')) { + // Temporary hack to avoid dependency conflicts + exclude group: 'io.micronaut' + exclude group: 'io.micronaut.flyway' + exclude group: 'io.micronaut.jaxrs' + exclude group: 'io.micronaut.security' + exclude group: 'io.micronaut.sql' + } implementation project(':airbyte-scheduler:scheduler-models') implementation project(':airbyte-api') + testAnnotationProcessor platform(libs.micronaut.bom) + testAnnotationProcessor libs.bundles.micronaut.test.annotation.processor + + integrationTestJavaAnnotationProcessor platform(libs.micronaut.bom) + integrationTestJavaAnnotationProcessor libs.bundles.micronaut.test.annotation.processor + + testImplementation libs.bundles.micronaut.test testImplementation 'io.temporal:temporal-testing:1.8.1' testImplementation 'com.jayway.jsonpath:json-path:2.7.0' - testImplementation libs.flyway.core - testImplementation 'org.mockito:mockito-inline:4.0.0' + testImplementation 'org.mockito:mockito-inline:4.7.0' testImplementation libs.postgresql testImplementation libs.platform.testcontainers testImplementation libs.platform.testcontainers.postgresql @@ -50,6 +78,7 @@ dependencies { testImplementation project(':airbyte-test-utils') integrationTestJavaImplementation project(':airbyte-workers') + integrationTestJavaImplementation libs.bundles.micronaut.test } jsonSchema2Pojo { @@ -66,13 +95,28 @@ jsonSchema2Pojo { includeSetters = true } -mainClassName = 'io.airbyte.workers.WorkerApp' +mainClassName = 'io.airbyte.workers.Application' application { + applicationName = project.name mainClass = mainClassName applicationDefaultJvmArgs = ['-XX:+ExitOnOutOfMemoryError', '-XX:MaxRAMPercentage=75.0'] } +Properties env = new Properties() +rootProject.file('.env.dev').withInputStream { env.load(it) } + +run { + // default for running on local machine. + env.each { entry -> + environment entry.getKey(), entry.getValue() + } + + environment 'AIRBYTE_ROLE', System.getenv('AIRBYTE_ROLE') + environment 'AIRBYTE_VERSION', env.VERSION + environment 'MICRONAUT_ENVIRONMENTS', 'local,control' +} + task cloudStorageIntegrationTest(type: Test) { useJUnitPlatform { includeTags cloudStorageTestTagName diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/Application.java b/airbyte-workers/src/main/java/io/airbyte/workers/Application.java new file mode 100644 index 0000000000000..025362fa7799d --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/Application.java @@ -0,0 +1,15 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers; + +import io.micronaut.runtime.Micronaut; + +public class Application { + + public static void main(final String[] args) { + Micronaut.run(Application.class, args); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/ApplicationInitializer.java b/airbyte-workers/src/main/java/io/airbyte/workers/ApplicationInitializer.java new file mode 100644 index 0000000000000..7d5325d13a5e8 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/ApplicationInitializer.java @@ -0,0 +1,357 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers; + +import io.airbyte.analytics.Deployment; +import io.airbyte.analytics.TrackingClientSingleton; +import io.airbyte.commons.version.AirbyteVersion; +import io.airbyte.config.Configs.DeploymentMode; +import io.airbyte.config.Configs.TrackingStrategy; +import io.airbyte.config.Configs.WorkerEnvironment; +import io.airbyte.config.Configs.WorkerPlane; +import io.airbyte.config.MaxWorkersConfig; +import io.airbyte.config.helpers.LogClientSingleton; +import io.airbyte.config.helpers.LogConfigs; +import io.airbyte.config.persistence.ConfigRepository; +import io.airbyte.db.check.DatabaseCheckException; +import io.airbyte.db.check.DatabaseMigrationCheck; +import io.airbyte.db.check.impl.JobsDatabaseAvailabilityCheck; +import io.airbyte.metrics.lib.MetricClientFactory; +import io.airbyte.metrics.lib.MetricEmittingApps; +import io.airbyte.scheduler.persistence.JobPersistence; +import io.airbyte.workers.process.KubePortManagerSingleton; +import io.airbyte.workers.temporal.TemporalJobType; +import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.check.connection.CheckConnectionWorkflowImpl; +import io.airbyte.workers.temporal.discover.catalog.DiscoverCatalogWorkflowImpl; +import io.airbyte.workers.temporal.scheduling.ConnectionManagerWorkflowImpl; +import io.airbyte.workers.temporal.spec.SpecWorkflowImpl; +import io.airbyte.workers.temporal.support.TemporalProxyHelper; +import io.airbyte.workers.temporal.sync.SyncWorkflowImpl; +import io.grpc.StatusRuntimeException; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import io.micronaut.context.env.Environment; +import io.micronaut.context.event.ApplicationEventListener; +import io.micronaut.core.util.StringUtils; +import io.micronaut.discovery.event.ServiceReadyEvent; +import io.micronaut.scheduling.TaskExecutors; +import io.temporal.api.workflowservice.v1.DescribeNamespaceRequest; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.worker.Worker; +import io.temporal.worker.WorkerFactory; +import io.temporal.worker.WorkerOptions; +import java.io.IOException; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.stream.Collectors; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +import lombok.extern.slf4j.Slf4j; + +/** + * Performs any required initialization logic on application context start. + */ +@Singleton +@Requires(notEnv = {Environment.TEST}) +@Slf4j +public class ApplicationInitializer implements ApplicationEventListener { + + @Value("${airbyte.role}") + private String airbyteRole; + @Inject + private AirbyteVersion airbyteVersion; + @Inject + @Named("checkConnectionActivities") + private Optional> checkConnectionActivities; + @Inject + @Named("configsDatabaseMigrationCheck") + private Optional configsDatabaseMigrationCheck; + @Inject + private Optional configRepository; + @Inject + @Named("connectionManagerActivities") + private Optional> connectionManagerActivities; + @Inject + private DeploymentMode deploymentMode; + @Inject + @Named("discoverActivities") + private Optional> discoverActivities; + @Inject + @Named(TaskExecutors.IO) + private ExecutorService executorService; + @Inject + @Named("jobsDatabaseMigrationCheck") + private Optional jobsDatabaseMigrationCheck; + @Inject + @Named("jobsDatabaseAvailabilityCheck") + private Optional jobsDatabaseAvailabilityCheck; + @Inject + private Optional jobPersistence; + @Inject + private Optional logConfigs; + @Value("${airbyte.worker.check.max-workers}") + private Integer maxCheckWorkers; + @Value("${airbyte.worker.discover.max-workers}") + private Integer maxDiscoverWorkers; + @Value("${airbyte.worker.spec.max-workers}") + private Integer maxSpecWorkers; + @Value("${airbyte.worker.sync.max-workers}") + private Integer maxSyncWorkers; + @Value("${airbyte.worker.check.enabled}") + private boolean shouldRunCheckConnectionWorkflows; + @Value("${airbyte.worker.connection.enabled}") + private boolean shouldRunConnectionManagerWorkflows; + @Value("${airbyte.worker.discover.enabled}") + private boolean shouldRunDiscoverWorkflows; + @Value("${airbyte.worker.spec.enabled}") + private boolean shouldRunGetSpecWorkflows; + @Value("${airbyte.worker.sync.enabled}") + private boolean shouldRunSyncWorkflows; + @Inject + @Named("specActivities") + private Optional> specActivities; + @Inject + @Named("syncActivities") + private Optional> syncActivities; + @Inject + private TemporalProxyHelper temporalProxyHelper; + @Inject + private WorkflowServiceStubs temporalService; + @Inject + private TemporalUtils temporalUtils; + @Value("${airbyte.temporal.worker.ports}") + private Set temporalWorkerPorts; + @Inject + private Optional trackingStrategy; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private WorkerFactory workerFactory; + @Inject + private WorkerPlane workerPlane; + @Value("${airbyte.workspace.root}") + private String workspaceRoot; + @Value("${temporal.cloud.namespace}") + private String temporalCloudNamespace; + @Value("${airbyte.data.sync.task-queue}") + private String syncTaskQueue; + + @Override + public void onApplicationEvent(final ServiceReadyEvent event) { + try { + initializeCommonDependencies(); + + if (WorkerPlane.CONTROL_PLANE.equals(workerPlane)) { + initializeControlPlaneDependencies(); + } else { + log.info("Skipping Control Plane dependency initialization."); + } + + registerWorkerFactory(workerFactory, new MaxWorkersConfig(maxCheckWorkers, maxDiscoverWorkers, maxSpecWorkers, maxSyncWorkers)); + + log.info("Starting worker factory..."); + workerFactory.start(); + + log.info("Application initialized (mode = {}).", workerPlane); + } catch (final DatabaseCheckException | ExecutionException | InterruptedException | IOException | TimeoutException e) { + log.error("Unable to initialize application.", e); + throw new IllegalStateException(e); + } + } + + private void initializeCommonDependencies() throws ExecutionException, InterruptedException, TimeoutException { + log.info("Initializing common worker dependencies."); + + // Initialize the metric client + MetricClientFactory.initialize(MetricEmittingApps.WORKER); + + // Configure logging client + LogClientSingleton.getInstance().setWorkspaceMdc(workerEnvironment, logConfigs.orElseThrow(), + LogClientSingleton.getInstance().getSchedulerLogsRoot(Path.of(workspaceRoot))); + + if (WorkerEnvironment.KUBERNETES.equals(workerEnvironment)) { + KubePortManagerSingleton.init(temporalWorkerPorts); + } + + configureTemporal(temporalUtils, temporalService); + } + + private void initializeControlPlaneDependencies() throws DatabaseCheckException, IOException { + // Ensure that the Configuration database has been migrated to the latest version + log.info("Checking config database flyway migration version..."); + configsDatabaseMigrationCheck.orElseThrow().check(); + + // Ensure that the Jobs database has been migrated to the latest version + log.info("Checking jobs database flyway migration version..."); + jobsDatabaseMigrationCheck.orElseThrow().check(); + + // Ensure that the Jobs database is available + log.info("Checking jobs database availability..."); + jobsDatabaseAvailabilityCheck.orElseThrow().check(); + + TrackingClientSingleton.initialize( + trackingStrategy.orElseThrow(), + new Deployment(deploymentMode, jobPersistence.orElseThrow().getDeployment().orElseThrow(), + workerEnvironment), + airbyteRole, + airbyteVersion, + configRepository.orElseThrow()); + } + + private void registerWorkerFactory(final WorkerFactory workerFactory, final MaxWorkersConfig maxWorkersConfiguration) { + log.info("Registering worker factories...."); + if (shouldRunGetSpecWorkflows) { + registerGetSpec(workerFactory, maxWorkersConfiguration); + } + + if (shouldRunCheckConnectionWorkflows) { + registerCheckConnection(workerFactory, maxWorkersConfiguration); + } + + if (shouldRunDiscoverWorkflows) { + registerDiscover(workerFactory, maxWorkersConfiguration); + } + + if (shouldRunSyncWorkflows) { + registerSync(workerFactory, maxWorkersConfiguration); + } + + if (shouldRunConnectionManagerWorkflows) { + registerConnectionManager(workerFactory, maxWorkersConfiguration); + } + } + + private void registerCheckConnection(final WorkerFactory factory, final MaxWorkersConfig maxWorkersConfig) { + final Worker checkConnectionWorker = + factory.newWorker(TemporalJobType.CHECK_CONNECTION.name(), getWorkerOptions(maxWorkersConfig.getMaxCheckWorkers())); + checkConnectionWorker + .registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(CheckConnectionWorkflowImpl.class)); + checkConnectionWorker.registerActivitiesImplementations(checkConnectionActivities.orElseThrow().toArray(new Object[] {})); + log.info("Check Connection Workflow registered."); + } + + private void registerConnectionManager(final WorkerFactory factory, final MaxWorkersConfig maxWorkersConfig) { + final Worker connectionUpdaterWorker = + factory.newWorker(TemporalJobType.CONNECTION_UPDATER.toString(), getWorkerOptions(maxWorkersConfig.getMaxSyncWorkers())); + connectionUpdaterWorker + .registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(ConnectionManagerWorkflowImpl.class)); + connectionUpdaterWorker.registerActivitiesImplementations(connectionManagerActivities.orElseThrow().toArray(new Object[] {})); + log.info("Connection Manager Workflow registered."); + } + + private void registerDiscover(final WorkerFactory factory, final MaxWorkersConfig maxWorkersConfig) { + final Worker discoverWorker = + factory.newWorker(TemporalJobType.DISCOVER_SCHEMA.name(), getWorkerOptions(maxWorkersConfig.getMaxDiscoverWorkers())); + discoverWorker + .registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(DiscoverCatalogWorkflowImpl.class)); + discoverWorker.registerActivitiesImplementations(discoverActivities.orElseThrow().toArray(new Object[] {})); + log.info("Discover Workflow registered."); + } + + private void registerGetSpec(final WorkerFactory factory, final MaxWorkersConfig maxWorkersConfig) { + final Worker specWorker = factory.newWorker(TemporalJobType.GET_SPEC.name(), getWorkerOptions(maxWorkersConfig.getMaxSpecWorkers())); + specWorker.registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(SpecWorkflowImpl.class)); + specWorker.registerActivitiesImplementations(specActivities.orElseThrow().toArray(new Object[] {})); + log.info("Get Spec Workflow registered."); + } + + private void registerSync(final WorkerFactory factory, final MaxWorkersConfig maxWorkersConfig) { + final Set taskQueues = getSyncTaskQueue(); + + // There should be a default value provided by the application framework. If not, do this + // as a safety check to ensure we don't attempt to register against no task queue. + if (taskQueues.isEmpty()) { + throw new IllegalStateException("Sync workflow task queue must be provided."); + } + + for (final String taskQueue : taskQueues) { + log.info("Registering sync workflow for task queue '{}'...", taskQueue); + final Worker syncWorker = factory.newWorker(taskQueue, getWorkerOptions(maxWorkersConfig.getMaxSyncWorkers())); + syncWorker.registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(SyncWorkflowImpl.class)); + syncWorker.registerActivitiesImplementations(syncActivities.orElseThrow().toArray(new Object[] {})); + } + log.info("Sync Workflow registered."); + } + + private WorkerOptions getWorkerOptions(final int max) { + return WorkerOptions.newBuilder() + .setMaxConcurrentActivityExecutionSize(max) + .build(); + } + + /** + * Performs additional configuration of the Temporal service/connection. + * + * @param temporalUtils A {@link TemporalUtils} instance. + * @param temporalService A {@link WorkflowServiceStubs} instance. + * @throws ExecutionException if unable to perform the additional configuration. + * @throws InterruptedException if unable to perform the additional configuration. + * @throws TimeoutException if unable to perform the additional configuration. + */ + private void configureTemporal(final TemporalUtils temporalUtils, final WorkflowServiceStubs temporalService) + throws ExecutionException, InterruptedException, TimeoutException { + log.info("Configuring Temporal...."); + // Create the default Temporal namespace + temporalUtils.configureTemporalNamespace(temporalService); + + // Ensure that the Temporal namespace exists before continuing. + // If it does not exist after 30 seconds, fail the startup. + executorService.submit(this::waitForTemporalNamespace).get(30, TimeUnit.SECONDS); + } + + /** + * Blocks until the Temporal {@link TemporalUtils#DEFAULT_NAMESPACE} has been created. This is + * necessary to avoid issues related to + * https://community.temporal.io/t/running-into-an-issue-when-creating-namespace-programmatically/2783/8. + */ + private void waitForTemporalNamespace() { + boolean namespaceExists = false; + final String temporalNamespace = getTemporalNamespace(); + while (!namespaceExists) { + try { + temporalService.blockingStub().describeNamespace(DescribeNamespaceRequest.newBuilder().setNamespace(temporalNamespace).build()); + namespaceExists = true; + // This is to allow the configured namespace to be available in the Temporal + // cache before continuing on with any additional configuration/bean creation. + Thread.sleep(TimeUnit.SECONDS.toMillis(5)); + } catch (final StatusRuntimeException e) { + log.debug("Namespace '{}' does not exist yet. Re-checking...", temporalNamespace); + } catch (final InterruptedException e) { + log.debug("Sleep interrupted. Exiting loop..."); + } + } + } + + /** + * Retrieve the Temporal namespace based on the configuration. + * + * @return The Temporal namespace. + */ + private String getTemporalNamespace() { + return StringUtils.isNotEmpty(temporalCloudNamespace) ? temporalCloudNamespace : TemporalUtils.DEFAULT_NAMESPACE; + } + + /** + * Retrieve and parse the sync workflow task queue configuration. + * + * @return A set of Temporal task queues for the sync workflow. + */ + private Set getSyncTaskQueue() { + if (StringUtils.isEmpty(syncTaskQueue)) { + return Set.of(); + } + return Arrays.stream(syncTaskQueue.split(",")).collect(Collectors.toSet()); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/ContainerOrchestratorConfig.java b/airbyte-workers/src/main/java/io/airbyte/workers/ContainerOrchestratorConfig.java new file mode 100644 index 0000000000000..bf0e68e5ee095 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/ContainerOrchestratorConfig.java @@ -0,0 +1,18 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers; + +import io.airbyte.workers.general.DocumentStoreClient; +import io.fabric8.kubernetes.client.KubernetesClient; + +public record ContainerOrchestratorConfig( + String namespace, + DocumentStoreClient documentStoreClient, + KubernetesClient kubernetesClient, + String secretName, + String secretMountPath, + String containerOrchestratorImage, + String containerOrchestratorImagePullPolicy, + String googleApplicationCredentials) {} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactory.java deleted file mode 100644 index 13ad099411b0a..0000000000000 --- a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactory.java +++ /dev/null @@ -1,13 +0,0 @@ -/* - * Copyright (c) 2022 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.workers; - -import io.airbyte.api.client.AirbyteApiClient; - -public interface WorkerApiClientFactory { - - AirbyteApiClient create(); - -} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactoryImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactoryImpl.java deleted file mode 100644 index 428d15668c1a9..0000000000000 --- a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApiClientFactoryImpl.java +++ /dev/null @@ -1,101 +0,0 @@ -/* - * Copyright (c) 2022 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.workers; - -import com.auth0.jwt.JWT; -import com.auth0.jwt.JWTCreator; -import com.auth0.jwt.algorithms.Algorithm; -import com.google.auth.oauth2.ServiceAccountCredentials; -import io.airbyte.api.client.AirbyteApiClient; -import io.airbyte.config.Configs; -import io.airbyte.config.Configs.WorkerPlane; -import java.io.FileInputStream; -import java.security.interfaces.RSAPrivateKey; -import java.util.Date; -import java.util.concurrent.TimeUnit; -import lombok.extern.slf4j.Slf4j; - -@Slf4j -public class WorkerApiClientFactoryImpl implements WorkerApiClientFactory { - - private static final int JWT_TTL_MINUTES = 5; - - private final AirbyteApiClient airbyteApiClient; - - public WorkerApiClientFactoryImpl(final Configs configs) { - final var authHeader = configs.getAirbyteApiAuthHeaderName(); - - // control plane workers communicate with the Airbyte API within their internal network, so https - // isn't needed - final var scheme = configs.getWorkerPlane().equals(WorkerPlane.CONTROL_PLANE) ? "http" : "https"; - - log.debug("Creating Airbyte Config Api Client with Scheme: {}, Host: {}, Port: {}, Auth-Header: {}", - scheme, configs.getAirbyteApiHost(), configs.getAirbyteApiPort(), authHeader); - - this.airbyteApiClient = new AirbyteApiClient( - new io.airbyte.api.client.invoker.generated.ApiClient() - .setScheme(scheme) - .setHost(configs.getAirbyteApiHost()) - .setPort(configs.getAirbyteApiPort()) - .setBasePath("/api") - .setRequestInterceptor(builder -> { - builder.setHeader("User-Agent", "WorkerApp"); - if (!authHeader.isBlank()) { - builder.setHeader(authHeader, generateAuthToken(configs)); - } - })); - } - - @Override - public AirbyteApiClient create() { - return this.airbyteApiClient; - } - - /** - * Generate an auth token based on configs. This is called by the Api Client's requestInterceptor - * for each request. - * - * For Data Plane workers, generate a signed JWT as described here: - * https://cloud.google.com/endpoints/docs/openapi/service-account-authentication - * - * Otherwise, use the AIRBYTE_API_AUTH_HEADER_VALUE from EnvConfigs. - */ - private static String generateAuthToken(final Configs configs) { - if (configs.getWorkerPlane().equals(WorkerPlane.CONTROL_PLANE)) { - // control plane workers communicate with the Airbyte API within their internal network, so a signed - // JWT isn't needed - return configs.getAirbyteApiAuthHeaderValue(); - } else if (configs.getWorkerPlane().equals(WorkerPlane.DATA_PLANE)) { - try { - final Date now = new Date(); - final Date expTime = new Date(System.currentTimeMillis() + TimeUnit.MINUTES.toMillis(JWT_TTL_MINUTES)); - final String saEmail = configs.getDataPlaneServiceAccountEmail(); - // Build the JWT payload - final JWTCreator.Builder token = JWT.create() - .withIssuedAt(now) - .withExpiresAt(expTime) - .withIssuer(saEmail) - .withAudience(configs.getControlPlaneAuthEndpoint()) - .withSubject(saEmail) - .withClaim("email", saEmail); - - // TODO multi-cloud phase 2: check performance of on-demand token generation in load testing. might - // need to pull some of this outside of this method which is called for every API request - final FileInputStream stream = new FileInputStream(configs.getDataPlaneServiceAccountCredentialsPath()); - final ServiceAccountCredentials cred = ServiceAccountCredentials.fromStream(stream); - final RSAPrivateKey key = (RSAPrivateKey) cred.getPrivateKey(); - final Algorithm algorithm = Algorithm.RSA256(null, key); - return "Bearer " + token.sign(algorithm); - } catch (final Exception e) { - log.warn("An issue occurred while generating a data plane auth token. Defaulting to empty string.", e); - return ""; - } - } else { - log.warn("Worker somehow wasn't a control plane or a data plane worker!"); - return ""; - } - } - -} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java deleted file mode 100644 index 6c66284919cca..0000000000000 --- a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java +++ /dev/null @@ -1,597 +0,0 @@ -/* - * Copyright (c) 2022 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.workers; - -import io.airbyte.analytics.Deployment; -import io.airbyte.analytics.TrackingClient; -import io.airbyte.analytics.TrackingClientSingleton; -import io.airbyte.api.client.AirbyteApiClient; -import io.airbyte.commons.features.EnvVariableFeatureFlags; -import io.airbyte.commons.features.FeatureFlags; -import io.airbyte.commons.lang.CloseableShutdownHook; -import io.airbyte.config.Configs; -import io.airbyte.config.Configs.WorkerEnvironment; -import io.airbyte.config.Configs.WorkerPlane; -import io.airbyte.config.EnvConfigs; -import io.airbyte.config.helpers.LogClientSingleton; -import io.airbyte.config.persistence.ConfigPersistence; -import io.airbyte.config.persistence.ConfigRepository; -import io.airbyte.config.persistence.DatabaseConfigPersistence; -import io.airbyte.config.persistence.StatePersistence; -import io.airbyte.config.persistence.StreamResetPersistence; -import io.airbyte.config.persistence.split_secrets.JsonSecretsProcessor; -import io.airbyte.config.persistence.split_secrets.SecretPersistence; -import io.airbyte.config.persistence.split_secrets.SecretsHydrator; -import io.airbyte.db.Database; -import io.airbyte.db.check.DatabaseCheckException; -import io.airbyte.db.check.impl.JobsDatabaseAvailabilityCheck; -import io.airbyte.db.factory.DSLContextFactory; -import io.airbyte.db.factory.DataSourceFactory; -import io.airbyte.db.factory.DatabaseCheckFactory; -import io.airbyte.db.factory.DatabaseDriver; -import io.airbyte.db.factory.FlywayFactory; -import io.airbyte.db.instance.DatabaseConstants; -import io.airbyte.db.instance.configs.ConfigsDatabaseMigrator; -import io.airbyte.db.instance.jobs.JobsDatabaseMigrator; -import io.airbyte.metrics.lib.MetricClientFactory; -import io.airbyte.metrics.lib.MetricEmittingApps; -import io.airbyte.scheduler.persistence.DefaultJobCreator; -import io.airbyte.scheduler.persistence.DefaultJobPersistence; -import io.airbyte.scheduler.persistence.JobNotifier; -import io.airbyte.scheduler.persistence.JobPersistence; -import io.airbyte.scheduler.persistence.WebUrlHelper; -import io.airbyte.scheduler.persistence.WorkspaceHelper; -import io.airbyte.scheduler.persistence.job_error_reporter.JobErrorReporter; -import io.airbyte.scheduler.persistence.job_error_reporter.JobErrorReportingClient; -import io.airbyte.scheduler.persistence.job_error_reporter.JobErrorReportingClientFactory; -import io.airbyte.scheduler.persistence.job_factory.DefaultSyncJobFactory; -import io.airbyte.scheduler.persistence.job_factory.OAuthConfigSupplier; -import io.airbyte.scheduler.persistence.job_factory.SyncJobFactory; -import io.airbyte.scheduler.persistence.job_tracker.JobTracker; -import io.airbyte.workers.general.DocumentStoreClient; -import io.airbyte.workers.helper.ConnectionHelper; -import io.airbyte.workers.normalization.NormalizationRunnerFactory; -import io.airbyte.workers.process.DockerProcessFactory; -import io.airbyte.workers.process.KubePortManagerSingleton; -import io.airbyte.workers.process.KubeProcessFactory; -import io.airbyte.workers.process.ProcessFactory; -import io.airbyte.workers.process.WorkerHeartbeatServer; -import io.airbyte.workers.run.TemporalWorkerRunFactory; -import io.airbyte.workers.storage.StateClients; -import io.airbyte.workers.temporal.TemporalClient; -import io.airbyte.workers.temporal.TemporalJobType; -import io.airbyte.workers.temporal.TemporalUtils; -import io.airbyte.workers.temporal.check.connection.CheckConnectionActivityImpl; -import io.airbyte.workers.temporal.check.connection.CheckConnectionWorkflowImpl; -import io.airbyte.workers.temporal.discover.catalog.DiscoverCatalogActivityImpl; -import io.airbyte.workers.temporal.discover.catalog.DiscoverCatalogWorkflowImpl; -import io.airbyte.workers.temporal.scheduling.ConnectionManagerWorkflowImpl; -import io.airbyte.workers.temporal.scheduling.activities.AutoDisableConnectionActivityImpl; -import io.airbyte.workers.temporal.scheduling.activities.ConfigFetchActivityImpl; -import io.airbyte.workers.temporal.scheduling.activities.ConnectionDeletionActivityImpl; -import io.airbyte.workers.temporal.scheduling.activities.GenerateInputActivityImpl; -import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivityImpl; -import io.airbyte.workers.temporal.scheduling.activities.StreamResetActivityImpl; -import io.airbyte.workers.temporal.spec.SpecActivityImpl; -import io.airbyte.workers.temporal.spec.SpecWorkflowImpl; -import io.airbyte.workers.temporal.sync.DbtTransformationActivityImpl; -import io.airbyte.workers.temporal.sync.NormalizationActivityImpl; -import io.airbyte.workers.temporal.sync.PersistStateActivityImpl; -import io.airbyte.workers.temporal.sync.ReplicationActivityImpl; -import io.airbyte.workers.temporal.sync.RouteToTaskQueueActivityImpl; -import io.airbyte.workers.temporal.sync.RouterService; -import io.airbyte.workers.temporal.sync.SyncWorkflowImpl; -import io.fabric8.kubernetes.client.DefaultKubernetesClient; -import io.fabric8.kubernetes.client.KubernetesClient; -import io.temporal.client.WorkflowClient; -import io.temporal.serviceclient.WorkflowServiceStubs; -import io.temporal.worker.Worker; -import io.temporal.worker.WorkerFactory; -import io.temporal.worker.WorkerOptions; -import java.io.IOException; -import java.net.InetAddress; -import java.nio.file.Path; -import java.time.Instant; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.Executors; -import javax.sql.DataSource; -import org.flywaydb.core.Flyway; -import org.jooq.DSLContext; -import org.jooq.SQLDialect; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; -import org.slf4j.MDC; - -@SuppressWarnings("PMD.AvoidCatchingThrowable") -public class WorkerApp { - - private static final Logger LOGGER = LoggerFactory.getLogger(WorkerApp.class); - public static final int KUBE_HEARTBEAT_PORT = 9000; - private static final String DRIVER_CLASS_NAME = DatabaseDriver.POSTGRESQL.getDriverClassName(); - - // IMPORTANT: Changing the storage location will orphan already existing kube pods when the new - // version is deployed! - public static final Path STATE_STORAGE_PREFIX = Path.of("/state"); - - private static Configs configs; - private static ProcessFactory defaultProcessFactory; - private static ProcessFactory specProcessFactory; - private static ProcessFactory checkProcessFactory; - private static ProcessFactory discoverProcessFactory; - private static ProcessFactory replicationProcessFactory; - private static SecretsHydrator secretsHydrator; - private static WorkflowClient workflowClient; - private static ConfigRepository configRepository; - private static WorkerConfigs defaultWorkerConfigs; - private static WorkerConfigs specWorkerConfigs; - private static WorkerConfigs checkWorkerConfigs; - private static WorkerConfigs discoverWorkerConfigs; - private static WorkerConfigs replicationWorkerConfigs; - private static SyncJobFactory jobFactory; - private static JobPersistence jobPersistence; - private static WorkflowServiceStubs temporalService; - private static TemporalWorkerRunFactory temporalWorkerRunFactory; - private static ConnectionHelper connectionHelper; - private static Optional containerOrchestratorConfig; - private static JobNotifier jobNotifier; - private static JobTracker jobTracker; - private static JobErrorReporter jobErrorReporter; - private static StreamResetPersistence streamResetPersistence; - private static FeatureFlags featureFlags; - private static DefaultJobCreator jobCreator; - // TODO (pmossman) the API client should be scoped down to only the clients necessary for the - // worker. - // can be done after the Config API is broken down into multiple smaller classes. - private static AirbyteApiClient airbyteApiClient; - private static RouterService routerService; - - private static void registerConnectionManager(final WorkerFactory factory) { - final FeatureFlags featureFlags = new EnvVariableFeatureFlags(); - - final Worker connectionUpdaterWorker = - factory.newWorker(TemporalJobType.CONNECTION_UPDATER.toString(), getWorkerOptions(configs.getMaxWorkers().getMaxSyncWorkers())); - connectionUpdaterWorker.registerWorkflowImplementationTypes(ConnectionManagerWorkflowImpl.class); - connectionUpdaterWorker.registerActivitiesImplementations( - new GenerateInputActivityImpl( - jobPersistence), - new JobCreationAndStatusUpdateActivityImpl( - jobFactory, - jobPersistence, - temporalWorkerRunFactory, - configs.getWorkerEnvironment(), - configs.getLogConfigs(), - jobNotifier, - jobTracker, - configRepository, - jobCreator, - streamResetPersistence, - jobErrorReporter), - new ConfigFetchActivityImpl(configRepository, jobPersistence, configs, () -> Instant.now().getEpochSecond()), - new ConnectionDeletionActivityImpl(connectionHelper), - new CheckConnectionActivityImpl( - checkWorkerConfigs, - checkProcessFactory, - secretsHydrator, - configs.getWorkspaceRoot(), - configs.getWorkerEnvironment(), - configs.getLogConfigs(), - airbyteApiClient, - configs.getAirbyteVersionOrWarning()), - new AutoDisableConnectionActivityImpl(configRepository, jobPersistence, featureFlags, configs, jobNotifier), - new StreamResetActivityImpl(streamResetPersistence, jobPersistence)); - } - - private static void registerSync(final WorkerFactory factory) { - registerWorkersForDataSyncTasks(factory); - registerSyncControlPlaneWorkers(factory); - } - - /** - * These workers process tasks related to the actual syncing and processing of data. Depending on - * configuration, these tasks can be handled by workers in the Control Plane, and/or workers in a - * Data Plane. - */ - private static void registerWorkersForDataSyncTasks(final WorkerFactory factory) { - if (!configs.getDataSyncTaskQueues().isEmpty()) { - final ReplicationActivityImpl replicationActivity = getReplicationActivityImpl(replicationWorkerConfigs, replicationProcessFactory); - // Note that the configuration injected here is for the normalization orchestrator, and not the - // normalization pod itself. - // Configuration for the normalization pod is injected via the SyncWorkflowImpl. - final NormalizationActivityImpl normalizationActivity = getNormalizationActivityImpl(defaultWorkerConfigs, defaultProcessFactory); - final DbtTransformationActivityImpl dbtTransformationActivity = getDbtActivityImpl( - defaultWorkerConfigs, - defaultProcessFactory); - final PersistStateActivityImpl persistStateActivity = new PersistStateActivityImpl(airbyteApiClient, featureFlags); - - for (final String taskQueue : configs.getDataSyncTaskQueues()) { - final Worker worker = factory.newWorker(taskQueue, getWorkerOptions(configs.getMaxWorkers().getMaxSyncWorkers())); - worker.registerActivitiesImplementations(replicationActivity, normalizationActivity, dbtTransformationActivity, persistStateActivity); - } - } - } - - /** - * Control Plane workers handle all workflow tasks for the SyncWorkflow, as well as the activity - * task to decide which task queue to use for Data Plane tasks. - */ - private static void registerSyncControlPlaneWorkers(final WorkerFactory factory) { - if (configs.getWorkerPlane().equals(WorkerPlane.CONTROL_PLANE)) { - final Worker syncWorker = factory.newWorker(TemporalJobType.SYNC.name(), getWorkerOptions(configs.getMaxWorkers().getMaxSyncWorkers())); - syncWorker.registerWorkflowImplementationTypes(SyncWorkflowImpl.class); - - final RouteToTaskQueueActivityImpl routeToTaskQueueActivity = getRouteToTaskQueueActivityImpl(); - syncWorker.registerActivitiesImplementations(routeToTaskQueueActivity); - } - } - - private static void registerDiscover(final WorkerFactory factory) { - final Worker discoverWorker = factory.newWorker(TemporalJobType.DISCOVER_SCHEMA.name(), - getWorkerOptions(configs.getMaxWorkers().getMaxDiscoverWorkers())); - discoverWorker.registerWorkflowImplementationTypes(DiscoverCatalogWorkflowImpl.class); - discoverWorker - .registerActivitiesImplementations( - new DiscoverCatalogActivityImpl(discoverWorkerConfigs, discoverProcessFactory, secretsHydrator, configs.getWorkspaceRoot(), - configs.getWorkerEnvironment(), - configs.getLogConfigs(), - airbyteApiClient, configs.getAirbyteVersionOrWarning())); - } - - private static void registerCheckConnection(final WorkerFactory factory) { - final Worker checkConnectionWorker = - factory.newWorker(TemporalJobType.CHECK_CONNECTION.name(), getWorkerOptions(configs.getMaxWorkers().getMaxCheckWorkers())); - checkConnectionWorker.registerWorkflowImplementationTypes(CheckConnectionWorkflowImpl.class); - checkConnectionWorker - .registerActivitiesImplementations( - new CheckConnectionActivityImpl(checkWorkerConfigs, checkProcessFactory, secretsHydrator, configs.getWorkspaceRoot(), - configs.getWorkerEnvironment(), configs.getLogConfigs(), - airbyteApiClient, configs.getAirbyteVersionOrWarning())); - } - - private static void registerGetSpec(final WorkerFactory factory) { - final Worker specWorker = factory.newWorker(TemporalJobType.GET_SPEC.name(), getWorkerOptions(configs.getMaxWorkers().getMaxSpecWorkers())); - specWorker.registerWorkflowImplementationTypes(SpecWorkflowImpl.class); - specWorker.registerActivitiesImplementations( - new SpecActivityImpl(specWorkerConfigs, specProcessFactory, configs.getWorkspaceRoot(), configs.getWorkerEnvironment(), - configs.getLogConfigs(), airbyteApiClient, - configs.getAirbyteVersionOrWarning())); - } - - private static ReplicationActivityImpl getReplicationActivityImpl(final WorkerConfigs workerConfigs, - final ProcessFactory jobProcessFactory) { - - return new ReplicationActivityImpl( - containerOrchestratorConfig, - workerConfigs, - jobProcessFactory, - secretsHydrator, - configs.getWorkspaceRoot(), - configs.getWorkerEnvironment(), - configs.getLogConfigs(), - airbyteApiClient, - configs.getAirbyteVersionOrWarning(), - featureFlags.useStreamCapableState()); - } - - private static NormalizationActivityImpl getNormalizationActivityImpl(final WorkerConfigs workerConfigs, - final ProcessFactory jobProcessFactory) { - - return new NormalizationActivityImpl( - containerOrchestratorConfig, - workerConfigs, - jobProcessFactory, - secretsHydrator, - configs.getWorkspaceRoot(), - configs.getWorkerEnvironment(), - configs.getLogConfigs(), - jobPersistence, - airbyteApiClient, - configs.getAirbyteVersionOrWarning()); - } - - private static DbtTransformationActivityImpl getDbtActivityImpl(final WorkerConfigs workerConfigs, - final ProcessFactory jobProcessFactory) { - - return new DbtTransformationActivityImpl( - containerOrchestratorConfig, - workerConfigs, - jobProcessFactory, - secretsHydrator, - configs.getWorkspaceRoot(), - configs.getWorkerEnvironment(), - configs.getLogConfigs(), - jobPersistence, - airbyteApiClient, - configs.getAirbyteVersionOrWarning()); - } - - private static RouteToTaskQueueActivityImpl getRouteToTaskQueueActivityImpl() { - return new RouteToTaskQueueActivityImpl(routerService); - } - - /** - * Return either a docker or kubernetes process factory depending on the environment in - * {@link WorkerConfigs} - * - * @param configs used to determine which process factory to create. - * @param workerConfigs used to create the process factory. - * @return either a {@link DockerProcessFactory} or a {@link KubeProcessFactory}. - * @throws IOException - */ - private static ProcessFactory getJobProcessFactory(final Configs configs, final WorkerConfigs workerConfigs) throws IOException { - if (configs.getWorkerEnvironment().equals(Configs.WorkerEnvironment.KUBERNETES)) { - final KubernetesClient fabricClient = new DefaultKubernetesClient(); - final String localIp = InetAddress.getLocalHost().getHostAddress(); - final String kubeHeartbeatUrl = localIp + ":" + KUBE_HEARTBEAT_PORT; - LOGGER.info("Using Kubernetes namespace: {}", configs.getJobKubeNamespace()); - return new KubeProcessFactory(workerConfigs, - configs.getJobKubeNamespace(), - fabricClient, - kubeHeartbeatUrl, - false); - } else { - return new DockerProcessFactory( - workerConfigs, - configs.getWorkspaceRoot(), - configs.getWorkspaceDockerMount(), - configs.getLocalDockerMount(), - configs.getDockerNetwork()); - } - } - - private static WorkerOptions getWorkerOptions(final int max) { - return WorkerOptions.newBuilder() - .setMaxConcurrentActivityExecutionSize(max) - .build(); - } - - public static record ContainerOrchestratorConfig( - String namespace, - DocumentStoreClient documentStoreClient, - KubernetesClient kubernetesClient, - String secretName, - String secretMountPath, - String containerOrchestratorImage, - String containerOrchestratorImagePullPolicy, - String googleApplicationCredentials) { - - } - - static Optional getContainerOrchestratorConfig(final Configs configs) { - if (configs.getContainerOrchestratorEnabled()) { - final var kubernetesClient = new DefaultKubernetesClient(); - - final DocumentStoreClient documentStoreClient = StateClients.create( - configs.getStateStorageCloudConfigs(), - STATE_STORAGE_PREFIX); - - return Optional.of(new ContainerOrchestratorConfig( - configs.getJobKubeNamespace(), - documentStoreClient, - kubernetesClient, - configs.getContainerOrchestratorSecretName(), - configs.getContainerOrchestratorSecretMountPath(), - configs.getContainerOrchestratorImage(), - configs.getJobKubeMainContainerImagePullPolicy(), - configs.getGoogleApplicationCredentials())); - } else { - return Optional.empty(); - } - } - - public static void start() { - final Map mdc = MDC.getCopyOfContextMap(); - Executors.newSingleThreadExecutor().submit( - () -> { - MDC.setContextMap(mdc); - try { - new WorkerHeartbeatServer(KUBE_HEARTBEAT_PORT).start(); - } catch (final Exception e) { - throw new RuntimeException(e); - } - }); - - final WorkerFactory factory = WorkerFactory.newInstance(workflowClient); - - if (configs.shouldRunGetSpecWorkflows()) { - registerGetSpec(factory); - } - - if (configs.shouldRunCheckConnectionWorkflows()) { - registerCheckConnection(factory); - } - - if (configs.shouldRunDiscoverWorkflows()) { - registerDiscover(factory); - } - - if (configs.shouldRunSyncWorkflows()) { - registerSync(factory); - } - - if (configs.shouldRunConnectionManagerWorkflows()) { - registerConnectionManager(factory); - } - - factory.start(); - } - - private static void initializeCommonDependencies() throws IOException { - LOGGER.info("Initializing common worker dependencies."); - configs = new EnvConfigs(); - LOGGER.info("workspaceRoot = " + configs.getWorkspaceRoot()); - - MetricClientFactory.initialize(MetricEmittingApps.WORKER); - - LogClientSingleton.getInstance().setWorkspaceMdc( - configs.getWorkerEnvironment(), - configs.getLogConfigs(), - LogClientSingleton.getInstance().getSchedulerLogsRoot(configs.getWorkspaceRoot())); - - if (configs.getWorkerEnvironment().equals(WorkerEnvironment.KUBERNETES)) { - KubePortManagerSingleton.init(configs.getTemporalWorkerPorts()); - } - - featureFlags = new EnvVariableFeatureFlags(); - defaultWorkerConfigs = new WorkerConfigs(configs); - temporalService = TemporalUtils.createTemporalService(); - workflowClient = TemporalUtils.createWorkflowClient(temporalService, TemporalUtils.getNamespace()); - TemporalUtils.configureTemporalNamespace(temporalService); - airbyteApiClient = new WorkerApiClientFactoryImpl(configs).create(); - - replicationWorkerConfigs = WorkerConfigs.buildReplicationWorkerConfigs(configs); - - defaultProcessFactory = getJobProcessFactory(configs, defaultWorkerConfigs); - replicationProcessFactory = getJobProcessFactory(configs, replicationWorkerConfigs); - - containerOrchestratorConfig = getContainerOrchestratorConfig(configs); - } - - /** - * These dependencies are only initialized by workers in the Control Plane, as they require database - * access. Workers in an external Data Plane - * - * @throws IOException - * @throws DatabaseCheckException - */ - private static void initializeControlPlaneDependencies() throws IOException, DatabaseCheckException { - LOGGER.info("Initializing control plane worker dependencies."); - - final DataSource configsDataSource = DataSourceFactory.create(configs.getConfigDatabaseUser(), configs.getConfigDatabasePassword(), - DRIVER_CLASS_NAME, configs.getConfigDatabaseUrl()); - final DataSource jobsDataSource = DataSourceFactory.create(configs.getDatabaseUser(), configs.getDatabasePassword(), - DRIVER_CLASS_NAME, configs.getDatabaseUrl()); - - // Manual configuration that will be replaced by Dependency Injection in the future - try (final DSLContext configsDslContext = DSLContextFactory.create(configsDataSource, SQLDialect.POSTGRES); - final DSLContext jobsDslContext = DSLContextFactory.create(jobsDataSource, SQLDialect.POSTGRES)) { - - // Ensure that the database resources are closed on application shutdown - CloseableShutdownHook.registerRuntimeShutdownHook(configsDataSource, jobsDataSource, configsDslContext, jobsDslContext); - - // TODO (pmossman) abstract all migration code here into a separate function. It's confusing to - // have migration code interwoven with dependency initialization as-is - final Flyway configsFlyway = FlywayFactory.create(configsDataSource, WorkerApp.class.getSimpleName(), - ConfigsDatabaseMigrator.DB_IDENTIFIER, ConfigsDatabaseMigrator.MIGRATION_FILE_LOCATION); - final Flyway jobsFlyway = FlywayFactory.create(jobsDataSource, WorkerApp.class.getSimpleName(), JobsDatabaseMigrator.DB_IDENTIFIER, - JobsDatabaseMigrator.MIGRATION_FILE_LOCATION); - - // Ensure that the Configuration database is available - DatabaseCheckFactory - .createConfigsDatabaseMigrationCheck(configsDslContext, configsFlyway, configs.getConfigsDatabaseMinimumFlywayMigrationVersion(), - configs.getConfigsDatabaseInitializationTimeoutMs()) - .check(); - - LOGGER.info("Checking jobs database flyway migration version.."); - DatabaseCheckFactory.createJobsDatabaseMigrationCheck(jobsDslContext, jobsFlyway, configs.getJobsDatabaseMinimumFlywayMigrationVersion(), - configs.getJobsDatabaseInitializationTimeoutMs()).check(); - - // Ensure that the Jobs database is available - new JobsDatabaseAvailabilityCheck(jobsDslContext, DatabaseConstants.DEFAULT_ASSERT_DATABASE_TIMEOUT_MS).check(); - - specWorkerConfigs = WorkerConfigs.buildSpecWorkerConfigs(configs); - checkWorkerConfigs = WorkerConfigs.buildCheckWorkerConfigs(configs); - discoverWorkerConfigs = WorkerConfigs.buildDiscoverWorkerConfigs(configs); - - specProcessFactory = getJobProcessFactory(configs, specWorkerConfigs); - checkProcessFactory = getJobProcessFactory(configs, checkWorkerConfigs); - discoverProcessFactory = getJobProcessFactory(configs, discoverWorkerConfigs); - - final Database configDatabase = new Database(configsDslContext); - final JsonSecretsProcessor jsonSecretsProcessor = JsonSecretsProcessor.builder() - .maskSecrets(!featureFlags.exposeSecretsInExport()) - .copySecrets(false) - .build(); - final ConfigPersistence configPersistence = DatabaseConfigPersistence.createWithValidation(configDatabase, jsonSecretsProcessor); - configRepository = new ConfigRepository(configPersistence, configDatabase); - - final Database jobDatabase = new Database(jobsDslContext); - - jobPersistence = new DefaultJobPersistence(jobDatabase); - final StatePersistence statePersistence = new StatePersistence(configDatabase); - jobCreator = new DefaultJobCreator( - jobPersistence, - defaultWorkerConfigs.getResourceRequirements(), - statePersistence); - - TrackingClientSingleton.initialize( - configs.getTrackingStrategy(), - new Deployment(configs.getDeploymentMode(), jobPersistence.getDeployment().orElseThrow(), configs.getWorkerEnvironment()), - configs.getAirbyteRole(), - configs.getAirbyteVersion(), - configRepository); - final TrackingClient trackingClient = TrackingClientSingleton.get(); - jobFactory = new DefaultSyncJobFactory( - configs.connectorSpecificResourceDefaultsEnabled(), - jobCreator, - configRepository, - new OAuthConfigSupplier(configRepository, trackingClient)); - - streamResetPersistence = new StreamResetPersistence(configDatabase); - - final TemporalClient temporalClient = new TemporalClient(workflowClient, configs.getWorkspaceRoot(), temporalService, streamResetPersistence); - - temporalWorkerRunFactory = new TemporalWorkerRunFactory( - temporalClient, - configs.getWorkspaceRoot(), - configs.getAirbyteVersionOrWarning(), - featureFlags); - - final WorkspaceHelper workspaceHelper = new WorkspaceHelper( - configRepository, - jobPersistence); - - connectionHelper = new ConnectionHelper(configRepository, workspaceHelper); - - final WebUrlHelper webUrlHelper = new WebUrlHelper(configs.getWebappUrl()); - - jobNotifier = new JobNotifier( - webUrlHelper, - configRepository, - workspaceHelper, - TrackingClientSingleton.get()); - - jobTracker = new JobTracker(configRepository, jobPersistence, trackingClient); - - final JobErrorReportingClient jobErrorReportingClient = JobErrorReportingClientFactory.getClient(configs.getJobErrorReportingStrategy(), - configs); - jobErrorReporter = new JobErrorReporter( - configRepository, - configs.getDeploymentMode(), - configs.getAirbyteVersionOrWarning(), - NormalizationRunnerFactory.BASE_NORMALIZATION_IMAGE_NAME, - NormalizationRunnerFactory.NORMALIZATION_VERSION, - webUrlHelper, - jobErrorReportingClient); - - // TODO (pmossman) refactor secretsHydrator instantiation to work better with Dependency Injection - // framework - secretsHydrator = SecretPersistence.getSecretsHydrator(configsDslContext, configs); - routerService = new RouterService(configs); - } - } - - public static void main(final String[] args) { - try { - initializeCommonDependencies(); - - if (configs.getWorkerPlane().equals(WorkerPlane.CONTROL_PLANE)) { - initializeControlPlaneDependencies(); - } else { - LOGGER.info("Skipping Control Plane dependency initialization."); - - // The SecretsHydrator is initialized with a database context for Control Plane workers. - // If the worker isn't in the Control Plane, we need to initialize the SecretsHydrator - // without a database context. - secretsHydrator = SecretPersistence.getSecretsHydrator(null, configs); - } - start(); - } catch (final Throwable t) { - LOGGER.error("Worker app failed", t); - System.exit(1); - } - } - -} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/ActivityBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/ActivityBeanFactory.java new file mode 100644 index 0000000000000..d39dc6179569c --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/ActivityBeanFactory.java @@ -0,0 +1,198 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.workers.exception.WorkerException; +import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.check.connection.CheckConnectionActivity; +import io.airbyte.workers.temporal.discover.catalog.DiscoverCatalogActivity; +import io.airbyte.workers.temporal.scheduling.activities.AutoDisableConnectionActivity; +import io.airbyte.workers.temporal.scheduling.activities.ConfigFetchActivity; +import io.airbyte.workers.temporal.scheduling.activities.ConnectionDeletionActivity; +import io.airbyte.workers.temporal.scheduling.activities.GenerateInputActivity; +import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivity; +import io.airbyte.workers.temporal.scheduling.activities.RecordMetricActivity; +import io.airbyte.workers.temporal.scheduling.activities.RouteToSyncTaskQueueActivity; +import io.airbyte.workers.temporal.scheduling.activities.StreamResetActivity; +import io.airbyte.workers.temporal.scheduling.activities.WorkflowConfigActivity; +import io.airbyte.workers.temporal.spec.SpecActivity; +import io.airbyte.workers.temporal.sync.DbtTransformationActivity; +import io.airbyte.workers.temporal.sync.NormalizationActivity; +import io.airbyte.workers.temporal.sync.PersistStateActivity; +import io.airbyte.workers.temporal.sync.ReplicationActivity; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Property; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import io.temporal.activity.ActivityCancellationType; +import io.temporal.activity.ActivityOptions; +import io.temporal.common.RetryOptions; +import java.time.Duration; +import java.util.List; +import javax.inject.Named; +import javax.inject.Singleton; + +/** + * Micronaut bean factory for activity-related singletons. + */ +@Factory +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class ActivityBeanFactory { + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("checkConnectionActivities") + public List checkConnectionActivities( + final CheckConnectionActivity checkConnectionActivity) { + return List.of(checkConnectionActivity); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("connectionManagerActivities") + public List connectionManagerActivities( + final GenerateInputActivity generateInputActivity, + final JobCreationAndStatusUpdateActivity jobCreationAndStatusUpdateActivity, + final ConfigFetchActivity configFetchActivity, + final ConnectionDeletionActivity connectionDeletionActivity, + final CheckConnectionActivity checkConnectionActivity, + final AutoDisableConnectionActivity autoDisableConnectionActivity, + final StreamResetActivity streamResetActivity, + final RecordMetricActivity recordMetricActivity, + final WorkflowConfigActivity workflowConfigActivity, + final RouteToSyncTaskQueueActivity routeToSyncTaskQueueActivity) { + return List.of(generateInputActivity, + jobCreationAndStatusUpdateActivity, + configFetchActivity, + connectionDeletionActivity, + checkConnectionActivity, + autoDisableConnectionActivity, + streamResetActivity, + recordMetricActivity, + workflowConfigActivity, + routeToSyncTaskQueueActivity); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("discoverActivities") + public List discoverActivities( + final DiscoverCatalogActivity discoverCatalogActivity) { + return List.of(discoverCatalogActivity); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("specActivities") + public List specActivities( + final SpecActivity specActivity) { + return List.of(specActivity); + } + + @Singleton + @Named("syncActivities") + public List syncActivities( + final ReplicationActivity replicationActivity, + final NormalizationActivity normalizationActivity, + final DbtTransformationActivity dbtTransformationActivity, + final PersistStateActivity persistStateActivity) { + return List.of(replicationActivity, normalizationActivity, dbtTransformationActivity, persistStateActivity); + } + + @Singleton + @Named("checkActivityOptions") + public ActivityOptions checkActivityOptions() { + return ActivityOptions.newBuilder() + .setScheduleToCloseTimeout(Duration.ofMinutes(5)) + .setRetryOptions(TemporalUtils.NO_RETRY) + .build(); + } + + @Singleton + @Named("discoveryActivityOptions") + public ActivityOptions discoveryActivityOptions() { + return ActivityOptions.newBuilder() + .setScheduleToCloseTimeout(Duration.ofHours(2)) + .setRetryOptions(TemporalUtils.NO_RETRY) + .build(); + } + + @Singleton + @Named("longRunActivityOptions") + public ActivityOptions longRunActivityOptions( + @Value("${airbyte.worker.sync.max-timeout}") final Long maxTimeout, + @Named("longRunActivityRetryOptions") final RetryOptions retryOptions) { + return ActivityOptions.newBuilder() + .setScheduleToCloseTimeout(Duration.ofDays(maxTimeout)) + .setStartToCloseTimeout(Duration.ofDays(maxTimeout)) + .setScheduleToStartTimeout(Duration.ofDays(maxTimeout)) + .setCancellationType(ActivityCancellationType.WAIT_CANCELLATION_COMPLETED) + .setRetryOptions(retryOptions) + .setHeartbeatTimeout(TemporalUtils.HEARTBEAT_TIMEOUT) + .build(); + } + + @Singleton + @Named("shortActivityOptions") + public ActivityOptions shortActivityOptions(@Property(name = "airbyte.activity.max-timeout", + defaultValue = "120") final Long maxTimeout, + @Named("shortRetryOptions") final RetryOptions shortRetryOptions) { + return ActivityOptions.newBuilder() + .setStartToCloseTimeout(Duration.ofSeconds(maxTimeout)) + .setCancellationType(ActivityCancellationType.WAIT_CANCELLATION_COMPLETED) + .setRetryOptions(shortRetryOptions) + .setHeartbeatTimeout(TemporalUtils.HEARTBEAT_TIMEOUT) + .build(); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("specActivityOptions") + public ActivityOptions specActivityOptions() { + return ActivityOptions.newBuilder() + .setScheduleToCloseTimeout(Duration.ofHours(1)) + .setRetryOptions(TemporalUtils.NO_RETRY) + .build(); + } + + @Singleton + @Requires(property = "airbyte.container.orchestrator.enabled", + value = "true") + @Named("longRunActivityRetryOptions") + public RetryOptions containerOrchestratorRetryOptions() { + return RetryOptions.newBuilder() + .setDoNotRetry(RuntimeException.class.getName(), WorkerException.class.getName()) + .build(); + } + + @Singleton + @Requires(property = "airbyte.container.orchestrator.enabled", + notEquals = "true") + @Named("longRunActivityRetryOptions") + public RetryOptions noRetryOptions() { + return TemporalUtils.NO_RETRY; + } + + @Singleton + @Named("shortRetryOptions") + public RetryOptions shortRetryOptions(@Property(name = "airbyte.activity.max-attempts", + defaultValue = "5") final Integer activityNumberOfAttempts, + @Property(name = "airbyte.activity.initial-delay", + defaultValue = "30") final Integer initialDelayBetweenActivityAttemptsSeconds, + @Property(name = "airbyte.activity.max-delay", + defaultValue = "600") final Integer maxDelayBetweenActivityAttemptsSeconds) { + return RetryOptions.newBuilder() + .setMaximumAttempts(activityNumberOfAttempts) + .setInitialInterval(Duration.ofSeconds(initialDelayBetweenActivityAttemptsSeconds)) + .setMaximumInterval(Duration.ofSeconds(maxDelayBetweenActivityAttemptsSeconds)) + .build(); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/ApiClientBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/ApiClientBeanFactory.java new file mode 100644 index 0000000000000..0b258712d4093 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/ApiClientBeanFactory.java @@ -0,0 +1,121 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import com.auth0.jwt.JWT; +import com.auth0.jwt.JWTCreator; +import com.auth0.jwt.algorithms.Algorithm; +import com.google.auth.oauth2.ServiceAccountCredentials; +import io.airbyte.api.client.AirbyteApiClient; +import io.airbyte.config.Configs.WorkerPlane; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Value; +import java.io.FileInputStream; +import java.security.interfaces.RSAPrivateKey; +import java.util.Date; +import java.util.concurrent.TimeUnit; +import javax.inject.Named; +import javax.inject.Singleton; +import lombok.extern.slf4j.Slf4j; + +/** + * Micronaut bean factory for API client singletons. + */ +@Factory +@Slf4j +public class ApiClientBeanFactory { + + private static final int JWT_TTL_MINUTES = 5; + + @Singleton + public AirbyteApiClient airbyteApiClient( + @Value("${airbyte.internal.api.auth-header.name}") final String airbyteApiAuthHeaderName, + @Value("${airbyte.internal.api.host}") final String airbyteApiHost, + @Named("internalApiAuthToken") final String internalApiAuthToken, + @Named("internalApiScheme") final String internalApiScheme) { + return new AirbyteApiClient( + new io.airbyte.api.client.invoker.generated.ApiClient() + .setScheme(internalApiScheme) + .setHost(parseHostName(airbyteApiHost)) + .setPort(parsePort(airbyteApiHost)) + .setBasePath("/api") + .setRequestInterceptor(builder -> { + builder.setHeader("User-Agent", "WorkerApp"); + if (!airbyteApiAuthHeaderName.isBlank()) { + builder.setHeader(airbyteApiAuthHeaderName, internalApiAuthToken); + } + })); + } + + @Singleton + @Named("internalApiScheme") + public String internalApiScheme(final WorkerPlane workerPlane) { + // control plane workers communicate with the Airbyte API within their internal network, so https + // isn't needed + return WorkerPlane.CONTROL_PLANE.equals(workerPlane) ? "http" : "https"; + } + + /** + * Generate an auth token based on configs. This is called by the Api Client's requestInterceptor + * for each request. + *

+ * For Data Plane workers, generate a signed JWT as described here: + * https://cloud.google.com/endpoints/docs/openapi/service-account-authentication + *

+ * Otherwise, use the AIRBYTE_API_AUTH_HEADER_VALUE from EnvConfigs. + */ + @Singleton + @Named("internalApiAuthToken") + public String internalApiAuthToken( + @Value("${airbyte.internal.api.auth-header.value}") final String airbyteApiAuthHeaderValue, + @Value("${airbyte.control.plane.auth-endpoint}") final String controlPlaneAuthEndpoint, + @Value("${airbyte.data.plane.service-account.email}") final String dataPlaneServiceAccountEmail, + @Value("${airbyte.data.plane.service-account.credentials-path}") final String dataPlaneServiceAccountCredentialsPath, + final WorkerPlane workerPlane) { + if (WorkerPlane.CONTROL_PLANE.equals(workerPlane)) { + // control plane workers communicate with the Airbyte API within their internal network, so a signed + // JWT isn't needed + return airbyteApiAuthHeaderValue; + } else if (WorkerPlane.DATA_PLANE.equals(workerPlane)) { + try { + final Date now = new Date(); + final Date expTime = new Date(System.currentTimeMillis() + TimeUnit.MINUTES.toMillis(JWT_TTL_MINUTES)); + // Build the JWT payload + final JWTCreator.Builder token = JWT.create() + .withIssuedAt(now) + .withExpiresAt(expTime) + .withIssuer(dataPlaneServiceAccountEmail) + .withAudience(controlPlaneAuthEndpoint) + .withSubject(dataPlaneServiceAccountEmail) + .withClaim("email", dataPlaneServiceAccountEmail); + + // TODO multi-cloud phase 2: check performance of on-demand token generation in load testing. might + // need to pull some of this outside of this method which is called for every API request + final FileInputStream stream = new FileInputStream(dataPlaneServiceAccountCredentialsPath); + final ServiceAccountCredentials cred = ServiceAccountCredentials.fromStream(stream); + final RSAPrivateKey key = (RSAPrivateKey) cred.getPrivateKey(); + final Algorithm algorithm = Algorithm.RSA256(null, key); + return "Bearer " + token.sign(algorithm); + } catch (final Exception e) { + log.warn( + "An issue occurred while generating a data plane auth token. Defaulting to empty string. Error Message: {}", + e.getMessage()); + return ""; + } + } else { + log.warn("Worker somehow wasn't a control plane or a data plane worker!"); + return ""; + } + } + + private String parseHostName(final String airbyteInternalApiHost) { + return airbyteInternalApiHost.split(":")[0]; + } + + private int parsePort(final String airbyteInternalApiHost) { + return Integer.parseInt(airbyteInternalApiHost.split(":")[1]); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/ApplicationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/ApplicationBeanFactory.java new file mode 100644 index 0000000000000..42fc15b1ccd21 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/ApplicationBeanFactory.java @@ -0,0 +1,175 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.analytics.TrackingClient; +import io.airbyte.commons.features.EnvVariableFeatureFlags; +import io.airbyte.commons.features.FeatureFlags; +import io.airbyte.commons.version.AirbyteVersion; +import io.airbyte.config.AirbyteConfigValidator; +import io.airbyte.config.Configs.DeploymentMode; +import io.airbyte.config.Configs.SecretPersistenceType; +import io.airbyte.config.Configs.TrackingStrategy; +import io.airbyte.config.Configs.WorkerEnvironment; +import io.airbyte.config.Configs.WorkerPlane; +import io.airbyte.config.persistence.ConfigRepository; +import io.airbyte.config.persistence.StatePersistence; +import io.airbyte.config.persistence.split_secrets.JsonSecretsProcessor; +import io.airbyte.metrics.lib.MetricClient; +import io.airbyte.metrics.lib.MetricClientFactory; +import io.airbyte.scheduler.persistence.DefaultJobCreator; +import io.airbyte.scheduler.persistence.JobNotifier; +import io.airbyte.scheduler.persistence.JobPersistence; +import io.airbyte.scheduler.persistence.WebUrlHelper; +import io.airbyte.scheduler.persistence.WorkspaceHelper; +import io.airbyte.scheduler.persistence.job_tracker.JobTracker; +import io.airbyte.workers.WorkerConfigs; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import io.micronaut.core.util.StringUtils; +import java.nio.file.Path; +import java.time.Instant; +import java.util.Locale; +import java.util.function.Function; +import java.util.function.Supplier; +import javax.inject.Named; +import javax.inject.Singleton; +import lombok.extern.slf4j.Slf4j; + +/** + * Micronaut bean factory for general singletons. + */ +@Factory +@Slf4j +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class ApplicationBeanFactory { + + @Singleton + public AirbyteVersion airbyteVersion(@Value("${airbyte.version}") final String airbyteVersion) { + return new AirbyteVersion(airbyteVersion); + } + + @Singleton + public DeploymentMode deploymentMode(@Value("${airbyte.deployment-mode}") final String deploymentMode) { + return convertToEnum(deploymentMode, DeploymentMode::valueOf, DeploymentMode.OSS); + } + + @Singleton + public SecretPersistenceType secretPersistenceType(@Value("${airbyte.secret.persistence}") final String secretPersistence) { + return convertToEnum(secretPersistence, SecretPersistenceType::valueOf, + SecretPersistenceType.TESTING_CONFIG_DB_TABLE); + } + + @Singleton + public TrackingStrategy trackingStrategy(@Value("${airbyte.tracking-strategy}") final String trackingStrategy) { + return convertToEnum(trackingStrategy, TrackingStrategy::valueOf, TrackingStrategy.LOGGING); + } + + @Singleton + public WorkerEnvironment workerEnvironment(@Value("${airbyte.worker.env}") final String workerEnv) { + return convertToEnum(workerEnv, WorkerEnvironment::valueOf, WorkerEnvironment.DOCKER); + } + + @Singleton + public WorkerPlane workerPlane(@Value("${airbyte.worker.plane}") final String workerPlane) { + return convertToEnum(workerPlane, WorkerPlane::valueOf, WorkerPlane.CONTROL_PLANE); + } + + @Singleton + @Named("workspaceRoot") + public Path workspaceRoot(@Value("${airbyte.workspace.root}") final String workspaceRoot) { + return Path.of(workspaceRoot); + } + + @Singleton + @Named("currentSecondsSupplier") + public Supplier currentSecondsSupplier() { + return () -> Instant.now().getEpochSecond(); + } + + @Singleton + public DefaultJobCreator defaultJobCreator(final JobPersistence jobPersistence, + @Named("defaultWorkerConfigs") final WorkerConfigs defaultWorkerConfigs, + final StatePersistence statePersistence) { + return new DefaultJobCreator( + jobPersistence, + defaultWorkerConfigs.getResourceRequirements(), + statePersistence); + } + + @Singleton + public FeatureFlags featureFlags() { + return new EnvVariableFeatureFlags(); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public JobNotifier jobNotifier( + final ConfigRepository configRepository, + final TrackingClient trackingClient, + final WebUrlHelper webUrlHelper, + final WorkspaceHelper workspaceHelper) { + return new JobNotifier( + webUrlHelper, + configRepository, + workspaceHelper, + trackingClient); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public JobTracker jobTracker( + final ConfigRepository configRepository, + final JobPersistence jobPersistence, + final TrackingClient trackingClient) { + return new JobTracker(configRepository, jobPersistence, trackingClient); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public JsonSecretsProcessor jsonSecretsProcessor(final FeatureFlags featureFlags) { + return JsonSecretsProcessor.builder() + .maskSecrets(!featureFlags.exposeSecretsInExport()) + .copySecrets(false) + .build(); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public WebUrlHelper webUrlHelper(@Value("${airbyte.web-app.url}") final String webAppUrl) { + return new WebUrlHelper(webAppUrl); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public WorkspaceHelper workspaceHelper( + final ConfigRepository configRepository, + final JobPersistence jobPersistence) { + return new WorkspaceHelper( + configRepository, + jobPersistence); + } + + @Singleton + public AirbyteConfigValidator airbyteConfigValidator() { + return new AirbyteConfigValidator(); + }; + + @Singleton + public MetricClient metricClient() { + return MetricClientFactory.getMetricClient(); + } + + private T convertToEnum(final String value, final Function creatorFunction, final T defaultValue) { + return StringUtils.isNotEmpty(value) ? creatorFunction.apply(value.toUpperCase(Locale.ROOT)) : defaultValue; + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/CloudStorageBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/CloudStorageBeanFactory.java new file mode 100644 index 0000000000000..0acaa60099ffa --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/CloudStorageBeanFactory.java @@ -0,0 +1,92 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.config.storage.CloudStorageConfigs; +import io.airbyte.config.storage.CloudStorageConfigs.GcsConfig; +import io.airbyte.config.storage.CloudStorageConfigs.MinioConfig; +import io.airbyte.config.storage.CloudStorageConfigs.S3Config; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import javax.inject.Named; +import javax.inject.Singleton; + +/** + * Micronaut bean factory for cloud storage-related singletons. + */ +@Factory +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class CloudStorageBeanFactory { + + @Singleton + @Requires(property = "airbyte.cloud.storage.logs.type", + value = "GCS") + @Named("logStorageConfigs") + public CloudStorageConfigs gcsLogStorageConfigs( + @Value("${airbyte.cloud.storage.logs.gcs.bucket}") final String gcsLogBucket, + @Value("${airbyte.cloud.storage.logs.gcs.application-credentials}") final String googleApplicationCredentials) { + return CloudStorageConfigs.gcs(new GcsConfig(gcsLogBucket, googleApplicationCredentials)); + } + + @Singleton + @Requires(property = "airbyte.cloud.storage.logs.type", + value = "MINIO") + @Named("logStorageConfigs") + public CloudStorageConfigs minioLogStorageConfigs( + @Value("${airbyte.cloud.storage.logs.minio.access-key}") final String awsAccessKeyId, + @Value("${airbyte.cloud.storage.logs.minio.secret-access-key}") final String awsSecretAccessKey, + @Value("${airbyte.cloud.storage.logs.minio.bucket}") final String s3LogBucket, + @Value("${airbyte.cloud.storage.logs.minio.endpoint}") final String s3MinioEndpoint) { + return CloudStorageConfigs.minio(new MinioConfig(s3LogBucket, awsAccessKeyId, awsSecretAccessKey, s3MinioEndpoint)); + } + + @Singleton + @Requires(property = "airbyte.cloud.storage.logs.type", + value = "S3") + @Named("logStorageConfigs") + public CloudStorageConfigs s3LogStorageConfigs( + @Value("${airbyte.cloud.storage.logs.s3.access-key}") final String awsAccessKeyId, + @Value("${airbyte.cloud.storage.logs.s3.secret-access-key}") final String awsSecretAccessKey, + @Value("${airbyte.cloud.storage.logs.s3.bucket}") final String s3LogBucket, + @Value("${airbyte.cloud.storage.logs.s3.region}") final String s3LogBucketRegion) { + return CloudStorageConfigs.s3(new S3Config(s3LogBucket, awsAccessKeyId, awsSecretAccessKey, s3LogBucketRegion)); + } + + @Singleton + @Requires(property = "airbyte.cloud.storage.state.type", + value = "GCS") + @Named("stateStorageConfigs") + public CloudStorageConfigs gcsStateStorageConfiguration( + @Value("${airbyte.cloud.storage.state.gcs.bucket}") final String gcsBucketName, + @Value("${airbyte.cloud.storage.state.gcs.application-credentials}") final String gcsApplicationCredentials) { + return CloudStorageConfigs.gcs(new GcsConfig(gcsBucketName, gcsApplicationCredentials)); + } + + @Singleton + @Requires(property = "airbyte.cloud.storage.state.type", + value = "MINIO") + @Named("stateStorageConfigs") + public CloudStorageConfigs minioStateStorageConfiguration( + @Value("${airbyte.cloud.storage.state.minio.bucket}") final String bucketName, + @Value("${airbyte.cloud.storage.state.minio.access-key}") final String awsAccessKey, + @Value("${airbyte.cloud.storage.state.minio.secret-access-key}") final String secretAccessKey, + @Value("${airbyte.cloud.storage.state.minio.endpoint}") final String endpoint) { + return CloudStorageConfigs.minio(new MinioConfig(bucketName, awsAccessKey, secretAccessKey, endpoint)); + } + + @Singleton + @Requires(property = "airbyte.cloud.storage.state.type", + value = "S3") + @Named("stateStorageConfigs") + public CloudStorageConfigs s3StateStorageConfiguration( + @Value("${airbyte.cloud.storage.state.s3.bucket}") final String bucketName, + @Value("${airbyte.cloud.storage.state.s3.access-key}") final String awsAccessKey, + @Value("${airbyte.cloud.storage.state.s3.secret-access-key}") final String secretAcessKey, + @Value("${airbyte.cloud.storage.state.s3.region}") final String s3Region) { + return CloudStorageConfigs.s3(new S3Config(bucketName, awsAccessKey, secretAcessKey, s3Region)); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/ContainerOrchestratorConfigBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/ContainerOrchestratorConfigBeanFactory.java new file mode 100644 index 0000000000000..aac3d7cf8e3d8 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/ContainerOrchestratorConfigBeanFactory.java @@ -0,0 +1,61 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.config.storage.CloudStorageConfigs; +import io.airbyte.workers.ContainerOrchestratorConfig; +import io.airbyte.workers.general.DocumentStoreClient; +import io.airbyte.workers.storage.StateClients; +import io.fabric8.kubernetes.client.DefaultKubernetesClient; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import io.micronaut.core.util.StringUtils; +import java.nio.file.Path; +import java.util.Optional; +import javax.inject.Named; +import javax.inject.Singleton; + +/** + * Micronaut bean factory for container orchestrator configuration-related singletons. + */ +@Factory +public class ContainerOrchestratorConfigBeanFactory { + + // IMPORTANT: Changing the storage location will orphan already existing kube pods when the new + // version is deployed! + private static final Path STATE_STORAGE_PREFIX = Path.of("/state"); + + @Singleton + @Requires(property = "airbyte.container.orchestrator.enabled", + value = "true") + @Named("containerOrchestratorConfig") + public ContainerOrchestratorConfig kubernetesContainerOrchestratorConfig( + @Named("stateStorageConfigs") final Optional cloudStateStorageConfiguration, + @Value("${airbyte.version}") final String airbyteVersion, + @Value("${airbyte.container.orchestrator.image}") final String containerOrchestratorImage, + @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String containerOrchestratorImagePullPolicy, + @Value("${airbyte.container.orchestrator.secret-mount-path}") final String containerOrchestratorSecretMountPath, + @Value("${airbyte.container.orchestrator.secret-name}") final String containerOrchestratorSecretName, + @Value("${google.application.credentials}") final String googleApplicationCredentials, + @Value("${airbyte.worker.job.kube.namespace}") final String namespace) { + final var kubernetesClient = new DefaultKubernetesClient(); + + final DocumentStoreClient documentStoreClient = StateClients.create( + cloudStateStorageConfiguration.orElse(null), + STATE_STORAGE_PREFIX); + + return new ContainerOrchestratorConfig( + namespace, + documentStoreClient, + kubernetesClient, + containerOrchestratorSecretName, + containerOrchestratorSecretMountPath, + StringUtils.isNotEmpty(containerOrchestratorImage) ? containerOrchestratorImage : "airbyte/container-orchestrator:" + airbyteVersion, + containerOrchestratorImagePullPolicy, + googleApplicationCredentials); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/DatabaseBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/DatabaseBeanFactory.java new file mode 100644 index 0000000000000..e99e0a09c8fb2 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/DatabaseBeanFactory.java @@ -0,0 +1,166 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.config.persistence.ConfigPersistence; +import io.airbyte.config.persistence.ConfigRepository; +import io.airbyte.config.persistence.DatabaseConfigPersistence; +import io.airbyte.config.persistence.StatePersistence; +import io.airbyte.config.persistence.StreamResetPersistence; +import io.airbyte.config.persistence.split_secrets.JsonSecretsProcessor; +import io.airbyte.db.Database; +import io.airbyte.db.check.DatabaseMigrationCheck; +import io.airbyte.db.check.impl.JobsDatabaseAvailabilityCheck; +import io.airbyte.db.factory.DatabaseCheckFactory; +import io.airbyte.db.instance.DatabaseConstants; +import io.airbyte.scheduler.persistence.DefaultJobPersistence; +import io.airbyte.scheduler.persistence.JobPersistence; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import io.micronaut.flyway.FlywayConfigurationProperties; +import java.io.IOException; +import javax.inject.Named; +import javax.inject.Singleton; +import javax.sql.DataSource; +import lombok.extern.slf4j.Slf4j; +import org.flywaydb.core.Flyway; +import org.jooq.DSLContext; + +/** + * Micronaut bean factory for database-related singletons. + */ +@Factory +@Slf4j +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class DatabaseBeanFactory { + + private static final String BASELINE_DESCRIPTION = "Baseline from file-based migration v1"; + private static final Boolean BASELINE_ON_MIGRATION = true; + private static final String INSTALLED_BY = "WorkerApp"; + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("configDatabase") + public Database configDatabase(@Named("config") final DSLContext dslContext) throws IOException { + return new Database(dslContext); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("jobsDatabase") + public Database jobsDatabase(@Named("jobs") final DSLContext dslContext) throws IOException { + return new Database(dslContext); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("configFlyway") + public Flyway configFlyway(@Named("config") final FlywayConfigurationProperties configFlywayConfigurationProperties, + @Named("config") final DataSource configDataSource, + @Value("${airbyte.flyway.configs.minimum-migration-version}") final String baselineVersion) { + return configFlywayConfigurationProperties.getFluentConfiguration() + .dataSource(configDataSource) + .baselineVersion(baselineVersion) + .baselineDescription(BASELINE_DESCRIPTION) + .baselineOnMigrate(BASELINE_ON_MIGRATION) + .installedBy(INSTALLED_BY) + .table(String.format("airbyte_%s_migrations", "configs")) + .load(); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("jobsFlyway") + public Flyway jobsFlyway(@Named("jobs") final FlywayConfigurationProperties jobsFlywayConfigurationProperties, + @Named("jobs") final DataSource jobsDataSource, + @Value("${airbyte.flyway.jobs.minimum-migration-version}") final String baselineVersion) { + return jobsFlywayConfigurationProperties.getFluentConfiguration() + .dataSource(jobsDataSource) + .baselineVersion(baselineVersion) + .baselineDescription(BASELINE_DESCRIPTION) + .baselineOnMigrate(BASELINE_ON_MIGRATION) + .installedBy(INSTALLED_BY) + .table(String.format("airbyte_%s_migrations", "jobs")) + .load(); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public ConfigPersistence configPersistence(@Named("configDatabase") final Database configDatabase, + final JsonSecretsProcessor jsonSecretsProcessor) { + return DatabaseConfigPersistence.createWithValidation(configDatabase, jsonSecretsProcessor); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public ConfigRepository configRepository(@Named("configPersistence") final ConfigPersistence configPersistence, + @Named("configDatabase") final Database configDatabase) { + return new ConfigRepository(configPersistence, configDatabase); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public JobPersistence jobPersistence(@Named("jobsDatabase") final Database jobDatabase) { + return new DefaultJobPersistence(jobDatabase); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public StatePersistence statePersistence(@Named("configDatabase") final Database configDatabase) { + return new StatePersistence(configDatabase); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public StreamResetPersistence streamResetPersistence(@Named("configDatabase") final Database configDatabase) { + return new StreamResetPersistence(configDatabase); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("configsDatabaseMigrationCheck") + public DatabaseMigrationCheck configsDatabaseMigrationCheck(@Named("config") final DSLContext dslContext, + @Named("configFlyway") final Flyway configsFlyway, + @Value("${airbyte.flyway.configs.minimum-migration-version}") final String configsDatabaseMinimumFlywayMigrationVersion, + @Value("${airbyte.flyway.configs.initialization-timeout-ms}") final Long configsDatabaseInitializationTimeoutMs) { + log.info("Configs database configuration: {} {}", configsDatabaseMinimumFlywayMigrationVersion, configsDatabaseInitializationTimeoutMs); + return DatabaseCheckFactory + .createConfigsDatabaseMigrationCheck(dslContext, configsFlyway, configsDatabaseMinimumFlywayMigrationVersion, + configsDatabaseInitializationTimeoutMs); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("jobsDatabaseMigrationCheck") + public DatabaseMigrationCheck jobsDatabaseMigrationCheck(@Named("jobs") final DSLContext dslContext, + @Named("jobsFlyway") final Flyway jobsFlyway, + @Value("${airbyte.flyway.jobs.minimum-migration-version}") final String jobsDatabaseMinimumFlywayMigrationVersion, + @Value("${airbyte.flyway.jobs.initialization-timeout-ms}") final Long jobsDatabaseInitializationTimeoutMs) { + return DatabaseCheckFactory + .createJobsDatabaseMigrationCheck(dslContext, jobsFlyway, jobsDatabaseMinimumFlywayMigrationVersion, + jobsDatabaseInitializationTimeoutMs); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("jobsDatabaseAvailabilityCheck") + public JobsDatabaseAvailabilityCheck jobsDatabaseAvailabilityCheck(@Named("jobs") final DSLContext dslContext) { + return new JobsDatabaseAvailabilityCheck(dslContext, DatabaseConstants.DEFAULT_ASSERT_DATABASE_TIMEOUT_MS); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/JobErrorReportingBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/JobErrorReportingBeanFactory.java new file mode 100644 index 0000000000000..63c561c6982d5 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/JobErrorReportingBeanFactory.java @@ -0,0 +1,79 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.config.Configs.DeploymentMode; +import io.airbyte.config.persistence.ConfigRepository; +import io.airbyte.scheduler.persistence.WebUrlHelper; +import io.airbyte.scheduler.persistence.job_error_reporter.JobErrorReporter; +import io.airbyte.scheduler.persistence.job_error_reporter.JobErrorReportingClient; +import io.airbyte.scheduler.persistence.job_error_reporter.LoggingJobErrorReportingClient; +import io.airbyte.scheduler.persistence.job_error_reporter.SentryExceptionHelper; +import io.airbyte.scheduler.persistence.job_error_reporter.SentryJobErrorReportingClient; +import io.airbyte.workers.normalization.NormalizationRunnerFactory; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import javax.inject.Named; +import javax.inject.Singleton; + +/** + * Micronaut bean factory for job error reporting-related singletons. + */ +@Factory +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class JobErrorReportingBeanFactory { + + @Singleton + @Requires(property = "airbyte.worker.job.error-reporting.strategy", + value = "SENTRY") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("jobErrorReportingClient") + public JobErrorReportingClient sentryJobErrorReportingClient( + @Value("${airbyte.worker.job.error-reporting.sentry.dsn}") final String sentryDsn) { + return new SentryJobErrorReportingClient(sentryDsn, new SentryExceptionHelper()); + } + + @Singleton + @Requires(property = "airbyte.worker.job.error-reporting.strategy", + value = "LOGGING") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("jobErrorReportingClient") + public JobErrorReportingClient loggingJobErrorReportingClient() { + return new LoggingJobErrorReportingClient(); + } + + @Singleton + @Requires(property = "airbyte.worker.job.error-reporting.strategy", + value = "") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("jobErrorReportingClient") + public JobErrorReportingClient defaultJobErrorReportingClient() { + return loggingJobErrorReportingClient(); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public JobErrorReporter jobErrorReporter( + @Value("${airbyte.version}") final String airbyteVersion, + final ConfigRepository configRepository, + final DeploymentMode deploymentMode, + @Named("jobErrorReportingClient") final JobErrorReportingClient jobErrorReportingClient, + final WebUrlHelper webUrlHelper) { + return new JobErrorReporter( + configRepository, + deploymentMode, + airbyteVersion, + NormalizationRunnerFactory.BASE_NORMALIZATION_IMAGE_NAME, + NormalizationRunnerFactory.NORMALIZATION_VERSION, + webUrlHelper, + jobErrorReportingClient); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/ProcessFactoryBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/ProcessFactoryBeanFactory.java new file mode 100644 index 0000000000000..199cc39ba03d7 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/ProcessFactoryBeanFactory.java @@ -0,0 +1,241 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.workers.WorkerConfigs; +import io.airbyte.workers.process.DockerProcessFactory; +import io.airbyte.workers.process.KubeProcessFactory; +import io.airbyte.workers.process.ProcessFactory; +import io.fabric8.kubernetes.client.DefaultKubernetesClient; +import io.fabric8.kubernetes.client.KubernetesClient; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import io.micronaut.core.util.StringUtils; +import java.net.InetAddress; +import java.net.UnknownHostException; +import java.nio.file.Path; +import javax.inject.Named; +import javax.inject.Singleton; + +/** + * Micronaut bean factory for process factory-related singletons. + */ +@Factory +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class ProcessFactoryBeanFactory { + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^(?!kubernetes$).*") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("checkProcessFactory") + public ProcessFactory checkDockerProcessFactory( + @Named("checkWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${docker.network}") final String dockerNetwork, + @Value("${airbyte.local.docker-mount}") final String localDockerMount, + @Value("${airbyte.local.root}") final String localRoot, + @Value("${airbyte.workspace.docker-mount}") final String workspaceDockerMount, + @Value("${airbyte.workspace.root}") final String workspaceRoot) { + return createDockerProcessFactory( + workerConfigs, + Path.of(workspaceRoot), + workspaceDockerMount, + localDockerMount, + localRoot, + dockerNetwork); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^kubernetes$") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("checkProcessFactory") + public ProcessFactory checkKubernetesProcessFactory( + @Named("checkWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${airbyte.worker.job.kube.namespace}") final String kubernetesNamespace, + @Value("${micronaut.server.port}") final Integer serverPort) + throws UnknownHostException { + return createKubernetesProcessFactory(workerConfigs, + kubernetesNamespace, + serverPort); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^(?!kubernetes$).*") + @Named("defaultProcessFactory") + public ProcessFactory defaultDockerProcessFactory( + @Named("defaultWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${docker.network}") final String dockerNetwork, + @Value("${airbyte.local.docker-mount}") final String localDockerMount, + @Value("${airbyte.local.root}") final String localRoot, + @Value("${airbyte.workspace.docker-mount}") final String workspaceDockerMount, + @Value("${airbyte.workspace.root}") final String workspaceRoot) { + return createDockerProcessFactory( + workerConfigs, + Path.of(workspaceRoot), + workspaceDockerMount, + localDockerMount, + localRoot, + dockerNetwork); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^kubernetes$") + @Named("defaultProcessFactory") + public ProcessFactory defaultKubernetesProcessFactory( + @Named("defaultWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${airbyte.worker.job.kube.namespace}") final String kubernetesNamespace, + @Value("${micronaut.server.port}") final Integer serverPort) + throws UnknownHostException { + return createKubernetesProcessFactory(workerConfigs, + kubernetesNamespace, + serverPort); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^(?!kubernetes$).*") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("discoverProcessFactory") + public ProcessFactory discoverDockerProcessFactory( + @Named("discoverWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${docker.network}") final String dockerNetwork, + @Value("${airbyte.local.docker-mount}") final String localDockerMount, + @Value("${airbyte.local.root}") final String localRoot, + @Value("${airbyte.workspace.docker-mount}") final String workspaceDockerMount, + @Value("${airbyte.workspace.root}") final String workspaceRoot) { + return createDockerProcessFactory( + workerConfigs, + Path.of(workspaceRoot), + workspaceDockerMount, + localDockerMount, + localRoot, + dockerNetwork); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^kubernetes$") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("discoverProcessFactory") + public ProcessFactory discoverKubernetesProcessFactory( + @Named("discoverWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${airbyte.worker.job.kube.namespace}") final String kubernetesNamespace, + @Value("${micronaut.server.port}") final Integer serverPort) + throws UnknownHostException { + return createKubernetesProcessFactory(workerConfigs, + kubernetesNamespace, + serverPort); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^(?!kubernetes$).*") + @Named("replicationProcessFactory") + public ProcessFactory replicationDockerProcessFactory( + @Named("replicationWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${docker.network}") final String dockerNetwork, + @Value("${airbyte.local.docker-mount}") final String localDockerMount, + @Value("${airbyte.local.root}") final String localRoot, + @Value("${airbyte.workspace.docker-mount}") final String workspaceDockerMount, + @Value("${airbyte.workspace.root}") final String workspaceRoot) { + return createDockerProcessFactory( + workerConfigs, + Path.of(workspaceRoot), + workspaceDockerMount, + localDockerMount, + localRoot, + dockerNetwork); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^kubernetes$") + @Named("replicationProcessFactory") + public ProcessFactory replicationKubernetesProcessFactory( + @Named("replicationWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${airbyte.worker.job.kube.namespace}") final String kubernetesNamespace, + @Value("${micronaut.server.port}") final Integer serverPort) + throws UnknownHostException { + return createKubernetesProcessFactory(workerConfigs, + kubernetesNamespace, + serverPort); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^(?!kubernetes$).*") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("specProcessFactory") + public ProcessFactory specDockerProcessFactory( + @Named("specWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${docker.network}") final String dockerNetwork, + @Value("${airbyte.local.docker-mount}") final String localDockerMount, + @Value("${airbyte.local.root}") final String localRoot, + @Value("${airbyte.workspace.docker-mount}") final String workspaceDockerMount, + @Value("${airbyte.workspace.root}") final String workspaceRoot) { + return createDockerProcessFactory( + workerConfigs, + Path.of(workspaceRoot), + workspaceDockerMount, + localDockerMount, + localRoot, + dockerNetwork); + } + + @Singleton + @Requires(property = "airbyte.worker.env", + pattern = "(?i)^kubernetes$") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("specProcessFactory") + public ProcessFactory specKubernetesProcessFactory( + @Named("specWorkerConfigs") final WorkerConfigs workerConfigs, + @Value("${airbyte.worker.job.kube.namespace}") final String kubernetesNamespace, + @Value("${micronaut.server.port}") final Integer serverPort) + throws UnknownHostException { + return createKubernetesProcessFactory(workerConfigs, + kubernetesNamespace, + serverPort); + } + + private ProcessFactory createDockerProcessFactory(final WorkerConfigs workerConfigs, + final Path workspaceRoot, + final String workspaceDockerMount, + final String localDockerMount, + final String localRoot, + final String dockerNetwork) { + return new DockerProcessFactory( + workerConfigs, + workspaceRoot, + StringUtils.isNotEmpty(workspaceDockerMount) ? workspaceDockerMount : workspaceRoot.toString(), + StringUtils.isNotEmpty(localDockerMount) ? localDockerMount : localRoot, + dockerNetwork); + } + + private ProcessFactory createKubernetesProcessFactory( + final WorkerConfigs workerConfigs, + final String kuberenetesNamespace, + final Integer serverPort) + throws UnknownHostException { + final KubernetesClient fabricClient = new DefaultKubernetesClient(); + final String localIp = InetAddress.getLocalHost().getHostAddress(); + final String kubeHeartbeatUrl = localIp + ":" + serverPort; + return new KubeProcessFactory(workerConfigs, + kuberenetesNamespace, + fabricClient, + kubeHeartbeatUrl, + false); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/SecretPersistenceBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/SecretPersistenceBeanFactory.java new file mode 100644 index 0000000000000..a01345ecf72fd --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/SecretPersistenceBeanFactory.java @@ -0,0 +1,88 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.config.persistence.split_secrets.GoogleSecretManagerPersistence; +import io.airbyte.config.persistence.split_secrets.LocalTestingSecretPersistence; +import io.airbyte.config.persistence.split_secrets.NoOpSecretsHydrator; +import io.airbyte.config.persistence.split_secrets.RealSecretsHydrator; +import io.airbyte.config.persistence.split_secrets.SecretPersistence; +import io.airbyte.config.persistence.split_secrets.SecretsHydrator; +import io.airbyte.config.persistence.split_secrets.VaultSecretPersistence; +import io.airbyte.db.Database; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import javax.inject.Named; +import javax.inject.Singleton; + +/** + * Micronaut bean factory for secret persistence-related singletons. + */ +@Factory +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class SecretPersistenceBeanFactory { + + @Singleton + @Requires(property = "airbyte.secret.persistence", + notEquals = "TESTING_CONFIG_DB_TABLE") + @Requires(property = "airbyte.secret.persistence", + notEquals = "GOOGLE_SECRET_MANAGER") + @Requires(property = "airbyte.secret.persistence", + notEquals = "VAULT") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("secretPersistence") + public SecretPersistence defaultSecretPersistence(@Named("configDatabase") final Database configDatabase) { + return localTestingSecretPersistence(configDatabase); + } + + @Singleton + @Requires(property = "airbyte.secret.persistence", + value = "TESTING_CONFIG_DB_TABLE") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("secretPersistence") + public SecretPersistence localTestingSecretPersistence(@Named("configDatabase") final Database configDatabase) { + return new LocalTestingSecretPersistence(configDatabase); + } + + @Singleton + @Requires(property = "airbyte.secret.persistence", + value = "GOOGLE_SECRET_MANAGER") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("secretPersistence") + public SecretPersistence googleSecretPersistence(@Value("${airbyte.secret.store.gcp.credentials}") final String credentials, + @Value("${airbyte.secret.store.gcp.project-id}") final String projectId) { + return GoogleSecretManagerPersistence.getLongLived(projectId, credentials); + } + + @Singleton + @Requires(property = "airbyte.secret.persistence", + value = "VAULT") + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("secretPersistence") + public SecretPersistence vaultSecretPersistence(@Value("${airbyte.secret.store.vault.address}") final String address, + @Value("${airbyte.secret.store.vault.prefix}") final String prefix, + @Value("${airbyte.secret.store.vault.token}") final String token) { + return new VaultSecretPersistence(address, prefix, token); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public SecretsHydrator secretsHydrator(@Named("secretPersistence") final SecretPersistence secretPersistence) { + return new RealSecretsHydrator(secretPersistence); + } + + @Singleton + @Requires(env = "data") + public SecretsHydrator secretsHydrator() { + return new NoOpSecretsHydrator(); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/TemporalBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/TemporalBeanFactory.java new file mode 100644 index 0000000000000..ea189c736f689 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/TemporalBeanFactory.java @@ -0,0 +1,92 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import io.airbyte.analytics.TrackingClient; +import io.airbyte.analytics.TrackingClientSingleton; +import io.airbyte.commons.features.FeatureFlags; +import io.airbyte.config.persistence.ConfigRepository; +import io.airbyte.scheduler.persistence.DefaultJobCreator; +import io.airbyte.scheduler.persistence.job_factory.DefaultSyncJobFactory; +import io.airbyte.scheduler.persistence.job_factory.OAuthConfigSupplier; +import io.airbyte.scheduler.persistence.job_factory.SyncJobFactory; +import io.airbyte.workers.run.TemporalWorkerRunFactory; +import io.airbyte.workers.temporal.TemporalClient; +import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.TemporalWorkflowUtils; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Property; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import io.temporal.client.WorkflowClient; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.worker.WorkerFactory; +import java.nio.file.Path; +import javax.inject.Singleton; + +/** + * Micronaut bean factory for Temporal-related singletons. + */ +@Factory +public class TemporalBeanFactory { + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public TrackingClient trackingClient() { + return TrackingClientSingleton.get(); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public SyncJobFactory jobFactory( + final ConfigRepository configRepository, + @Property(name = "airbyte.connector.specific-resource-defaults-enabled", + defaultValue = "false") final boolean connectorSpecificResourceDefaultsEnabled, + final DefaultJobCreator jobCreator, + final TrackingClient trackingClient) { + return new DefaultSyncJobFactory( + connectorSpecificResourceDefaultsEnabled, + jobCreator, + configRepository, + new OAuthConfigSupplier(configRepository, trackingClient)); + } + + @Singleton + public WorkflowServiceStubs temporalService(final TemporalUtils temporalUtils) { + return temporalUtils.createTemporalService(); + } + + @Singleton + public WorkflowClient workflowClient( + final TemporalUtils temporalUtils, + final WorkflowServiceStubs temporalService) { + return TemporalWorkflowUtils.createWorkflowClient(temporalService, temporalUtils.getNamespace()); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + public TemporalWorkerRunFactory temporalWorkerRunFactory( + @Value("${airbyte.version}") final String airbyteVersion, + final FeatureFlags featureFlags, + final TemporalClient temporalClient, + final TemporalUtils temporalUtils, + final WorkflowServiceStubs temporalService, + @Value("${airbyte.workspace.root}") final String workspaceRoot) { + return new TemporalWorkerRunFactory( + temporalClient, + Path.of(workspaceRoot), + airbyteVersion, + featureFlags); + } + + @Singleton + public WorkerFactory workerFactory(final WorkflowClient workflowClient) { + return WorkerFactory.newInstance(workflowClient); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java new file mode 100644 index 0000000000000..e5564644b1113 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/config/WorkerConfigurationBeanFactory.java @@ -0,0 +1,392 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.config; + +import com.google.common.base.Splitter; +import com.google.common.base.Strings; +import io.airbyte.commons.map.MoreMaps; +import io.airbyte.config.Configs.DeploymentMode; +import io.airbyte.config.Configs.WorkerEnvironment; +import io.airbyte.config.ResourceRequirements; +import io.airbyte.config.TolerationPOJO; +import io.airbyte.workers.WorkerConfigs; +import io.micronaut.context.annotation.Factory; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.inject.Named; +import javax.inject.Singleton; +import lombok.extern.slf4j.Slf4j; + +/** + * Micronaut bean factory for worker configuration-related singletons. + */ +@Factory +@Slf4j +@SuppressWarnings("PMD.AvoidDuplicateLiterals") +public class WorkerConfigurationBeanFactory { + + private static final String AIRBYTE_ROLE = "AIRBYTE_ROLE"; + private static final String AIRBYTE_VERSION = "AIRBYTE_VERSION"; + private static final String DEPLOYMENT_MODE = "DEPLOYMENT_MODE"; + private static final String JOB_DEFAULT_ENV_PREFIX = "JOB_DEFAULT_ENV_"; + private static final String WORKER_ENVIRONMENT = "WORKER_ENVIRONMENT"; + + @Singleton + @Named("checkJobKubeAnnotations") + public Map checkJobKubeAnnotations(@Value("${airbyte.worker.check.kube.annotations}") final String kubeAnnotations) { + return splitKVPairsFromEnvString(kubeAnnotations); + } + + @Singleton + @Named("checkJobKubeNodeSelectors") + public Map checkJobKubeNodeSelectors(@Value("${airbyte.worker.check.kube.node-selectors}") final String kubeNodeSelectors) { + return splitKVPairsFromEnvString(kubeNodeSelectors); + } + + @Singleton + @Named("discoverJobKubeAnnotations") + public Map discoverJobKubeAnnotations(@Value("${airbyte.worker.discover.kube.annotations}") final String kubeAnnotations) { + return splitKVPairsFromEnvString(kubeAnnotations); + } + + @Singleton + @Named("discoverJobKubeNodeSelectors") + public Map discoverJobKubeNodeSelectors(@Value("${airbyte.worker.discover.kube.node-selectors}") final String kubeNodeSelectors) { + return splitKVPairsFromEnvString(kubeNodeSelectors); + } + + @Singleton + @Named("defaultJobKubeAnnotations") + public Map jobKubeAnnotations(@Value("${airbyte.worker.job.kube.annotations}") final String kubeAnnotations) { + return splitKVPairsFromEnvString(kubeAnnotations); + } + + @Singleton + @Named("defaultJobKubeNodeSelectors") + public Map jobKubeNodeSelectors(@Value("${airbyte.worker.job.kube.node-selectors}") final String kubeNodeSelectors) { + return splitKVPairsFromEnvString(kubeNodeSelectors); + } + + @Singleton + @Named("specJobKubeAnnotations") + public Map specJobKubeAnnotations(@Value("${airbyte.worker.spec.kube.annotations}") final String kubeAnnotations) { + return splitKVPairsFromEnvString(kubeAnnotations); + } + + @Singleton + @Named("specJobKubeNodeSelectors") + public Map specJobKubeNodeSelectors(@Value("${airbyte.worker.spec.kube.node-selectors}") final String kubeNodeSelectors) { + return splitKVPairsFromEnvString(kubeNodeSelectors); + } + + @Singleton + @Named("jobDefaultEnvMap") + public Map jobDefaultEnvMap( + @Value("${airbyte.role}") final String airbyteRole, + @Value("${airbyte.version}") final String airbyteVersion, + final DeploymentMode deploymentMode, + final WorkerEnvironment workerEnvironment) { + final Map envMap = System.getenv(); + final Map jobPrefixedEnvMap = envMap.keySet().stream() + .filter(key -> key.startsWith(JOB_DEFAULT_ENV_PREFIX)) + .collect(Collectors.toMap(key -> key.replace(JOB_DEFAULT_ENV_PREFIX, ""), envMap::get)); + final Map jobSharedEnvMap = Map.of(AIRBYTE_ROLE, airbyteRole, + AIRBYTE_VERSION, airbyteVersion, + DEPLOYMENT_MODE, deploymentMode.name(), + WORKER_ENVIRONMENT, workerEnvironment.name()); + return MoreMaps.merge(jobPrefixedEnvMap, jobSharedEnvMap); + } + + /** + * Returns worker pod tolerations parsed from its own environment variable. The value of the env is + * a string that represents one or more tolerations. + *

    + *
  • Tolerations are separated by a `;` + *
  • Each toleration contains k=v pairs mentioning some/all of key, effect, operator and value and + * separated by `,` + *
+ *

+ * For example:- The following represents two tolerations, one checking existence and another + * matching a value + *

+ * key=airbyte-server,operator=Exists,effect=NoSchedule;key=airbyte-server,operator=Equals,value=true,effect=NoSchedule + * + * @return list of WorkerKubeToleration parsed from env + */ + @Singleton + public List jobKubeTolerations(@Value("${airbyte.worker.job.kube.tolerations}") final String jobKubeTolerations) { + final Stream tolerations = Strings.isNullOrEmpty(jobKubeTolerations) ? Stream.of() + : Splitter.on(";") + .splitToStream(jobKubeTolerations) + .filter(tolerationStr -> !Strings.isNullOrEmpty(tolerationStr)); + + return tolerations + .map(this::parseToleration) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + } + + @Singleton + @Named("checkResourceRequirements") + public ResourceRequirements checkResourceRequirements( + @Value("${airbyte.worker.check.main.container.cpu.request}") final String cpuRequest, + @Value("${airbyte.worker.check.main.container.cpu.limit}") final String cpuLimit, + @Value("${airbyte.worker.check.main.container.memory.request}") final String memoryRequest, + @Value("${airbyte.worker.check.main.container.memory.limit}") final String memoryLimit) { + return new ResourceRequirements() + .withCpuRequest(cpuRequest) + .withCpuLimit(cpuLimit) + .withMemoryRequest(memoryRequest) + .withMemoryLimit(memoryLimit); + } + + @Singleton + @Named("defaultResourceRequirements") + public ResourceRequirements defaultResourceRequirements( + @Value("${airbyte.worker.job.main.container.cpu.request}") final String cpuRequest, + @Value("${airbyte.worker.job.main.container.cpu.limit}") final String cpuLimit, + @Value("${airbyte.worker.job.main.container.memory.request}") final String memoryRequest, + @Value("${airbyte.worker.job.main.container.memory.limit}") final String memoryLimit) { + return new ResourceRequirements() + .withCpuRequest(cpuRequest) + .withCpuLimit(cpuLimit) + .withMemoryRequest(memoryRequest) + .withMemoryLimit(memoryLimit); + } + + @Singleton + @Named("normalizationResourceRequirements") + public ResourceRequirements normalizationResourceRequirements( + @Value("${airbyte.worker.normalization.main.container.cpu.request}") final String cpuRequest, + @Value("${airbyte.worker.normalization.main.container.cpu.limit}") final String cpuLimit, + @Value("${airbyte.worker.normalization.main.container.memory.request}") final String memoryRequest, + @Value("${airbyte.worker.normalization.main.container.memory.limit}") final String memoryLimit, + @Value("${airbyte.worker.job.main.container.cpu.request}") final String defaultCpuRequest, + @Value("${airbyte.worker.job.main.container.cpu.limit}") final String defaultCpuLimit, + @Value("${airbyte.worker.job.main.container.memory.request}") final String defaultMemoryRequest, + @Value("${airbyte.worker.job.main.container.memory.limit}") final String defaultMemoryLimit) { + return new ResourceRequirements() + .withCpuRequest(Optional.ofNullable(cpuRequest).orElse(defaultCpuRequest)) + .withCpuLimit(Optional.ofNullable(cpuLimit).orElse(defaultCpuLimit)) + .withMemoryRequest(Optional.ofNullable(memoryRequest).orElse(defaultMemoryRequest)) + .withMemoryLimit(Optional.ofNullable(memoryLimit).orElse(defaultMemoryLimit)); + } + + @Singleton + @Named("replicationResourceRequirements") + public ResourceRequirements replicationResourceRequirements( + @Value("${airbyte.worker.replication.orchestrator.cpu.request}") final String cpuRequest, + @Value("${airbyte.worker.replication.orchestrator.cpu.limit}") final String cpuLimit, + @Value("${airbyte.worker.replication.orchestrator.memory.request}") final String memoryRequest, + @Value("${airbyte.worker.replication.orchestrator.memory.limit}") final String memoryLimit) { + return new ResourceRequirements() + .withCpuRequest(cpuRequest) + .withCpuLimit(cpuLimit) + .withMemoryRequest(memoryRequest) + .withMemoryLimit(memoryLimit); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("checkWorkerConfigs") + public WorkerConfigs checkWorkerConfigs( + final WorkerEnvironment workerEnvironment, + @Named("checkResourceRequirements") final ResourceRequirements resourceRequirements, + final List jobKubeTolerations, + @Named("checkJobKubeNodeSelectors") final Map nodeSelectors, + @Named("checkJobKubeAnnotations") final Map annotations, + @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, + @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.images.socat}") final String socatImage, + @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, + @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, + @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + return new WorkerConfigs( + workerEnvironment, + resourceRequirements, + jobKubeTolerations, + nodeSelectors, + annotations, + mainContainerImagePullSecret, + mainContainerImagePullPolicy, + sidecarContainerImagePullPolicy, + socatImage, + busyboxImage, + curlImage, + jobDefaultEnvMap); + } + + @Singleton + @Named("defaultWorkerConfigs") + public WorkerConfigs defaultWorkerConfigs( + final WorkerEnvironment workerEnvironment, + @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, + final List jobKubeTolerations, + @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, + @Named("defaultJobKubeAnnotations") final Map annotations, + @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, + @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.images.socat}") final String socatImage, + @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, + @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, + @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + return new WorkerConfigs( + workerEnvironment, + resourceRequirements, + jobKubeTolerations, + nodeSelectors, + annotations, + mainContainerImagePullSecret, + mainContainerImagePullPolicy, + sidecarContainerImagePullPolicy, + socatImage, + busyboxImage, + curlImage, + jobDefaultEnvMap); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("discoverWorkerConfigs") + public WorkerConfigs discoverWorkerConfigs( + final WorkerEnvironment workerEnvironment, + @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, + final List jobKubeTolerations, + @Named("discoverJobKubeNodeSelectors") final Map nodeSelectors, + @Named("discoverJobKubeAnnotations") final Map annotations, + @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, + @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.images.socat}") final String socatImage, + @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, + @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, + @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + return new WorkerConfigs( + workerEnvironment, + resourceRequirements, + jobKubeTolerations, + nodeSelectors, + annotations, + mainContainerImagePullSecret, + mainContainerImagePullPolicy, + sidecarContainerImagePullPolicy, + socatImage, + busyboxImage, + curlImage, + jobDefaultEnvMap); + } + + @Singleton + @Named("replicationWorkerConfigs") + public WorkerConfigs replicationWorkerConfigs( + final WorkerEnvironment workerEnvironment, + @Named("replicationResourceRequirements") final ResourceRequirements resourceRequirements, + final List jobKubeTolerations, + @Named("defaultJobKubeNodeSelectors") final Map nodeSelectors, + @Named("defaultJobKubeAnnotations") final Map annotations, + @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, + @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.images.socat}") final String socatImage, + @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, + @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, + @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + return new WorkerConfigs( + workerEnvironment, + resourceRequirements, + jobKubeTolerations, + nodeSelectors, + annotations, + mainContainerImagePullSecret, + mainContainerImagePullPolicy, + sidecarContainerImagePullPolicy, + socatImage, + busyboxImage, + curlImage, + jobDefaultEnvMap); + } + + @Singleton + @Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") + @Named("specWorkerConfigs") + public WorkerConfigs specWorkerConfigs( + final WorkerEnvironment workerEnvironment, + @Named("defaultResourceRequirements") final ResourceRequirements resourceRequirements, + final List jobKubeTolerations, + @Named("specJobKubeNodeSelectors") final Map nodeSelectors, + @Named("specJobKubeAnnotations") final Map annotations, + @Value("${airbyte.worker.job.kube.main.container.image-pull-secret}") final String mainContainerImagePullSecret, + @Value("${airbyte.worker.job.kube.main.container.image-pull-policy}") final String mainContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.sidecar.container.image-pull-policy}") final String sidecarContainerImagePullPolicy, + @Value("${airbyte.worker.job.kube.images.socat}") final String socatImage, + @Value("${airbyte.worker.job.kube.images.busybox}") final String busyboxImage, + @Value("${airbyte.worker.job.kube.images.curl}") final String curlImage, + @Named("jobDefaultEnvMap") final Map jobDefaultEnvMap) { + return new WorkerConfigs( + workerEnvironment, + resourceRequirements, + jobKubeTolerations, + nodeSelectors, + annotations, + mainContainerImagePullSecret, + mainContainerImagePullPolicy, + sidecarContainerImagePullPolicy, + socatImage, + busyboxImage, + curlImage, + jobDefaultEnvMap); + } + + /** + * Splits key value pairs from the input string into a map. Each kv-pair is separated by a ','. The + * key and the value are separated by '='. + *

+ * For example:- The following represents two map entries + *

+ * key1=value1,key2=value2 + * + * @param input string + * @return map containing kv pairs + */ + private Map splitKVPairsFromEnvString(final String input) { + return Splitter.on(",") + .splitToStream(input) + .filter(s -> !Strings.isNullOrEmpty(s) && s.contains("=")) + .map(s -> s.split("=")) + .collect(Collectors.toMap(s -> s[0].trim(), s -> s[1].trim())); + } + + private TolerationPOJO parseToleration(final String tolerationStr) { + final Map tolerationMap = Splitter.on(",") + .splitToStream(tolerationStr) + .map(s -> s.split("=")) + .collect(Collectors.toMap(s -> s[0], s -> s[1])); + + if (tolerationMap.containsKey("key") && tolerationMap.containsKey("effect") && tolerationMap.containsKey("operator")) { + return new TolerationPOJO( + tolerationMap.get("key"), + tolerationMap.get("effect"), + tolerationMap.get("value"), + tolerationMap.get("operator")); + } else { + log.warn( + "Ignoring toleration {}, missing one of key,effect or operator", + tolerationStr); + return null; + } + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/controller/HeartbeatController.java b/airbyte-workers/src/main/java/io/airbyte/workers/controller/HeartbeatController.java new file mode 100644 index 0000000000000..5dd650554f5d8 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/controller/HeartbeatController.java @@ -0,0 +1,51 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.controller; + +import io.micronaut.http.HttpHeaders; +import io.micronaut.http.HttpResponse; +import io.micronaut.http.MediaType; +import io.micronaut.http.MutableHttpResponse; +import io.micronaut.http.annotation.Controller; +import io.micronaut.http.annotation.Get; +import io.micronaut.http.annotation.Options; +import io.micronaut.http.annotation.Post; +import java.util.Map; + +/** + * Heartbeat controller + */ +@Controller("/") +public class HeartbeatController { + + private static final Map CORS_FILTER_MAP = Map.of( + HttpHeaders.ACCESS_CONTROL_ALLOW_ORIGIN, "*", + HttpHeaders.ACCESS_CONTROL_ALLOW_HEADERS, "Origin, Content-Type, Accept, Content-Encoding", + HttpHeaders.ACCESS_CONTROL_ALLOW_METHODS, "GET, POST, PUT, DELETE, OPTIONS, HEAD"); + + private final static Map DEFAULT_RESPONSE_BODY = Map.of("up", true); + + @Get(produces = MediaType.APPLICATION_JSON) + @Post(produces = MediaType.APPLICATION_JSON) + public HttpResponse> heartbeat() { + final MutableHttpResponse> response = HttpResponse.ok(DEFAULT_RESPONSE_BODY); + addCorsHeaders(response); + return response; + } + + @Options + public HttpResponse> emptyHeartbeat() { + final MutableHttpResponse> response = HttpResponse.ok(); + addCorsHeaders(response); + return response; + } + + private void addCorsHeaders(final MutableHttpResponse response) { + for (final Map.Entry entry : CORS_FILTER_MAP.entrySet()) { + response.header(entry.getKey(), entry.getValue()); + } + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/helper/ConnectionHelper.java b/airbyte-workers/src/main/java/io/airbyte/workers/helper/ConnectionHelper.java index e2280ab8edf15..821ec9921b643 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/helper/ConnectionHelper.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/helper/ConnectionHelper.java @@ -17,18 +17,26 @@ import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.scheduler.persistence.WorkspaceHelper; import io.airbyte.validation.json.JsonValidationException; +import io.micronaut.context.annotation.Requires; import java.io.IOException; import java.util.HashSet; import java.util.Set; import java.util.UUID; +import javax.inject.Inject; +import javax.inject.Singleton; import lombok.AllArgsConstructor; // todo (cgardens) - we are not getting any value out of instantiating this class. we should just // use it as statics. not doing it now, because already in the middle of another refactor. @AllArgsConstructor +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class ConnectionHelper { - private final ConfigRepository configRepository; + @Inject + private ConfigRepository configRepository; + @Inject private final WorkspaceHelper workspaceHelper; public void deleteConnection(final UUID connectionId) throws JsonValidationException, ConfigNotFoundException, IOException { @@ -140,7 +148,7 @@ public static void validateWorkspace(final WorkspaceHelper workspaceHelper, } // Helper method to convert between TimeUnit enums for old and new schedule schemas. - private static BasicSchedule.TimeUnit convertTimeUnitSchema(Schedule.TimeUnit timeUnit) { + private static BasicSchedule.TimeUnit convertTimeUnitSchema(final Schedule.TimeUnit timeUnit) { switch (timeUnit) { case MINUTES: return BasicSchedule.TimeUnit.MINUTES; diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/process/AsyncOrchestratorPodProcess.java b/airbyte-workers/src/main/java/io/airbyte/workers/process/AsyncOrchestratorPodProcess.java index b17583e613fab..9edb75417bc05 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/process/AsyncOrchestratorPodProcess.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/process/AsyncOrchestratorPodProcess.java @@ -10,7 +10,6 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.ResourceRequirements; import io.airbyte.config.helpers.LogClientSingleton; -import io.airbyte.workers.WorkerApp; import io.airbyte.workers.general.DocumentStoreClient; import io.fabric8.kubernetes.api.model.ContainerBuilder; import io.fabric8.kubernetes.api.model.ContainerPort; @@ -62,6 +61,7 @@ public class AsyncOrchestratorPodProcess implements KubePod { private final String googleApplicationCredentials; private final AtomicReference> cachedExitValue; private final boolean useStreamCapableState; + private final Integer serverPort; public AsyncOrchestratorPodProcess( final KubePodInfo kubePodInfo, @@ -70,7 +70,8 @@ public AsyncOrchestratorPodProcess( final String secretName, final String secretMountPath, final String googleApplicationCredentials, - final boolean useStreamCapableState) { + final boolean useStreamCapableState, + final Integer serverPort) { this.kubePodInfo = kubePodInfo; this.documentStoreClient = documentStoreClient; this.kubernetesClient = kubernetesClient; @@ -79,6 +80,7 @@ public AsyncOrchestratorPodProcess( this.googleApplicationCredentials = googleApplicationCredentials; this.cachedExitValue = new AtomicReference<>(Optional.empty()); this.useStreamCapableState = useStreamCapableState; + this.serverPort = serverPort; } public Optional getOutput() { @@ -287,7 +289,7 @@ public void create(final Map allLabels, envVars.add(new EnvVar(EnvConfigs.PUBLISH_METRICS, Boolean.toString(envConfigs.getPublishMetrics()), null)); envVars.add(new EnvVar(EnvVariableFeatureFlags.USE_STREAM_CAPABLE_STATE, Boolean.toString(useStreamCapableState), null)); final List containerPorts = KubePodProcess.createContainerPortList(portMap); - containerPorts.add(new ContainerPort(WorkerApp.KUBE_HEARTBEAT_PORT, null, null, null, null)); + containerPorts.add(new ContainerPort(serverPort, null, null, null, null)); final var mainContainer = new ContainerBuilder() .withName(KubePodProcess.MAIN_CONTAINER_NAME) diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/CancellationHandler.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/CancellationHandler.java index 676369650b0c8..aad52655f85b1 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/CancellationHandler.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/CancellationHandler.java @@ -19,7 +19,7 @@ class TemporalCancellationHandler implements CancellationHandler { private final ActivityExecutionContext activityContext; - public TemporalCancellationHandler(ActivityExecutionContext activityContext) { + public TemporalCancellationHandler(final ActivityExecutionContext activityContext) { this.activityContext = activityContext; } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/ConnectionManagerUtils.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/ConnectionManagerUtils.java index a6cdb57e1d084..f468ea5275fe1 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/ConnectionManagerUtils.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/ConnectionManagerUtils.java @@ -22,22 +22,18 @@ import java.util.Optional; import java.util.UUID; import java.util.function.Function; +import javax.inject.Singleton; +import lombok.NoArgsConstructor; import lombok.extern.slf4j.Slf4j; /** * Encapsulates logic specific to retrieving, starting, and signaling the ConnectionManagerWorkflow. */ +@NoArgsConstructor +@Singleton @Slf4j public class ConnectionManagerUtils { - private static final ConnectionManagerUtils instance = new ConnectionManagerUtils(); - - private ConnectionManagerUtils() {} - - public static ConnectionManagerUtils getInstance() { - return instance; - } - /** * Attempts to send a signal to the existing ConnectionManagerWorkflow for the provided connection. * @@ -52,9 +48,9 @@ public static ConnectionManagerUtils getInstance() { * @return the healthy connection manager workflow that was signaled * @throws DeletedWorkflowException if the connection manager workflow was deleted */ - ConnectionManagerWorkflow signalWorkflowAndRepairIfNecessary(final WorkflowClient client, - final UUID connectionId, - final Function signalMethod) + public ConnectionManagerWorkflow signalWorkflowAndRepairIfNecessary(final WorkflowClient client, + final UUID connectionId, + final Function signalMethod) throws DeletedWorkflowException { return signalWorkflowAndRepairIfNecessary(client, connectionId, signalMethod, Optional.empty()); } @@ -74,10 +70,10 @@ ConnectionManagerWorkflow signalWorkflowAndRepairIfNecessary(final WorkflowClien * @return the healthy connection manager workflow that was signaled * @throws DeletedWorkflowException if the connection manager workflow was deleted */ - ConnectionManagerWorkflow signalWorkflowAndRepairIfNecessary(final WorkflowClient client, - final UUID connectionId, - final Function> signalMethod, - final T signalArgument) + public ConnectionManagerWorkflow signalWorkflowAndRepairIfNecessary(final WorkflowClient client, + final UUID connectionId, + final Function> signalMethod, + final T signalArgument) throws DeletedWorkflowException { return signalWorkflowAndRepairIfNecessary(client, connectionId, signalMethod, Optional.of(signalArgument)); } @@ -116,7 +112,7 @@ private ConnectionManagerWorkflow signalWorkflowAndRepairIfNecessary(final W safeTerminateWorkflow(client, connectionId, "Terminating workflow in unreachable state before starting a new workflow for this connection"); final ConnectionManagerWorkflow connectionManagerWorkflow = newConnectionManagerWorkflowStub(client, connectionId); - final ConnectionUpdaterInput startWorkflowInput = buildStartWorkflowInput(connectionId); + final ConnectionUpdaterInput startWorkflowInput = TemporalWorkflowUtils.buildStartWorkflowInput(connectionId); final BatchRequest batchRequest = client.newSignalWithStartRequest(); batchRequest.add(connectionManagerWorkflow::run, startWorkflowInput); @@ -149,13 +145,13 @@ void safeTerminateWorkflow(final WorkflowClient client, final String workflowId, } } - void safeTerminateWorkflow(final WorkflowClient client, final UUID connectionId, final String reason) { + public void safeTerminateWorkflow(final WorkflowClient client, final UUID connectionId, final String reason) { safeTerminateWorkflow(client, getConnectionManagerName(connectionId), reason); } - ConnectionManagerWorkflow startConnectionManagerNoSignal(final WorkflowClient client, final UUID connectionId) { + public ConnectionManagerWorkflow startConnectionManagerNoSignal(final WorkflowClient client, final UUID connectionId) { final ConnectionManagerWorkflow connectionManagerWorkflow = newConnectionManagerWorkflowStub(client, connectionId); - final ConnectionUpdaterInput input = buildStartWorkflowInput(connectionId); + final ConnectionUpdaterInput input = TemporalWorkflowUtils.buildStartWorkflowInput(connectionId); WorkflowClient.start(connectionManagerWorkflow::run, input); return connectionManagerWorkflow; @@ -169,7 +165,7 @@ ConnectionManagerWorkflow startConnectionManagerNoSignal(final WorkflowClient cl * @throws DeletedWorkflowException if the workflow was deleted, according to the workflow state * @throws UnreachableWorkflowException if the workflow is in an unreachable state */ - ConnectionManagerWorkflow getConnectionManagerWorkflow(final WorkflowClient client, final UUID connectionId) + public ConnectionManagerWorkflow getConnectionManagerWorkflow(final WorkflowClient client, final UUID connectionId) throws DeletedWorkflowException, UnreachableWorkflowException { final ConnectionManagerWorkflow connectionManagerWorkflow; @@ -215,7 +211,7 @@ boolean isWorkflowStateRunning(final WorkflowClient client, final UUID connectio } } - WorkflowExecutionStatus getConnectionManagerWorkflowStatus(final WorkflowClient workflowClient, final UUID connectionId) { + public WorkflowExecutionStatus getConnectionManagerWorkflowStatus(final WorkflowClient workflowClient, final UUID connectionId) { final DescribeWorkflowExecutionRequest describeWorkflowExecutionRequest = DescribeWorkflowExecutionRequest.newBuilder() .setExecution(WorkflowExecution.newBuilder() .setWorkflowId(getConnectionManagerName(connectionId)) @@ -228,7 +224,7 @@ WorkflowExecutionStatus getConnectionManagerWorkflowStatus(final WorkflowClient return describeWorkflowExecutionResponse.getWorkflowExecutionInfo().getStatus(); } - long getCurrentJobId(final WorkflowClient client, final UUID connectionId) { + public long getCurrentJobId(final WorkflowClient client, final UUID connectionId) { try { final ConnectionManagerWorkflow connectionManagerWorkflow = getConnectionManagerWorkflow(client, connectionId); return connectionManagerWorkflow.getJobInformation().getJobId(); @@ -237,26 +233,13 @@ long getCurrentJobId(final WorkflowClient client, final UUID connectionId) { } } - ConnectionManagerWorkflow newConnectionManagerWorkflowStub(final WorkflowClient client, final UUID connectionId) { + public ConnectionManagerWorkflow newConnectionManagerWorkflowStub(final WorkflowClient client, final UUID connectionId) { return client.newWorkflowStub(ConnectionManagerWorkflow.class, - TemporalUtils.getWorkflowOptionsWithWorkflowId(TemporalJobType.CONNECTION_UPDATER, getConnectionManagerName(connectionId))); + TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.CONNECTION_UPDATER, getConnectionManagerName(connectionId))); } - String getConnectionManagerName(final UUID connectionId) { + public String getConnectionManagerName(final UUID connectionId) { return "connection_manager_" + connectionId; } - public ConnectionUpdaterInput buildStartWorkflowInput(final UUID connectionId) { - return ConnectionUpdaterInput.builder() - .connectionId(connectionId) - .jobId(null) - .attemptId(null) - .fromFailure(false) - .attemptNumber(1) - .workflowState(null) - .resetConnection(false) - .fromJobResetFailure(false) - .build(); - } - } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/StreamResetRecordsHelper.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/StreamResetRecordsHelper.java new file mode 100644 index 0000000000000..da2bb97b5d08e --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/StreamResetRecordsHelper.java @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal; + +import io.airbyte.config.JobConfig.ConfigType; +import io.airbyte.config.persistence.StreamResetPersistence; +import io.airbyte.protocol.models.StreamDescriptor; +import io.airbyte.scheduler.models.Job; +import io.airbyte.scheduler.persistence.JobPersistence; +import io.airbyte.workers.temporal.exception.RetryableException; +import java.io.IOException; +import java.util.List; +import java.util.UUID; +import javax.inject.Inject; +import javax.inject.Singleton; +import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +/** + * Helper class that provides methods for dealing with stream reset records. + */ +@Singleton +@AllArgsConstructor +@NoArgsConstructor +@Slf4j +public class StreamResetRecordsHelper { + + @Inject + private JobPersistence jobPersistence; + + @Inject + private StreamResetPersistence streamResetPersistence; + + /** + * Deletes all stream reset records related to the provided job and connection. + * + * @param jobId The job ID. + * @param connectionId the connection ID. + */ + public void deleteStreamResetRecordsForJob(final Long jobId, final UUID connectionId) { + if (jobId == null) { + log.info("deleteStreamResetRecordsForJob was called with a null job id; returning."); + return; + } + + try { + final Job job = jobPersistence.getJob(jobId); + final ConfigType configType = job.getConfig().getConfigType(); + if (!ConfigType.RESET_CONNECTION.equals(configType)) { + log.info("deleteStreamResetRecordsForJob was called for job {} with config type {}. Returning, as config type is not {}.", + jobId, + configType, + ConfigType.RESET_CONNECTION); + return; + } + + final List resetStreams = job.getConfig().getResetConnection().getResetSourceConfiguration().getStreamsToReset(); + log.info("Deleting the following streams for reset job {} from the stream_reset table: {}", jobId, resetStreams); + streamResetPersistence.deleteStreamResets(connectionId, resetStreams); + } catch (final IOException e) { + throw new RetryableException(e); + } + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalClient.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalClient.java index 595ee5480433f..021d95170b2d3 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalClient.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalClient.java @@ -29,6 +29,7 @@ import io.airbyte.workers.temporal.scheduling.ConnectionManagerWorkflow; import io.airbyte.workers.temporal.spec.SpecWorkflow; import io.airbyte.workers.temporal.sync.SyncWorkflow; +import io.micronaut.context.annotation.Requires; import io.temporal.api.common.v1.WorkflowType; import io.temporal.api.enums.v1.WorkflowExecutionStatus; import io.temporal.api.workflowservice.v1.ListOpenWorkflowExecutionsRequest; @@ -52,47 +53,45 @@ import java.util.concurrent.TimeoutException; import java.util.function.Supplier; import java.util.stream.Collectors; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +import lombok.AllArgsConstructor; import lombok.Builder; +import lombok.NoArgsConstructor; import lombok.Value; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.time.StopWatch; +@AllArgsConstructor +@NoArgsConstructor @Slf4j +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class TemporalClient { - private final Path workspaceRoot; - private final WorkflowClient client; - private final WorkflowServiceStubs service; - private final StreamResetPersistence streamResetPersistence; - private final ConnectionManagerUtils connectionManagerUtils; - /** - * This is use to sleep between 2 temporal queries. The query are needed to ensure that the cancel + * This is used to sleep between 2 temporal queries. The query is needed to ensure that the cancel * and start manual sync methods wait before returning. Since temporal signals are async, we need to * use the queries to make sure that we are in a state in which we want to continue with. */ private static final int DELAY_BETWEEN_QUERY_MS = 10; - public TemporalClient(final WorkflowClient client, - final Path workspaceRoot, - final WorkflowServiceStubs workflowServiceStubs, - final StreamResetPersistence streamResetPersistence) { - this(client, workspaceRoot, workflowServiceStubs, streamResetPersistence, ConnectionManagerUtils.getInstance()); - } - - @VisibleForTesting - TemporalClient(final WorkflowClient client, - final Path workspaceRoot, - final WorkflowServiceStubs workflowServiceStubs, - final StreamResetPersistence streamResetPersistence, - final ConnectionManagerUtils connectionManagerUtils) { - this.client = client; - this.workspaceRoot = workspaceRoot; - this.service = workflowServiceStubs; - this.streamResetPersistence = streamResetPersistence; - this.connectionManagerUtils = connectionManagerUtils; - } + @Inject + @Named("workspaceRoot") + private Path workspaceRoot; + @Inject + private WorkflowClient client; + @Inject + private WorkflowServiceStubs service; + @Inject + private StreamResetPersistence streamResetPersistence; + @Inject + private ConnectionManagerUtils connectionManagerUtils; + @Inject + private StreamResetRecordsHelper streamResetRecordsHelper; /** * Direct termination of Temporal Workflows should generally be avoided. This method exists for some @@ -104,7 +103,7 @@ public void dangerouslyTerminateWorkflow(final String workflowId, final String r } public TemporalResponse submitGetSpec(final UUID jobId, final int attempt, final JobGetSpecConfig config) { - final JobRunConfig jobRunConfig = TemporalUtils.createJobRunConfig(jobId, attempt); + final JobRunConfig jobRunConfig = TemporalWorkflowUtils.createJobRunConfig(jobId, attempt); final IntegrationLauncherConfig launcherConfig = new IntegrationLauncherConfig() .withJobId(jobId.toString()) @@ -118,7 +117,7 @@ public TemporalResponse submitGetSpec(final UUID jobId, fina public TemporalResponse submitCheckConnection(final UUID jobId, final int attempt, final JobCheckConnectionConfig config) { - final JobRunConfig jobRunConfig = TemporalUtils.createJobRunConfig(jobId, attempt); + final JobRunConfig jobRunConfig = TemporalWorkflowUtils.createJobRunConfig(jobId, attempt); final IntegrationLauncherConfig launcherConfig = new IntegrationLauncherConfig() .withJobId(jobId.toString()) .withAttemptId((long) attempt) @@ -132,7 +131,7 @@ public TemporalResponse submitCheckConnection(final UUID job public TemporalResponse submitDiscoverSchema(final UUID jobId, final int attempt, final JobDiscoverCatalogConfig config) { - final JobRunConfig jobRunConfig = TemporalUtils.createJobRunConfig(jobId, attempt); + final JobRunConfig jobRunConfig = TemporalWorkflowUtils.createJobRunConfig(jobId, attempt); final IntegrationLauncherConfig launcherConfig = new IntegrationLauncherConfig() .withJobId(jobId.toString()) .withAttemptId((long) attempt) @@ -144,7 +143,7 @@ public TemporalResponse submitDiscoverSchema(final UUID jobI } public TemporalResponse submitSync(final long jobId, final int attempt, final JobSyncConfig config, final UUID connectionId) { - final JobRunConfig jobRunConfig = TemporalUtils.createJobRunConfig(jobId, attempt); + final JobRunConfig jobRunConfig = TemporalWorkflowUtils.createJobRunConfig(jobId, attempt); final IntegrationLauncherConfig sourceLauncherConfig = new IntegrationLauncherConfig() .withJobId(String.valueOf(jobId)) @@ -373,6 +372,8 @@ public ManualOperationResult startNewCancellation(final UUID connectionId) { } } while (connectionManagerUtils.isWorkflowStateRunning(client, connectionId)); + streamResetRecordsHelper.deleteStreamResetRecordsForJob(jobId, connectionId); + log.info("end of manual cancellation"); return new ManualOperationResult( @@ -531,7 +532,7 @@ Set filterOutRunningWorkspaceId(final Set workflowIds) { } private T getWorkflowStub(final Class workflowClass, final TemporalJobType jobType) { - return client.newWorkflowStub(workflowClass, TemporalUtils.getWorkflowOptions(jobType)); + return client.newWorkflowStub(workflowClass, TemporalWorkflowUtils.buildWorkflowOptions(jobType)); } private boolean getConnectorJobSucceeded(final ConnectorJobOutput output) { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalUtils.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalUtils.java index a19a4f186c4c5..5c9d4add2adc1 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalUtils.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalUtils.java @@ -4,16 +4,14 @@ package io.airbyte.workers.temporal; -import com.google.common.annotations.VisibleForTesting; import com.uber.m3.tally.RootScopeBuilder; import com.uber.m3.tally.Scope; import com.uber.m3.tally.StatsReporter; import io.airbyte.commons.lang.Exceptions; -import io.airbyte.config.Configs; -import io.airbyte.config.EnvConfigs; import io.airbyte.metrics.lib.MetricClientFactory; -import io.airbyte.scheduler.models.JobRunConfig; import io.micrometer.core.instrument.MeterRegistry; +import io.micronaut.context.annotation.Property; +import io.micronaut.context.annotation.Value; import io.temporal.activity.ActivityExecutionContext; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.namespace.v1.NamespaceConfig; @@ -22,8 +20,6 @@ import io.temporal.api.workflowservice.v1.UpdateNamespaceRequest; import io.temporal.client.ActivityCompletionException; import io.temporal.client.WorkflowClient; -import io.temporal.client.WorkflowClientOptions; -import io.temporal.client.WorkflowOptions; import io.temporal.client.WorkflowStub; import io.temporal.common.RetryOptions; import io.temporal.common.reporter.MicrometerClientStatsReporter; @@ -44,34 +40,50 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; +import javax.inject.Singleton; import javax.net.ssl.SSLException; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.time.DurationFormatUtils; import org.apache.commons.lang3.tuple.ImmutablePair; @Slf4j +@Singleton public class TemporalUtils { - private static final Configs configs = new EnvConfigs(); - private static final Duration WORKFLOW_EXECUTION_TTL = Duration.ofDays(configs.getTemporalRetentionInDays()); private static final Duration WAIT_INTERVAL = Duration.ofSeconds(2); private static final Duration MAX_TIME_TO_CONNECT = Duration.ofMinutes(2); private static final Duration WAIT_TIME_AFTER_CONNECT = Duration.ofSeconds(5); - private static final String HUMAN_READABLE_WORKFLOW_EXECUTION_TTL = - DurationFormatUtils.formatDurationWords(WORKFLOW_EXECUTION_TTL.toMillis(), true, true); - public static final String DEFAULT_NAMESPACE = "default"; public static final Duration SEND_HEARTBEAT_INTERVAL = Duration.ofSeconds(10); public static final Duration HEARTBEAT_TIMEOUT = Duration.ofSeconds(30); public static final RetryOptions NO_RETRY = RetryOptions.newBuilder().setMaximumAttempts(1).build(); - public static final RetryOptions RETRY = RetryOptions.newBuilder() - .setMaximumAttempts(configs.getActivityNumberOfAttempt()) - .setInitialInterval(Duration.ofSeconds(configs.getInitialDelayBetweenActivityAttemptsSeconds())) - .setMaximumInterval(Duration.ofSeconds(configs.getMaxDelayBetweenActivityAttemptsSeconds())) - .build(); private static final double REPORT_INTERVAL_SECONDS = 120.0; - public static WorkflowServiceStubs createTemporalService(final WorkflowServiceStubsOptions options, final String namespace) { + private final String temporalCloudClientCert; + private final String temporalCloudClientKey; + private final Boolean temporalCloudEnabled; + private final String temporalCloudHost; + private final String temporalCloudNamespace; + private final String temporalHost; + private final Integer temporalRetentionInDays; + + public TemporalUtils(@Property(name = "temporal.cloud.client.cert") final String temporalCloudClientCert, + @Property(name = "temporal.cloud.client.key") final String temporalCloudClientKey, + @Property(name = "temporal.cloud.enabled", defaultValue = "false") final Boolean temporalCloudEnabled, + @Value("${temporal.cloud.host}") final String temporalCloudHost, + @Value("${temporal.cloud.namespace}") final String temporalCloudNamespace, + @Value("${temporal.host}") final String temporalHost, + @Value("${temporal.retention}") final Integer temporalRetentionInDays) { + this.temporalCloudClientCert = temporalCloudClientCert; + this.temporalCloudClientKey = temporalCloudClientKey; + this.temporalCloudEnabled = temporalCloudEnabled; + this.temporalCloudHost = temporalCloudHost; + this.temporalCloudNamespace = temporalCloudNamespace; + this.temporalHost = temporalHost; + this.temporalRetentionInDays = temporalRetentionInDays; + } + + public WorkflowServiceStubs createTemporalService(final WorkflowServiceStubsOptions options, final String namespace) { return getTemporalClientWhenConnected( WAIT_INTERVAL, MAX_TIME_TO_CONNECT, @@ -83,25 +95,25 @@ public static WorkflowServiceStubs createTemporalService(final WorkflowServiceSt // TODO consider consolidating this method's logic into createTemporalService() after the Temporal // Cloud migration is complete. // The Temporal Migration migrator is the only reason this public method exists. - public static WorkflowServiceStubs createTemporalService(final boolean isCloud) { - final WorkflowServiceStubsOptions options = isCloud ? getCloudTemporalOptions() : getAirbyteTemporalOptions(configs.getTemporalHost()); - final String namespace = isCloud ? configs.getTemporalCloudNamespace() : DEFAULT_NAMESPACE; + public WorkflowServiceStubs createTemporalService(final boolean isCloud) { + final WorkflowServiceStubsOptions options = isCloud ? getCloudTemporalOptions() : TemporalWorkflowUtils.getAirbyteTemporalOptions(temporalHost); + final String namespace = isCloud ? temporalCloudNamespace : DEFAULT_NAMESPACE; return createTemporalService(options, namespace); } - public static WorkflowServiceStubs createTemporalService() { - return createTemporalService(configs.temporalCloudEnabled()); + public WorkflowServiceStubs createTemporalService() { + return createTemporalService(temporalCloudEnabled); } - private static WorkflowServiceStubsOptions getCloudTemporalOptions() { - final InputStream clientCert = new ByteArrayInputStream(configs.getTemporalCloudClientCert().getBytes(StandardCharsets.UTF_8)); - final InputStream clientKey = new ByteArrayInputStream(configs.getTemporalCloudClientKey().getBytes(StandardCharsets.UTF_8)); - WorkflowServiceStubsOptions.Builder optionBuilder; + private WorkflowServiceStubsOptions getCloudTemporalOptions() { + final InputStream clientCert = new ByteArrayInputStream(temporalCloudClientCert.getBytes(StandardCharsets.UTF_8)); + final InputStream clientKey = new ByteArrayInputStream(temporalCloudClientKey.getBytes(StandardCharsets.UTF_8)); + final WorkflowServiceStubsOptions.Builder optionBuilder; try { optionBuilder = WorkflowServiceStubsOptions.newBuilder() .setSslContext(SimpleSslContextBuilder.forPKCS8(clientCert, clientKey).build()) - .setTarget(configs.getTemporalCloudHost()); + .setTarget(temporalCloudHost); } catch (final SSLException e) { log.error("SSL Exception occurred attempting to establish Temporal Cloud options."); throw new RuntimeException(e); @@ -111,34 +123,19 @@ private static WorkflowServiceStubsOptions getCloudTemporalOptions() { return optionBuilder.build(); } - private static void configureTemporalMeterRegistry(WorkflowServiceStubsOptions.Builder optionalBuilder) { - MeterRegistry registry = MetricClientFactory.getMeterRegistry(); + private void configureTemporalMeterRegistry(final WorkflowServiceStubsOptions.Builder optionalBuilder) { + final MeterRegistry registry = MetricClientFactory.getMeterRegistry(); if (registry != null) { - StatsReporter reporter = new MicrometerClientStatsReporter(registry); - Scope scope = new RootScopeBuilder() + final StatsReporter reporter = new MicrometerClientStatsReporter(registry); + final Scope scope = new RootScopeBuilder() .reporter(reporter) .reportEvery(com.uber.m3.util.Duration.ofSeconds(REPORT_INTERVAL_SECONDS)); optionalBuilder.setMetricsScope(scope); } } - @VisibleForTesting - public static WorkflowServiceStubsOptions getAirbyteTemporalOptions(final String temporalHost) { - return WorkflowServiceStubsOptions.newBuilder() - .setTarget(temporalHost) - .build(); - } - - public static WorkflowClient createWorkflowClient(final WorkflowServiceStubs workflowServiceStubs, final String namespace) { - return WorkflowClient.newInstance( - workflowServiceStubs, - WorkflowClientOptions.newBuilder() - .setNamespace(namespace) - .build()); - } - - public static String getNamespace() { - return configs.temporalCloudEnabled() ? configs.getTemporalCloudNamespace() : DEFAULT_NAMESPACE; + public String getNamespace() { + return temporalCloudEnabled ? temporalCloudNamespace : DEFAULT_NAMESPACE; } /** @@ -146,8 +143,8 @@ public static String getNamespace() { * This should not be called when using Temporal Cloud, because Temporal Cloud does not allow * programmatic modification of workflow execution retention TTL. */ - public static void configureTemporalNamespace(final WorkflowServiceStubs temporalService) { - if (configs.temporalCloudEnabled()) { + public void configureTemporalNamespace(final WorkflowServiceStubs temporalService) { + if (temporalCloudEnabled) { log.info("Skipping Temporal Namespace configuration because Temporal Cloud is in use."); return; } @@ -156,17 +153,19 @@ public static void configureTemporalNamespace(final WorkflowServiceStubs tempora final var describeNamespaceRequest = DescribeNamespaceRequest.newBuilder().setNamespace(DEFAULT_NAMESPACE).build(); final var currentRetentionGrpcDuration = client.describeNamespace(describeNamespaceRequest).getConfig().getWorkflowExecutionRetentionTtl(); final var currentRetention = Duration.ofSeconds(currentRetentionGrpcDuration.getSeconds()); + final var workflowExecutionTtl = Duration.ofDays(temporalRetentionInDays); + final var humanReadableWorkflowExecutionTtl = DurationFormatUtils.formatDurationWords(workflowExecutionTtl.toMillis(), true, true); - if (currentRetention.equals(WORKFLOW_EXECUTION_TTL)) { + if (currentRetention.equals(workflowExecutionTtl)) { log.info("Workflow execution TTL already set for namespace " + DEFAULT_NAMESPACE + ". Remains unchanged as: " - + HUMAN_READABLE_WORKFLOW_EXECUTION_TTL); + + humanReadableWorkflowExecutionTtl); } else { - final var newGrpcDuration = com.google.protobuf.Duration.newBuilder().setSeconds(WORKFLOW_EXECUTION_TTL.getSeconds()).build(); + final var newGrpcDuration = com.google.protobuf.Duration.newBuilder().setSeconds(workflowExecutionTtl.getSeconds()).build(); final var humanReadableCurrentRetention = DurationFormatUtils.formatDurationWords(currentRetention.toMillis(), true, true); final var namespaceConfig = NamespaceConfig.newBuilder().setWorkflowExecutionRetentionTtl(newGrpcDuration).build(); final var updateNamespaceRequest = UpdateNamespaceRequest.newBuilder().setNamespace(DEFAULT_NAMESPACE).setConfig(namespaceConfig).build(); log.info("Workflow execution TTL differs for namespace " + DEFAULT_NAMESPACE + ". Changing from (" + humanReadableCurrentRetention + ") to (" - + HUMAN_READABLE_WORKFLOW_EXECUTION_TTL + "). "); + + humanReadableWorkflowExecutionTtl + "). "); client.updateNamespace(updateNamespaceRequest); } } @@ -178,39 +177,6 @@ public interface TemporalJobCreator { } - public static WorkflowOptions getWorkflowOptionsWithWorkflowId(final TemporalJobType jobType, final String workflowId) { - - return WorkflowOptions.newBuilder() - .setWorkflowId(workflowId) - .setRetryOptions(NO_RETRY) - .setTaskQueue(jobType.name()) - .build(); - } - - public static WorkflowOptions getWorkflowOptions(final TemporalJobType jobType) { - return WorkflowOptions.newBuilder() - .setTaskQueue(jobType.name()) - .setWorkflowTaskTimeout(Duration.ofSeconds(27)) // TODO parker - temporarily increasing this to a recognizable number to see if it changes - // error I'm seeing - // todo (cgardens) we do not leverage Temporal retries. - .setRetryOptions(RetryOptions.newBuilder().setMaximumAttempts(1).build()) - .build(); - } - - public static JobRunConfig createJobRunConfig(final UUID jobId, final int attemptId) { - return createJobRunConfig(String.valueOf(jobId), attemptId); - } - - public static JobRunConfig createJobRunConfig(final long jobId, final int attemptId) { - return createJobRunConfig(String.valueOf(jobId), attemptId); - } - - public static JobRunConfig createJobRunConfig(final String jobId, final int attemptId) { - return new JobRunConfig() - .withJobId(jobId) - .withAttemptId((long) attemptId); - } - /** * Allows running a given temporal workflow stub asynchronously. This method only works for * workflows that take one argument. Because of the iface that Temporal supplies, in order to handle @@ -228,10 +194,10 @@ public static JobRunConfig createJobRunConfig(final String jobId, final int atte * @return pair of the workflow execution (contains metadata on the asynchronously running job) and * future that can be used to await the result of the workflow stub's function */ - public static ImmutablePair> asyncExecute(final STUB workflowStub, - final Functions.Func1 function, - final A1 arg1, - final Class outputType) { + public ImmutablePair> asyncExecute(final STUB workflowStub, + final Functions.Func1 function, + final A1 arg1, + final Class outputType) { final WorkflowStub untyped = WorkflowStub.fromTyped(workflowStub); final WorkflowExecution workflowExecution = WorkflowClient.start(function, arg1); final CompletableFuture resultAsync = untyped.getResultAsync(outputType); @@ -244,12 +210,12 @@ public static ImmutablePair temporalServiceSupplier, - final String namespace) { + public WorkflowServiceStubs getTemporalClientWhenConnected( + final Duration waitInterval, + final Duration maxTimeToConnect, + final Duration waitAfterConnection, + final Supplier temporalServiceSupplier, + final String namespace) { log.info("Waiting for temporal server..."); boolean temporalNamespaceInitialized = false; @@ -283,7 +249,7 @@ public static WorkflowServiceStubs getTemporalClientWhenConnected( return temporalService; } - protected static NamespaceInfo getNamespaceInfo(final WorkflowServiceStubs temporalService, final String namespace) { + protected NamespaceInfo getNamespaceInfo(final WorkflowServiceStubs temporalService, final String namespace) { return temporalService.blockingStub() .describeNamespace(DescribeNamespaceRequest.newBuilder().setNamespace(namespace).build()) .getNamespaceInfo(); @@ -293,8 +259,8 @@ protected static NamespaceInfo getNamespaceInfo(final WorkflowServiceStubs tempo * Runs the code within the supplier while heartbeating in the backgroud. Also makes sure to shut * down the heartbeat server after the fact. */ - public static T withBackgroundHeartbeat(final Callable callable, - final Supplier activityContext) { + public T withBackgroundHeartbeat(final Callable callable, + final Supplier activityContext) { final ScheduledExecutorService scheduledExecutor = Executors.newSingleThreadScheduledExecutor(); try { @@ -314,9 +280,9 @@ public static T withBackgroundHeartbeat(final Callable callable, } } - public static T withBackgroundHeartbeat(final AtomicReference cancellationCallbackRef, - final Callable callable, - final Supplier activityContext) { + public T withBackgroundHeartbeat(final AtomicReference cancellationCallbackRef, + final Callable callable, + final Supplier activityContext) { final ScheduledExecutorService scheduledExecutor = Executors.newSingleThreadScheduledExecutor(); try { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalWorkflowUtils.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalWorkflowUtils.java new file mode 100644 index 0000000000000..8ff088e33206d --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalWorkflowUtils.java @@ -0,0 +1,91 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal; + +import com.google.common.annotations.VisibleForTesting; +import io.airbyte.scheduler.models.JobRunConfig; +import io.airbyte.workers.temporal.scheduling.ConnectionUpdaterInput; +import io.temporal.client.WorkflowClient; +import io.temporal.client.WorkflowClientOptions; +import io.temporal.client.WorkflowOptions; +import io.temporal.common.RetryOptions; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.serviceclient.WorkflowServiceStubsOptions; +import java.time.Duration; +import java.util.UUID; + +/** + * Collection of Temporal workflow related utility methods. + * + * N.B: These methods should not store any state or depend on any other objects/singletons + * managed by the application framework. + */ +public class TemporalWorkflowUtils { + + public static final RetryOptions NO_RETRY = RetryOptions.newBuilder().setMaximumAttempts(1).build(); + + private TemporalWorkflowUtils() {} + + public static ConnectionUpdaterInput buildStartWorkflowInput(final UUID connectionId) { + return ConnectionUpdaterInput.builder() + .connectionId(connectionId) + .jobId(null) + .attemptId(null) + .fromFailure(false) + .attemptNumber(1) + .workflowState(null) + .resetConnection(false) + .fromJobResetFailure(false) + .build(); + } + + public static WorkflowOptions buildWorkflowOptions(final TemporalJobType jobType, final String workflowId) { + return WorkflowOptions.newBuilder() + .setWorkflowId(workflowId) + .setRetryOptions(NO_RETRY) + .setTaskQueue(jobType.name()) + .build(); + } + + public static WorkflowOptions buildWorkflowOptions(final TemporalJobType jobType) { + return WorkflowOptions.newBuilder() + .setTaskQueue(jobType.name()) + .setWorkflowTaskTimeout(Duration.ofSeconds(27)) // TODO parker - temporarily increasing this to a recognizable number to see if it changes + // error I'm seeing + // todo (cgardens) we do not leverage Temporal retries. + .setRetryOptions(RetryOptions.newBuilder().setMaximumAttempts(1).build()) + .build(); + } + + public static JobRunConfig createJobRunConfig(final UUID jobId, final int attemptId) { + return createJobRunConfig(String.valueOf(jobId), attemptId); + } + + public static JobRunConfig createJobRunConfig(final long jobId, final int attemptId) { + return createJobRunConfig(String.valueOf(jobId), attemptId); + } + + public static JobRunConfig createJobRunConfig(final String jobId, final int attemptId) { + return new JobRunConfig() + .withJobId(jobId) + .withAttemptId((long) attemptId); + } + + @VisibleForTesting + public static WorkflowServiceStubsOptions getAirbyteTemporalOptions(final String temporalHost) { + return WorkflowServiceStubsOptions.newBuilder() + .setTarget(temporalHost) + .build(); + } + + public static WorkflowClient createWorkflowClient(final WorkflowServiceStubs workflowServiceStubs, final String namespace) { + return WorkflowClient.newInstance( + workflowServiceStubs, + WorkflowClientOptions.newBuilder() + .setNamespace(namespace) + .build()); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/annotations/TemporalActivityStub.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/annotations/TemporalActivityStub.java new file mode 100644 index 0000000000000..ee402796bbadc --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/annotations/TemporalActivityStub.java @@ -0,0 +1,30 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.annotations; + +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +/** + * Denotes a field in a Temporal workflow that represents a Temporal activity stub. Fields marked + * with this annotation will automatically have a Temporal activity stub created, if not already + * initialized when execution of the Temporal workflow starts. + */ +@Retention(RetentionPolicy.RUNTIME) +@Target({ElementType.FIELD}) +public @interface TemporalActivityStub { + + /** + * The name of the singleton bean that holds the Temporal activity options for the Temporal activity + * stub annotated by this annotation. This bean must exist in the application context. + * + * @return The name of the singleton bean that holds the Temporal activity options for that Temporal + * activity stub. + */ + String activityOptionsBeanName(); + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java index cf2b63af6016c..98e30baff4891 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionActivityImpl.java @@ -23,38 +23,39 @@ import io.airbyte.workers.process.ProcessFactory; import io.airbyte.workers.temporal.CancellationHandler; import io.airbyte.workers.temporal.TemporalAttemptExecution; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; import io.temporal.activity.Activity; import io.temporal.activity.ActivityExecutionContext; import java.nio.file.Path; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class CheckConnectionActivityImpl implements CheckConnectionActivity { - private final WorkerConfigs workerConfigs; - private final ProcessFactory processFactory; - private final SecretsHydrator secretsHydrator; - private final Path workspaceRoot; - private final WorkerEnvironment workerEnvironment; - private final LogConfigs logConfigs; - private final AirbyteApiClient airbyteApiClient; - private final String airbyteVersion; - - public CheckConnectionActivityImpl(final WorkerConfigs workerConfigs, - final ProcessFactory processFactory, - final SecretsHydrator secretsHydrator, - final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final LogConfigs logConfigs, - final AirbyteApiClient airbyteApiClient, - final String airbyteVersion) { - this.workerConfigs = workerConfigs; - this.processFactory = processFactory; - this.secretsHydrator = secretsHydrator; - this.workspaceRoot = workspaceRoot; - this.workerEnvironment = workerEnvironment; - this.logConfigs = logConfigs; - this.airbyteApiClient = airbyteApiClient; - this.airbyteVersion = airbyteVersion; - } + @Inject + @Named("checkWorkerConfigs") + private WorkerConfigs workerConfigs; + @Inject + @Named("checkProcessFactory") + private ProcessFactory processFactory; + @Inject + private SecretsHydrator secretsHydrator; + @Inject + @Named("workspaceRoot") + private Path workspaceRoot; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private LogConfigs logConfigs; + @Inject + private AirbyteApiClient airbyteApiClient; + @Value("${airbyte.version}") + private String airbyteVersion; @Override public ConnectorJobOutput runWithJobOutput(final CheckConnectionInput args) { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowImpl.java index 4728ba07f9d39..c673732ebcda5 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowImpl.java @@ -10,22 +10,25 @@ import io.airbyte.config.StandardCheckConnectionOutput; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; import io.airbyte.workers.temporal.check.connection.CheckConnectionActivity.CheckConnectionInput; -import io.airbyte.workers.temporal.scheduling.shared.ActivityConfiguration; import io.temporal.workflow.Workflow; +import javax.inject.Singleton; +@Singleton public class CheckConnectionWorkflowImpl implements CheckConnectionWorkflow { - private final CheckConnectionActivity activity = - Workflow.newActivityStub(CheckConnectionActivity.class, ActivityConfiguration.CHECK_ACTIVITY_OPTIONS); - private static final String CHECK_JOB_OUTPUT_TAG = "check_job_output"; private static final int CHECK_JOB_OUTPUT_TAG_CURRENT_VERSION = 1; + @TemporalActivityStub(activityOptionsBeanName = "checkActivityOptions") + private CheckConnectionActivity activity; + @Override public ConnectorJobOutput run(final JobRunConfig jobRunConfig, final IntegrationLauncherConfig launcherConfig, final StandardCheckConnectionInput connectionConfiguration) { + final CheckConnectionInput checkInput = new CheckConnectionInput(jobRunConfig, launcherConfig, connectionConfiguration); final int jobOutputVersion = diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java index 9011c7222b97a..8e6ed061ea9a5 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogActivityImpl.java @@ -24,38 +24,41 @@ import io.airbyte.workers.process.ProcessFactory; import io.airbyte.workers.temporal.CancellationHandler; import io.airbyte.workers.temporal.TemporalAttemptExecution; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; import io.temporal.activity.Activity; import io.temporal.activity.ActivityExecutionContext; import java.nio.file.Path; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +import lombok.extern.slf4j.Slf4j; +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") +@Slf4j public class DiscoverCatalogActivityImpl implements DiscoverCatalogActivity { - private final WorkerConfigs workerConfigs; - private final ProcessFactory processFactory; - private final SecretsHydrator secretsHydrator; - private final Path workspaceRoot; - private final WorkerEnvironment workerEnvironment; - private final LogConfigs logConfigs; - private final AirbyteApiClient airbyteApiClient; - private final String airbyteVersion; - - public DiscoverCatalogActivityImpl(final WorkerConfigs workerConfigs, - final ProcessFactory processFactory, - final SecretsHydrator secretsHydrator, - final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final LogConfigs logConfigs, - final AirbyteApiClient airbyteApiClient, - final String airbyteVersion) { - this.workerConfigs = workerConfigs; - this.processFactory = processFactory; - this.secretsHydrator = secretsHydrator; - this.workspaceRoot = workspaceRoot; - this.workerEnvironment = workerEnvironment; - this.logConfigs = logConfigs; - this.airbyteApiClient = airbyteApiClient; - this.airbyteVersion = airbyteVersion; - } + @Inject + @Named("discoverWorkerConfigs") + private WorkerConfigs workerConfigs; + @Inject + @Named("discoverProcessFactory") + private ProcessFactory processFactory; + @Inject + private SecretsHydrator secretsHydrator; + @Inject + @Named("workspaceRoot") + private Path workspaceRoot; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private LogConfigs logConfigs; + @Inject + private AirbyteApiClient airbyteApiClient;; + @Value("${airbyte.version}") + private String airbyteVersion; @Override public ConnectorJobOutput run(final JobRunConfig jobRunConfig, @@ -68,6 +71,8 @@ public ConnectorJobOutput run(final JobRunConfig jobRunConfig, final ActivityExecutionContext context = Activity.getExecutionContext(); + log.info("Fetching catalog data {}", fullConfig); + final TemporalAttemptExecution temporalAttemptExecution = new TemporalAttemptExecution<>( workspaceRoot, diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogWorkflowImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogWorkflowImpl.java index d0c68d7470bf9..ba35f5094cc79 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogWorkflowImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/discover/catalog/DiscoverCatalogWorkflowImpl.java @@ -8,18 +8,14 @@ import io.airbyte.config.StandardDiscoverCatalogInput; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.workers.temporal.TemporalUtils; -import io.temporal.activity.ActivityOptions; -import io.temporal.workflow.Workflow; -import java.time.Duration; +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; +import javax.inject.Singleton; +@Singleton public class DiscoverCatalogWorkflowImpl implements DiscoverCatalogWorkflow { - final ActivityOptions options = ActivityOptions.newBuilder() - .setScheduleToCloseTimeout(Duration.ofHours(2)) - .setRetryOptions(TemporalUtils.NO_RETRY) - .build(); - private final DiscoverCatalogActivity activity = Workflow.newActivityStub(DiscoverCatalogActivity.class, options); + @TemporalActivityStub(activityOptionsBeanName = "discoveryActivityOptions") + private DiscoverCatalogActivity activity; @Override public ConnectorJobOutput run(final JobRunConfig jobRunConfig, diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java index 67ffdc7f27bde..8f24e13d7c04c 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowImpl.java @@ -7,7 +7,6 @@ import com.fasterxml.jackson.databind.JsonNode; import io.airbyte.config.ConnectorJobOutput; import io.airbyte.config.ConnectorJobOutput.OutputType; -import io.airbyte.config.EnvConfigs; import io.airbyte.config.FailureReason; import io.airbyte.config.FailureReason.FailureType; import io.airbyte.config.NormalizationSummary; @@ -17,17 +16,14 @@ import io.airbyte.config.StandardSyncOutput; import io.airbyte.config.StandardSyncSummary; import io.airbyte.config.StandardSyncSummary.ReplicationStatus; -import io.airbyte.metrics.lib.MetricAttribute; -import io.airbyte.metrics.lib.MetricClient; -import io.airbyte.metrics.lib.MetricClientFactory; -import io.airbyte.metrics.lib.MetricTags; import io.airbyte.metrics.lib.OssMetricsRegistry; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; import io.airbyte.workers.WorkerConstants; import io.airbyte.workers.helper.FailureHelper; -import io.airbyte.workers.temporal.ConnectionManagerUtils; import io.airbyte.workers.temporal.TemporalJobType; +import io.airbyte.workers.temporal.TemporalWorkflowUtils; +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; import io.airbyte.workers.temporal.check.connection.CheckConnectionActivity; import io.airbyte.workers.temporal.check.connection.CheckConnectionActivity.CheckConnectionInput; import io.airbyte.workers.temporal.exception.RetryableException; @@ -58,9 +54,15 @@ import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivity.JobSuccessInput; import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivity.JobSuccessInputWithAttemptNumber; import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivity.ReportJobStartInput; +import io.airbyte.workers.temporal.scheduling.activities.RecordMetricActivity; +import io.airbyte.workers.temporal.scheduling.activities.RecordMetricActivity.FailureCause; +import io.airbyte.workers.temporal.scheduling.activities.RecordMetricActivity.RecordMetricInput; +import io.airbyte.workers.temporal.scheduling.activities.RouteToSyncTaskQueueActivity; +import io.airbyte.workers.temporal.scheduling.activities.RouteToSyncTaskQueueActivity.RouteToSyncTaskQueueInput; +import io.airbyte.workers.temporal.scheduling.activities.RouteToSyncTaskQueueActivity.RouteToSyncTaskQueueOutput; import io.airbyte.workers.temporal.scheduling.activities.StreamResetActivity; import io.airbyte.workers.temporal.scheduling.activities.StreamResetActivity.DeleteStreamResetRecordsForJobInput; -import io.airbyte.workers.temporal.scheduling.shared.ActivityConfiguration; +import io.airbyte.workers.temporal.scheduling.activities.WorkflowConfigActivity; import io.airbyte.workers.temporal.scheduling.state.WorkflowInternalState; import io.airbyte.workers.temporal.scheduling.state.WorkflowState; import io.airbyte.workers.temporal.scheduling.state.listener.NoopStateListener; @@ -74,17 +76,15 @@ import io.temporal.workflow.Workflow; import java.time.Duration; import java.time.Instant; -import java.util.ArrayList; -import java.util.List; +import java.util.Optional; import java.util.Set; import java.util.UUID; import java.util.function.Consumer; import java.util.function.Function; -import java.util.stream.Collectors; -import java.util.stream.Stream; import lombok.extern.slf4j.Slf4j; @Slf4j +@SuppressWarnings("PMD.AvoidDuplicateLiterals") public class ConnectionManagerWorkflowImpl implements ConnectionManagerWorkflow { public static final long NON_RUNNING_JOB_ID = -1; @@ -107,40 +107,51 @@ public class ConnectionManagerWorkflowImpl implements ConnectionManagerWorkflow private static final String DELETE_RESET_JOB_STREAMS_TAG = "delete_reset_job_streams"; private static final int DELETE_RESET_JOB_STREAMS_CURRENT_VERSION = 1; + private static final String RECORD_METRIC_TAG = "record_metric"; + private static final int RECORD_METRIC_CURRENT_VERSION = 1; + private static final String WORKFLOW_CONFIG_TAG = "workflow_config"; + private static final int WORKFLOW_CONFIG_CURRENT_VERSION = 1; - static final Duration WORKFLOW_FAILURE_RESTART_DELAY = Duration.ofSeconds(new EnvConfigs().getWorkflowFailureRestartDelaySeconds()); + private static final String ROUTE_ACTIVITY_TAG = "route_activity"; + private static final int ROUTE_ACTIVITY_CURRENT_VERSION = 1; private WorkflowState workflowState = new WorkflowState(UUID.randomUUID(), new NoopStateListener()); private final WorkflowInternalState workflowInternalState = new WorkflowInternalState(); - private final GenerateInputActivity getSyncInputActivity = - Workflow.newActivityStub(GenerateInputActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - private final JobCreationAndStatusUpdateActivity jobCreationAndStatusUpdateActivity = - Workflow.newActivityStub(JobCreationAndStatusUpdateActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - private final ConfigFetchActivity configFetchActivity = - Workflow.newActivityStub(ConfigFetchActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - private final ConnectionDeletionActivity connectionDeletionActivity = - Workflow.newActivityStub(ConnectionDeletionActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - private final AutoDisableConnectionActivity autoDisableConnectionActivity = - Workflow.newActivityStub(AutoDisableConnectionActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - private final CheckConnectionActivity checkActivity = - Workflow.newActivityStub(CheckConnectionActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - private final StreamResetActivity streamResetActivity = - Workflow.newActivityStub(StreamResetActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private GenerateInputActivity getSyncInputActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private JobCreationAndStatusUpdateActivity jobCreationAndStatusUpdateActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private ConfigFetchActivity configFetchActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private ConnectionDeletionActivity connectionDeletionActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private AutoDisableConnectionActivity autoDisableConnectionActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private CheckConnectionActivity checkActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private StreamResetActivity streamResetActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private RecordMetricActivity recordMetricActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private WorkflowConfigActivity workflowConfigActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private RouteToSyncTaskQueueActivity routeToSyncTaskQueueActivity; private CancellationScope cancellableSyncWorkflow; private UUID connectionId; - private final MetricClient metricClient = MetricClientFactory.getMetricClient(); - - public ConnectionManagerWorkflowImpl() {} + private Duration workflowDelay; @Override public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws RetryableException { - recordWorkflowCountMetric(connectionUpdaterInput, OssMetricsRegistry.TEMPORAL_WORKFLOW_ATTEMPT); try { + recordMetric(new RecordMetricInput(connectionUpdaterInput, Optional.empty(), OssMetricsRegistry.TEMPORAL_WORKFLOW_ATTEMPT, null)); + workflowDelay = getWorkflowRestartDelaySeconds(); + try { cancellableSyncWorkflow = generateSyncWorkflowRunnable(connectionUpdaterInput); cancellableSyncWorkflow.run(); @@ -148,7 +159,8 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr // When a scope is cancelled temporal will throw a CanceledFailure as you can see here: // https://github.com/temporalio/sdk-java/blob/master/temporal-sdk/src/main/java/io/temporal/workflow/CancellationScope.java#L72 // The naming is very misleading, it is not a failure but the expected behavior... - recordWorkflowFailureCountMetric(connectionUpdaterInput, FailureCause.CANCELED); + recordMetric( + new RecordMetricInput(connectionUpdaterInput, Optional.of(FailureCause.CANCELED), OssMetricsRegistry.TEMPORAL_WORKFLOW_FAILURE, null)); } if (workflowState.isDeleted()) { @@ -168,7 +180,6 @@ public void run(final ConnectionUpdaterInput connectionUpdaterInput) throws Retr // "Cancel" button was pressed on a job if (workflowState.isCancelled()) { - deleteResetJobStreams(); reportCancelledAndContinueWith(false, connectionUpdaterInput); } @@ -246,8 +257,8 @@ private CancellationScope generateSyncWorkflowRunnable(final ConnectionUpdaterIn // exception since we expect it, we just // silently ignore it. if (childWorkflowFailure.getCause() instanceof CanceledFailure) { + log.debug("Ignoring canceled failure as it is handled by the cancellation scope."); // do nothing, cancellation handled by cancellationScope - recordWorkflowFailureCountMetric(connectionUpdaterInput, FailureCause.CANCELED); } else if (childWorkflowFailure.getCause()instanceof final ActivityFailure af) { // Allows us to classify unhandled failures from the sync workflow. e.g. If the normalization // activity throws an exception, for @@ -291,7 +302,7 @@ private void reportSuccess(final ConnectionUpdaterInput connectionUpdaterInput, deleteResetJobStreams(); // Record the success metric - recordWorkflowCountMetric(connectionUpdaterInput, OssMetricsRegistry.TEMPORAL_WORKFLOW_SUCCESS); + recordMetric(new RecordMetricInput(connectionUpdaterInput, Optional.empty(), OssMetricsRegistry.TEMPORAL_WORKFLOW_SUCCESS, null)); resetNewConnectionInput(connectionUpdaterInput); } @@ -345,7 +356,7 @@ private void reportFailure(final ConnectionUpdaterInput connectionUpdaterInput, } // Record the failure metric - recordWorkflowFailureCountMetric(connectionUpdaterInput, failureCause); + recordMetric(new RecordMetricInput(connectionUpdaterInput, Optional.of(failureCause), OssMetricsRegistry.TEMPORAL_WORKFLOW_FAILURE, null)); resetNewConnectionInput(connectionUpdaterInput); } @@ -539,8 +550,8 @@ private OUTPUT runMandatoryActivityWithOutput(final Function OUTPUT runMandatoryActivityWithOutput(final Function workflowState.isRetryFailedActivity()); + Workflow.await(workflowDelay, () -> workflowState.isRetryFailedActivity()); // Accept a manual signal to retry the failed activity during this window if (workflowState.isRetryFailedActivity()) { @@ -561,7 +572,7 @@ private OUTPUT runMandatoryActivityWithOutput(final Function baseMetricAttributes = generateMetricAttributes(connectionUpdaterInput); - if (metricAttributes != null) { - baseMetricAttributes.addAll(Stream.of(metricAttributes).collect(Collectors.toList())); + if (workflowConfigVersion < WORKFLOW_CONFIG_CURRENT_VERSION) { + return Duration.ofMinutes(10L); } - metricClient.count(metricName, 1L, baseMetricAttributes.toArray(new MetricAttribute[] {})); - } - - /** - * Generates the list of {@link MetricAttribute}s to be included when recording a metric. - * - * @param connectionUpdaterInput The {@link ConnectionUpdaterInput} that represents the workflow to - * be executed. - * @return The list of {@link MetricAttribute}s to be included when recording a metric. - */ - private List generateMetricAttributes(final ConnectionUpdaterInput connectionUpdaterInput) { - final List metricAttributes = new ArrayList<>(); - metricAttributes.add(new MetricAttribute(MetricTags.CONNECTION_ID, String.valueOf(connectionUpdaterInput.getConnectionId()))); - return metricAttributes; - } - /** - * Enumeration of workflow failure causes. - */ - private enum FailureCause { - ACTIVITY, - CANCELED, - CONNECTION, - UNKNOWN, - WORKFLOW + return workflowConfigActivity.getWorkflowRestartDelaySeconds(); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityImpl.java index 9761aee66dbed..fc8be792b70ed 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityImpl.java @@ -9,8 +9,8 @@ import static io.airbyte.scheduler.persistence.JobNotifier.CONNECTION_DISABLED_WARNING_NOTIFICATION; import static java.time.temporal.ChronoUnit.DAYS; +import com.google.common.annotations.VisibleForTesting; import io.airbyte.commons.features.FeatureFlags; -import io.airbyte.config.Configs; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSync.Status; import io.airbyte.config.persistence.ConfigRepository; @@ -21,19 +21,31 @@ import io.airbyte.scheduler.persistence.JobPersistence; import io.airbyte.validation.json.JsonValidationException; import io.airbyte.workers.temporal.exception.RetryableException; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; import java.io.IOException; import java.util.List; import java.util.Optional; import java.util.concurrent.TimeUnit; -import lombok.AllArgsConstructor; +import javax.inject.Inject; +import javax.inject.Singleton; -@AllArgsConstructor +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class AutoDisableConnectionActivityImpl implements AutoDisableConnectionActivity { + @Inject private ConfigRepository configRepository; + @Inject private JobPersistence jobPersistence; + @Inject private FeatureFlags featureFlags; - private Configs configs; + @Value("${airbyte.worker.job.failed.max-days}") + private Integer maxDaysOfOnlyFailedJobsBeforeConnectionDisable; + @Value("${airbyte.worker.job.failed.max-jobs}") + private Integer maxFailedJobsInARowBeforeConnectionDisable; + @Inject private JobNotifier jobNotifier; // Given a connection id and current timestamp, this activity will set a connection to INACTIVE if @@ -54,9 +66,9 @@ public AutoDisableConnectionOutput autoDisableFailingConnection(final AutoDisabl return new AutoDisableConnectionOutput(false); } - final int maxDaysOfOnlyFailedJobs = configs.getMaxDaysOfOnlyFailedJobsBeforeConnectionDisable(); + final int maxDaysOfOnlyFailedJobs = maxDaysOfOnlyFailedJobsBeforeConnectionDisable; final int maxDaysOfOnlyFailedJobsBeforeWarning = maxDaysOfOnlyFailedJobs / 2; - final int maxFailedJobsInARowBeforeConnectionDisableWarning = configs.getMaxFailedJobsInARowBeforeConnectionDisable() / 2; + final int maxFailedJobsInARowBeforeConnectionDisableWarning = maxFailedJobsInARowBeforeConnectionDisable / 2; final long currTimestampInSeconds = input.getCurrTimestamp().getEpochSecond(); final Job lastJob = jobPersistence.getLastReplicationJob(input.getConnectionId()) .orElseThrow(() -> new Exception("Auto-Disable Connection should not have been attempted if can't get latest replication job.")); @@ -84,7 +96,7 @@ public AutoDisableConnectionOutput autoDisableFailingConnection(final AutoDisabl if (numFailures == 0) { return new AutoDisableConnectionOutput(false); - } else if (numFailures >= configs.getMaxFailedJobsInARowBeforeConnectionDisable()) { + } else if (numFailures >= maxFailedJobsInARowBeforeConnectionDisable) { // disable connection if max consecutive failed jobs limit has been hit disableConnection(standardSync, lastJob); return new AutoDisableConnectionOutput(true); @@ -186,4 +198,34 @@ private void disableConnection(final StandardSync standardSync, final Job lastJo jobNotifier.notifyJobByEmail(null, CONNECTION_DISABLED_NOTIFICATION, lastJob); } + @VisibleForTesting + void setConfigRepository(final ConfigRepository configRepository) { + this.configRepository = configRepository; + } + + @VisibleForTesting + void setJobPersistence(final JobPersistence jobPersistence) { + this.jobPersistence = jobPersistence; + } + + @VisibleForTesting + void setFeatureFlags(final FeatureFlags featureFlags) { + this.featureFlags = featureFlags; + } + + @VisibleForTesting + void setMaxDaysOfOnlyFailedJobsBeforeConnectionDisable(final Integer maxDaysOfOnlyFailedJobsBeforeConnectionDisable) { + this.maxDaysOfOnlyFailedJobsBeforeConnectionDisable = maxDaysOfOnlyFailedJobsBeforeConnectionDisable; + } + + @VisibleForTesting + void setMaxFailedJobsInARowBeforeConnectionDisable(final Integer maxFailedJobsInARowBeforeConnectionDisable) { + this.maxFailedJobsInARowBeforeConnectionDisable = maxFailedJobsInARowBeforeConnectionDisable; + } + + @VisibleForTesting + void setJobNotifier(final JobNotifier jobNotifier) { + this.jobNotifier = jobNotifier; + } + } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityImpl.java index 19ee4b0289581..b9a8efa492cd1 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityImpl.java @@ -4,7 +4,7 @@ package io.airbyte.workers.temporal.scheduling.activities; -import io.airbyte.config.Configs; +import com.google.common.annotations.VisibleForTesting; import io.airbyte.config.Cron; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSync.ScheduleType; @@ -16,6 +16,8 @@ import io.airbyte.scheduler.persistence.JobPersistence; import io.airbyte.validation.json.JsonValidationException; import io.airbyte.workers.temporal.exception.RetryableException; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; import java.io.IOException; import java.text.ParseException; import java.time.DateTimeException; @@ -25,21 +27,30 @@ import java.util.TimeZone; import java.util.UUID; import java.util.function.Supplier; -import lombok.AllArgsConstructor; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; import lombok.extern.slf4j.Slf4j; import org.joda.time.DateTimeZone; import org.quartz.CronExpression; -@AllArgsConstructor @Slf4j +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class ConfigFetchActivityImpl implements ConfigFetchActivity { private final static long MS_PER_SECOND = 1000L; private final static long MIN_CRON_INTERVAL_SECONDS = 60; + @Inject private ConfigRepository configRepository; + @Inject private JobPersistence jobPersistence; - private Configs configs; + @Value("${airbyte.worker.sync.max-attempts}") + private Integer syncJobMaxAttempts; + @Inject + @Named("currentSecondsSupplier") private Supplier currentSecondsSupplier; @Override @@ -64,7 +75,7 @@ public ScheduleRetrieverOutput getTimeToWait(final ScheduleRetrieverInput input) * * This method consumes the `scheduleType` and `scheduleData` fields. */ - private ScheduleRetrieverOutput getTimeToWaitFromScheduleType(final StandardSync standardSync, UUID connectionId) throws IOException { + private ScheduleRetrieverOutput getTimeToWaitFromScheduleType(final StandardSync standardSync, final UUID connectionId) throws IOException { if (standardSync.getScheduleType() == ScheduleType.MANUAL || standardSync.getStatus() != Status.ACTIVE) { // Manual syncs wait for their first run return new ScheduleRetrieverOutput(Duration.ofDays(100 * 365)); @@ -103,7 +114,7 @@ private ScheduleRetrieverOutput getTimeToWaitFromScheduleType(final StandardSync final Duration timeToWait = Duration.ofSeconds( Math.max(0, nextRunStart.getTime() / MS_PER_SECOND - currentSecondsSupplier.get())); return new ScheduleRetrieverOutput(timeToWait); - } catch (ParseException e) { + } catch (final ParseException e) { throw (DateTimeException) new DateTimeException(e.getMessage()).initCause(e); } } @@ -117,7 +128,7 @@ private ScheduleRetrieverOutput getTimeToWaitFromScheduleType(final StandardSync * * This method consumes the `schedule` field. */ - private ScheduleRetrieverOutput getTimeToWaitFromLegacy(final StandardSync standardSync, UUID connectionId) throws IOException { + private ScheduleRetrieverOutput getTimeToWaitFromLegacy(final StandardSync standardSync, final UUID connectionId) throws IOException { if (standardSync.getSchedule() == null || standardSync.getStatus() != Status.ACTIVE) { // Manual syncs wait for their first run return new ScheduleRetrieverOutput(Duration.ofDays(100 * 365)); @@ -144,7 +155,27 @@ private ScheduleRetrieverOutput getTimeToWaitFromLegacy(final StandardSync stand @Override public GetMaxAttemptOutput getMaxAttempt() { - return new GetMaxAttemptOutput(configs.getSyncJobMaxAttempts()); + return new GetMaxAttemptOutput(syncJobMaxAttempts); + } + + @VisibleForTesting + void setConfigRepository(final ConfigRepository configRepository) { + this.configRepository = configRepository; + } + + @VisibleForTesting + void setJobPersistence(final JobPersistence jobPersistence) { + this.jobPersistence = jobPersistence; + } + + @VisibleForTesting + void setSyncJobMaxAttempts(final Integer syncJobMaxAttempts) { + this.syncJobMaxAttempts = syncJobMaxAttempts; + } + + @VisibleForTesting + void setCurrentSecondsSupplier(final Supplier currentSecondsSupplier) { + this.currentSecondsSupplier = currentSecondsSupplier; } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConnectionDeletionActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConnectionDeletionActivityImpl.java index a047a2729c95b..528d7de9e6a88 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConnectionDeletionActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/ConnectionDeletionActivityImpl.java @@ -8,13 +8,22 @@ import io.airbyte.validation.json.JsonValidationException; import io.airbyte.workers.helper.ConnectionHelper; import io.airbyte.workers.temporal.exception.RetryableException; +import io.micronaut.context.annotation.Requires; import java.io.IOException; +import javax.inject.Inject; +import javax.inject.Singleton; import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; @AllArgsConstructor +@NoArgsConstructor +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class ConnectionDeletionActivityImpl implements ConnectionDeletionActivity { - private final ConnectionHelper connectionHelper; + @Inject + private ConnectionHelper connectionHelper; @Override public void deleteConnection(final ConnectionDeletionInput input) { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/GenerateInputActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/GenerateInputActivityImpl.java index 2276054781eea..7cb67c093c997 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/GenerateInputActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/GenerateInputActivityImpl.java @@ -15,14 +15,23 @@ import io.airbyte.scheduler.models.JobRunConfig; import io.airbyte.scheduler.persistence.JobPersistence; import io.airbyte.workers.WorkerConstants; -import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.TemporalWorkflowUtils; import io.airbyte.workers.temporal.exception.RetryableException; +import io.micronaut.context.annotation.Requires; import java.util.List; +import javax.inject.Inject; +import javax.inject.Singleton; import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; @AllArgsConstructor +@NoArgsConstructor +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class GenerateInputActivityImpl implements GenerateInputActivity { + @Inject private JobPersistence jobPersistence; @Override @@ -61,7 +70,7 @@ public GeneratedJobInput getSyncWorkflowInput(final SyncInput input) { List.of(ConfigType.SYNC, ConfigType.RESET_CONNECTION))); } - final JobRunConfig jobRunConfig = TemporalUtils.createJobRunConfig(jobId, attempt); + final JobRunConfig jobRunConfig = TemporalWorkflowUtils.createJobRunConfig(jobId, attempt); final IntegrationLauncherConfig sourceLauncherConfig = new IntegrationLauncherConfig() .withJobId(String.valueOf(jobId)) diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityImpl.java index 8edf65e79f597..19efb47b09e54 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/JobCreationAndStatusUpdateActivityImpl.java @@ -44,29 +44,48 @@ import io.airbyte.workers.run.TemporalWorkerRunFactory; import io.airbyte.workers.run.WorkerRun; import io.airbyte.workers.temporal.exception.RetryableException; +import io.micronaut.context.annotation.Requires; import java.io.IOException; import java.nio.file.Path; import java.util.List; import java.util.Optional; import java.util.UUID; +import javax.inject.Inject; +import javax.inject.Singleton; import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; import lombok.extern.slf4j.Slf4j; @AllArgsConstructor +@NoArgsConstructor @Slf4j +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class JobCreationAndStatusUpdateActivityImpl implements JobCreationAndStatusUpdateActivity { - private final SyncJobFactory jobFactory; - private final JobPersistence jobPersistence; - private final TemporalWorkerRunFactory temporalWorkerRunFactory; - private final WorkerEnvironment workerEnvironment; - private final LogConfigs logConfigs; - private final JobNotifier jobNotifier; - private final JobTracker jobTracker; - private final ConfigRepository configRepository; - private final JobCreator jobCreator; - private final StreamResetPersistence streamResetPersistence; - private final JobErrorReporter jobErrorReporter; + @Inject + private SyncJobFactory jobFactory; + @Inject + private JobPersistence jobPersistence; + @Inject + private TemporalWorkerRunFactory temporalWorkerRunFactory; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private LogConfigs logConfigs; + @Inject + private JobNotifier jobNotifier; + @Inject + private JobTracker jobTracker; + @Inject + private ConfigRepository configRepository; + @Inject + private JobCreator jobCreator; + @Inject + private StreamResetPersistence streamResetPersistence; + @Inject + private JobErrorReporter jobErrorReporter; @Override public JobCreationOutput createNewJob(final JobCreationInput input) { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivity.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivity.java new file mode 100644 index 0000000000000..e8eac26d0725c --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivity.java @@ -0,0 +1,51 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import io.airbyte.metrics.lib.MetricAttribute; +import io.airbyte.metrics.lib.OssMetricsRegistry; +import io.airbyte.workers.temporal.scheduling.ConnectionUpdaterInput; +import io.temporal.activity.ActivityInterface; +import io.temporal.activity.ActivityMethod; +import java.util.Optional; +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.NoArgsConstructor; + +/** + * Custom Temporal activity that records metrics. + */ +@ActivityInterface +public interface RecordMetricActivity { + + enum FailureCause { + ACTIVITY, + CANCELED, + CONNECTION, + UNKNOWN, + WORKFLOW + } + + @Data + @NoArgsConstructor + @AllArgsConstructor + class RecordMetricInput { + + private ConnectionUpdaterInput connectionUpdaterInput; + private Optional failureCause; + private OssMetricsRegistry metricName; + private MetricAttribute[] metricAttributes; + + } + + /** + * Records a counter metric. + * + * @param metricInput The metric information. + */ + @ActivityMethod + void recordWorkflowCountMetric(final RecordMetricInput metricInput); + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImpl.java new file mode 100644 index 0000000000000..d3ee4c5d7614d --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImpl.java @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import io.airbyte.metrics.lib.MetricAttribute; +import io.airbyte.metrics.lib.MetricClient; +import io.airbyte.metrics.lib.MetricTags; +import io.airbyte.workers.temporal.scheduling.ConnectionUpdaterInput; +import io.micronaut.context.annotation.Requires; +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import javax.inject.Inject; +import javax.inject.Singleton; +import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; +import lombok.extern.slf4j.Slf4j; + +/** + * Implementation of the {@link RecordMetricActivity} that is managed by the application framework + * and therefore has access to other singletons managed by the framework. + */ +@AllArgsConstructor +@NoArgsConstructor +@Slf4j +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") +public class RecordMetricActivityImpl implements RecordMetricActivity { + + @Inject + private MetricClient metricClient; + + /** + * Records a workflow counter for the specified metric. + * + * @param metricInput The information about the metric to record. + */ + @Override + public void recordWorkflowCountMetric(final RecordMetricInput metricInput) { + final List baseMetricAttributes = generateMetricAttributes(metricInput.getConnectionUpdaterInput()); + if (metricInput.getMetricAttributes() != null) { + baseMetricAttributes.addAll(Stream.of(metricInput.getMetricAttributes()).collect(Collectors.toList())); + } + metricInput.getFailureCause().ifPresent(fc -> baseMetricAttributes.add(new MetricAttribute(MetricTags.RESET_WORKFLOW_FAILURE_CAUSE, fc.name()))); + metricClient.count(metricInput.getMetricName(), 1L, baseMetricAttributes.toArray(new MetricAttribute[] {})); + } + + /** + * Generates the list of {@link MetricAttribute}s to be included when recording a metric. + * + * @param connectionUpdaterInput The {@link ConnectionUpdaterInput} that represents the workflow to + * be executed. + * @return The list of {@link MetricAttribute}s to be included when recording a metric. + */ + private List generateMetricAttributes(final ConnectionUpdaterInput connectionUpdaterInput) { + final List metricAttributes = new ArrayList<>(); + metricAttributes.add(new MetricAttribute(MetricTags.CONNECTION_ID, String.valueOf(connectionUpdaterInput.getConnectionId()))); + return metricAttributes; + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivity.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivity.java new file mode 100644 index 0000000000000..ce7835cb4865e --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivity.java @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import io.temporal.activity.ActivityInterface; +import io.temporal.activity.ActivityMethod; +import java.util.UUID; +import lombok.AllArgsConstructor; +import lombok.Data; +import lombok.NoArgsConstructor; + +@ActivityInterface +public interface RouteToSyncTaskQueueActivity { + + @ActivityMethod + RouteToSyncTaskQueueOutput route(RouteToSyncTaskQueueInput input); + + @Data + @NoArgsConstructor + @AllArgsConstructor + class RouteToSyncTaskQueueInput { + + private UUID connectionId; + + } + + @Data + @NoArgsConstructor + @AllArgsConstructor + class RouteToSyncTaskQueueOutput { + + private String taskQueue; + + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivityImpl.java new file mode 100644 index 0000000000000..6bb3a3babca52 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/RouteToSyncTaskQueueActivityImpl.java @@ -0,0 +1,24 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import io.airbyte.workers.temporal.sync.RouterService; +import javax.inject.Inject; +import javax.inject.Singleton; + +@Singleton +public class RouteToSyncTaskQueueActivityImpl implements RouteToSyncTaskQueueActivity { + + @Inject + private RouterService routerService; + + @Override + public RouteToSyncTaskQueueOutput route(final RouteToSyncTaskQueueInput input) { + final String taskQueueForConnectionId = routerService.getTaskQueue(input.getConnectionId()); + + return new RouteToSyncTaskQueueOutput(taskQueueForConnectionId); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityImpl.java index e48d057891584..d3bd511696f0e 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityImpl.java @@ -4,49 +4,28 @@ package io.airbyte.workers.temporal.scheduling.activities; -import io.airbyte.config.JobConfig.ConfigType; -import io.airbyte.config.persistence.StreamResetPersistence; -import io.airbyte.protocol.models.StreamDescriptor; -import io.airbyte.scheduler.models.Job; -import io.airbyte.scheduler.persistence.JobPersistence; -import io.airbyte.workers.temporal.exception.RetryableException; -import java.io.IOException; -import java.util.List; +import io.airbyte.workers.temporal.StreamResetRecordsHelper; +import io.micronaut.context.annotation.Requires; +import javax.inject.Inject; +import javax.inject.Singleton; import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; import lombok.extern.slf4j.Slf4j; @AllArgsConstructor +@NoArgsConstructor @Slf4j +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class StreamResetActivityImpl implements StreamResetActivity { - private StreamResetPersistence streamResetPersistence; - private JobPersistence jobPersistence; + @Inject + private StreamResetRecordsHelper streamResetRecordsHelper; @Override public void deleteStreamResetRecordsForJob(final DeleteStreamResetRecordsForJobInput input) { - // if there is no job, there is nothing to delete - if (input.getJobId() == null) { - log.info("deleteStreamResetRecordsForJob was called with a null job id; returning."); - return; - } - - try { - final Job job = jobPersistence.getJob(input.getJobId()); - final ConfigType configType = job.getConfig().getConfigType(); - if (!ConfigType.RESET_CONNECTION.equals(configType)) { - log.info("deleteStreamResetRecordsForJob was called for job {} with config type {}. Returning, as config type is not {}.", - input.getJobId(), - configType, - ConfigType.RESET_CONNECTION); - return; - } - - final List resetStreams = job.getConfig().getResetConnection().getResetSourceConfiguration().getStreamsToReset(); - log.info("Deleting the following streams for reset job {} from the stream_reset table: {}", input.getJobId(), resetStreams); - streamResetPersistence.deleteStreamResets(input.getConnectionId(), resetStreams); - } catch (final IOException e) { - throw new RetryableException(e); - } + streamResetRecordsHelper.deleteStreamResetRecordsForJob(input.getJobId(), input.getConnectionId()); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivity.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivity.java new file mode 100644 index 0000000000000..fb9bb00cb3349 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivity.java @@ -0,0 +1,26 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import io.temporal.activity.ActivityInterface; +import io.temporal.activity.ActivityMethod; +import java.time.Duration; + +/** + * Custom Temporal activity that can be used to retrieve configuration values managed by the + * application framework from the application context. + */ +@ActivityInterface +public interface WorkflowConfigActivity { + + /** + * Fetches the configured workflow restart delay in seconds from the application context. + * + * @return The workflow restart delay in seconds. + */ + @ActivityMethod + Duration getWorkflowRestartDelaySeconds(); + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImpl.java new file mode 100644 index 0000000000000..3e8e46798d2ff --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImpl.java @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import com.google.common.annotations.VisibleForTesting; +import io.micronaut.context.annotation.Property; +import io.micronaut.context.annotation.Requires; +import java.time.Duration; +import javax.inject.Singleton; +import lombok.extern.slf4j.Slf4j; + +/** + * Implementation of the {@link WorkflowConfigActivity} that is managed by the application framework + * and therefore has access to the configuration loaded by the framework. + */ +@Slf4j +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") +public class WorkflowConfigActivityImpl implements WorkflowConfigActivity { + + @Property(name = "airbyte.workflow.failure.restart-delay", + defaultValue = "600") + private Long workflowRestartDelaySeconds; + + @Override + public Duration getWorkflowRestartDelaySeconds() { + return Duration.ofSeconds(workflowRestartDelaySeconds); + } + + @VisibleForTesting + void setWorkflowRestartDelaySeconds(final Long workflowRestartDelaySeconds) { + this.workflowRestartDelaySeconds = workflowRestartDelaySeconds; + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/shared/ActivityConfiguration.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/shared/ActivityConfiguration.java deleted file mode 100644 index 212c679e9221f..0000000000000 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/scheduling/shared/ActivityConfiguration.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * Copyright (c) 2022 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.workers.temporal.scheduling.shared; - -import io.airbyte.config.Configs; -import io.airbyte.config.EnvConfigs; -import io.airbyte.workers.exception.WorkerException; -import io.airbyte.workers.temporal.TemporalUtils; -import io.temporal.activity.ActivityCancellationType; -import io.temporal.activity.ActivityOptions; -import io.temporal.common.RetryOptions; -import java.time.Duration; - -/** - * Shared temporal workflow configuration in order to ensure that - * {@link io.airbyte.workers.temporal.scheduling.ConnectionManagerWorkflow} and - * {@link io.airbyte.workers.temporal.sync.SyncWorkflow} configurations are on sync, especially for - * the grace period. - */ -public class ActivityConfiguration { - - private static final Configs configs = new EnvConfigs(); - - private static final int MAX_SYNC_TIMEOUT_DAYS = configs.getSyncJobMaxTimeoutDays(); - private static final Duration DB_INTERACTION_TIMEOUT = Duration.ofSeconds(configs.getMaxActivityTimeoutSecond()); - - // retry infinitely if the worker is killed without exceptions and dies due to timeouts - // but fail for everything thrown by the call itself which is rethrown as runtime exceptions - private static final RetryOptions ORCHESTRATOR_RETRY = RetryOptions.newBuilder() - .setDoNotRetry(RuntimeException.class.getName(), WorkerException.class.getName()) - .build(); - - private static final RetryOptions RETRY_POLICY = new EnvConfigs().getContainerOrchestratorEnabled() ? ORCHESTRATOR_RETRY : TemporalUtils.NO_RETRY; - - public static final ActivityOptions LONG_RUN_OPTIONS = ActivityOptions.newBuilder() - .setScheduleToCloseTimeout(Duration.ofDays(MAX_SYNC_TIMEOUT_DAYS)) - .setStartToCloseTimeout(Duration.ofDays(MAX_SYNC_TIMEOUT_DAYS)) - .setScheduleToStartTimeout(Duration.ofDays(MAX_SYNC_TIMEOUT_DAYS)) - .setCancellationType(ActivityCancellationType.WAIT_CANCELLATION_COMPLETED) - .setRetryOptions(RETRY_POLICY) - .setHeartbeatTimeout(TemporalUtils.HEARTBEAT_TIMEOUT) - .build(); - - public static final ActivityOptions CHECK_ACTIVITY_OPTIONS = ActivityOptions.newBuilder() - .setScheduleToCloseTimeout(Duration.ofMinutes(5)) - .setRetryOptions(TemporalUtils.NO_RETRY) - .build(); - - public static final ActivityOptions SHORT_ACTIVITY_OPTIONS = ActivityOptions.newBuilder() - .setStartToCloseTimeout(DB_INTERACTION_TIMEOUT) - .setCancellationType(ActivityCancellationType.WAIT_CANCELLATION_COMPLETED) - .setRetryOptions(TemporalUtils.RETRY) - .setHeartbeatTimeout(TemporalUtils.HEARTBEAT_TIMEOUT) - .build(); - -} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java index 0ede06fa40551..b8b607bc59e60 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecActivityImpl.java @@ -20,36 +20,38 @@ import io.airbyte.workers.process.ProcessFactory; import io.airbyte.workers.temporal.CancellationHandler; import io.airbyte.workers.temporal.TemporalAttemptExecution; +import io.micronaut.context.annotation.Requires; +import io.micronaut.context.annotation.Value; import io.temporal.activity.Activity; import io.temporal.activity.ActivityExecutionContext; import java.nio.file.Path; import java.util.function.Supplier; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +@Singleton +@Requires(property = "airbyte.worker.plane", + notEquals = "DATA_PLANE") public class SpecActivityImpl implements SpecActivity { - private final WorkerConfigs workerConfigs; - private final ProcessFactory processFactory; - private final Path workspaceRoot; - private final WorkerEnvironment workerEnvironment; - private final LogConfigs logConfigs; - private final AirbyteApiClient airbyteApiClient; - private final String airbyteVersion; - - public SpecActivityImpl(final WorkerConfigs workerConfigs, - final ProcessFactory processFactory, - final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final LogConfigs logConfigs, - final AirbyteApiClient airbyteApiClient, - final String airbyteVersion) { - this.workerConfigs = workerConfigs; - this.processFactory = processFactory; - this.workspaceRoot = workspaceRoot; - this.workerEnvironment = workerEnvironment; - this.logConfigs = logConfigs; - this.airbyteApiClient = airbyteApiClient; - this.airbyteVersion = airbyteVersion; - } + @Inject + @Named("specWorkerConfigs") + private WorkerConfigs workerConfigs; + @Inject + @Named("specProcessFactory") + private ProcessFactory processFactory; + @Inject + @Named("workspaceRoot") + private Path workspaceRoot; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private LogConfigs logConfigs; + @Inject + private AirbyteApiClient airbyteApiClient; + @Value("${airbyte.version}") + private String airbyteVersion; @Override public ConnectorJobOutput run(final JobRunConfig jobRunConfig, final IntegrationLauncherConfig launcherConfig) { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecWorkflowImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecWorkflowImpl.java index 2ed95b062c1f8..a2ee201da0189 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecWorkflowImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/spec/SpecWorkflowImpl.java @@ -7,18 +7,14 @@ import io.airbyte.config.ConnectorJobOutput; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.workers.temporal.TemporalUtils; -import io.temporal.activity.ActivityOptions; -import io.temporal.workflow.Workflow; -import java.time.Duration; +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; +import javax.inject.Singleton; +@Singleton public class SpecWorkflowImpl implements SpecWorkflow { - final ActivityOptions options = ActivityOptions.newBuilder() - .setScheduleToCloseTimeout(Duration.ofHours(1)) - .setRetryOptions(TemporalUtils.NO_RETRY) - .build(); - private final SpecActivity activity = Workflow.newActivityStub(SpecActivity.class, options); + @TemporalActivityStub(activityOptionsBeanName = "specActivityOptions") + private SpecActivity activity; @Override public ConnectorJobOutput run(final JobRunConfig jobRunConfig, final IntegrationLauncherConfig launcherConfig) { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubGeneratorFunction.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubGeneratorFunction.java new file mode 100644 index 0000000000000..79365ef202936 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubGeneratorFunction.java @@ -0,0 +1,21 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.support; + +import io.temporal.activity.ActivityOptions; + +/** + * Functional interface that defines the function used to generate a Temporal activity stub. + * + * @param The Temporal activity stub class. + * @param The {@link ActivityOptions} for the Temporal activity stub. + * @param The Temporal activity stub object. + */ +@FunctionalInterface +public interface TemporalActivityStubGeneratorFunction, A extends ActivityOptions, O> { + + O apply(C c, A a); + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptor.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptor.java new file mode 100644 index 0000000000000..0ab20eb9f29b3 --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptor.java @@ -0,0 +1,171 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.support; + +import com.google.common.annotations.VisibleForTesting; +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; +import io.micronaut.context.BeanRegistration; +import io.temporal.activity.ActivityOptions; +import io.temporal.workflow.Workflow; +import java.lang.reflect.Field; +import java.util.Collection; +import java.util.Optional; +import java.util.concurrent.Callable; +import lombok.extern.slf4j.Slf4j; +import net.bytebuddy.implementation.bind.annotation.RuntimeType; +import net.bytebuddy.implementation.bind.annotation.SuperCall; +import net.bytebuddy.implementation.bind.annotation.This; +import org.springframework.util.ReflectionUtils; + +/** + * Custom interceptor that handles invocations of Temporal workflow implementations to ensure that + * any and all Temporal activity stubs are created prior to the first execution of the workflow. + * This class is used in conjunction with {@link TemporalProxyHelper}. This approach is inspired by + * https://github.com/applicaai/spring-boot-starter-temporal. + * + * @param The type of the Temporal workflow. + */ +@Slf4j +public class TemporalActivityStubInterceptor { + + /** + * Function that generates Temporal activity stubs. + * + * Replace this value for unit testing. + */ + private TemporalActivityStubGeneratorFunction, ActivityOptions, Object> activityStubGenerator = Workflow::newActivityStub; + + /** + * The collection of configured {@link ActivityOptions} beans provided by the application framework. + */ + private final Collection> availableActivityOptions; + + /** + * The type of the workflow implementation to be proxied. + */ + private final Class workflowImplClass; + + /** + * Constructs a new interceptor for the provided workflow implementation class. + * + * @param workflowImplClass The Temporal workflow implementation class that will be intercepted. + * @param availableActivityOptions The collection of configured {@link ActivityOptions} beans + * provided by the application framework. + */ + public TemporalActivityStubInterceptor(final Class workflowImplClass, + final Collection> availableActivityOptions) { + this.workflowImplClass = workflowImplClass; + this.availableActivityOptions = availableActivityOptions; + } + + /** + * Main interceptor method that will be invoked by the proxy. + * + * @param workflowImplInstance The actual workflow implementation object invoked on the proxy + * Temporal workflow instance. + * @param call A {@link Callable} used to invoke the proxied method. + * @return The result of the proxied method execution. + * @throws Exception if the proxied method throws a checked exception + * @throws IllegalStateException if the Temporal activity stubs associated with the workflow cannot + * be initialized. + */ + @RuntimeType + public Object execute(@This final T workflowImplInstance, @SuperCall final Callable call) + throws Exception { + // Initialize the activity stubs, if not already done, before execution of the workflow method + initializeActivityStubs(workflowImplClass, workflowImplInstance); + return call.call(); + } + + /** + * Initializes all Temporal activity stubs present on the provided workflow instance. A Temporal + * activity stub is denoted by the use of the {@link TemporalActivityStub} annotation on the field. + * + * @param workflowImplClass The target class of the proxy. + * @param workflowInstance The workflow instance that may contain Temporal activity stub fields. + */ + private void initializeActivityStubs(final Class workflowImplClass, + final T workflowInstance) { + for (final Field field : workflowImplClass.getDeclaredFields()) { + if (field.isAnnotationPresent(TemporalActivityStub.class)) { + initializeActivityStub(workflowInstance, field); + } + } + } + + /** + * Initializes the Temporal activity stub represented by the provided field on the provided object, + * if not already set. + * + * @param workflowInstance The Temporal workflow instance that contains the Temporal activity stub + * field. + * @param activityStubField The field that represents the Temporal activity stub. + */ + private void initializeActivityStub(final T workflowInstance, + final Field activityStubField) { + try { + log.debug("Attempting to initialize Temporal activity stub for activity '{}' on workflow '{}'...", activityStubField.getType(), + workflowInstance.getClass().getName()); + ReflectionUtils.makeAccessible(activityStubField); + if (activityStubField.get(workflowInstance) == null) { + final ActivityOptions activityOptions = getActivityOptions(activityStubField); + final Object activityStub = generateActivityStub(activityStubField, activityOptions); + activityStubField.set(workflowInstance, activityStub); + log.debug("Initialized Temporal activity stub for activity '{}' for workflow '{}'.", activityStubField.getType(), + workflowInstance.getClass().getName()); + } else { + log.debug("Temporal activity stub '{}' is already initialized for Temporal workflow '{}'.", + activityStubField.getType(), + workflowInstance.getClass().getName()); + } + } catch (final IllegalArgumentException | IllegalAccessException | IllegalStateException e) { + log.error("Unable to initialize Temporal activity stub for activity '{}' for workflow '{}'.", activityStubField.getType(), + workflowInstance.getClass().getName(), e); + throw new RuntimeException(e); + } + } + + /** + * Extracts the Temporal {@link ActivityOptions} from the {@link Field} on the provided target + * instance object. + * + * @param activityStubField The field that represents the Temporal activity stub. + * @return The Temporal {@link ActivityOptions} from the {@link Field} on the provided Temporal + * workflow instance object. + * @throws IllegalStateException if the referenced Temporal {@link ActivityOptions} bean cannot be + * located. + */ + private ActivityOptions getActivityOptions(final Field activityStubField) { + final TemporalActivityStub annotation = activityStubField.getAnnotation(TemporalActivityStub.class); + final String activityOptionsBeanName = annotation.activityOptionsBeanName(); + final Optional selectedActivityOptions = + availableActivityOptions.stream().filter(b -> b.getIdentifier().getName().equalsIgnoreCase(activityOptionsBeanName)).map(b -> b.getBean()) + .findFirst(); + if (selectedActivityOptions.isPresent()) { + return selectedActivityOptions.get(); + } else { + throw new IllegalStateException("No activity options bean of name '" + activityOptionsBeanName + "' exists."); + } + } + + /** + * Retrieve the activity stub generator function associated with the Temporal activity stub. + * + * @param activityStubField The field that represents the Temporal activity stub. + * @return The {@link TemporalActivityStubGeneratorFunction} associated with the Temporal activity + * stub. + * @throws IllegalStateException if the referenced {@link TemporalActivityStubGeneratorFunction} + * bean cannot be located. + */ + private Object generateActivityStub(final Field activityStubField, final ActivityOptions activityOptions) { + return activityStubGenerator.apply(activityStubField.getType(), activityOptions); + } + + @VisibleForTesting + void setActivityStubGenerator(final TemporalActivityStubGeneratorFunction, ActivityOptions, Object> activityStubGenerator) { + this.activityStubGenerator = activityStubGenerator; + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalProxyHelper.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalProxyHelper.java new file mode 100644 index 0000000000000..722bdaba9a05e --- /dev/null +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/support/TemporalProxyHelper.java @@ -0,0 +1,136 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.support; + +import com.google.common.annotations.VisibleForTesting; +import io.micronaut.context.BeanRegistration; +import io.temporal.activity.ActivityOptions; +import io.temporal.workflow.QueryMethod; +import io.temporal.workflow.SignalMethod; +import io.temporal.workflow.WorkflowMethod; +import java.lang.annotation.Annotation; +import java.lang.reflect.Method; +import java.util.Collection; +import java.util.HashSet; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; +import javax.inject.Inject; +import javax.inject.Singleton; +import lombok.NoArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.TypeCache; +import net.bytebuddy.dynamic.loading.ClassLoadingStrategy; +import net.bytebuddy.implementation.MethodDelegation; +import net.bytebuddy.matcher.ElementMatchers; +import org.springframework.core.MethodIntrospector; +import org.springframework.core.annotation.AnnotationUtils; +import org.springframework.util.ReflectionUtils; + +/** + * Generates proxy classes which can be registered with Temporal. These proxies delegate all methods + * to the provided bean/singleton to allow for the dependency injection framework to manage the + * lifecycle of a Temporal workflow implementation. This approach is inspired by + * https://github.com/applicaai/spring-boot-starter-temporal. + */ +@Singleton +@NoArgsConstructor +@Slf4j +public class TemporalProxyHelper { + + /** + * Cache of already generated proxies to reduce the cost of creating and loading the proxies. + */ + private final TypeCache> WORKFLOW_PROXY_CACHE = new TypeCache<>(); + + /** + * Collection of available {@link ActivityOptions} beans which will be used to initialize Temporal + * activity stubs in each registered Temporal workflow. + */ + @Inject + private Collection> availableActivityOptions; + + private Optional, ActivityOptions, Object>> activityStubGenerator = Optional.empty(); + + public TemporalProxyHelper(final Collection> availableActivityOptions) { + this.availableActivityOptions = availableActivityOptions; + } + + /** + * Creates a proxy class for the given workflow class implementation and instance. + * + * @param workflowImplClass The workflow implementation class to proxy. proxy. + * @return A proxied workflow implementation class that can be registered with Temporal. + * @param The type of the workflow implementation class. + */ + @SuppressWarnings("PMD.UnnecessaryCast") + public Class proxyWorkflowClass(final Class workflowImplClass) { + log.debug("Creating a Temporal proxy for worker class '{}' with interface '{}'...", workflowImplClass.getName(), + workflowImplClass.getInterfaces()[0]); + return (Class) WORKFLOW_PROXY_CACHE.findOrInsert(workflowImplClass.getClassLoader(), workflowImplClass, () -> { + final Set workflowMethods = findAnnotatedMethods(workflowImplClass, WorkflowMethod.class); + final Set signalMethods = findAnnotatedMethods(workflowImplClass, SignalMethod.class); + final Set queryMethods = findAnnotatedMethods(workflowImplClass, QueryMethod.class); + + final Set proxiedMethods = new HashSet<>(); + proxiedMethods.add((Method) workflowMethods.toArray()[0]); + proxiedMethods.addAll(signalMethods.stream().collect(Collectors.toList())); + proxiedMethods.addAll(queryMethods.stream().collect(Collectors.toList())); + + final Class type = (Class) new ByteBuddy() + .subclass(workflowImplClass) + .name(workflowImplClass.getSimpleName() + "Proxy") + .implement(workflowImplClass.getInterfaces()[0]) + .method(ElementMatchers.anyOf(proxiedMethods.toArray(new Method[] {}))) + .intercept( + MethodDelegation.to(generateInterceptor(workflowImplClass, availableActivityOptions))) + .make() + .load(workflowImplClass.getClassLoader(), ClassLoadingStrategy.Default.WRAPPER) + .getLoaded(); + + log.debug("Temporal workflow proxy '{}' created for worker class '{}' with interface '{}'.", type.getName(), workflowImplClass.getName(), + workflowImplClass.getInterfaces()[0]); + return type; + }); + } + + /** + * Finds the methods annotated with the provided annotation type in the given class. + * + * @param workflowImplClass The workflow implementation class. + * @param annotationClass The annotation. + * @return The set of methods annotated with the provided annotation. + * @param The type of the annotation. + */ + private Set findAnnotatedMethods(final Class workflowImplClass, final Class annotationClass) { + return MethodIntrospector.selectMethods( + workflowImplClass, + (ReflectionUtils.MethodFilter) method -> AnnotationUtils.findAnnotation(method, annotationClass) != null); + } + + /** + * Generates a {@link TemporalActivityStubInterceptor} instance for use with the generated proxy + * workflow implementation. + * + * @param workflowImplClass The workflow implementation class. + * @param activityOptions The collection of available {@link ActivityOptions} beans which will be + * used to initialize Temporal activity stubs in each registered Temporal workflow. + * @return The generated {@link TemporalActivityStubInterceptor} instance. + * @param The workflow implementation type. + */ + private TemporalActivityStubInterceptor generateInterceptor(final Class workflowImplClass, + final Collection> activityOptions) { + final TemporalActivityStubInterceptor interceptor = new TemporalActivityStubInterceptor(workflowImplClass, activityOptions); + activityStubGenerator.ifPresent(a -> interceptor.setActivityStubGenerator(a)); + return interceptor; + } + + @VisibleForTesting + void setActivityStubGenerator(final TemporalActivityStubGeneratorFunction, ActivityOptions, Object> activityStubGenerator) { + this.activityStubGenerator = Optional.ofNullable(activityStubGenerator); + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtLauncherWorker.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtLauncherWorker.java index be4e2584d4369..e8b8a963c8af9 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtLauncherWorker.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtLauncherWorker.java @@ -8,8 +8,9 @@ import io.airbyte.config.OperatorDbtInput; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.workers.WorkerApp; +import io.airbyte.workers.ContainerOrchestratorConfig; import io.airbyte.workers.WorkerConfigs; +import io.airbyte.workers.temporal.TemporalUtils; import io.temporal.activity.ActivityExecutionContext; import java.util.Map; import java.util.UUID; @@ -25,8 +26,10 @@ public DbtLauncherWorker(final UUID connectionId, final IntegrationLauncherConfig destinationLauncherConfig, final JobRunConfig jobRunConfig, final WorkerConfigs workerConfigs, - final WorkerApp.ContainerOrchestratorConfig containerOrchestratorConfig, - final Supplier activityContext) { + final ContainerOrchestratorConfig containerOrchestratorConfig, + final Supplier activityContext, + final Integer serverPort, + final TemporalUtils temporalUtils) { super( connectionId, DBT, @@ -37,7 +40,9 @@ public DbtLauncherWorker(final UUID connectionId, containerOrchestratorConfig, workerConfigs.getResourceRequirements(), Void.class, - activityContext); + activityContext, + serverPort, + temporalUtils); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtTransformationActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtTransformationActivityImpl.java index 1815502aa6865..726145bf86d93 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtTransformationActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/DbtTransformationActivityImpl.java @@ -5,6 +5,8 @@ package io.airbyte.workers.temporal.sync; import io.airbyte.api.client.AirbyteApiClient; +import io.airbyte.api.client.invoker.generated.ApiException; +import io.airbyte.api.client.model.generated.JobIdRequestBody; import io.airbyte.commons.functional.CheckedSupplier; import io.airbyte.commons.json.Jsons; import io.airbyte.config.AirbyteConfigValidator; @@ -16,9 +18,8 @@ import io.airbyte.config.persistence.split_secrets.SecretsHydrator; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.scheduler.persistence.JobPersistence; +import io.airbyte.workers.ContainerOrchestratorConfig; import io.airbyte.workers.Worker; -import io.airbyte.workers.WorkerApp; import io.airbyte.workers.WorkerConfigs; import io.airbyte.workers.general.DbtTransformationRunner; import io.airbyte.workers.general.DbtTransformationWorker; @@ -27,51 +28,48 @@ import io.airbyte.workers.temporal.CancellationHandler; import io.airbyte.workers.temporal.TemporalAttemptExecution; import io.airbyte.workers.temporal.TemporalUtils; +import io.micronaut.context.annotation.Value; import io.temporal.activity.Activity; import io.temporal.activity.ActivityExecutionContext; -import java.io.IOException; import java.nio.file.Path; import java.util.Optional; import java.util.UUID; import java.util.function.Supplier; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +@Singleton public class DbtTransformationActivityImpl implements DbtTransformationActivity { - private final WorkerConfigs workerConfigs; - private final ProcessFactory jobProcessFactory; - private final SecretsHydrator secretsHydrator; - private final Path workspaceRoot; - private final AirbyteConfigValidator validator; - private final WorkerEnvironment workerEnvironment; - private final LogConfigs logConfigs; - private final JobPersistence jobPersistence; - - private final AirbyteApiClient airbyteApiClient; - private final String airbyteVersion; - private final Optional containerOrchestratorConfig; - - public DbtTransformationActivityImpl(final Optional containerOrchestratorConfig, - final WorkerConfigs workerConfigs, - final ProcessFactory jobProcessFactory, - final SecretsHydrator secretsHydrator, - final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final LogConfigs logConfigs, - final JobPersistence jobPersistence, - final AirbyteApiClient airbyteApiClient, - final String airbyteVersion) { - this.containerOrchestratorConfig = containerOrchestratorConfig; - this.workerConfigs = workerConfigs; - this.jobProcessFactory = jobProcessFactory; - this.secretsHydrator = secretsHydrator; - this.workspaceRoot = workspaceRoot; - this.validator = new AirbyteConfigValidator(); - this.workerEnvironment = workerEnvironment; - this.logConfigs = logConfigs; - this.jobPersistence = jobPersistence; - this.airbyteApiClient = airbyteApiClient; - this.airbyteVersion = airbyteVersion; - } + @Inject + @Named("containerOrchestratorConfig") + private Optional containerOrchestratorConfig; + @Inject + @Named("defaultWorkerConfigs") + private WorkerConfigs workerConfigs; + @Inject + @Named("defaultProcessFactory") + private ProcessFactory processFactory; + @Inject + private SecretsHydrator secretsHydrator; + @Inject + @Named("workspaceRoot") + private Path workspaceRoot; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private LogConfigs logConfigs; + @Value("${airbyte.version}") + private String airbyteVersion; + @Value("${micronaut.server.port}") + private Integer serverPort; + @Inject + private AirbyteConfigValidator airbyteConfigValidator; + @Inject + private TemporalUtils temporalUtils; + @Inject + private AirbyteApiClient airbyteApiClient; @Override public Void run(final JobRunConfig jobRunConfig, @@ -79,13 +77,13 @@ public Void run(final JobRunConfig jobRunConfig, final ResourceRequirements resourceRequirements, final OperatorDbtInput input) { final ActivityExecutionContext context = Activity.getExecutionContext(); - return TemporalUtils.withBackgroundHeartbeat( + return temporalUtils.withBackgroundHeartbeat( () -> { final var fullDestinationConfig = secretsHydrator.hydrate(input.getDestinationConfiguration()); final var fullInput = Jsons.clone(input).withDestinationConfiguration(fullDestinationConfig); final Supplier inputSupplier = () -> { - validator.ensureAsRuntime(ConfigSchema.OPERATOR_DBT_INPUT, Jsons.jsonNode(fullInput)); + airbyteConfigValidator.ensureAsRuntime(ConfigSchema.OPERATOR_DBT_INPUT, Jsons.jsonNode(fullInput)); return fullInput; }; @@ -124,10 +122,10 @@ private CheckedSupplier, Exception> getLegacyWork resourceRequirements, new DbtTransformationRunner( workerConfigs, - jobProcessFactory, NormalizationRunnerFactory.create( + processFactory, NormalizationRunnerFactory.create( workerConfigs, destinationLauncherConfig.getDockerImage(), - jobProcessFactory, + processFactory, NormalizationRunnerFactory.NORMALIZATION_VERSION))); } @@ -136,8 +134,10 @@ private CheckedSupplier, Exception> getContainerL final IntegrationLauncherConfig destinationLauncherConfig, final JobRunConfig jobRunConfig, final Supplier activityContext) - throws IOException { - final var jobScope = jobPersistence.getJob(Long.parseLong(jobRunConfig.getJobId())).getScope(); + throws ApiException { + final JobIdRequestBody id = new JobIdRequestBody(); + id.setId(Long.valueOf(jobRunConfig.getJobId())); + final var jobScope = airbyteApiClient.getJobsApi().getJobInfo(id).getJob().getConfigId(); final var connectionId = UUID.fromString(jobScope); return () -> new DbtLauncherWorker( @@ -146,7 +146,9 @@ private CheckedSupplier, Exception> getContainerL jobRunConfig, workerConfigs, containerOrchestratorConfig.get(), - activityContext); + activityContext, + serverPort, + temporalUtils); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/LauncherWorker.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/LauncherWorker.java index d58951414a136..a03bd74d4fedb 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/LauncherWorker.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/LauncherWorker.java @@ -10,8 +10,8 @@ import io.airbyte.commons.lang.Exceptions; import io.airbyte.config.ResourceRequirements; import io.airbyte.scheduler.models.JobRunConfig; +import io.airbyte.workers.ContainerOrchestratorConfig; import io.airbyte.workers.Worker; -import io.airbyte.workers.WorkerApp; import io.airbyte.workers.exception.WorkerException; import io.airbyte.workers.process.AsyncKubePodStatus; import io.airbyte.workers.process.AsyncOrchestratorPodProcess; @@ -56,10 +56,12 @@ public class LauncherWorker implements Worker { private final String podNamePrefix; private final JobRunConfig jobRunConfig; private final Map additionalFileMap; - private final WorkerApp.ContainerOrchestratorConfig containerOrchestratorConfig; + private final ContainerOrchestratorConfig containerOrchestratorConfig; private final ResourceRequirements resourceRequirements; private final Class outputClass; private final Supplier activityContext; + private final Integer serverPort; + private final TemporalUtils temporalUtils; private final AtomicBoolean cancelled = new AtomicBoolean(false); private AsyncOrchestratorPodProcess process; @@ -69,10 +71,13 @@ public LauncherWorker(final UUID connectionId, final String podNamePrefix, final JobRunConfig jobRunConfig, final Map additionalFileMap, - final WorkerApp.ContainerOrchestratorConfig containerOrchestratorConfig, + final ContainerOrchestratorConfig containerOrchestratorConfig, final ResourceRequirements resourceRequirements, final Class outputClass, - final Supplier activityContext) { + final Supplier activityContext, + final Integer serverPort, + final TemporalUtils temporalUtils) { + this.connectionId = connectionId; this.application = application; this.podNamePrefix = podNamePrefix; @@ -82,13 +87,15 @@ public LauncherWorker(final UUID connectionId, this.resourceRequirements = resourceRequirements; this.outputClass = outputClass; this.activityContext = activityContext; + this.serverPort = serverPort; + this.temporalUtils = temporalUtils; } @Override public OUTPUT run(final INPUT input, final Path jobRoot) throws WorkerException { final AtomicBoolean isCanceled = new AtomicBoolean(false); final AtomicReference cancellationCallback = new AtomicReference<>(null); - return TemporalUtils.withBackgroundHeartbeat(cancellationCallback, () -> { + return temporalUtils.withBackgroundHeartbeat(cancellationCallback, () -> { try { final Map envMap = System.getenv().entrySet().stream() .filter(entry -> OrchestratorConstants.ENV_VARS_TO_TRANSFER.contains(entry.getKey())) @@ -102,7 +109,7 @@ public OUTPUT run(final INPUT input, final Path jobRoot) throws WorkerException OrchestratorConstants.INIT_FILE_ENV_MAP, Jsons.serialize(envMap))); final Map portMap = Map.of( - WorkerApp.KUBE_HEARTBEAT_PORT, WorkerApp.KUBE_HEARTBEAT_PORT, + serverPort, serverPort, OrchestratorConstants.PORT1, OrchestratorConstants.PORT1, OrchestratorConstants.PORT2, OrchestratorConstants.PORT2, OrchestratorConstants.PORT3, OrchestratorConstants.PORT3, @@ -129,7 +136,8 @@ public OUTPUT run(final INPUT input, final Path jobRoot) throws WorkerException containerOrchestratorConfig.secretName(), containerOrchestratorConfig.secretMountPath(), containerOrchestratorConfig.googleApplicationCredentials(), - featureFlag.useStreamCapableState()); + featureFlag.useStreamCapableState(), + serverPort); cancellationCallback.set(() -> { // When cancelled, try to set to true. diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivity.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivity.java index 3ed1af99ac401..02de62c27ba6b 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivity.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivity.java @@ -6,6 +6,8 @@ import io.airbyte.config.NormalizationInput; import io.airbyte.config.NormalizationSummary; +import io.airbyte.config.StandardSyncInput; +import io.airbyte.config.StandardSyncOutput; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; import io.temporal.activity.ActivityInterface; @@ -19,4 +21,7 @@ NormalizationSummary normalize(JobRunConfig jobRunConfig, IntegrationLauncherConfig destinationLauncherConfig, NormalizationInput input); + @ActivityMethod + NormalizationInput generateNormalizationInput(final StandardSyncInput syncInput, final StandardSyncOutput syncOutput); + } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivityImpl.java index cc1e6b2037e36..a675ede1550e3 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationActivityImpl.java @@ -5,6 +5,8 @@ package io.airbyte.workers.temporal.sync; import io.airbyte.api.client.AirbyteApiClient; +import io.airbyte.api.client.invoker.generated.ApiException; +import io.airbyte.api.client.model.generated.JobIdRequestBody; import io.airbyte.commons.functional.CheckedSupplier; import io.airbyte.commons.json.Jsons; import io.airbyte.config.AirbyteConfigValidator; @@ -12,13 +14,15 @@ import io.airbyte.config.Configs.WorkerEnvironment; import io.airbyte.config.NormalizationInput; import io.airbyte.config.NormalizationSummary; +import io.airbyte.config.ResourceRequirements; +import io.airbyte.config.StandardSyncInput; +import io.airbyte.config.StandardSyncOutput; import io.airbyte.config.helpers.LogConfigs; import io.airbyte.config.persistence.split_secrets.SecretsHydrator; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.scheduler.persistence.JobPersistence; +import io.airbyte.workers.ContainerOrchestratorConfig; import io.airbyte.workers.Worker; -import io.airbyte.workers.WorkerApp; import io.airbyte.workers.WorkerConfigs; import io.airbyte.workers.general.DefaultNormalizationWorker; import io.airbyte.workers.normalization.NormalizationRunnerFactory; @@ -26,63 +30,63 @@ import io.airbyte.workers.temporal.CancellationHandler; import io.airbyte.workers.temporal.TemporalAttemptExecution; import io.airbyte.workers.temporal.TemporalUtils; +import io.micronaut.context.annotation.Value; import io.temporal.activity.Activity; import io.temporal.activity.ActivityExecutionContext; -import java.io.IOException; import java.nio.file.Path; import java.util.Optional; import java.util.UUID; import java.util.function.Supplier; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; +@Singleton public class NormalizationActivityImpl implements NormalizationActivity { - private final WorkerConfigs workerConfigs; - private final ProcessFactory jobProcessFactory; - private final SecretsHydrator secretsHydrator; - private final Path workspaceRoot; - private final AirbyteConfigValidator validator; - private final WorkerEnvironment workerEnvironment; - private final LogConfigs logConfigs; - private final JobPersistence jobPersistence; - - private final AirbyteApiClient airbyteApiClient; - private final String airbyteVersion; - private final Optional containerOrchestratorConfig; - - public NormalizationActivityImpl(final Optional containerOrchestratorConfig, - final WorkerConfigs workerConfigs, - final ProcessFactory jobProcessFactory, - final SecretsHydrator secretsHydrator, - final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final LogConfigs logConfigs, - final JobPersistence jobPersistence, - final AirbyteApiClient airbyteApiClient, - final String airbyteVersion) { - this.containerOrchestratorConfig = containerOrchestratorConfig; - this.workerConfigs = workerConfigs; - this.jobProcessFactory = jobProcessFactory; - this.secretsHydrator = secretsHydrator; - this.workspaceRoot = workspaceRoot; - this.validator = new AirbyteConfigValidator(); - this.workerEnvironment = workerEnvironment; - this.logConfigs = logConfigs; - this.jobPersistence = jobPersistence; - this.airbyteApiClient = airbyteApiClient; - this.airbyteVersion = airbyteVersion; - } + @Inject + @Named("containerOrchestratorConfig") + private Optional containerOrchestratorConfig; + @Inject + @Named("defaultWorkerConfigs") + private WorkerConfigs workerConfigs; + @Inject + @Named("defaultProcessFactory") + private ProcessFactory processFactory; + @Inject + private SecretsHydrator secretsHydrator; + @Inject + @Named("workspaceRoot") + private Path workspaceRoot; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private LogConfigs logConfigs; + @Value("${airbyte.version}") + private String airbyteVersion; + @Value("${micronaut.server.port}") + private Integer serverPort; + @Inject + private AirbyteConfigValidator airbyteConfigValidator; + @Inject + private TemporalUtils temporalUtils; + @Inject + @Named("normalizationResourceRequirements") + private ResourceRequirements normalizationResourceRequirements; + @Inject + private AirbyteApiClient airbyteApiClient; @Override public NormalizationSummary normalize(final JobRunConfig jobRunConfig, final IntegrationLauncherConfig destinationLauncherConfig, final NormalizationInput input) { final ActivityExecutionContext context = Activity.getExecutionContext(); - return TemporalUtils.withBackgroundHeartbeat(() -> { + return temporalUtils.withBackgroundHeartbeat(() -> { final var fullDestinationConfig = secretsHydrator.hydrate(input.getDestinationConfiguration()); final var fullInput = Jsons.clone(input).withDestinationConfiguration(fullDestinationConfig); final Supplier inputSupplier = () -> { - validator.ensureAsRuntime(ConfigSchema.NORMALIZATION_INPUT, Jsons.jsonNode(fullInput)); + airbyteConfigValidator.ensureAsRuntime(ConfigSchema.NORMALIZATION_INPUT, Jsons.jsonNode(fullInput)); return fullInput; }; @@ -110,6 +114,14 @@ public NormalizationSummary normalize(final JobRunConfig jobRunConfig, () -> context); } + @Override + public NormalizationInput generateNormalizationInput(final StandardSyncInput syncInput, final StandardSyncOutput syncOutput) { + return new NormalizationInput() + .withDestinationConfiguration(syncInput.getDestinationConfiguration()) + .withCatalog(syncOutput.getOutputCatalog()) + .withResourceRequirements(normalizationResourceRequirements); + } + private CheckedSupplier, Exception> getLegacyWorkerFactory( final WorkerConfigs workerConfigs, final IntegrationLauncherConfig destinationLauncherConfig, @@ -120,7 +132,7 @@ private CheckedSupplier, Except NormalizationRunnerFactory.create( workerConfigs, destinationLauncherConfig.getDockerImage(), - jobProcessFactory, + processFactory, NormalizationRunnerFactory.NORMALIZATION_VERSION), workerEnvironment); } @@ -130,8 +142,10 @@ private CheckedSupplier, Except final IntegrationLauncherConfig destinationLauncherConfig, final JobRunConfig jobRunConfig, final Supplier activityContext) - throws IOException { - final var jobScope = jobPersistence.getJob(Long.parseLong(jobRunConfig.getJobId())).getScope(); + throws ApiException { + final JobIdRequestBody id = new JobIdRequestBody(); + id.setId(Long.valueOf(jobRunConfig.getJobId())); + final var jobScope = airbyteApiClient.getJobsApi().getJobInfo(id).getJob().getConfigId(); final var connectionId = UUID.fromString(jobScope); return () -> new NormalizationLauncherWorker( connectionId, @@ -139,7 +153,9 @@ private CheckedSupplier, Except jobRunConfig, workerConfigs, containerOrchestratorConfig.get(), - activityContext); + activityContext, + serverPort, + temporalUtils); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationLauncherWorker.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationLauncherWorker.java index 40ca5f2245c82..7039eed6d58a4 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationLauncherWorker.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/NormalizationLauncherWorker.java @@ -9,8 +9,9 @@ import io.airbyte.config.NormalizationSummary; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.workers.WorkerApp; +import io.airbyte.workers.ContainerOrchestratorConfig; import io.airbyte.workers.WorkerConfigs; +import io.airbyte.workers.temporal.TemporalUtils; import io.temporal.activity.ActivityExecutionContext; import java.util.Map; import java.util.UUID; @@ -26,8 +27,10 @@ public NormalizationLauncherWorker(final UUID connectionId, final IntegrationLauncherConfig destinationLauncherConfig, final JobRunConfig jobRunConfig, final WorkerConfigs workerConfigs, - final WorkerApp.ContainerOrchestratorConfig containerOrchestratorConfig, - final Supplier activityContext) { + final ContainerOrchestratorConfig containerOrchestratorConfig, + final Supplier activityContext, + final Integer serverPort, + final TemporalUtils temporalUtils) { super( connectionId, NORMALIZATION, @@ -38,7 +41,9 @@ public NormalizationLauncherWorker(final UUID connectionId, containerOrchestratorConfig, workerConfigs.getResourceRequirements(), NormalizationSummary.class, - activityContext); + activityContext, + serverPort, + temporalUtils); } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/PersistStateActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/PersistStateActivityImpl.java index e9ebc9265e6df..1983d1e1ebdf8 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/PersistStateActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/PersistStateActivityImpl.java @@ -26,13 +26,20 @@ import java.util.List; import java.util.Optional; import java.util.UUID; +import javax.inject.Inject; +import javax.inject.Singleton; import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; @AllArgsConstructor +@NoArgsConstructor +@Singleton public class PersistStateActivityImpl implements PersistStateActivity { - private final AirbyteApiClient airbyteApiClient; - private final FeatureFlags featureFlags; + @Inject + private AirbyteApiClient airbyteApiClient; + @Inject + private FeatureFlags featureFlags; @Override public boolean persist(final UUID connectionId, final StandardSyncOutput syncOutput, final ConfiguredAirbyteCatalog configuredCatalog) { diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationActivityImpl.java index 0647c8966f78c..1a3129f42c6de 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationActivityImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationActivityImpl.java @@ -4,10 +4,10 @@ package io.airbyte.workers.temporal.sync; -import com.google.common.annotations.VisibleForTesting; import io.airbyte.api.client.AirbyteApiClient; import io.airbyte.api.client.invoker.generated.ApiException; import io.airbyte.api.client.model.generated.JobIdRequestBody; +import io.airbyte.commons.features.FeatureFlags; import io.airbyte.commons.functional.CheckedSupplier; import io.airbyte.commons.json.Jsons; import io.airbyte.config.AirbyteConfigValidator; @@ -25,10 +25,9 @@ import io.airbyte.metrics.lib.MetricEmittingApps; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; +import io.airbyte.workers.ContainerOrchestratorConfig; import io.airbyte.workers.RecordSchemaValidator; import io.airbyte.workers.Worker; -import io.airbyte.workers.WorkerApp; -import io.airbyte.workers.WorkerApp.ContainerOrchestratorConfig; import io.airbyte.workers.WorkerConfigs; import io.airbyte.workers.WorkerConstants; import io.airbyte.workers.WorkerMetricReporter; @@ -46,70 +45,54 @@ import io.airbyte.workers.temporal.CancellationHandler; import io.airbyte.workers.temporal.TemporalAttemptExecution; import io.airbyte.workers.temporal.TemporalUtils; +import io.micronaut.context.annotation.Value; import io.temporal.activity.Activity; import io.temporal.activity.ActivityExecutionContext; import java.nio.file.Path; import java.util.Optional; import java.util.UUID; import java.util.function.Supplier; +import javax.inject.Inject; +import javax.inject.Named; +import javax.inject.Singleton; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +@Singleton public class ReplicationActivityImpl implements ReplicationActivity { private static final Logger LOGGER = LoggerFactory.getLogger(ReplicationActivityImpl.class); - private final Optional containerOrchestratorConfig; - private final WorkerConfigs workerConfigs; - private final ProcessFactory processFactory; - private final SecretsHydrator secretsHydrator; - private final Path workspaceRoot; - private final AirbyteConfigValidator validator; - private final WorkerEnvironment workerEnvironment; - private final LogConfigs logConfigs; - - private final AirbyteApiClient airbyteApiClient; - private final String airbyteVersion; - private final boolean useStreamCapableState; - - public ReplicationActivityImpl(final Optional containerOrchestratorConfig, - final WorkerConfigs workerConfigs, - final ProcessFactory processFactory, - final SecretsHydrator secretsHydrator, - final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final LogConfigs logConfigs, - final AirbyteApiClient airbyteApiClient, - final String airbyteVersion, - final boolean useStreamCapableState) { - this(containerOrchestratorConfig, workerConfigs, processFactory, secretsHydrator, workspaceRoot, workerEnvironment, logConfigs, - new AirbyteConfigValidator(), airbyteApiClient, airbyteVersion, useStreamCapableState); - } - - @VisibleForTesting - ReplicationActivityImpl(final Optional containerOrchestratorConfig, - final WorkerConfigs workerConfigs, - final ProcessFactory processFactory, - final SecretsHydrator secretsHydrator, - final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final LogConfigs logConfigs, - final AirbyteConfigValidator validator, - final AirbyteApiClient airbyteApiClient, - final String airbyteVersion, - final boolean useStreamCapableState) { - this.containerOrchestratorConfig = containerOrchestratorConfig; - this.workerConfigs = workerConfigs; - this.processFactory = processFactory; - this.secretsHydrator = secretsHydrator; - this.workspaceRoot = workspaceRoot; - this.validator = validator; - this.workerEnvironment = workerEnvironment; - this.logConfigs = logConfigs; - this.airbyteApiClient = airbyteApiClient; - this.airbyteVersion = airbyteVersion; - this.useStreamCapableState = useStreamCapableState; - } + @Inject + @Named("containerOrchestratorConfig") + private Optional containerOrchestratorConfig; + @Inject + @Named("replicationWorkerConfigs") + private WorkerConfigs workerConfigs; + @Inject + @Named("replicationProcessFactory") + private ProcessFactory processFactory; + @Inject + private SecretsHydrator secretsHydrator; + @Inject + @Named("workspaceRoot") + private Path workspaceRoot; + @Inject + private WorkerEnvironment workerEnvironment; + @Inject + private LogConfigs logConfigs; + @Value("${airbyte.version}") + private String airbyteVersion; + @Inject + private FeatureFlags featureFlags; + @Value("${micronaut.server.port}") + private Integer serverPort; + @Inject + private AirbyteConfigValidator airbyteConfigValidator; + @Inject + private TemporalUtils temporalUtils; + @Inject + private AirbyteApiClient airbyteApiClient; @Override public StandardSyncOutput replicate(final JobRunConfig jobRunConfig, @@ -117,7 +100,7 @@ public StandardSyncOutput replicate(final JobRunConfig jobRunConfig, final IntegrationLauncherConfig destinationLauncherConfig, final StandardSyncInput syncInput) { final ActivityExecutionContext context = Activity.getExecutionContext(); - return TemporalUtils.withBackgroundHeartbeat( + return temporalUtils.withBackgroundHeartbeat( () -> { final var fullSourceConfig = secretsHydrator.hydrate(syncInput.getSourceConfiguration()); @@ -128,7 +111,7 @@ public StandardSyncOutput replicate(final JobRunConfig jobRunConfig, .withDestinationConfiguration(fullDestinationConfig); final Supplier inputSupplier = () -> { - validator.ensureAsRuntime(ConfigSchema.STANDARD_SYNC_INPUT, Jsons.jsonNode(fullSyncInput)); + airbyteConfigValidator.ensureAsRuntime(ConfigSchema.STANDARD_SYNC_INPUT, Jsons.jsonNode(fullSyncInput)); return fullSyncInput; }; @@ -213,7 +196,7 @@ private CheckedSupplier, Exception> // reset jobs use an empty source to induce resetting all data in destination. final AirbyteSource airbyteSource = WorkerConstants.RESET_JOB_SOURCE_DOCKER_IMAGE_STUB.equals(sourceLauncherConfig.getDockerImage()) - ? new EmptyAirbyteSource(useStreamCapableState) + ? new EmptyAirbyteSource(featureFlags.useStreamCapableState()) : new DefaultAirbyteSource(workerConfigs, sourceLauncher); MetricClientFactory.initialize(MetricEmittingApps.WORKER); final MetricClient metricClient = MetricClientFactory.getMetricClient(); @@ -252,7 +235,9 @@ private CheckedSupplier, Exception> destinationLauncherConfig, jobRunConfig, resourceRequirements, - activityContext); + activityContext, + serverPort, + temporalUtils); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationLauncherWorker.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationLauncherWorker.java index 2e39139f12a04..0a2cba19bde4d 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationLauncherWorker.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/ReplicationLauncherWorker.java @@ -10,7 +10,8 @@ import io.airbyte.config.StandardSyncInput; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.workers.WorkerApp; +import io.airbyte.workers.ContainerOrchestratorConfig; +import io.airbyte.workers.temporal.TemporalUtils; import io.temporal.activity.ActivityExecutionContext; import java.util.Map; import java.util.UUID; @@ -29,12 +30,14 @@ public class ReplicationLauncherWorker extends LauncherWorker activityContext) { + final Supplier activityContext, + final Integer serverPort, + final TemporalUtils temporalUtils) { super( connectionId, REPLICATION, @@ -46,7 +49,9 @@ public ReplicationLauncherWorker(final UUID connectionId, containerOrchestratorConfig, resourceRequirements, ReplicationOutput.class, - activityContext); + activityContext, + serverPort, + temporalUtils); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivityImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivityImpl.java deleted file mode 100644 index c1821f35ab2c5..0000000000000 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivityImpl.java +++ /dev/null @@ -1,22 +0,0 @@ -/* - * Copyright (c) 2022 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.workers.temporal.sync; - -import java.util.UUID; - -public class RouteToTaskQueueActivityImpl implements RouteToTaskQueueActivity { - - private final RouterService routerService; - - public RouteToTaskQueueActivityImpl(final RouterService routerService) { - this.routerService = routerService; - } - - @Override - public String routeToTaskQueue(final UUID connectionId) { - return routerService.getTaskQueue(connectionId); - } - -} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouterService.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouterService.java index 3b2e05ddc45ee..84017b83e7785 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouterService.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouterService.java @@ -4,27 +4,45 @@ package io.airbyte.workers.temporal.sync; -import io.airbyte.config.Configs; +import com.google.common.annotations.VisibleForTesting; import io.airbyte.workers.temporal.TemporalJobType; +import io.micronaut.context.annotation.Value; +import io.micronaut.core.util.StringUtils; +import java.util.Arrays; +import java.util.Set; import java.util.UUID; -import lombok.AllArgsConstructor; +import java.util.stream.Collectors; +import javax.inject.Singleton; -@AllArgsConstructor +@Singleton public class RouterService { - private static final String MVP_DATA_PLANE_TASK_QUEUE = "MVP_DATA_PLANE"; + static final String MVP_DATA_PLANE_TASK_QUEUE = "MVP_DATA_PLANE"; - private final Configs configs; + @Value("${airbyte.data.plane.connection-ids-mvp}") + private String connectionIdsForMvpDataPlane; /** * For now, returns a Task Queue by checking to see if the connectionId is on the env var list for * usage in the MVP Data Plane. This will be replaced by a proper Router Service in the future. */ public String getTaskQueue(final UUID connectionId) { - if (configs.connectionIdsForMvpDataPlane().contains(connectionId.toString())) { - return MVP_DATA_PLANE_TASK_QUEUE; + if (connectionId != null) { + if (getConnectionIdsForMvpDataPlane().contains(connectionId.toString())) { + return MVP_DATA_PLANE_TASK_QUEUE; + } } return TemporalJobType.SYNC.name(); } + private Set getConnectionIdsForMvpDataPlane() { + return StringUtils.isNotEmpty(connectionIdsForMvpDataPlane) ? Arrays.stream(connectionIdsForMvpDataPlane.split(",")).collect(Collectors.toSet()) + : Set.of(); + } + + @VisibleForTesting + void setConnectionIdsForMvpDataPlane(final String connectionIdsForMvpDataPlane) { + this.connectionIdsForMvpDataPlane = connectionIdsForMvpDataPlane; + } + } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/SyncWorkflowImpl.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/SyncWorkflowImpl.java index 22b5d0282e02b..29d7cfa29277e 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/SyncWorkflowImpl.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/SyncWorkflowImpl.java @@ -4,12 +4,9 @@ package io.airbyte.workers.temporal.sync; -import io.airbyte.config.Configs; -import io.airbyte.config.EnvConfigs; import io.airbyte.config.NormalizationInput; import io.airbyte.config.NormalizationSummary; import io.airbyte.config.OperatorDbtInput; -import io.airbyte.config.ResourceRequirements; import io.airbyte.config.StandardSyncInput; import io.airbyte.config.StandardSyncOperation; import io.airbyte.config.StandardSyncOperation.OperatorType; @@ -17,19 +14,29 @@ import io.airbyte.protocol.models.ConfiguredAirbyteCatalog; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; -import io.airbyte.workers.temporal.scheduling.shared.ActivityConfiguration; -import io.temporal.activity.ActivityOptions; +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; import io.temporal.workflow.Workflow; import java.util.UUID; +import javax.inject.Singleton; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +@Singleton +@SuppressWarnings("PMD.AvoidDuplicateLiterals") public class SyncWorkflowImpl implements SyncWorkflow { private static final Logger LOGGER = LoggerFactory.getLogger(SyncWorkflowImpl.class); private static final String VERSION_LABEL = "sync-workflow"; private static final int CURRENT_VERSION = 2; - private static final int PREV_VERSION = 1; + + @TemporalActivityStub(activityOptionsBeanName = "longRunActivityOptions") + private ReplicationActivity replicationActivity; + @TemporalActivityStub(activityOptionsBeanName = "longRunActivityOptions") + private NormalizationActivity normalizationActivity; + @TemporalActivityStub(activityOptionsBeanName = "longRunActivityOptions") + private DbtTransformationActivity dbtTransformationActivity; + @TemporalActivityStub(activityOptionsBeanName = "shortActivityOptions") + private PersistStateActivity persistActivity; @Override public StandardSyncOutput run(final JobRunConfig jobRunConfig, @@ -39,38 +46,6 @@ public StandardSyncOutput run(final JobRunConfig jobRunConfig, final UUID connectionId) { final int version = Workflow.getVersion(VERSION_LABEL, Workflow.DEFAULT_VERSION, CURRENT_VERSION); - - final ReplicationActivity replicationActivity; - final NormalizationActivity normalizationActivity; - final DbtTransformationActivity dbtTransformationActivity; - final PersistStateActivity persistActivity; - - /** - * The current version calls a new activity to determine which Task Queue to use for other - * activities. The previous version doesn't call this new activity, and instead lets each activity - * inherit the workflow's Task Queue. - */ - if (version > PREV_VERSION) { - final RouteToTaskQueueActivity routeToTaskQueueActivity = - Workflow.newActivityStub(RouteToTaskQueueActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - - final String dataPlaneTaskQueue = routeToTaskQueueActivity.routeToTaskQueue(connectionId); - - replicationActivity = - Workflow.newActivityStub(ReplicationActivity.class, setTaskQueue(ActivityConfiguration.LONG_RUN_OPTIONS, dataPlaneTaskQueue)); - persistActivity = - Workflow.newActivityStub(PersistStateActivity.class, setTaskQueue(ActivityConfiguration.SHORT_ACTIVITY_OPTIONS, dataPlaneTaskQueue)); - normalizationActivity = - Workflow.newActivityStub(NormalizationActivity.class, setTaskQueue(ActivityConfiguration.LONG_RUN_OPTIONS, dataPlaneTaskQueue)); - dbtTransformationActivity = - Workflow.newActivityStub(DbtTransformationActivity.class, setTaskQueue(ActivityConfiguration.LONG_RUN_OPTIONS, dataPlaneTaskQueue)); - } else { - replicationActivity = Workflow.newActivityStub(ReplicationActivity.class, ActivityConfiguration.LONG_RUN_OPTIONS); - normalizationActivity = Workflow.newActivityStub(NormalizationActivity.class, ActivityConfiguration.LONG_RUN_OPTIONS); - dbtTransformationActivity = Workflow.newActivityStub(DbtTransformationActivity.class, ActivityConfiguration.LONG_RUN_OPTIONS); - persistActivity = Workflow.newActivityStub(PersistStateActivity.class, ActivityConfiguration.SHORT_ACTIVITY_OPTIONS); - } - StandardSyncOutput syncOutput = replicationActivity.replicate(jobRunConfig, sourceLauncherConfig, destinationLauncherConfig, syncInput); if (version > Workflow.DEFAULT_VERSION) { @@ -84,9 +59,7 @@ public StandardSyncOutput run(final JobRunConfig jobRunConfig, if (syncInput.getOperationSequence() != null && !syncInput.getOperationSequence().isEmpty()) { for (final StandardSyncOperation standardSyncOperation : syncInput.getOperationSequence()) { if (standardSyncOperation.getOperatorType() == OperatorType.NORMALIZATION) { - final Configs configs = new EnvConfigs(); - final NormalizationInput normalizationInput = generateNormalizationInput(syncInput, syncOutput, configs); - + final NormalizationInput normalizationInput = generateNormalizationInput(syncInput, syncOutput); final NormalizationSummary normalizationSummary = normalizationActivity.normalize(jobRunConfig, destinationLauncherConfig, normalizationInput); syncOutput = syncOutput.withNormalizationSummary(normalizationSummary); @@ -108,22 +81,9 @@ public StandardSyncOutput run(final JobRunConfig jobRunConfig, } private NormalizationInput generateNormalizationInput(final StandardSyncInput syncInput, - final StandardSyncOutput syncOutput, - final Configs configs) { - final ResourceRequirements resourceReqs = new ResourceRequirements() - .withCpuRequest(configs.getNormalizationJobMainContainerCpuRequest()) - .withCpuLimit(configs.getNormalizationJobMainContainerCpuLimit()) - .withMemoryRequest(configs.getNormalizationJobMainContainerMemoryRequest()) - .withMemoryLimit(configs.getNormalizationJobMainContainerMemoryLimit()); - - return new NormalizationInput() - .withDestinationConfiguration(syncInput.getDestinationConfiguration()) - .withCatalog(syncOutput.getOutputCatalog()) - .withResourceRequirements(resourceReqs); - } + final StandardSyncOutput syncOutput) { - private ActivityOptions setTaskQueue(final ActivityOptions activityOptions, final String taskQueue) { - return ActivityOptions.newBuilder(activityOptions).setTaskQueue(taskQueue).build(); + return normalizationActivity.generateNormalizationInput(syncInput, syncOutput); } } diff --git a/airbyte-workers/src/main/resources/application-control.yml b/airbyte-workers/src/main/resources/application-control.yml new file mode 100644 index 0000000000000..752bc85a5f112 --- /dev/null +++ b/airbyte-workers/src/main/resources/application-control.yml @@ -0,0 +1,40 @@ +datasources: + config: + connection-test-query: SELECT 1 + connection-timeout: 30000 + idle-timeout: 600000 + maximum-pool-size: 10 + url: ${DATABASE_URL} + driverClassName: org.postgresql.Driver + username: ${DATABASE_USER} + password: ${DATABASE_PASSWORD} + jobs: + connection-test-query: SELECT 1 + connection-timeout: 30000 + idle-timeout: 600000 + maximum-pool-size: 10 + url: ${DATABASE_URL} + driverClassName: org.postgresql.Driver + username: ${DATABASE_USER} + password: ${DATABASE_PASSWORD} + +flyway: + enabled: true + datasources: + config: + enabled: false + locations: + - 'classpath:io/airbyte/db/instance/configs/migrations' + jobs: + enabled: false + locations: + - 'classpath:io/airbyte/db/instance/jobs/migrations' + +jooq: + datasources: + config: + jackson-converter-enabled: true + sql-dialect: POSTGRES + jobs: + jackson-converter-enabled: true + sql-dialect: POSTGRES \ No newline at end of file diff --git a/airbyte-workers/src/main/resources/application.yml b/airbyte-workers/src/main/resources/application.yml new file mode 100644 index 0000000000000..6961c8a3801c6 --- /dev/null +++ b/airbyte-workers/src/main/resources/application.yml @@ -0,0 +1,219 @@ +micronaut: + application: + name: airbyte-workers + security: + intercept-url-map: + - pattern: /** + httpMethod: GET + access: + - isAnonymous() + server: + port: 9000 + +airbyte: + activity: + initial-delay: ${ACTIVITY_INITIAL_DELAY_BETWEEN_ATTEMPTS_SECONDS:30} + max-attempts: ${ACTIVITY_MAX_ATTEMPT:5} + max-delay: ${ACTIVITY_MAX_DELAY_BETWEEN_ATTEMPTS_SECONDS:600} + max-timeout: ${ACTIVITY_MAX_TIMEOUT_SECOND:120} + cloud: + storage: + logs: + type: ${WORKER_LOGS_STORAGE_TYPE:} + gcs: + application-credentials: ${GOOGLE_APPLICATION_CREDENTIALS:} + bucket: ${GCS_LOG_BUCKET:} + minio: + access-key: ${AWS_ACCESS_KEY_ID:} + bucket: ${S3_LOG_BUCKET:} + endpoint: ${S3_MINIO_ENDPOINT:} + secret-access-key: ${AWS_SECRET_ACCESS_KEY:} + s3: + access-key: ${AWS_ACCESS_KEY_ID:} + bucket: ${S3_LOG_BUCKET:} + region: ${S3_LOG_BUCKET_REGION:} + secret-access-key: ${AWS_SECRET_ACCESS_KEY:} + state: + type: ${WORKER_STATE_STORAGE_TYPE:} + gcs: + application-credentials: ${STATE_STORAGE_GCS_APPLICATION_CREDENTIALS:} + bucket: ${STATE_STORAGE_GCS_BUCKET_NAME:} + minio: + access-key: ${STATE_STORAGE_MINIO_ACCESS_KEY:} + bucket: ${STATE_STORAGE_MINIO_BUCKET_NAME:} + endpoint: ${STATE_STORAGE_MINIO_ENDPOINT:} + secret-access-key: ${STATE_STORAGE_MINIO_SECRET_ACCESS_KEY:} + s3: + access-key: ${STATE_STORAGE_S3_ACCESS_KEY:} + bucket: ${STATE_STORAGE_S3_BUCKET_NAME:} + region: ${STATE_STORAGE_S3_BUCKET_REGION:} + secret-access-key: ${STATE_STORAGE_S3_SECRET_ACCESS_KEY:} + connector: + specific-resource-defaults-enabled: ${CONNECTOR_SPECIFIC_RESOURCE_DEFAULTS_ENABLED:false} + container: + orchestrator: + enabled: ${CONTAINER_ORCHESTRATOR_ENABLED:false} + image: ${CONTAINER_ORCHESTRATOR_IMAGE:} + secret-mount-path: ${CONTAINER_ORCHESTRATOR_SECRET_MOUNT_PATH:} + secret-name: ${CONTAINER_ORCHESTRATOR_SECRET_NAME:} + control: + plane: + auth-endpoint: ${CONTROL_PLANE_AUTH_ENDPOINT:} + data: + sync: + task-queue: ${DATA_SYNC_TASK_QUEUES:SYNC} + plane: + connection-ids-mvp: ${CONNECTION_IDS_FOR_MVP_DATA_PLANE:} + service-account: + credentials-path: ${DATA_PLANE_SERVICE_ACCOUNT_CREDENTIALS_PATH:} + email: ${DATA_PLANE_SERVICE_ACCOUNT_EMAIL:} + deployment-mode: ${DEPLOYMENT_MODE:OSS} + flyway: + configs: + initialization-timeout-ms: ${CONFIGS_DATABASE_INITIALIZATION_TIMEOUT_MS:60000} + minimum-migration-version: ${CONFIGS_DATABASE_MINIMUM_FLYWAY_MIGRATION_VERSION} + jobs: + initialization-timeout-ms: ${JOBS_DATABASE_INITIALIZATION_TIMEOUT_MS:60000} + minimum-migration-version: ${JOBS_DATABASE_MINIMUM_FLYWAY_MIGRATION_VERSION} + internal: + api: + auth-header: + name: ${AIRBYTE_API_AUTH_HEADER_NAME:} + value: ${AIRBYTE_API_AUTH_HEADER_VALUE:} + host: ${INTERNAL_API_HOST} + local: + docker-mount: ${LOCAL_DOCKER_MOUNT} + root: ${LOCAL_ROOT} + worker: + env: ${WORKER_ENVIRONMENT:DOCKER} + check: + enabled: ${SHOULD_RUN_CHECK_CONNECTION_WORKFLOWS:true} + kube: + annotations: ${CHECK_JOB_KUBE_ANNOTATION:} + node-selectors: ${CHECK_JOB_KUBE_NODE_SELECTORS:} + max-workers: ${MAX_CHECK_WORKERS:5} + main: + container: + cpu: + limit: ${CHECK_JOB_MAIN_CONTAINER_CPU_LIMIT:} + request: ${CHECK_JOB_MAIN_CONTAINER_CPU_REQUEST:} + memory: + limit: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_LIMIT:} + request: ${CHECK_JOB_MAIN_CONTAINER_MEMORY_REQUEST:} + connection: + enabled: ${SHOULD_RUN_CONNECTION_MANAGER_WORKFLOWS:true} + discover: + enabled: ${SHOULD_RUN_DISCOVER_WORKFLOWS:true} + kube: + annotations: ${DISCOVER_JOB_KUBE_ANNOTATIONS:} + node-selectors: ${DISCOVER_JOB_KUBE_NODE_SELECTORS:} + max-workers: ${MAX_DISCOVER_WORKERS:5} + job: + error-reporting: + sentry: + dsn: ${JOB_ERROR_REPORTING_SENTRY_DSN} + strategy: ${JOB_ERROR_REPORTING_STRATEGY:LOGGING} + failed: + max-days: ${MAX_DAYS_OF_ONLY_FAILED_JOBS_BEFORE_CONNECTION_DISABLE:14} + max-jobs: ${MAX_FAILED_JOBS_IN_A_ROW_BEFORE_CONNECTION_DISABLE:100} + kube: + annotations: ${JOB_KUBE_ANNOTATIONS:} + images: + busybox: ${JOB_KUBE_BUSYBOX_IMAGE:`busybox:1.28`} + curl: ${JOB_KUBE_CURL_IMAGE:`curlimages/curl:7.83.1`} + socat: ${JOB_KUBE_SOCAT_IMAGE:`alpine/socat:1.7.4.3-r0`} + main: + container: + image-pull-policy: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_POLICY:IfNotPresent} + image-pull-secret: ${JOB_KUBE_MAIN_CONTAINER_IMAGE_PULL_SECRET:} + namespace: ${JOB_KUBE_NAMESPACE:default} + node-selectors: ${JOB_KUBE_NODE_SELECTORS:} + sidecar: + container: + image-pull-policy: ${JOB_KUBE_SIDECAR_CONTAINER_IMAGE_PULL_POLICY:IfNotPresent} + tolerations: ${JOB_KUBE_TOLERATIONS:} + main: + container: + cpu: + limit: ${JOB_MAIN_CONTAINER_CPU_LIMIT:} + request: ${JOB_MAIN_CONTAINER_CPU_REQUEST:} + memory: + limit: ${JOB_MAIN_CONTAINER_MEMORY_LIMIT:} + request: ${JOB_MAIN_CONTAINER_MEMORY_REQUEST:} + normalization: + main: + container: + cpu: + limit: ${NORMALIZATION_JOB_MAIN_CONTAINER_CPU_LIMIT:} + request: ${NORMALIZATION_JOB_MAIN_CONTAINER_CPU_REQUEST:} + memory: + limit: ${NORMALIZATION_JOB_MAIN_CONTAINER_MEMORY_LIMIT:} + request: ${NORMALIZATION_JOB_MAIN_CONTAINER_MEMORY_REQUEST:} + plane: ${WORKER_PLANE:CONTROL_PLANE} + replication: + orchestrator: + cpu: + limit: ${REPLICATION_ORCHESTRATOR_CPU_LIMIT:} + request: ${REPLICATION_ORCHESTRATOR_CPU_REQUEST:} + memory: + limit: ${REPLICATION_ORCHESTRATOR_MEMORY_LIMIT:} + request: ${REPLICATION_ORCHESTRATOR_MEMORY_REQUEST:} + spec: + enabled: ${SHOULD_RUN_GET_SPEC_WORKFLOWS:true} + kube: + annotations: ${SPEC_JOB_KUBE_ANNOTATIONS:} + node-selectors: ${SPEC_JOB_KUBE_NODE_SELECTORS:} + max-workers: ${MAX_SPEC_WORKERS:5} + sync: + enabled: ${SHOULD_RUN_SYNC_WORKFLOWS:true} + max-workers: ${MAX_SYNC_WORKERS:5} + max-attempts: ${SYNC_JOB_MAX_ATTEMPTS:3} + max-timeout: ${SYNC_JOB_MAX_TIMEOUT_DAYS:3} + role: ${AIRBYTE_ROLE:} + secret: + persistence: ${SECRET_PERSISTENCE:TESTING_CONFIG_DB_TABLE} + store: + gcp: + credentials: ${SECRET_STORE_GCP_CREDENTIALS:} + project-id: ${SECRET_STORE_GCP_PROJECT_ID:} + vault: + address: ${VAULT_ADDRESS:} + prefix: ${VAULT_PREFIX:} + token: ${VAULT_AUTH_TOKEN:} + temporal: + worker: + ports: ${TEMPORAL_WORKER_PORTS:} + tracking-strategy: ${TRACKING_STRATEGY:LOGGING} + version: ${AIRBYTE_VERSION} + web-app: + url: ${WEBAPP_URL:} + workflow: + failure: + restart-delay: ${WORKFLOW_FAILURE_RESTART_DELAY_SECONDS:600} + workspace: + docker-mount: ${WORKSPACE_DOCKER_MOUNT:} + root: ${WORKSPACE_ROOT} + +docker: + network: ${DOCKER_NETWORK:host} + +endpoints: + all: + enabled: true + +temporal: + cloud: + client: + cert: ${TEMPORAL_CLOUD_CLIENT_CERT:} + key: ${TEMPORAL_CLOUD_CLIENT_KEY:} + enabled: ${TEMPORAL_CLOUD_ENABLED:false} + host: ${TEMPORAL_CLOUD_HOST:} + namespace: ${TEMPORAL_CLOUD_NAMESPACE:} + host: ${TEMPORAL_HOST:`airbyte-temporal:7233`} + retention: ${TEMPORAL_HISTORY_RETENTION_IN_DAYS:30} + +logger: + levels: + io.airbyte.bootloader: DEBUG +# Uncomment to help resolve issues with conditional beans +# io.micronaut.context.condition: DEBUG diff --git a/airbyte-workers/src/main/resources/micronaut-banner.txt b/airbyte-workers/src/main/resources/micronaut-banner.txt new file mode 100644 index 0000000000000..f5939528f6ddd --- /dev/null +++ b/airbyte-workers/src/main/resources/micronaut-banner.txt @@ -0,0 +1,8 @@ + + ___ _ __ __ + / | (_)____/ /_ __ __/ /____ + / /| | / / ___/ __ \/ / / / __/ _ \ + / ___ |/ / / / /_/ / /_/ / /_/ __/ +/_/ |_/_/_/ /_.___/\__, /\__/\___/ + /____/ + : airbyte-workers : \ No newline at end of file diff --git a/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/AsyncOrchestratorPodProcessIntegrationTest.java b/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/AsyncOrchestratorPodProcessIntegrationTest.java index eec896b4e3994..d3b794315811e 100644 --- a/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/AsyncOrchestratorPodProcessIntegrationTest.java +++ b/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/AsyncOrchestratorPodProcessIntegrationTest.java @@ -11,7 +11,6 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.storage.CloudStorageConfigs; import io.airbyte.config.storage.MinioS3ClientFactory; -import io.airbyte.workers.WorkerApp; import io.airbyte.workers.WorkerConfigs; import io.airbyte.workers.general.DocumentStoreClient; import io.airbyte.workers.storage.S3DocumentStoreClient; @@ -23,6 +22,7 @@ import io.fabric8.kubernetes.api.model.PodBuilder; import io.fabric8.kubernetes.client.DefaultKubernetesClient; import io.fabric8.kubernetes.client.KubernetesClient; +import io.fabric8.kubernetes.client.KubernetesClientException; import java.nio.file.Path; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -99,7 +99,7 @@ public static void init() throws Exception { @ValueSource(strings = {"IfNotPresent", " Always"}) @ParameterizedTest public void testAsyncOrchestratorPodProcess(final String pullPolicy) throws InterruptedException { - + final var serverPort = 8080; final var podName = "test-async-" + RandomStringUtils.randomAlphabetic(10).toLowerCase(); final var mainContainerInfo = new KubeContainerInfo("airbyte/container-orchestrator:dev", pullPolicy); // make kubepodinfo @@ -113,10 +113,11 @@ public void testAsyncOrchestratorPodProcess(final String pullPolicy) throws Inte null, null, null, - true); + true, + serverPort); final Map portMap = Map.of( - WorkerApp.KUBE_HEARTBEAT_PORT, WorkerApp.KUBE_HEARTBEAT_PORT, + serverPort, serverPort, OrchestratorConstants.PORT1, OrchestratorConstants.PORT1, OrchestratorConstants.PORT2, OrchestratorConstants.PORT2, OrchestratorConstants.PORT3, OrchestratorConstants.PORT3, @@ -143,18 +144,9 @@ public void testAsyncOrchestratorPodProcess(final String pullPolicy) throws Inte } @AfterAll - public static void teardown() { - try { - portForwardProcess.destroyForcibly(); - } catch (final Exception e) { - e.printStackTrace(); - } - - try { - kubernetesClient.pods().delete(); - } catch (final Exception e) { - e.printStackTrace(); - } + public static void teardown() throws KubernetesClientException { + portForwardProcess.destroyForcibly(); + kubernetesClient.pods().delete(); } } diff --git a/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java b/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java index 21f9957408d98..55035ad738b02 100644 --- a/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java +++ b/airbyte-workers/src/test-integration/java/io/airbyte/workers/process/KubePodProcessIntegrationTest.java @@ -22,6 +22,8 @@ import io.fabric8.kubernetes.client.DefaultKubernetesClient; import io.fabric8.kubernetes.client.KubernetesClient; import io.fabric8.kubernetes.client.informers.SharedIndexInformer; +import io.micronaut.context.annotation.Value; +import io.micronaut.test.extensions.junit5.annotation.MicronautTest; import java.io.IOException; import java.net.Inet4Address; import java.net.ServerSocket; @@ -43,7 +45,6 @@ import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import org.apache.commons.lang3.RandomStringUtils; -import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.RepeatedTest; @@ -69,6 +70,7 @@ */ @Timeout(value = 6, unit = TimeUnit.MINUTES) +@MicronautTest public class KubePodProcessIntegrationTest { private static final Logger LOGGER = LoggerFactory.getLogger(KubePodProcessIntegrationTest.class); @@ -77,26 +79,26 @@ public class KubePodProcessIntegrationTest { private static final boolean IS_MINIKUBE = Boolean.parseBoolean(Optional.ofNullable(System.getenv("IS_MINIKUBE")).orElse("false")); private static List openPorts; - private static int heartbeatPort; - private static String heartbeatUrl; - private static KubernetesClient fabricClient; - private static KubeProcessFactory processFactory; + @Value("${micronaut.server.port}") + private Integer heartbeatPort; + private String heartbeatUrl; + private KubernetesClient fabricClient; + private KubeProcessFactory processFactory; private static final ResourceRequirements DEFAULT_RESOURCE_REQUIREMENTS = new WorkerConfigs(new EnvConfigs()).getResourceRequirements(); - private WorkerHeartbeatServer server; - @BeforeAll public static void init() throws Exception { - openPorts = new ArrayList<>(getOpenPorts(30)); // todo: should we offer port pairs to prevent deadlock? can create test here with fewer to get - // this + // todo: should we offer port pairs to prevent deadlock? can create test here with fewer to get this + openPorts = new ArrayList<>(getOpenPorts(30)); + KubePortManagerSingleton.init(new HashSet<>(openPorts.subList(1, openPorts.size() - 1))); + } - heartbeatPort = openPorts.get(0); + @BeforeEach + public void setup() throws Exception { heartbeatUrl = getHost() + ":" + heartbeatPort; fabricClient = new DefaultKubernetesClient(); - KubePortManagerSingleton.init(new HashSet<>(openPorts.subList(1, openPorts.size() - 1))); - final WorkerConfigs workerConfigs = spy(new WorkerConfigs(new EnvConfigs())); when(workerConfigs.getEnvMap()).thenReturn(Map.of("ENV_VAR_1", "ENV_VALUE_1")); @@ -110,17 +112,6 @@ public static void init() throws Exception { false); } - @BeforeEach - public void setup() throws Exception { - server = new WorkerHeartbeatServer(heartbeatPort); - server.startBackground(); - } - - @AfterEach - public void teardown() throws Exception { - server.stop(); - } - /** * In the past we've had some issues with transient / stuck pods. The idea here is to run a few at * once, and check that they are all running in hopes of identifying regressions that introduce @@ -344,13 +335,26 @@ public void testMissingEntrypoint() throws WorkerException, InterruptedException @RetryingTest(3) public void testKillingWithoutHeartbeat() throws Exception { + heartbeatUrl = "invalid_host"; + + fabricClient = new DefaultKubernetesClient(); + + final WorkerConfigs workerConfigs = spy(new WorkerConfigs(new EnvConfigs())); + when(workerConfigs.getEnvMap()).thenReturn(Map.of("ENV_VAR_1", "ENV_VALUE_1")); + + processFactory = + new KubeProcessFactory( + workerConfigs, + "default", + fabricClient, + heartbeatUrl, + getHost(), + false); + // start an infinite process final var availablePortsBefore = KubePortManagerSingleton.getInstance().getNumAvailablePorts(); final Process process = getProcess("while true; do echo hi; sleep 1; done"); - // kill the heartbeat server - server.stop(); - // waiting for process process.waitFor(); diff --git a/airbyte-workers/src/test-integration/resources/application-test.yml b/airbyte-workers/src/test-integration/resources/application-test.yml new file mode 100644 index 0000000000000..064a9883c0c6c --- /dev/null +++ b/airbyte-workers/src/test-integration/resources/application-test.yml @@ -0,0 +1,11 @@ +datasources: + config: + url: jdbc:h2:mem:default + driverClassName: org.h2.Driver + username: sa + password: "" + jobs: + url: jdbc:h2:mem:default + driverClassName: org.h2.Driver + username: sa + password: "" \ No newline at end of file diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/CancellationHandlerTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/CancellationHandlerTest.java index 963d4fbc32dca..56a37adcb3d83 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/CancellationHandlerTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/CancellationHandlerTest.java @@ -6,6 +6,7 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import io.airbyte.workers.temporal.stubs.HeartbeatWorkflow; import io.temporal.activity.Activity; import io.temporal.activity.ActivityExecutionContext; import io.temporal.client.WorkflowClient; @@ -27,7 +28,7 @@ void testCancellationHandler() { final WorkflowClient client = testEnv.getWorkflowClient(); worker.registerActivitiesImplementations(new HeartbeatWorkflow.HeartbeatActivityImpl(() -> { - ActivityExecutionContext context = Activity.getExecutionContext(); + final ActivityExecutionContext context = Activity.getExecutionContext(); new CancellationHandler.TemporalCancellationHandler(context).checkAndHandleCancellation(() -> {}); })); diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/StreamResetRecordsHelperTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/StreamResetRecordsHelperTest.java new file mode 100644 index 0000000000000..f46768898dab5 --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/StreamResetRecordsHelperTest.java @@ -0,0 +1,72 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal; + +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.when; + +import io.airbyte.config.JobConfig.ConfigType; +import io.airbyte.config.persistence.StreamResetPersistence; +import io.airbyte.protocol.models.StreamDescriptor; +import io.airbyte.scheduler.models.Job; +import io.airbyte.scheduler.persistence.JobPersistence; +import java.io.IOException; +import java.util.List; +import java.util.UUID; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.InjectMocks; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; + +/** + * Test suite for the {@link StreamResetRecordsHelper} class. + */ +@ExtendWith(MockitoExtension.class) +class StreamResetRecordsHelperTest { + + private static final UUID CONNECTION_ID = UUID.randomUUID(); + private static final Long JOB_ID = Long.valueOf("123"); + + @Mock + private JobPersistence jobPersistence; + @Mock + private StreamResetPersistence streamResetPersistence; + @InjectMocks + private StreamResetRecordsHelper streamResetRecordsHelper; + + @Test + void testDeleteStreamResetRecordsForJob() throws IOException { + final Job jobMock = mock(Job.class, RETURNS_DEEP_STUBS); + when(jobPersistence.getJob(JOB_ID)).thenReturn(jobMock); + + when(jobMock.getConfig().getConfigType()).thenReturn(ConfigType.RESET_CONNECTION); + final List streamsToDelete = List.of(new StreamDescriptor().withName("streamname").withNamespace("namespace")); + when(jobMock.getConfig().getResetConnection().getResetSourceConfiguration().getStreamsToReset()).thenReturn(streamsToDelete); + streamResetRecordsHelper.deleteStreamResetRecordsForJob(JOB_ID, CONNECTION_ID); + Mockito.verify(streamResetPersistence).deleteStreamResets(CONNECTION_ID, streamsToDelete); + } + + @Test + void testIncorrectConfigType() throws IOException { + final Job jobMock = mock(Job.class, RETURNS_DEEP_STUBS); + when(jobPersistence.getJob(JOB_ID)).thenReturn(jobMock); + + when(jobMock.getConfig().getConfigType()).thenReturn(ConfigType.SYNC); + streamResetRecordsHelper.deleteStreamResetRecordsForJob(JOB_ID, CONNECTION_ID); + Mockito.verify(streamResetPersistence, never()).deleteStreamResets(Mockito.any(UUID.class), Mockito.anyList()); + } + + @Test + void testNoJobId() throws IOException { + streamResetRecordsHelper.deleteStreamResetRecordsForJob(null, CONNECTION_ID); + Mockito.verify(jobPersistence, never()).getJob(Mockito.anyLong()); + Mockito.verify(streamResetPersistence, never()).deleteStreamResets(Mockito.any(UUID.class), Mockito.anyList()); + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalClientTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalClientTest.java index c89a829256972..e6be802ade59b 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalClientTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalClientTest.java @@ -102,6 +102,7 @@ class TemporalClientTest { private WorkflowServiceBlockingStub workflowServiceBlockingStub; private StreamResetPersistence streamResetPersistence; private ConnectionManagerUtils connectionManagerUtils; + private StreamResetRecordsHelper streamResetRecordsHelper; @BeforeEach void setup() throws IOException { @@ -115,8 +116,11 @@ void setup() throws IOException { when(workflowServiceStubs.blockingStub()).thenReturn(workflowServiceBlockingStub); streamResetPersistence = mock(StreamResetPersistence.class); mockWorkflowStatus(WorkflowExecutionStatus.WORKFLOW_EXECUTION_STATUS_RUNNING); - connectionManagerUtils = spy(ConnectionManagerUtils.class); - temporalClient = spy(new TemporalClient(workflowClient, workspaceRoot, workflowServiceStubs, streamResetPersistence, connectionManagerUtils)); + connectionManagerUtils = spy(new ConnectionManagerUtils()); + streamResetRecordsHelper = mock(StreamResetRecordsHelper.class); + temporalClient = + spy(new TemporalClient(workspaceRoot, workflowClient, workflowServiceStubs, streamResetPersistence, connectionManagerUtils, + streamResetRecordsHelper)); } @Nested @@ -178,19 +182,21 @@ class TestJobSubmission { @Test void testSubmitGetSpec() { final SpecWorkflow specWorkflow = mock(SpecWorkflow.class); - when(workflowClient.newWorkflowStub(SpecWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.GET_SPEC))).thenReturn(specWorkflow); + when(workflowClient.newWorkflowStub(SpecWorkflow.class, TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.GET_SPEC))) + .thenReturn(specWorkflow); final JobGetSpecConfig getSpecConfig = new JobGetSpecConfig().withDockerImage(IMAGE_NAME1); temporalClient.submitGetSpec(JOB_UUID, ATTEMPT_ID, getSpecConfig); specWorkflow.run(JOB_RUN_CONFIG, UUID_LAUNCHER_CONFIG); - verify(workflowClient).newWorkflowStub(SpecWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.GET_SPEC)); + verify(workflowClient).newWorkflowStub(SpecWorkflow.class, TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.GET_SPEC)); } @Test void testSubmitCheckConnection() { final CheckConnectionWorkflow checkConnectionWorkflow = mock(CheckConnectionWorkflow.class); - when(workflowClient.newWorkflowStub(CheckConnectionWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.CHECK_CONNECTION))) - .thenReturn(checkConnectionWorkflow); + when( + workflowClient.newWorkflowStub(CheckConnectionWorkflow.class, TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.CHECK_CONNECTION))) + .thenReturn(checkConnectionWorkflow); final JobCheckConnectionConfig checkConnectionConfig = new JobCheckConnectionConfig() .withDockerImage(IMAGE_NAME1) .withConnectionConfiguration(Jsons.emptyObject()); @@ -199,13 +205,14 @@ void testSubmitCheckConnection() { temporalClient.submitCheckConnection(JOB_UUID, ATTEMPT_ID, checkConnectionConfig); checkConnectionWorkflow.run(JOB_RUN_CONFIG, UUID_LAUNCHER_CONFIG, input); - verify(workflowClient).newWorkflowStub(CheckConnectionWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.CHECK_CONNECTION)); + verify(workflowClient).newWorkflowStub(CheckConnectionWorkflow.class, + TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.CHECK_CONNECTION)); } @Test void testSubmitDiscoverSchema() { final DiscoverCatalogWorkflow discoverCatalogWorkflow = mock(DiscoverCatalogWorkflow.class); - when(workflowClient.newWorkflowStub(DiscoverCatalogWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.DISCOVER_SCHEMA))) + when(workflowClient.newWorkflowStub(DiscoverCatalogWorkflow.class, TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.DISCOVER_SCHEMA))) .thenReturn(discoverCatalogWorkflow); final JobDiscoverCatalogConfig checkConnectionConfig = new JobDiscoverCatalogConfig() .withDockerImage(IMAGE_NAME1) @@ -215,13 +222,14 @@ void testSubmitDiscoverSchema() { temporalClient.submitDiscoverSchema(JOB_UUID, ATTEMPT_ID, checkConnectionConfig); discoverCatalogWorkflow.run(JOB_RUN_CONFIG, UUID_LAUNCHER_CONFIG, input); - verify(workflowClient).newWorkflowStub(DiscoverCatalogWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.DISCOVER_SCHEMA)); + verify(workflowClient).newWorkflowStub(DiscoverCatalogWorkflow.class, + TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.DISCOVER_SCHEMA)); } @Test void testSubmitSync() { final SyncWorkflow discoverCatalogWorkflow = mock(SyncWorkflow.class); - when(workflowClient.newWorkflowStub(SyncWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.SYNC))) + when(workflowClient.newWorkflowStub(SyncWorkflow.class, TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.SYNC))) .thenReturn(discoverCatalogWorkflow); final JobSyncConfig syncConfig = new JobSyncConfig() .withSourceDockerImage(IMAGE_NAME1) @@ -247,7 +255,7 @@ void testSubmitSync() { temporalClient.submitSync(JOB_ID, ATTEMPT_ID, syncConfig, CONNECTION_ID); discoverCatalogWorkflow.run(JOB_RUN_CONFIG, LAUNCHER_CONFIG, destinationLauncherConfig, input, CONNECTION_ID); - verify(workflowClient).newWorkflowStub(SyncWorkflow.class, TemporalUtils.getWorkflowOptions(TemporalJobType.SYNC)); + verify(workflowClient).newWorkflowStub(SyncWorkflow.class, TemporalWorkflowUtils.buildWorkflowOptions(TemporalJobType.SYNC)); } @Test @@ -291,10 +299,8 @@ void migrateCalled() { final UUID nonMigratedId = UUID.randomUUID(); final UUID migratedId = UUID.randomUUID(); - doReturn(false) - .when(temporalClient).isInRunningWorkflowCache(ConnectionManagerUtils.getInstance().getConnectionManagerName(nonMigratedId)); - doReturn(true) - .when(temporalClient).isInRunningWorkflowCache(ConnectionManagerUtils.getInstance().getConnectionManagerName(migratedId)); + when(temporalClient.isInRunningWorkflowCache(connectionManagerUtils.getConnectionManagerName(nonMigratedId))).thenReturn(false); + when(temporalClient.isInRunningWorkflowCache(connectionManagerUtils.getConnectionManagerName(migratedId))).thenReturn(true); doNothing() .when(temporalClient).refreshRunningWorkflow(); @@ -572,6 +578,7 @@ void testStartNewCancellationSuccess() { assertEquals(JOB_ID, result.getJobId().get()); assertFalse(result.getFailingReason().isPresent()); verify(mConnectionManagerWorkflow).cancelJob(); + verify(streamResetRecordsHelper).deleteStreamResetRecordsForJob(JOB_ID, CONNECTION_ID); } @Test @@ -744,7 +751,9 @@ public void init() throws IOException { mConnectionManagerUtils = mock(ConnectionManagerUtils.class); final Path workspaceRoot = Files.createTempDirectory(Path.of("/tmp"), "temporal_client_test"); - temporalClient = spy(new TemporalClient(workflowClient, workspaceRoot, workflowServiceStubs, streamResetPersistence, mConnectionManagerUtils)); + temporalClient = spy( + new TemporalClient(workspaceRoot, workflowClient, workflowServiceStubs, streamResetPersistence, mConnectionManagerUtils, + streamResetRecordsHelper)); } @Test diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalUtilsTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalUtilsTest.java index 37bdaa4405154..b4fd33ce429fc 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalUtilsTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalUtilsTest.java @@ -4,7 +4,6 @@ package io.airbyte.workers.temporal; -import static io.airbyte.workers.temporal.TemporalUtils.getTemporalClientWhenConnected; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -15,6 +14,7 @@ import io.airbyte.commons.concurrency.VoidCallable; import io.airbyte.workers.exception.WorkerException; +import io.airbyte.workers.temporal.stubs.HeartbeatWorkflow; import io.temporal.activity.Activity; import io.temporal.activity.ActivityCancellationType; import io.temporal.activity.ActivityExecutionContext; @@ -56,6 +56,7 @@ class TemporalUtilsTest { @Test void testAsyncExecute() throws Exception { + final TemporalUtils temporalUtils = new TemporalUtils(null, null, null, null, null, null, null); final CountDownLatch countDownLatch = new CountDownLatch(1); final VoidCallable callable = mock(VoidCallable.class); @@ -75,7 +76,7 @@ void testAsyncExecute() throws Exception { testEnv.start(); final TestWorkflow workflowStub = client.newWorkflowStub(TestWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build()); - final ImmutablePair> pair = TemporalUtils.asyncExecute( + final ImmutablePair> pair = temporalUtils.asyncExecute( workflowStub, workflowStub::run, "whatever", @@ -98,6 +99,7 @@ void testAsyncExecute() throws Exception { @Test void testWaitForTemporalServerAndLogThrowsException() { + final TemporalUtils temporalUtils = new TemporalUtils(null, null, null, null, null, null, null); final WorkflowServiceStubs workflowServiceStubs = mock(WorkflowServiceStubs.class, Mockito.RETURNS_DEEP_STUBS); final DescribeNamespaceResponse describeNamespaceResponse = mock(DescribeNamespaceResponse.class); final NamespaceInfo namespaceInfo = mock(NamespaceInfo.class); @@ -113,11 +115,12 @@ void testWaitForTemporalServerAndLogThrowsException() { when(workflowServiceStubs.blockingStub().describeNamespace(any())) .thenThrow(RuntimeException.class) .thenReturn(describeNamespaceResponse); - getTemporalClientWhenConnected(Duration.ofMillis(10), Duration.ofSeconds(1), Duration.ofSeconds(0), serviceSupplier, namespace); + temporalUtils.getTemporalClientWhenConnected(Duration.ofMillis(10), Duration.ofSeconds(1), Duration.ofSeconds(0), serviceSupplier, namespace); } @Test void testWaitThatTimesOut() { + final TemporalUtils temporalUtils = new TemporalUtils(null, null, null, null, null, null, null); final WorkflowServiceStubs workflowServiceStubs = mock(WorkflowServiceStubs.class, Mockito.RETURNS_DEEP_STUBS); final DescribeNamespaceResponse describeNamespaceResponse = mock(DescribeNamespaceResponse.class); final NamespaceInfo namespaceInfo = mock(NamespaceInfo.class); @@ -133,7 +136,7 @@ void testWaitThatTimesOut() { .thenThrow(RuntimeException.class) .thenReturn(List.of(describeNamespaceResponse)); assertThrows(RuntimeException.class, () -> { - getTemporalClientWhenConnected(Duration.ofMillis(100), Duration.ofMillis(10), Duration.ofSeconds(0), serviceSupplier, namespace); + temporalUtils.getTemporalClientWhenConnected(Duration.ofMillis(100), Duration.ofMillis(10), Duration.ofSeconds(0), serviceSupplier, namespace); }); } @@ -181,7 +184,7 @@ void testWorkerExceptionOnHeartbeatWrapper() { @Test void testHeartbeatWithContext() throws InterruptedException { - + final TemporalUtils temporalUtils = new TemporalUtils(null, null, null, null, null, null, null); final TestWorkflowEnvironment testEnv = TestWorkflowEnvironment.newInstance(); final Worker worker = testEnv.newWorker(TASK_QUEUE); @@ -193,7 +196,7 @@ void testHeartbeatWithContext() throws InterruptedException { worker.registerActivitiesImplementations(new HeartbeatWorkflow.HeartbeatActivityImpl(() -> { final ActivityExecutionContext context = Activity.getExecutionContext(); - TemporalUtils.withBackgroundHeartbeat( + temporalUtils.withBackgroundHeartbeat( // TODO (itaseski) figure out how to decrease heartbeat intervals using reflection () -> { latch.await(); @@ -222,7 +225,7 @@ void testHeartbeatWithContext() throws InterruptedException { @Test void testHeartbeatWithContextAndCallbackRef() throws InterruptedException { - + final TemporalUtils temporalUtils = new TemporalUtils(null, null, null, null, null, null, null); final TestWorkflowEnvironment testEnv = TestWorkflowEnvironment.newInstance(); final Worker worker = testEnv.newWorker(TASK_QUEUE); @@ -234,7 +237,7 @@ void testHeartbeatWithContextAndCallbackRef() throws InterruptedException { worker.registerActivitiesImplementations(new HeartbeatWorkflow.HeartbeatActivityImpl(() -> { final ActivityExecutionContext context = Activity.getExecutionContext(); - TemporalUtils.withBackgroundHeartbeat( + temporalUtils.withBackgroundHeartbeat( // TODO (itaseski) figure out how to decrease heartbeat intervals using reflection new AtomicReference<>(() -> {}), () -> { @@ -376,6 +379,8 @@ class Activity1Impl implements Activity1 { private final AtomicInteger timesReachedEnd; + private final TemporalUtils temporalUtils = new TemporalUtils(null, null, null, null, null, null, null); + public Activity1Impl(final AtomicInteger timesReachedEnd) { this.timesReachedEnd = timesReachedEnd; } @@ -384,7 +389,7 @@ public Activity1Impl(final AtomicInteger timesReachedEnd) { public void activity(final String arg) { LOGGER.info(BEFORE, ACTIVITY1); final ActivityExecutionContext context = Activity.getExecutionContext(); - TemporalUtils.withBackgroundHeartbeat( + temporalUtils.withBackgroundHeartbeat( new AtomicReference<>(null), () -> { if (timesReachedEnd.get() == 0) { diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowTest.java index e9f44295d7bc9..52d21fe19d222 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/check/connection/CheckConnectionWorkflowTest.java @@ -4,19 +4,49 @@ package io.airbyte.workers.temporal.check.connection; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.support.TemporalProxyHelper; +import io.micronaut.context.BeanRegistration; +import io.micronaut.inject.BeanIdentifier; +import io.temporal.activity.ActivityOptions; import io.temporal.testing.WorkflowReplayer; +import java.time.Duration; +import java.util.List; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert") class CheckConnectionWorkflowTest { + private ActivityOptions activityOptions; + private TemporalProxyHelper temporalProxyHelper; + + @BeforeEach + void setUp() { + activityOptions = ActivityOptions.newBuilder() + .setScheduleToCloseTimeout(Duration.ofMinutes(5)) + .setRetryOptions(TemporalUtils.NO_RETRY) + .build(); + + final BeanIdentifier activityOptionsBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration activityOptionsBeanRegistration = mock(BeanRegistration.class); + when(activityOptionsBeanIdentifier.getName()).thenReturn("checkActivityOptions"); + when(activityOptionsBeanRegistration.getIdentifier()).thenReturn(activityOptionsBeanIdentifier); + when(activityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + temporalProxyHelper = new TemporalProxyHelper(List.of(activityOptionsBeanRegistration)); + } + @Test void replayOldWorkflow() throws Exception { // This test ensures that a new version of the workflow doesn't break an in-progress execution // This JSON file is exported from Temporal directly (e.g. // `http://${temporal-ui}/namespaces/default/workflows/${uuid}/${uuid}/history`) and export - WorkflowReplayer.replayWorkflowExecutionFromResource("checkWorkflowHistory.json", CheckConnectionWorkflowImpl.class); + WorkflowReplayer.replayWorkflowExecutionFromResource("checkWorkflowHistory.json", + temporalProxyHelper.proxyWorkflowClass(CheckConnectionWorkflowImpl.class)); } } diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java index 567e219159b3a..1cf42673dbc54 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/ConnectionManagerWorkflowTest.java @@ -5,6 +5,8 @@ package io.airbyte.workers.temporal.scheduling; import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import io.airbyte.config.ConnectorJobOutput; import io.airbyte.config.ConnectorJobOutput.OutputType; @@ -35,8 +37,11 @@ import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivity.JobCancelledInputWithAttemptNumber; import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivity.JobCreationOutput; import io.airbyte.workers.temporal.scheduling.activities.JobCreationAndStatusUpdateActivity.JobSuccessInputWithAttemptNumber; +import io.airbyte.workers.temporal.scheduling.activities.RecordMetricActivity; +import io.airbyte.workers.temporal.scheduling.activities.RouteToSyncTaskQueueActivity; +import io.airbyte.workers.temporal.scheduling.activities.RouteToSyncTaskQueueActivity.RouteToSyncTaskQueueOutput; import io.airbyte.workers.temporal.scheduling.activities.StreamResetActivity; -import io.airbyte.workers.temporal.scheduling.activities.StreamResetActivity.DeleteStreamResetRecordsForJobInput; +import io.airbyte.workers.temporal.scheduling.activities.WorkflowConfigActivity; import io.airbyte.workers.temporal.scheduling.state.WorkflowState; import io.airbyte.workers.temporal.scheduling.state.listener.TestStateListener; import io.airbyte.workers.temporal.scheduling.state.listener.WorkflowStateChangedListener.ChangedStateEvent; @@ -49,13 +54,18 @@ import io.airbyte.workers.temporal.scheduling.testsyncworkflow.ReplicateFailureSyncWorkflow; import io.airbyte.workers.temporal.scheduling.testsyncworkflow.SleepingSyncWorkflow; import io.airbyte.workers.temporal.scheduling.testsyncworkflow.SourceAndDestinationFailureSyncWorkflow; +import io.airbyte.workers.temporal.support.TemporalProxyHelper; import io.airbyte.workers.temporal.sync.SyncWorkflow; +import io.micronaut.context.BeanRegistration; +import io.micronaut.inject.BeanIdentifier; +import io.temporal.activity.ActivityOptions; import io.temporal.api.enums.v1.WorkflowExecutionStatus; import io.temporal.api.filter.v1.WorkflowExecutionFilter; import io.temporal.api.workflowservice.v1.ListClosedWorkflowExecutionsRequest; import io.temporal.api.workflowservice.v1.ListClosedWorkflowExecutionsResponse; import io.temporal.client.WorkflowClient; import io.temporal.client.WorkflowOptions; +import io.temporal.common.RetryOptions; import io.temporal.failure.ApplicationFailure; import io.temporal.testing.TestWorkflowEnvironment; import io.temporal.worker.Worker; @@ -95,25 +105,35 @@ class ConnectionManagerWorkflowTest { private static final Duration SCHEDULE_WAIT = Duration.ofMinutes(20L); private static final String WORKFLOW_ID = "workflow-id"; + private static final Duration WORKFLOW_FAILURE_RESTART_DELAY = Duration.ofSeconds(600); + private final ConfigFetchActivity mConfigFetchActivity = - Mockito.mock(ConfigFetchActivity.class, Mockito.withSettings().withoutAnnotations()); + mock(ConfigFetchActivity.class, Mockito.withSettings().withoutAnnotations()); private final CheckConnectionActivity mCheckConnectionActivity = - Mockito.mock(CheckConnectionActivity.class, Mockito.withSettings().withoutAnnotations()); + mock(CheckConnectionActivity.class, Mockito.withSettings().withoutAnnotations()); private static final ConnectionDeletionActivity mConnectionDeletionActivity = - Mockito.mock(ConnectionDeletionActivity.class, Mockito.withSettings().withoutAnnotations()); + mock(ConnectionDeletionActivity.class, Mockito.withSettings().withoutAnnotations()); private static final GenerateInputActivityImpl mGenerateInputActivityImpl = - Mockito.mock(GenerateInputActivityImpl.class, Mockito.withSettings().withoutAnnotations()); + mock(GenerateInputActivityImpl.class, Mockito.withSettings().withoutAnnotations()); private static final JobCreationAndStatusUpdateActivity mJobCreationAndStatusUpdateActivity = - Mockito.mock(JobCreationAndStatusUpdateActivity.class, Mockito.withSettings().withoutAnnotations()); + mock(JobCreationAndStatusUpdateActivity.class, Mockito.withSettings().withoutAnnotations()); private static final AutoDisableConnectionActivity mAutoDisableConnectionActivity = - Mockito.mock(AutoDisableConnectionActivity.class, Mockito.withSettings().withoutAnnotations()); + mock(AutoDisableConnectionActivity.class, Mockito.withSettings().withoutAnnotations()); private static final StreamResetActivity mStreamResetActivity = - Mockito.mock(StreamResetActivity.class, Mockito.withSettings().withoutAnnotations()); + mock(StreamResetActivity.class, Mockito.withSettings().withoutAnnotations()); + private static final RecordMetricActivity mRecordMetricActivity = + mock(RecordMetricActivity.class, Mockito.withSettings().withoutAnnotations()); + private static final WorkflowConfigActivity mWorkflowConfigActivity = + mock(WorkflowConfigActivity.class, Mockito.withSettings().withoutAnnotations()); + private static final RouteToSyncTaskQueueActivity mRouteToSyncTaskQueueActivity = + mock(RouteToSyncTaskQueueActivity.class, Mockito.withSettings().withoutAnnotations()); private static final String EVENT = "event = "; private TestWorkflowEnvironment testEnv; private WorkflowClient client; private ConnectionManagerWorkflow workflow; + private ActivityOptions activityOptions; + private TemporalProxyHelper temporalProxyHelper; static Stream getMaxAttemptForResetRetry() { return Stream.of( @@ -134,20 +154,23 @@ void setUp() { Mockito.reset(mJobCreationAndStatusUpdateActivity); Mockito.reset(mAutoDisableConnectionActivity); Mockito.reset(mStreamResetActivity); + Mockito.reset(mRecordMetricActivity); + Mockito.reset(mWorkflowConfigActivity); + Mockito.reset(mRouteToSyncTaskQueueActivity); // default is to wait "forever" - Mockito.when(mConfigFetchActivity.getTimeToWait(Mockito.any())).thenReturn(new ScheduleRetrieverOutput( + when(mConfigFetchActivity.getTimeToWait(Mockito.any())).thenReturn(new ScheduleRetrieverOutput( Duration.ofDays(100 * 365))); - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenReturn(new JobCreationOutput( 1L)); - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) .thenReturn(new AttemptNumberCreationOutput( 1)); - Mockito.when(mGenerateInputActivityImpl.getSyncWorkflowInputWithAttemptNumber(Mockito.any(SyncInputWithAttemptNumber.class))) + when(mGenerateInputActivityImpl.getSyncWorkflowInputWithAttemptNumber(Mockito.any(SyncInputWithAttemptNumber.class))) .thenReturn( new GeneratedJobInput( new JobRunConfig(), @@ -155,12 +178,38 @@ void setUp() { new IntegrationLauncherConfig(), new StandardSyncInput())); - Mockito.when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) + when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) .thenReturn(new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION) .withCheckConnection(new StandardCheckConnectionOutput().withStatus(Status.SUCCEEDED).withMessage("check worked"))); - Mockito.when(mAutoDisableConnectionActivity.autoDisableFailingConnection(Mockito.any())) + when(mAutoDisableConnectionActivity.autoDisableFailingConnection(Mockito.any())) .thenReturn(new AutoDisableConnectionOutput(false)); + + when(mWorkflowConfigActivity.getWorkflowRestartDelaySeconds()) + .thenReturn(WORKFLOW_FAILURE_RESTART_DELAY); + + // TODO: for now, always route to default 'SYNC' task queue. Add test for routing + // to a different task queue + when(mRouteToSyncTaskQueueActivity.route(Mockito.any())) + .thenReturn(new RouteToSyncTaskQueueOutput(TemporalJobType.SYNC.name())); + + activityOptions = ActivityOptions.newBuilder() + .setHeartbeatTimeout(Duration.ofSeconds(30)) + .setStartToCloseTimeout(Duration.ofSeconds(120)) + .setRetryOptions(RetryOptions.newBuilder() + .setMaximumAttempts(5) + .setInitialInterval(Duration.ofSeconds(30)) + .setMaximumInterval(Duration.ofSeconds(600)) + .build()) + + .build(); + + final BeanIdentifier activityOptionsBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration activityOptionsBeanRegistration = mock(BeanRegistration.class); + when(activityOptionsBeanIdentifier.getName()).thenReturn("shortActivityOptions"); + when(activityOptionsBeanRegistration.getIdentifier()).thenReturn(activityOptionsBeanIdentifier); + when(activityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + temporalProxyHelper = new TemporalProxyHelper(List.of(activityOptionsBeanRegistration)); } @AfterEach @@ -170,7 +219,7 @@ void tearDown() { } private void mockResetJobInput() { - Mockito.when(mGenerateInputActivityImpl.getSyncWorkflowInputWithAttemptNumber(Mockito.any(SyncInputWithAttemptNumber.class))) + when(mGenerateInputActivityImpl.getSyncWorkflowInputWithAttemptNumber(Mockito.any(SyncInputWithAttemptNumber.class))) .thenReturn( new GeneratedJobInput( new JobRunConfig(), @@ -193,7 +242,7 @@ void setup() { unit = TimeUnit.SECONDS) @DisplayName("Test that a successful workflow retries and waits") void runSuccess() throws InterruptedException { - Mockito.when(mConfigFetchActivity.getTimeToWait(Mockito.any())) + when(mConfigFetchActivity.getTimeToWait(Mockito.any())) .thenReturn(new ScheduleRetrieverOutput(SCHEDULE_WAIT)); final UUID testId = UUID.randomUUID(); @@ -217,11 +266,11 @@ void runSuccess() throws InterruptedException { Assertions.assertThat(events) .filteredOn(changedStateEvent -> changedStateEvent.getField() == StateField.RUNNING && changedStateEvent.isValue()) - .hasSize(2); + .hasSize(3); Assertions.assertThat(events) .filteredOn(changedStateEvent -> changedStateEvent.getField() == StateField.DONE_WAITING && changedStateEvent.isValue()) - .hasSize(2); + .hasSize(3); Assertions.assertThat(events) .filteredOn(changedStateEvent -> (changedStateEvent.getField() != StateField.RUNNING @@ -236,7 +285,7 @@ void runSuccess() throws InterruptedException { unit = TimeUnit.SECONDS) @DisplayName("Test workflow does not wait to run after a failure") void retryAfterFail() throws InterruptedException { - Mockito.when(mConfigFetchActivity.getTimeToWait(Mockito.any())) + when(mConfigFetchActivity.getTimeToWait(Mockito.any())) .thenReturn(new ScheduleRetrieverOutput(SCHEDULE_WAIT)); final UUID testId = UUID.randomUUID(); @@ -258,11 +307,11 @@ void retryAfterFail() throws InterruptedException { Assertions.assertThat(events) .filteredOn(changedStateEvent -> changedStateEvent.getField() == StateField.RUNNING && changedStateEvent.isValue()) - .hasSize(1); + .hasSize(3); Assertions.assertThat(events) .filteredOn(changedStateEvent -> changedStateEvent.getField() == StateField.DONE_WAITING && changedStateEvent.isValue()) - .hasSize(1); + .hasSize(3); Assertions.assertThat(events) .filteredOn(changedStateEvent -> (changedStateEvent.getField() != StateField.RUNNING @@ -496,7 +545,7 @@ void setup() { unit = TimeUnit.SECONDS) @DisplayName("Test workflow which receives a manual sync while running a scheduled sync does nothing") void manualRun() throws InterruptedException { - Mockito.when(mConfigFetchActivity.getTimeToWait(Mockito.any())) + when(mConfigFetchActivity.getTimeToWait(Mockito.any())) .thenReturn(new ScheduleRetrieverOutput(SCHEDULE_WAIT)); final UUID testId = UUID.randomUUID(); @@ -710,35 +759,6 @@ void resetCancelRunningWorkflow() throws InterruptedException { } - @Test - @Timeout(value = 60, - unit = TimeUnit.SECONDS) - @DisplayName("Test that cancelling a reset deletes streamsToReset from stream_resets table") - void cancelResetRemovesStreamsToReset() throws InterruptedException { - final UUID connectionId = UUID.randomUUID(); - final UUID testId = UUID.randomUUID(); - final TestStateListener testStateListener = new TestStateListener(); - final WorkflowState workflowState = new WorkflowState(testId, testStateListener); - - final ConnectionUpdaterInput input = Mockito.spy(ConnectionUpdaterInput.builder() - .connectionId(connectionId) - .jobId(JOB_ID) - .attemptId(ATTEMPT_ID) - .fromFailure(false) - .attemptNumber(1) - .workflowState(workflowState) - .skipScheduling(true) - .build()); - - startWorkflowAndWaitUntilReady(workflow, input); - - testEnv.sleep(Duration.ofSeconds(30L)); - workflow.cancelJob(); - Thread.sleep(500); - - Mockito.verify(mStreamResetActivity).deleteStreamResetRecordsForJob(new DeleteStreamResetRecordsForJobInput(connectionId, JOB_ID)); - } - @Test @DisplayName("Test that running workflow which receives an update signal waits for the current run and reports the job status") void updatedSignalReceivedWhileRunning() throws InterruptedException { @@ -807,15 +827,16 @@ void setup() { testEnv = TestWorkflowEnvironment.newInstance(); final Worker managerWorker = testEnv.newWorker(TemporalJobType.CONNECTION_UPDATER.name()); - managerWorker.registerWorkflowImplementationTypes(ConnectionManagerWorkflowImpl.class); + managerWorker.registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(ConnectionManagerWorkflowImpl.class)); managerWorker.registerActivitiesImplementations(mConfigFetchActivity, mCheckConnectionActivity, mConnectionDeletionActivity, - mGenerateInputActivityImpl, mJobCreationAndStatusUpdateActivity, mAutoDisableConnectionActivity); + mGenerateInputActivityImpl, mJobCreationAndStatusUpdateActivity, mAutoDisableConnectionActivity, mRecordMetricActivity, + mWorkflowConfigActivity, mRouteToSyncTaskQueueActivity); client = testEnv.getWorkflowClient(); workflow = client.newWorkflowStub(ConnectionManagerWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TemporalJobType.CONNECTION_UPDATER.name()).build()); - Mockito.when(mConfigFetchActivity.getMaxAttempt()).thenReturn(new GetMaxAttemptOutput(1)); + when(mConfigFetchActivity.getMaxAttempt()).thenReturn(new GetMaxAttemptOutput(1)); } @Test @@ -902,15 +923,16 @@ void setup() { testEnv = TestWorkflowEnvironment.newInstance(); final Worker managerWorker = testEnv.newWorker(TemporalJobType.CONNECTION_UPDATER.name()); - managerWorker.registerWorkflowImplementationTypes(ConnectionManagerWorkflowImpl.class); + managerWorker.registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(ConnectionManagerWorkflowImpl.class)); managerWorker.registerActivitiesImplementations(mConfigFetchActivity, mCheckConnectionActivity, mConnectionDeletionActivity, - mGenerateInputActivityImpl, mJobCreationAndStatusUpdateActivity, mAutoDisableConnectionActivity); + mGenerateInputActivityImpl, mJobCreationAndStatusUpdateActivity, mAutoDisableConnectionActivity, mRecordMetricActivity, + mWorkflowConfigActivity, mRouteToSyncTaskQueueActivity); client = testEnv.getWorkflowClient(); workflow = client.newWorkflowStub(ConnectionManagerWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TemporalJobType.CONNECTION_UPDATER.name()).build()); - Mockito.when(mConfigFetchActivity.getMaxAttempt()).thenReturn(new GetMaxAttemptOutput(1)); + when(mConfigFetchActivity.getMaxAttempt()).thenReturn(new GetMaxAttemptOutput(1)); } @Test @@ -918,11 +940,11 @@ void setup() { unit = TimeUnit.SECONDS) @DisplayName("Test that Source CHECK failures are recorded") void testSourceCheckFailuresRecorded() throws InterruptedException { - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenReturn(new JobCreationOutput(JOB_ID)); - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) .thenReturn(new AttemptNumberCreationOutput(ATTEMPT_ID)); - Mockito.when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) + when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) .thenReturn(new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION) .withCheckConnection(new StandardCheckConnectionOutput().withStatus(Status.FAILED).withMessage("nope"))); @@ -957,11 +979,11 @@ void testSourceCheckFailuresRecorded() throws InterruptedException { unit = TimeUnit.SECONDS) @DisplayName("Test that Source CHECK failure reasons are recorded") void testSourceCheckFailureReasonsRecorded() throws InterruptedException { - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenReturn(new JobCreationOutput(JOB_ID)); - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) .thenReturn(new AttemptNumberCreationOutput(ATTEMPT_ID)); - Mockito.when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) + when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) .thenReturn(new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION) .withFailureReason(new FailureReason().withFailureType(FailureType.SYSTEM_ERROR))); @@ -996,11 +1018,11 @@ void testSourceCheckFailureReasonsRecorded() throws InterruptedException { unit = TimeUnit.SECONDS) @DisplayName("Test that Destination CHECK failures are recorded") void testDestinationCheckFailuresRecorded() throws InterruptedException { - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenReturn(new JobCreationOutput(JOB_ID)); - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) .thenReturn(new AttemptNumberCreationOutput(ATTEMPT_ID)); - Mockito.when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) + when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) // First call (source) succeeds .thenReturn(new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION) .withCheckConnection(new StandardCheckConnectionOutput().withStatus(Status.SUCCEEDED).withMessage("all good"))) @@ -1040,11 +1062,11 @@ void testDestinationCheckFailuresRecorded() throws InterruptedException { unit = TimeUnit.SECONDS) @DisplayName("Test that Destination CHECK failure reasons are recorded") void testDestinationCheckFailureReasonsRecorded() throws InterruptedException { - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenReturn(new JobCreationOutput(JOB_ID)); - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) .thenReturn(new AttemptNumberCreationOutput(ATTEMPT_ID)); - Mockito.when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) + when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) // First call (source) succeeds .thenReturn(new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION) .withCheckConnection(new StandardCheckConnectionOutput().withStatus(Status.SUCCEEDED).withMessage("all good"))) @@ -1084,12 +1106,12 @@ void testDestinationCheckFailureReasonsRecorded() throws InterruptedException { unit = TimeUnit.SECONDS) @DisplayName("Test that reset workflows do not CHECK the source") void testSourceCheckSkippedWhenReset() throws InterruptedException { - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenReturn(new JobCreationOutput(JOB_ID)); - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) .thenReturn(new AttemptNumberCreationOutput(ATTEMPT_ID)); mockResetJobInput(); - Mockito.when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) + when(mCheckConnectionActivity.runWithJobOutput(Mockito.any())) // first call, but should fail destination because source check is skipped .thenReturn(new ConnectorJobOutput().withOutputType(OutputType.CHECK_CONNECTION) .withCheckConnection(new StandardCheckConnectionOutput().withStatus(Status.FAILED).withMessage("nope"))); @@ -1339,14 +1361,14 @@ void setup() { static Stream getSetupFailingActivity() { return Stream.of( - Arguments.of(new Thread(() -> Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + Arguments.of(new Thread(() -> when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenThrow(ApplicationFailure.newNonRetryableFailure("", "")))), - Arguments.of(new Thread(() -> Mockito.when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) + Arguments.of(new Thread(() -> when(mJobCreationAndStatusUpdateActivity.createNewAttemptNumber(Mockito.any())) .thenThrow(ApplicationFailure.newNonRetryableFailure("", "")))), Arguments.of(new Thread(() -> Mockito.doThrow(ApplicationFailure.newNonRetryableFailure("", "")) .when(mJobCreationAndStatusUpdateActivity).reportJobStart(Mockito.any()))), Arguments.of(new Thread( - () -> Mockito.when(mGenerateInputActivityImpl.getSyncWorkflowInputWithAttemptNumber(Mockito.any(SyncInputWithAttemptNumber.class))) + () -> when(mGenerateInputActivityImpl.getSyncWorkflowInputWithAttemptNumber(Mockito.any(SyncInputWithAttemptNumber.class))) .thenThrow(ApplicationFailure.newNonRetryableFailure("", ""))))); } @@ -1354,7 +1376,7 @@ static Stream getSetupFailingActivity() { @MethodSource("getSetupFailingActivity") void testWorkflowRestartedAfterFailedActivity(final Thread mockSetup) throws InterruptedException { mockSetup.run(); - Mockito.when(mConfigFetchActivity.getTimeToWait(Mockito.any())).thenReturn(new ScheduleRetrieverOutput( + when(mConfigFetchActivity.getTimeToWait(Mockito.any())).thenReturn(new ScheduleRetrieverOutput( Duration.ZERO)); final UUID testId = UUID.randomUUID(); @@ -1375,7 +1397,7 @@ void testWorkflowRestartedAfterFailedActivity(final Thread mockSetup) throws Int // Sleep test env for restart delay, plus a small buffer to ensure that the workflow executed the // logic after the delay - testEnv.sleep(ConnectionManagerWorkflowImpl.WORKFLOW_FAILURE_RESTART_DELAY.plus(Duration.ofSeconds(10))); + testEnv.sleep(WORKFLOW_FAILURE_RESTART_DELAY.plus(Duration.ofSeconds(10))); final Queue events = testStateListener.events(testId); @@ -1388,11 +1410,11 @@ void testWorkflowRestartedAfterFailedActivity(final Thread mockSetup) throws Int @Test void testCanRetryFailedActivity() throws InterruptedException { - Mockito.when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) + when(mJobCreationAndStatusUpdateActivity.createNewJob(Mockito.any())) .thenThrow(ApplicationFailure.newNonRetryableFailure("", "")) .thenReturn(new JobCreationOutput(1l)); - Mockito.when(mConfigFetchActivity.getTimeToWait(Mockito.any())).thenReturn(new ScheduleRetrieverOutput( + when(mConfigFetchActivity.getTimeToWait(Mockito.any())).thenReturn(new ScheduleRetrieverOutput( Duration.ZERO)); final UUID testId = UUID.randomUUID(); @@ -1411,10 +1433,9 @@ void testCanRetryFailedActivity() throws InterruptedException { startWorkflowAndWaitUntilReady(workflow, input); // Sleep test env for half of restart delay, so that we know we are in the middle of the delay - testEnv.sleep(ConnectionManagerWorkflowImpl.WORKFLOW_FAILURE_RESTART_DELAY.dividedBy(2)); + testEnv.sleep(WORKFLOW_FAILURE_RESTART_DELAY.dividedBy(2)); workflow.retryFailedActivity(); Thread.sleep(500); // any time after no-waiting manual run - final Queue events = testStateListener.events(testId); Assertions.assertThat(events) @@ -1498,9 +1519,10 @@ private void setupSpecificChildWorkflow(final Class syncWorker.registerWorkflowImplementationTypes(mockedSyncedWorkflow); final Worker managerWorker = testEnv.newWorker(TemporalJobType.CONNECTION_UPDATER.name()); - managerWorker.registerWorkflowImplementationTypes(ConnectionManagerWorkflowImpl.class); + managerWorker.registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(ConnectionManagerWorkflowImpl.class)); managerWorker.registerActivitiesImplementations(mConfigFetchActivity, mCheckConnectionActivity, mConnectionDeletionActivity, - mGenerateInputActivityImpl, mJobCreationAndStatusUpdateActivity, mAutoDisableConnectionActivity, mStreamResetActivity); + mGenerateInputActivityImpl, mJobCreationAndStatusUpdateActivity, mAutoDisableConnectionActivity, mRecordMetricActivity, + mWorkflowConfigActivity, mRouteToSyncTaskQueueActivity); client = testEnv.getWorkflowClient(); testEnv.start(); diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/WorkflowReplayingTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/WorkflowReplayingTest.java index b2f0b9504ce55..99b149f60a8a9 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/WorkflowReplayingTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/WorkflowReplayingTest.java @@ -4,15 +4,49 @@ package io.airbyte.workers.temporal.scheduling; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.airbyte.workers.temporal.support.TemporalProxyHelper; +import io.micronaut.context.BeanRegistration; +import io.micronaut.inject.BeanIdentifier; +import io.temporal.activity.ActivityOptions; +import io.temporal.common.RetryOptions; import io.temporal.testing.WorkflowReplayer; import java.io.File; import java.net.URL; +import java.time.Duration; +import java.util.List; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; // TODO: Auto generation of the input and more scenario coverage @SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert") class WorkflowReplayingTest { + private ActivityOptions activityOptions; + private TemporalProxyHelper temporalProxyHelper; + + @BeforeEach + void setUp() { + activityOptions = ActivityOptions.newBuilder() + .setHeartbeatTimeout(Duration.ofSeconds(30)) + .setStartToCloseTimeout(Duration.ofSeconds(120)) + .setRetryOptions(RetryOptions.newBuilder() + .setMaximumAttempts(5) + .setInitialInterval(Duration.ofSeconds(30)) + .setMaximumInterval(Duration.ofSeconds(600)) + .build()) + .build(); + + final BeanIdentifier shortActivitiesBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration shortActivityOptionsBeanRegistration = mock(BeanRegistration.class); + when(shortActivitiesBeanIdentifier.getName()).thenReturn("shortActivityOptions"); + when(shortActivityOptionsBeanRegistration.getIdentifier()).thenReturn(shortActivitiesBeanIdentifier); + when(shortActivityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + temporalProxyHelper = new TemporalProxyHelper(List.of(shortActivityOptionsBeanRegistration)); + } + @Test void replaySimpleSuccessfulWorkflow() throws Exception { // This test ensures that a new version of the workflow doesn't break an in-progress execution @@ -23,7 +57,7 @@ void replaySimpleSuccessfulWorkflow() throws Exception { final File historyFile = new File(historyPath.toURI()); - WorkflowReplayer.replayWorkflowExecution(historyFile, ConnectionManagerWorkflowImpl.class); + WorkflowReplayer.replayWorkflowExecution(historyFile, temporalProxyHelper.proxyWorkflowClass(ConnectionManagerWorkflowImpl.class)); } } diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityTest.java index a53acdcc4512b..2f60f6202474b 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/AutoDisableConnectionActivityTest.java @@ -10,7 +10,6 @@ import static org.assertj.core.api.AssertionsForClassTypes.assertThat; import io.airbyte.commons.features.FeatureFlags; -import io.airbyte.config.Configs; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSync.Status; import io.airbyte.config.persistence.ConfigNotFoundException; @@ -36,7 +35,6 @@ import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; @@ -53,16 +51,12 @@ class AutoDisableConnectionActivityTest { @Mock private JobPersistence mJobPersistence; - @Mock - private Configs mConfigs; - @Mock private JobNotifier mJobNotifier; @Mock private Job mJob; - @InjectMocks private AutoDisableConnectionActivityImpl autoDisableActivity; private static final UUID CONNECTION_ID = UUID.randomUUID(); @@ -86,9 +80,16 @@ void setUp() throws IOException, JsonValidationException, ConfigNotFoundExceptio Mockito.when(mConfigRepository.getStandardSync(CONNECTION_ID)).thenReturn(standardSync); standardSync.setStatus(Status.ACTIVE); Mockito.when(mFeatureFlags.autoDisablesFailingConnections()).thenReturn(true); - Mockito.when(mConfigs.getMaxDaysOfOnlyFailedJobsBeforeConnectionDisable()).thenReturn(MAX_DAYS_OF_ONLY_FAILED_JOBS); Mockito.when(mJobPersistence.getLastReplicationJob(CONNECTION_ID)).thenReturn(Optional.of(mJob)); Mockito.when(mJobPersistence.getFirstReplicationJob(CONNECTION_ID)).thenReturn(Optional.of(mJob)); + + autoDisableActivity = new AutoDisableConnectionActivityImpl(); + autoDisableActivity.setConfigRepository(mConfigRepository); + autoDisableActivity.setJobPersistence(mJobPersistence); + autoDisableActivity.setFeatureFlags(mFeatureFlags); + autoDisableActivity.setMaxDaysOfOnlyFailedJobsBeforeConnectionDisable(MAX_DAYS_OF_ONLY_FAILED_JOBS); + autoDisableActivity.setMaxFailedJobsInARowBeforeConnectionDisable(MAX_FAILURE_JOBS_IN_A_ROW); + autoDisableActivity.setJobNotifier(mJobNotifier); } // test warnings @@ -96,11 +97,11 @@ void setUp() throws IOException, JsonValidationException, ConfigNotFoundExceptio @Test @DisplayName("Test that a notification warning is sent for connections that have failed `MAX_FAILURE_JOBS_IN_A_ROW / 2` times") void testWarningNotificationsForAutoDisablingMaxNumFailures() throws IOException { + // from most recent to least recent: MAX_FAILURE_JOBS_IN_A_ROW/2 and 1 success final List jobs = new ArrayList<>(Collections.nCopies(MAX_FAILURE_JOBS_IN_A_ROW / 2, FAILED_JOB)); jobs.add(SUCCEEDED_JOB); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))).thenReturn(jobs); @@ -114,12 +115,10 @@ void testWarningNotificationsForAutoDisablingMaxNumFailures() throws IOException @Test @DisplayName("Test that a notification warning is sent after only failed jobs in last `MAX_DAYS_OF_STRAIGHT_FAILURE / 2` days") void testWarningNotificationsForAutoDisablingMaxDaysOfFailure() throws IOException { - Mockito.when(mConfigs.getMaxDaysOfOnlyFailedJobsBeforeConnectionDisable()).thenReturn(MAX_DAYS_OF_ONLY_FAILED_JOBS); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))) .thenReturn(Collections.singletonList(FAILED_JOB)); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJob.getCreatedAtInSecond()).thenReturn( CURR_INSTANT.getEpochSecond() - TimeUnit.DAYS.toSeconds(MAX_DAYS_OF_ONLY_FAILED_JOBS_BEFORE_WARNING)); @@ -136,11 +135,9 @@ void testWarningNotificationsDoesNotSpam() throws IOException { final List jobs = new ArrayList<>(Collections.nCopies(2, FAILED_JOB)); final long mJobCreateOrUpdatedInSeconds = CURR_INSTANT.getEpochSecond() - TimeUnit.DAYS.toSeconds(MAX_DAYS_OF_ONLY_FAILED_JOBS_BEFORE_WARNING); - Mockito.when(mConfigs.getMaxDaysOfOnlyFailedJobsBeforeConnectionDisable()).thenReturn(MAX_DAYS_OF_ONLY_FAILED_JOBS); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))).thenReturn(jobs); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJob.getCreatedAtInSecond()).thenReturn(mJobCreateOrUpdatedInSeconds); Mockito.when(mJob.getUpdatedAtInSecond()).thenReturn(mJobCreateOrUpdatedInSeconds); @@ -157,11 +154,9 @@ void testWarningNotificationsDoesNotSpamAfterConsecutiveFailures() throws IOExce final List jobs = new ArrayList<>(Collections.nCopies(MAX_FAILURE_JOBS_IN_A_ROW - 1, FAILED_JOB)); final long mJobCreateOrUpdatedInSeconds = CURR_INSTANT.getEpochSecond() - TimeUnit.DAYS.toSeconds(MAX_DAYS_OF_ONLY_FAILED_JOBS_BEFORE_WARNING); - Mockito.when(mConfigs.getMaxDaysOfOnlyFailedJobsBeforeConnectionDisable()).thenReturn(MAX_DAYS_OF_ONLY_FAILED_JOBS); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))).thenReturn(jobs); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJob.getCreatedAtInSecond()).thenReturn(mJobCreateOrUpdatedInSeconds); Mockito.when(mJob.getUpdatedAtInSecond()).thenReturn(mJobCreateOrUpdatedInSeconds); @@ -175,12 +170,10 @@ void testWarningNotificationsDoesNotSpamAfterConsecutiveFailures() throws IOExce @Test @DisplayName("Test that the connection is _not_ disabled and no warning is sent after only failed jobs and oldest job is less than `MAX_DAYS_OF_STRAIGHT_FAILURE / 2 `days old") void testOnlyFailuresButFirstJobYoungerThanMaxDaysWarning() throws IOException { - Mockito.when(mConfigs.getMaxDaysOfOnlyFailedJobsBeforeConnectionDisable()).thenReturn(MAX_DAYS_OF_ONLY_FAILED_JOBS); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))) .thenReturn(Collections.singletonList(FAILED_JOB)); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJob.getCreatedAtInSecond()).thenReturn(CURR_INSTANT.getEpochSecond()); final AutoDisableConnectionOutput output = autoDisableActivity.autoDisableFailingConnection(ACTIVITY_INPUT); @@ -199,7 +192,6 @@ void testMaxFailuresInARow() throws IOException, JsonValidationException, Config final List jobs = new ArrayList<>(Collections.nCopies(MAX_FAILURE_JOBS_IN_A_ROW, FAILED_JOB)); jobs.add(SUCCEEDED_JOB); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))).thenReturn(jobs); Mockito.when(mConfigRepository.getStandardSync(CONNECTION_ID)).thenReturn(standardSync); @@ -218,7 +210,6 @@ void testLessThanMaxFailuresInARow() throws IOException { final List jobs = new ArrayList<>(Collections.nCopies(MAX_FAILURE_JOBS_IN_A_ROW - 1, FAILED_JOB)); jobs.add(SUCCEEDED_JOB); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))).thenReturn(jobs); Mockito.when(mJob.getCreatedAtInSecond()).thenReturn( @@ -250,12 +241,10 @@ void testNoRuns() throws IOException { @Test @DisplayName("Test that the connection is disabled after only failed jobs in last MAX_DAYS_OF_STRAIGHT_FAILURE days") void testOnlyFailuresInMaxDays() throws IOException, JsonValidationException, ConfigNotFoundException { - Mockito.when(mConfigs.getMaxDaysOfOnlyFailedJobsBeforeConnectionDisable()).thenReturn(MAX_DAYS_OF_ONLY_FAILED_JOBS); Mockito.when(mJobPersistence.listJobStatusAndTimestampWithConnection(CONNECTION_ID, REPLICATION_TYPES, CURR_INSTANT.minus(MAX_DAYS_OF_ONLY_FAILED_JOBS, ChronoUnit.DAYS))) .thenReturn(Collections.singletonList(FAILED_JOB)); - Mockito.when(mConfigs.getMaxFailedJobsInARowBeforeConnectionDisable()).thenReturn(MAX_FAILURE_JOBS_IN_A_ROW); Mockito.when(mJob.getCreatedAtInSecond()).thenReturn( CURR_INSTANT.getEpochSecond() - TimeUnit.DAYS.toSeconds(MAX_DAYS_OF_ONLY_FAILED_JOBS)); Mockito.when(mConfigRepository.getStandardSync(CONNECTION_ID)).thenReturn(standardSync); diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityTest.java index a5bab67db82b7..f485e0a8e06b0 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/ConfigFetchActivityTest.java @@ -5,7 +5,6 @@ package io.airbyte.workers.temporal.scheduling.activities; import io.airbyte.config.BasicSchedule; -import io.airbyte.config.Configs; import io.airbyte.config.Cron; import io.airbyte.config.Schedule; import io.airbyte.config.ScheduleData; @@ -26,6 +25,7 @@ import java.util.TimeZone; import java.util.UUID; import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; @@ -37,19 +37,18 @@ @ExtendWith(MockitoExtension.class) class ConfigFetchActivityTest { + private static final Integer SYNC_JOB_MAX_ATTEMPTS = 3; + @Mock private ConfigRepository mConfigRepository; @Mock private JobPersistence mJobPersistence; - @Mock - private Configs mConfigs; - @Mock private Job mJob; - private ConfigFetchActivity configFetchActivity; + private ConfigFetchActivityImpl configFetchActivity; private final static UUID connectionId = UUID.randomUUID(); private final static StandardSync standardSyncWithLegacySchedule = new StandardSync() @@ -91,13 +90,21 @@ class ConfigFetchActivityTest { .withStatus(Status.DEPRECATED); private static final StandardSync standardSyncWithoutSchedule = new StandardSync(); + @BeforeEach + void setup() { + configFetchActivity = new ConfigFetchActivityImpl(); + configFetchActivity.setConfigRepository(mConfigRepository); + configFetchActivity.setJobPersistence(mJobPersistence); + configFetchActivity.setSyncJobMaxAttempts(SYNC_JOB_MAX_ATTEMPTS); + configFetchActivity.setCurrentSecondsSupplier(() -> Instant.now().getEpochSecond()); + } + @Nested class TimeToWaitTest { @Test @DisplayName("Test that the job gets scheduled if it is not manual and if it is the first run with legacy schedule schema") void testFirstJobNonManual() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> Instant.now().getEpochSecond()); Mockito.when(mJobPersistence.getLastReplicationJob(connectionId)) .thenReturn(Optional.empty()); @@ -115,8 +122,6 @@ void testFirstJobNonManual() throws IOException, JsonValidationException, Config @Test @DisplayName("Test that the job will wait for a long time if it is manual in the legacy schedule schema") void testManual() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> Instant.now().getEpochSecond()); - Mockito.when(mConfigRepository.getStandardSync(connectionId)) .thenReturn(standardSyncWithoutSchedule); @@ -131,8 +136,6 @@ void testManual() throws IOException, JsonValidationException, ConfigNotFoundExc @Test @DisplayName("Test that the job will wait for a long time if it is disabled") void testDisable() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> Instant.now().getEpochSecond()); - Mockito.when(mConfigRepository.getStandardSync(connectionId)) .thenReturn(standardSyncWithScheduleDisable); @@ -147,8 +150,6 @@ void testDisable() throws IOException, JsonValidationException, ConfigNotFoundEx @Test @DisplayName("Test that the connection will wait for a long time if it is deleted") void testDeleted() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> Instant.now().getEpochSecond()); - Mockito.when(mConfigRepository.getStandardSync(connectionId)) .thenReturn(standardSyncWithScheduleDeleted); @@ -163,7 +164,7 @@ void testDeleted() throws IOException, JsonValidationException, ConfigNotFoundEx @Test @DisplayName("Test we will wait the required amount of time with legacy config") void testWait() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> 60L * 3); + configFetchActivity.setCurrentSecondsSupplier(() -> 60L * 3); Mockito.when(mJob.getStartedAtInSecond()) .thenReturn(Optional.of(60L)); @@ -185,7 +186,7 @@ void testWait() throws IOException, JsonValidationException, ConfigNotFoundExcep @Test @DisplayName("Test we will not wait if we are late in the legacy schedule schema") void testNotWaitIfLate() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> 60L * 10); + configFetchActivity.setCurrentSecondsSupplier(() -> 60L * 10); Mockito.when(mJob.getStartedAtInSecond()) .thenReturn(Optional.of(60L)); @@ -209,8 +210,6 @@ void testNotWaitIfLate() throws IOException, JsonValidationException, ConfigNotF @Test @DisplayName("Test that the job will wait a long time if it is MANUAL scheduleType") void testManualScheduleType() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> Instant.now().getEpochSecond()); - Mockito.when(mConfigRepository.getStandardSync(connectionId)) .thenReturn(standardSyncWithManualScheduleType); @@ -225,7 +224,6 @@ void testManualScheduleType() throws IOException, JsonValidationException, Confi @Test @DisplayName("Test that the job will be immediately scheduled if it is a BASIC_SCHEDULE type on the first run") void testBasicScheduleTypeFirstRun() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> Instant.now().getEpochSecond()); Mockito.when(mJobPersistence.getLastReplicationJob(connectionId)) .thenReturn(Optional.empty()); @@ -243,7 +241,7 @@ void testBasicScheduleTypeFirstRun() throws IOException, JsonValidationException @Test @DisplayName("Test that we will wait the required amount of time with a BASIC_SCHEDULE type on a subsequent run") void testBasicScheduleSubsequentRun() throws IOException, JsonValidationException, ConfigNotFoundException { - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> 60L * 3); + configFetchActivity.setCurrentSecondsSupplier(() -> 60L * 3); Mockito.when(mJob.getStartedAtInSecond()) .thenReturn(Optional.of(60L)); @@ -265,12 +263,13 @@ void testBasicScheduleSubsequentRun() throws IOException, JsonValidationExceptio @Test @DisplayName("Test that the job will wait to be scheduled if it is a CRON type") void testCronScheduleSubsequentRun() throws IOException, JsonValidationException, ConfigNotFoundException { - Calendar mockRightNow = Calendar.getInstance(TimeZone.getTimeZone("UTC")); + final Calendar mockRightNow = Calendar.getInstance(TimeZone.getTimeZone("UTC")); mockRightNow.set(Calendar.HOUR_OF_DAY, 0); mockRightNow.set(Calendar.MINUTE, 0); mockRightNow.set(Calendar.SECOND, 0); mockRightNow.set(Calendar.MILLISECOND, 0); - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> mockRightNow.getTimeInMillis() / 1000L); + + configFetchActivity.setCurrentSecondsSupplier(() -> mockRightNow.getTimeInMillis() / 1000L); Mockito.when(mJobPersistence.getLastReplicationJob(connectionId)) .thenReturn(Optional.of(mJob)); @@ -289,12 +288,12 @@ void testCronScheduleSubsequentRun() throws IOException, JsonValidationException @Test @DisplayName("Test that the job will only be scheduled once per minimum cron interval") void testCronScheduleMinimumInterval() throws IOException, JsonValidationException, ConfigNotFoundException { - Calendar mockRightNow = Calendar.getInstance(TimeZone.getTimeZone("UTC")); + final Calendar mockRightNow = Calendar.getInstance(TimeZone.getTimeZone("UTC")); mockRightNow.set(Calendar.HOUR_OF_DAY, 12); mockRightNow.set(Calendar.MINUTE, 0); mockRightNow.set(Calendar.SECOND, 0); mockRightNow.set(Calendar.MILLISECOND, 0); - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> mockRightNow.getTimeInMillis() / 1000L); + configFetchActivity.setCurrentSecondsSupplier(() -> mockRightNow.getTimeInMillis() / 1000L); Mockito.when(mJob.getStartedAtInSecond()).thenReturn(Optional.of(mockRightNow.getTimeInMillis() / 1000L)); Mockito.when(mJobPersistence.getLastReplicationJob(connectionId)) @@ -315,14 +314,10 @@ void testCronScheduleMinimumInterval() throws IOException, JsonValidationExcepti class TestGetMaxAttempt { @Test - @DisplayName("Test that we are using to right service to get the maximum amout of attempt") + @DisplayName("Test that we are using to right service to get the maximum amount of attempt") void testGetMaxAttempt() { final int maxAttempt = 15031990; - Mockito.when(mConfigs.getSyncJobMaxAttempts()) - .thenReturn(15031990); - - configFetchActivity = new ConfigFetchActivityImpl(mConfigRepository, mJobPersistence, mConfigs, () -> Instant.now().getEpochSecond()); - + configFetchActivity.setSyncJobMaxAttempts(maxAttempt); Assertions.assertThat(configFetchActivity.getMaxAttempt().getMaxAttempt()) .isEqualTo(maxAttempt); } diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImplTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImplTest.java new file mode 100644 index 0000000000000..da8786a54de89 --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/RecordMetricActivityImplTest.java @@ -0,0 +1,86 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.airbyte.metrics.lib.MetricAttribute; +import io.airbyte.metrics.lib.MetricClient; +import io.airbyte.metrics.lib.MetricTags; +import io.airbyte.metrics.lib.OssMetricsRegistry; +import io.airbyte.workers.temporal.scheduling.ConnectionUpdaterInput; +import io.airbyte.workers.temporal.scheduling.activities.RecordMetricActivity.FailureCause; +import io.airbyte.workers.temporal.scheduling.activities.RecordMetricActivity.RecordMetricInput; +import java.util.Optional; +import java.util.UUID; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Test suite for the {@link RecordMetricActivityImpl} class. + */ +class RecordMetricActivityImplTest { + + private MetricClient metricClient; + + private RecordMetricActivityImpl activity; + + @BeforeEach + void setup() { + metricClient = mock(MetricClient.class); + activity = new RecordMetricActivityImpl(metricClient); + } + + @Test + void testRecordingMetricCounter() { + final UUID connectionId = UUID.randomUUID(); + final OssMetricsRegistry metricName = OssMetricsRegistry.TEMPORAL_WORKFLOW_ATTEMPT; + final ConnectionUpdaterInput connectionUpdaterInput = mock(ConnectionUpdaterInput.class); + final RecordMetricInput metricInput = new RecordMetricInput(connectionUpdaterInput, Optional.empty(), metricName, null); + + when(connectionUpdaterInput.getConnectionId()).thenReturn(connectionId); + + activity.recordWorkflowCountMetric(metricInput); + + verify(metricClient).count(eq(metricName), eq(1L), eq(new MetricAttribute(MetricTags.CONNECTION_ID, String.valueOf(connectionId)))); + } + + @Test + void testRecordingMetricCounterWithAdditionalAttributes() { + final UUID connectionId = UUID.randomUUID(); + final OssMetricsRegistry metricName = OssMetricsRegistry.TEMPORAL_WORKFLOW_ATTEMPT; + final ConnectionUpdaterInput connectionUpdaterInput = mock(ConnectionUpdaterInput.class); + final MetricAttribute additionalAttribute = new MetricAttribute(MetricTags.JOB_STATUS, "test"); + final RecordMetricInput metricInput = + new RecordMetricInput(connectionUpdaterInput, Optional.empty(), metricName, new MetricAttribute[] {additionalAttribute}); + + when(connectionUpdaterInput.getConnectionId()).thenReturn(connectionId); + + activity.recordWorkflowCountMetric(metricInput); + + verify(metricClient).count(eq(metricName), eq(1L), eq(new MetricAttribute(MetricTags.CONNECTION_ID, String.valueOf(connectionId))), + eq(additionalAttribute)); + } + + @Test + void testRecordingMetricCounterWithFailureCause() { + final UUID connectionId = UUID.randomUUID(); + final OssMetricsRegistry metricName = OssMetricsRegistry.TEMPORAL_WORKFLOW_ATTEMPT; + final ConnectionUpdaterInput connectionUpdaterInput = mock(ConnectionUpdaterInput.class); + final FailureCause failureCause = FailureCause.CANCELED; + final RecordMetricInput metricInput = new RecordMetricInput(connectionUpdaterInput, Optional.of(failureCause), metricName, null); + + when(connectionUpdaterInput.getConnectionId()).thenReturn(connectionId); + + activity.recordWorkflowCountMetric(metricInput); + + verify(metricClient).count(eq(metricName), eq(1L), eq(new MetricAttribute(MetricTags.CONNECTION_ID, String.valueOf(connectionId))), + eq(new MetricAttribute(MetricTags.RESET_WORKFLOW_FAILURE_CAUSE, failureCause.name()))); + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityTest.java index cd7f41d18ebfe..d2fc66d761847 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/StreamResetActivityTest.java @@ -4,66 +4,31 @@ package io.airbyte.workers.temporal.scheduling.activities; -import static org.mockito.Mockito.RETURNS_DEEP_STUBS; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.when; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.verify; -import io.airbyte.config.JobConfig.ConfigType; -import io.airbyte.config.persistence.StreamResetPersistence; -import io.airbyte.protocol.models.StreamDescriptor; -import io.airbyte.scheduler.models.Job; -import io.airbyte.scheduler.persistence.JobPersistence; +import io.airbyte.workers.temporal.StreamResetRecordsHelper; import io.airbyte.workers.temporal.scheduling.activities.StreamResetActivity.DeleteStreamResetRecordsForJobInput; -import java.io.IOException; -import java.util.List; import java.util.UUID; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.InjectMocks; import org.mockito.Mock; -import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; @ExtendWith(MockitoExtension.class) class StreamResetActivityTest { @Mock - private StreamResetPersistence streamResetPersistence; - @Mock - private JobPersistence jobPersistence; + private StreamResetRecordsHelper streamResetRecordsHelper; @InjectMocks private StreamResetActivityImpl streamResetActivity; - private final DeleteStreamResetRecordsForJobInput input = new DeleteStreamResetRecordsForJobInput(UUID.randomUUID(), Long.valueOf("123")); - private final DeleteStreamResetRecordsForJobInput noJobIdInput = new DeleteStreamResetRecordsForJobInput(UUID.randomUUID(), null); - - @Test - void testDeleteStreamResetRecordsForJob() throws IOException { - final Job jobMock = mock(Job.class, RETURNS_DEEP_STUBS); - when(jobPersistence.getJob(input.getJobId())).thenReturn(jobMock); - - when(jobMock.getConfig().getConfigType()).thenReturn(ConfigType.RESET_CONNECTION); - final List streamsToDelete = List.of(new StreamDescriptor().withName("streamname").withNamespace("namespace")); - when(jobMock.getConfig().getResetConnection().getResetSourceConfiguration().getStreamsToReset()).thenReturn(streamsToDelete); - streamResetActivity.deleteStreamResetRecordsForJob(input); - Mockito.verify(streamResetPersistence).deleteStreamResets(input.getConnectionId(), streamsToDelete); - } @Test - void testIncorrectConfigType() throws IOException { - final Job jobMock = mock(Job.class, RETURNS_DEEP_STUBS); - when(jobPersistence.getJob(input.getJobId())).thenReturn(jobMock); - - when(jobMock.getConfig().getConfigType()).thenReturn(ConfigType.SYNC); + void testDeleteStreamResetRecordsForJob() { + final DeleteStreamResetRecordsForJobInput input = new DeleteStreamResetRecordsForJobInput(UUID.randomUUID(), Long.valueOf("123")); streamResetActivity.deleteStreamResetRecordsForJob(input); - Mockito.verify(streamResetPersistence, never()).deleteStreamResets(Mockito.any(UUID.class), Mockito.anyList()); - } - - @Test - void testNoJobId() throws IOException { - streamResetActivity.deleteStreamResetRecordsForJob(noJobIdInput); - Mockito.verify(jobPersistence, never()).getJob(Mockito.anyLong()); - Mockito.verify(streamResetPersistence, never()).deleteStreamResets(Mockito.any(UUID.class), Mockito.anyList()); + verify(streamResetRecordsHelper).deleteStreamResetRecordsForJob(eq(input.getJobId()), eq(input.getConnectionId())); } } diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImplTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImplTest.java new file mode 100644 index 0000000000000..faece206fd3a3 --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/scheduling/activities/WorkflowConfigActivityImplTest.java @@ -0,0 +1,23 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.scheduling.activities; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Test suite for the {@link WorkflowConfigActivityImpl} class. + */ +class WorkflowConfigActivityImplTest { + + @Test + void testFetchingWorkflowRestartDelayInSeconds() { + final Long workflowRestartDelaySeconds = 30L; + final WorkflowConfigActivityImpl activity = new WorkflowConfigActivityImpl(); + activity.setWorkflowRestartDelaySeconds(workflowRestartDelaySeconds); + Assertions.assertEquals(workflowRestartDelaySeconds, activity.getWorkflowRestartDelaySeconds().getSeconds()); + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ErrorTestWorkflowImpl.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ErrorTestWorkflowImpl.java new file mode 100644 index 0000000000000..ed419ba481f96 --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ErrorTestWorkflowImpl.java @@ -0,0 +1,24 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.stubs; + +import io.airbyte.workers.temporal.exception.RetryableException; + +public class ErrorTestWorkflowImpl implements TestWorkflow { + + @Override + public void run() throws RetryableException { + throw new RetryableException(new NullPointerException("test")); + } + + @Override + public void cancel() {} + + @Override + public Integer getState() { + return 1; + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/HeartbeatWorkflow.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/HeartbeatWorkflow.java similarity index 90% rename from airbyte-workers/src/test/java/io/airbyte/workers/temporal/HeartbeatWorkflow.java rename to airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/HeartbeatWorkflow.java index 3fc2c9feaf749..306b6b9612aa0 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/HeartbeatWorkflow.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/HeartbeatWorkflow.java @@ -2,8 +2,9 @@ * Copyright (c) 2022 Airbyte, Inc., all rights reserved. */ -package io.airbyte.workers.temporal; +package io.airbyte.workers.temporal.stubs; +import io.airbyte.workers.temporal.TemporalUtils; import io.temporal.activity.ActivityCancellationType; import io.temporal.activity.ActivityInterface; import io.temporal.activity.ActivityMethod; @@ -48,7 +49,7 @@ class HeartbeatActivityImpl implements HeartbeatActivity { private final Runnable runnable; - public HeartbeatActivityImpl(Runnable runnable) { + public HeartbeatActivityImpl(final Runnable runnable) { this.runnable = runnable; } diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/InvalidTestWorkflowImpl.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/InvalidTestWorkflowImpl.java new file mode 100644 index 0000000000000..2a749650d91cb --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/InvalidTestWorkflowImpl.java @@ -0,0 +1,27 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.stubs; + +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; +import io.airbyte.workers.temporal.exception.RetryableException; + +@SuppressWarnings("PMD.UnusedPrivateField") +public class InvalidTestWorkflowImpl implements TestWorkflow { + + @TemporalActivityStub(activityOptionsBeanName = "missingActivityOptions") + private TestActivity testActivity; + + @Override + public void run() throws RetryableException {} + + @Override + public void cancel() {} + + @Override + public Integer getState() { + return 1; + } + +} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivity.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/TestActivity.java similarity index 55% rename from airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivity.java rename to airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/TestActivity.java index c6b41d68c3915..8a66474ac833e 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/sync/RouteToTaskQueueActivity.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/TestActivity.java @@ -2,16 +2,15 @@ * Copyright (c) 2022 Airbyte, Inc., all rights reserved. */ -package io.airbyte.workers.temporal.sync; +package io.airbyte.workers.temporal.stubs; import io.temporal.activity.ActivityInterface; import io.temporal.activity.ActivityMethod; -import java.util.UUID; @ActivityInterface -public interface RouteToTaskQueueActivity { +public interface TestActivity { @ActivityMethod - String routeToTaskQueue(final UUID connectionId); + String getValue(); } diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/TestWorkflow.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/TestWorkflow.java new file mode 100644 index 0000000000000..d8aa6acc0a95c --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/TestWorkflow.java @@ -0,0 +1,24 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.stubs; + +import io.temporal.workflow.QueryMethod; +import io.temporal.workflow.SignalMethod; +import io.temporal.workflow.WorkflowInterface; +import io.temporal.workflow.WorkflowMethod; + +@WorkflowInterface +public interface TestWorkflow { + + @WorkflowMethod + void run(); + + @SignalMethod + void cancel(); + + @QueryMethod + Integer getState(); + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ValidTestWorkflowImpl.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ValidTestWorkflowImpl.java new file mode 100644 index 0000000000000..d620c521be0d7 --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/stubs/ValidTestWorkflowImpl.java @@ -0,0 +1,41 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.stubs; + +import io.airbyte.workers.temporal.annotations.TemporalActivityStub; + +public class ValidTestWorkflowImpl implements TestWorkflow { + + private boolean cancelled = false; + private boolean hasRun = false; + + @TemporalActivityStub(activityOptionsBeanName = "activityOptions") + private TestActivity testActivity; + + @Override + public void run() { + testActivity.getValue(); + hasRun = true; + } + + @Override + public void cancel() { + cancelled = true; + } + + @Override + public Integer getState() { + return 1; + } + + public boolean isCancelled() { + return cancelled; + } + + public boolean isHasRun() { + return hasRun; + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptorTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptorTest.java new file mode 100644 index 0000000000000..bbf5c75f017a9 --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalActivityStubInterceptorTest.java @@ -0,0 +1,142 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.support; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.airbyte.workers.temporal.exception.RetryableException; +import io.airbyte.workers.temporal.stubs.ErrorTestWorkflowImpl; +import io.airbyte.workers.temporal.stubs.InvalidTestWorkflowImpl; +import io.airbyte.workers.temporal.stubs.TestActivity; +import io.airbyte.workers.temporal.stubs.ValidTestWorkflowImpl; +import io.micronaut.context.BeanRegistration; +import io.micronaut.inject.BeanIdentifier; +import io.temporal.activity.ActivityOptions; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Test suite for the {@link TemporalActivityStubInterceptor} class. + */ +class TemporalActivityStubInterceptorTest { + + private static final String ACTIVITY_OPTIONS = "activityOptions"; + + @Test + void testExecutionOfValidWorkflowWithActivities() throws Exception { + final ActivityOptions activityOptions = mock(ActivityOptions.class); + final TestActivity testActivity = mock(TestActivity.class); + + final BeanIdentifier activityOptionsBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration activityOptionsBeanRegistration = mock(BeanRegistration.class); + when(activityOptionsBeanIdentifier.getName()).thenReturn(ACTIVITY_OPTIONS); + when(activityOptionsBeanRegistration.getIdentifier()).thenReturn(activityOptionsBeanIdentifier); + when(activityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + + final TemporalActivityStubInterceptor interceptor = + new TemporalActivityStubInterceptor(ValidTestWorkflowImpl.class, List.of(activityOptionsBeanRegistration)); + interceptor.setActivityStubGenerator((c, a) -> testActivity); + + final ValidTestWorkflowImpl validTestWorklowImpl = new ValidTestWorkflowImpl(); + final Callable callable = () -> { + validTestWorklowImpl.run(); + return null; + }; + + interceptor.execute(validTestWorklowImpl, callable); + Assertions.assertTrue(validTestWorklowImpl.isHasRun()); + } + + @Test + void testExecutionOfValidWorkflowWithActivitiesThatThrows() { + final ActivityOptions activityOptions = mock(ActivityOptions.class); + final TestActivity testActivity = mock(TestActivity.class); + + final BeanIdentifier activityOptionsBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration activityOptionsBeanRegistration = mock(BeanRegistration.class); + when(activityOptionsBeanIdentifier.getName()).thenReturn(ACTIVITY_OPTIONS); + when(activityOptionsBeanRegistration.getIdentifier()).thenReturn(activityOptionsBeanIdentifier); + when(activityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + + final TemporalActivityStubInterceptor interceptor = + new TemporalActivityStubInterceptor(ErrorTestWorkflowImpl.class, List.of(activityOptionsBeanRegistration)); + interceptor.setActivityStubGenerator((c, a) -> testActivity); + + final ErrorTestWorkflowImpl errorTestWorkflowImpl = new ErrorTestWorkflowImpl(); + final Callable callable = () -> { + errorTestWorkflowImpl.run(); + return null; + }; + + Assertions.assertThrows(RetryableException.class, () -> { + interceptor.execute(errorTestWorkflowImpl, callable); + }); + } + + @Test + void testActivityStubsAreOnlyInitializedOnce() throws Exception { + final AtomicInteger activityStubInitializationCounter = new AtomicInteger(0); + final ActivityOptions activityOptions = mock(ActivityOptions.class); + final TestActivity testActivity = mock(TestActivity.class); + final TemporalActivityStubGeneratorFunction, ActivityOptions, Object> activityStubFunction = (c, a) -> { + activityStubInitializationCounter.incrementAndGet(); + return testActivity; + }; + + final BeanIdentifier activityOptionsBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration activityOptionsBeanRegistration = mock(BeanRegistration.class); + when(activityOptionsBeanIdentifier.getName()).thenReturn(ACTIVITY_OPTIONS); + when(activityOptionsBeanRegistration.getIdentifier()).thenReturn(activityOptionsBeanIdentifier); + when(activityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + + final TemporalActivityStubInterceptor interceptor = + new TemporalActivityStubInterceptor(ValidTestWorkflowImpl.class, List.of(activityOptionsBeanRegistration)); + interceptor.setActivityStubGenerator(activityStubFunction); + + final ValidTestWorkflowImpl validTestWorklowImpl = new ValidTestWorkflowImpl(); + final Callable callable = () -> { + validTestWorklowImpl.run(); + return null; + }; + interceptor.execute(validTestWorklowImpl, callable); + interceptor.execute(validTestWorklowImpl, callable); + interceptor.execute(validTestWorklowImpl, callable); + interceptor.execute(validTestWorklowImpl, callable); + + Assertions.assertEquals(1, activityStubInitializationCounter.get()); + } + + @Test + void testExecutionOfInvalidWorkflowWithActivityWithMissingActivityOptions() { + final ActivityOptions activityOptions = mock(ActivityOptions.class); + final TestActivity testActivity = mock(TestActivity.class); + + final BeanIdentifier activityOptionsBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration activityOptionsBeanRegistration = mock(BeanRegistration.class); + when(activityOptionsBeanIdentifier.getName()).thenReturn(ACTIVITY_OPTIONS); + when(activityOptionsBeanRegistration.getIdentifier()).thenReturn(activityOptionsBeanIdentifier); + when(activityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + + final TemporalActivityStubInterceptor interceptor = + new TemporalActivityStubInterceptor(InvalidTestWorkflowImpl.class, List.of(activityOptionsBeanRegistration)); + interceptor.setActivityStubGenerator((c, a) -> testActivity); + + final InvalidTestWorkflowImpl invalidTestWorklowImpl = new InvalidTestWorkflowImpl(); + final Callable callable = () -> { + invalidTestWorklowImpl.run(); + return null; + }; + + final RuntimeException exception = Assertions.assertThrows(RuntimeException.class, () -> { + interceptor.execute(invalidTestWorklowImpl, callable); + }); + Assertions.assertEquals(IllegalStateException.class, exception.getCause().getClass()); + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalProxyHelperTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalProxyHelperTest.java new file mode 100644 index 0000000000000..2733cb662d0ec --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/support/TemporalProxyHelperTest.java @@ -0,0 +1,57 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.support; + +import static org.junit.Assert.assertNotNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.airbyte.workers.temporal.stubs.ValidTestWorkflowImpl; +import io.micronaut.context.BeanRegistration; +import io.micronaut.inject.BeanIdentifier; +import io.temporal.activity.ActivityOptions; +import io.temporal.common.RetryOptions; +import java.lang.reflect.InvocationTargetException; +import java.time.Duration; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Test suite for the {@link TemporalProxyHelper} class. + */ +class TemporalProxyHelperTest { + + @Test + void testProxyToImplementation() throws NoSuchMethodException, InvocationTargetException, InstantiationException, IllegalAccessException { + final ActivityOptions activityOptions = ActivityOptions.newBuilder() + .setHeartbeatTimeout(Duration.ofSeconds(30)) + .setStartToCloseTimeout(Duration.ofSeconds(120)) + .setRetryOptions(RetryOptions.newBuilder() + .setMaximumAttempts(5) + .setInitialInterval(Duration.ofSeconds(30)) + .setMaximumInterval(Duration.ofSeconds(600)) + .build()) + .build(); + + final BeanIdentifier activityOptionsBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration activityOptionsBeanRegistration = mock(BeanRegistration.class); + when(activityOptionsBeanIdentifier.getName()).thenReturn("activityOptions"); + when(activityOptionsBeanRegistration.getIdentifier()).thenReturn(activityOptionsBeanIdentifier); + when(activityOptionsBeanRegistration.getBean()).thenReturn(activityOptions); + + final TemporalProxyHelper temporalProxyHelper = new TemporalProxyHelper(List.of(activityOptionsBeanRegistration)); + temporalProxyHelper.setActivityStubGenerator((c, a) -> mock(c)); + + final Class proxy = temporalProxyHelper.proxyWorkflowClass(ValidTestWorkflowImpl.class); + + assertNotNull(proxy); + + final ValidTestWorkflowImpl proxyImplementation = proxy.getDeclaredConstructor().newInstance(); + proxyImplementation.run(); + Assertions.assertTrue(proxyImplementation.isHasRun()); + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/PersistStateActivityTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/PersistStateActivityTest.java index f729d8fce8606..63075a57b16f8 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/PersistStateActivityTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/PersistStateActivityTest.java @@ -26,8 +26,8 @@ import io.airbyte.protocol.models.ConfiguredAirbyteStream; import io.airbyte.protocol.models.SyncMode; import java.util.List; +import java.util.Map; import java.util.UUID; -import org.elasticsearch.common.collect.Map; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/RouterServiceTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/RouterServiceTest.java new file mode 100644 index 0000000000000..9eecf2dcab731 --- /dev/null +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/RouterServiceTest.java @@ -0,0 +1,60 @@ +/* + * Copyright (c) 2022 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.workers.temporal.sync; + +import static io.airbyte.workers.temporal.sync.RouterService.MVP_DATA_PLANE_TASK_QUEUE; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import io.airbyte.workers.temporal.TemporalJobType; +import java.util.UUID; +import org.junit.jupiter.api.Test; + +/** + * Test suite for the {@link RouterService} class. + */ +class RouterServiceTest { + + @Test + void testSelectionOfTaskQueueForDataPlane() { + final UUID connectionId = UUID.randomUUID(); + final String connectionIdsForMvpDataPlane = connectionId.toString(); + final RouterService routerService = new RouterService(); + routerService.setConnectionIdsForMvpDataPlane(connectionIdsForMvpDataPlane); + + final String taskQueue = routerService.getTaskQueue(connectionId); + assertEquals(MVP_DATA_PLANE_TASK_QUEUE, taskQueue); + } + + @Test + void testSelectionOfTaskQueueForNonMatchingConnectionId() { + final UUID connectionId = UUID.randomUUID(); + final String connectionIdsForMvpDataPlane = "1,2,3,4,5"; + final RouterService routerService = new RouterService(); + routerService.setConnectionIdsForMvpDataPlane(connectionIdsForMvpDataPlane); + + final String taskQueue = routerService.getTaskQueue(connectionId); + assertEquals(TemporalJobType.SYNC.name(), taskQueue); + } + + @Test + void testSelectionOfTaskQueueForNullConnectionId() { + final String connectionIdsForMvpDataPlane = "1,2,3,4,5"; + final RouterService routerService = new RouterService(); + routerService.setConnectionIdsForMvpDataPlane(connectionIdsForMvpDataPlane); + + final String taskQueue = routerService.getTaskQueue(null); + assertEquals(TemporalJobType.SYNC.name(), taskQueue); + } + + @Test + void testSelectionOfTaskQueueForBlankConnectionIdSet() { + final UUID connectionId = UUID.randomUUID(); + final RouterService routerService = new RouterService(); + + final String taskQueue = routerService.getTaskQueue(connectionId); + assertEquals(TemporalJobType.SYNC.name(), taskQueue); + } + +} diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/SyncWorkflowTest.java similarity index 77% rename from airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java rename to airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/SyncWorkflowTest.java index bf1621bb1bffa..e1df428166f5a 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/sync/SyncWorkflowTest.java @@ -2,16 +2,18 @@ * Copyright (c) 2022 Airbyte, Inc., all rights reserved. */ -package io.airbyte.workers.temporal; +package io.airbyte.workers.temporal.sync; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; import io.airbyte.config.NormalizationInput; import io.airbyte.config.NormalizationSummary; @@ -24,27 +26,23 @@ import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; import io.airbyte.workers.TestConfigHelpers; -import io.airbyte.workers.temporal.sync.DbtTransformationActivity; -import io.airbyte.workers.temporal.sync.DbtTransformationActivityImpl; -import io.airbyte.workers.temporal.sync.NormalizationActivity; -import io.airbyte.workers.temporal.sync.NormalizationActivityImpl; -import io.airbyte.workers.temporal.sync.PersistStateActivity; -import io.airbyte.workers.temporal.sync.PersistStateActivityImpl; -import io.airbyte.workers.temporal.sync.ReplicationActivity; -import io.airbyte.workers.temporal.sync.ReplicationActivityImpl; -import io.airbyte.workers.temporal.sync.RouteToTaskQueueActivity; -import io.airbyte.workers.temporal.sync.RouteToTaskQueueActivityImpl; -import io.airbyte.workers.temporal.sync.SyncWorkflow; -import io.airbyte.workers.temporal.sync.SyncWorkflowImpl; +import io.airbyte.workers.temporal.TemporalUtils; +import io.airbyte.workers.temporal.support.TemporalProxyHelper; +import io.micronaut.context.BeanRegistration; +import io.micronaut.inject.BeanIdentifier; +import io.temporal.activity.ActivityCancellationType; +import io.temporal.activity.ActivityOptions; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.workflowservice.v1.RequestCancelWorkflowExecutionRequest; import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc.WorkflowServiceBlockingStub; import io.temporal.client.WorkflowClient; import io.temporal.client.WorkflowFailedException; import io.temporal.client.WorkflowOptions; +import io.temporal.common.RetryOptions; import io.temporal.testing.TestWorkflowEnvironment; import io.temporal.worker.Worker; -import java.util.UUID; +import java.time.Duration; +import java.util.List; import org.apache.commons.lang3.tuple.ImmutablePair; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -55,16 +53,14 @@ class SyncWorkflowTest { // TEMPORAL private TestWorkflowEnvironment testEnv; - private Worker syncControlPlaneWorker; - private Worker syncDataPlaneWorker; + private Worker syncWorker; private WorkflowClient client; private ReplicationActivityImpl replicationActivity; private NormalizationActivityImpl normalizationActivity; private DbtTransformationActivityImpl dbtTransformationActivity; private PersistStateActivityImpl persistStateActivity; - private RouteToTaskQueueActivityImpl routeToTaskQueueActivity; - private static final String DATA_PLANE_TASK_QUEUE = "SYNC_DATA_PLANE"; + private static final String SYNC_TASK_QUEUE = "SYNC_TASK_QUEUE"; // AIRBYTE CONFIGURATION private static final long JOB_ID = 11L; @@ -87,18 +83,16 @@ class SyncWorkflowTest { private StandardSyncInput syncInput; private NormalizationInput normalizationInput; private OperatorDbtInput operatorDbtInput; - private StandardSyncOutput replicationSuccessOutput; private NormalizationSummary normalizationSummary; + private ActivityOptions longActivityOptions; + private ActivityOptions shortActivityOptions; + private TemporalProxyHelper temporalProxyHelper; @BeforeEach - public void setUp() { + void setUp() { testEnv = TestWorkflowEnvironment.newInstance(); - syncControlPlaneWorker = testEnv.newWorker(TemporalJobType.SYNC.name()); - syncControlPlaneWorker.registerWorkflowImplementationTypes(SyncWorkflowImpl.class); - - syncDataPlaneWorker = testEnv.newWorker(DATA_PLANE_TASK_QUEUE); - + syncWorker = testEnv.newWorker(SYNC_TASK_QUEUE); client = testEnv.getWorkflowClient(); final ImmutablePair syncPair = TestConfigHelpers.createSyncConfig(); @@ -120,18 +114,52 @@ public void setUp() { normalizationActivity = mock(NormalizationActivityImpl.class); dbtTransformationActivity = mock(DbtTransformationActivityImpl.class); persistStateActivity = mock(PersistStateActivityImpl.class); - routeToTaskQueueActivity = mock(RouteToTaskQueueActivityImpl.class); - doReturn(DATA_PLANE_TASK_QUEUE).when(routeToTaskQueueActivity).routeToTaskQueue(sync.getConnectionId()); + when(normalizationActivity.generateNormalizationInput(any(), any())).thenReturn(normalizationInput); + + longActivityOptions = ActivityOptions.newBuilder() + .setScheduleToCloseTimeout(Duration.ofDays(3)) + .setStartToCloseTimeout(Duration.ofDays(3)) + .setScheduleToStartTimeout(Duration.ofDays(3)) + .setCancellationType(ActivityCancellationType.WAIT_CANCELLATION_COMPLETED) + .setRetryOptions(TemporalUtils.NO_RETRY) + .setHeartbeatTimeout(TemporalUtils.HEARTBEAT_TIMEOUT) + .build(); + shortActivityOptions = ActivityOptions.newBuilder() + .setStartToCloseTimeout(Duration.ofSeconds(120)) + .setRetryOptions(RetryOptions.newBuilder() + .setMaximumAttempts(5) + .setInitialInterval(Duration.ofSeconds(30)) + .setMaximumInterval(Duration.ofSeconds(600)) + .build()) + .build(); + + final BeanIdentifier longActivitiesBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration longActivityOptionsBeanRegistration = mock(BeanRegistration.class); + when(longActivitiesBeanIdentifier.getName()).thenReturn("longRunActivityOptions"); + when(longActivityOptionsBeanRegistration.getIdentifier()).thenReturn(longActivitiesBeanIdentifier); + when(longActivityOptionsBeanRegistration.getBean()).thenReturn(longActivityOptions); + final BeanIdentifier shortActivitiesBeanIdentifier = mock(BeanIdentifier.class); + final BeanRegistration shortActivityOptionsBeanRegistration = mock(BeanRegistration.class); + when(shortActivitiesBeanIdentifier.getName()).thenReturn("shortActivityOptions"); + when(shortActivityOptionsBeanRegistration.getIdentifier()).thenReturn(shortActivitiesBeanIdentifier); + when(shortActivityOptionsBeanRegistration.getBean()).thenReturn(shortActivityOptions); + temporalProxyHelper = new TemporalProxyHelper(List.of(longActivityOptionsBeanRegistration, shortActivityOptionsBeanRegistration)); + + syncWorker.registerWorkflowImplementationTypes(temporalProxyHelper.proxyWorkflowClass(SyncWorkflowImpl.class)); + } + + @AfterEach + public void tearDown() { + testEnv.close(); } // bundle up all the temporal worker setup / execution into one method. private StandardSyncOutput execute() { - syncControlPlaneWorker.registerActivitiesImplementations(routeToTaskQueueActivity); - syncDataPlaneWorker.registerActivitiesImplementations(replicationActivity, normalizationActivity, dbtTransformationActivity, + syncWorker.registerActivitiesImplementations(replicationActivity, normalizationActivity, dbtTransformationActivity, persistStateActivity); testEnv.start(); final SyncWorkflow workflow = - client.newWorkflowStub(SyncWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TemporalJobType.SYNC.name()).build()); + client.newWorkflowStub(SyncWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(SYNC_TASK_QUEUE).build()); return workflow.run(JOB_RUN_CONFIG, SOURCE_LAUNCHER_CONFIG, DESTINATION_LAUNCHER_CONFIG, syncInput, sync.getConnectionId()); } @@ -151,7 +179,6 @@ void testSuccess() { final StandardSyncOutput actualOutput = execute(); - verifyRouteToTaskQueue(routeToTaskQueueActivity, sync.getConnectionId()); verifyReplication(replicationActivity, syncInput); verifyPersistState(persistStateActivity, sync, replicationSuccessOutput, syncInput.getCatalog()); verifyNormalize(normalizationActivity, normalizationInput); @@ -256,11 +283,6 @@ private void cancelWorkflow() { testEnv.getWorkflowService().blockingStub().requestCancelWorkflowExecution(cancelRequest); } - private static void verifyRouteToTaskQueue(final RouteToTaskQueueActivity routeToTaskQueueActivity, - final UUID connectionId) { - verify(routeToTaskQueueActivity).routeToTaskQueue(connectionId); - } - private static void verifyReplication(final ReplicationActivity replicationActivity, final StandardSyncInput syncInput) { verify(replicationActivity).replicate( JOB_RUN_CONFIG, @@ -296,9 +318,4 @@ private static void verifyDbtTransform(final DbtTransformationActivity dbtTransf operatorDbtInput); } - @AfterEach - public void tearDown() { - testEnv.close(); - } - } diff --git a/charts/airbyte-worker/templates/deployment.yaml b/charts/airbyte-worker/templates/deployment.yaml index bc6afeb7cd601..fc3f08a0fbe32 100644 --- a/charts/airbyte-worker/templates/deployment.yaml +++ b/charts/airbyte-worker/templates/deployment.yaml @@ -294,6 +294,22 @@ spec: configMapKeyRef: name: {{ .Release.Name }}-airbyte-env key: USE_STREAM_CAPABLE_STATE + - name: MICRONAUT_ENVIRONMENTS + valueFrom: + configMapKeyRef: + name: {{ .Release.Name }}-airbyte-env + key: WORKERS_MICRONAUT_ENVIRONMENTS + - name: WORKER_LOGS_STORAGE_TYPE + valueFrom: + configMapKeyRef: + name: {{ .Release.Name }}-airbyte-env + key: WORKER_LOGS_STORAGE_TYPE + - name: WORKER_STATE_STORAGE_TYPE + valueFrom: + configMapKeyRef: + name: {{ .Release.Name }}-airbyte-env + key: WORKER_STATE_STORAGE_TYPE + {{- end }} # Values from secret {{- if .Values.secrets }} diff --git a/charts/airbyte/templates/env-configmap.yaml b/charts/airbyte/templates/env-configmap.yaml index 2ba1c67f9cb75..992f8beb4d022 100644 --- a/charts/airbyte/templates/env-configmap.yaml +++ b/charts/airbyte/templates/env-configmap.yaml @@ -63,4 +63,7 @@ data: USE_STREAM_CAPABLE_STATE: "true" CONTAINER_ORCHESTRATOR_ENABLED: {{ .Values.worker.containerOrchestrator.enabled | quote }} CONTAINER_ORCHESTRATOR_IMAGE: {{ .Values.worker.containerOrchestrator.image | quote }} + WORKERS_MICRONAUT_ENVIRONMENTS: "control" + WORKER_LOGS_STORAGE_TYPE: {{ .Values.global.logs.storage.type | quote }} + WORKER_STATE_STORAGE_TYPE: {{ .Values.global.state.storage.type | quote }} {{- end }} diff --git a/charts/airbyte/values.yaml b/charts/airbyte/values.yaml index 740b928721243..0da7b7525f9cb 100644 --- a/charts/airbyte/values.yaml +++ b/charts/airbyte/values.yaml @@ -20,6 +20,10 @@ global: secretValue: "" host: "example.com" port: "5432" + state: + ## state.storage.type Determines which state storage will be utilized. One of "MINIO", "S3" or "GCS" + storage: + type: "MINIO" logs: ## logs.accessKey.password Logs Access Key ## logs.accessKey.existingSecret @@ -36,6 +40,11 @@ global: existingSecret: "" existingSecretKey: "" + ## logs.storage.type Determines which log storage will be utilized. One of "MINIO", "S3" or "GCS" + ## Used in conjunction with logs.minio.*, logs.s3.* or logs.gcs.* + storage: + type: "MINIO" + ## logs.minio.enabled Switch to enable or disable the Minio helm chart minio: enabled: true diff --git a/deps.toml b/deps.toml index 9c4bc88b157f4..da8e4a490cd1c 100644 --- a/deps.toml +++ b/deps.toml @@ -1,15 +1,15 @@ [versions] -fasterxml_version = "2.13.0" +fasterxml_version = "2.13.3" flyway = "7.14.0" glassfish_version = "2.31" hikaricp = "5.0.1" commons_io = "2.7" -log4j = "2.17.1" -slf4j = "1.7.30" -lombok = "1.18.22" +log4j = "2.17.2" +slf4j = "1.7.36" +lombok = "1.18.24" jooq = "3.13.4" -junit-jupiter = "5.8.2" -micronaut = "3.6.0" +junit-jupiter = "5.9.0" +micronaut = "3.6.3" micronaut-test = "3.5.0" postgresql = "42.3.5" connectors-testcontainers = "1.15.3" @@ -21,7 +21,7 @@ connectors-testcontainers-tidb = "1.16.3" connectors-destination-testcontainers-clickhouse = "1.17.3" connectors-destination-testcontainers-oracle-xe = "1.17.3" connectors-source-testcontainers-clickhouse = "1.17.3" -platform-testcontainers = "1.17.1" +platform-testcontainers = "1.17.3" [libraries] fasterxml = { module = "com.fasterxml.jackson:jackson-bom", version.ref = "fasterxml_version" } @@ -34,7 +34,7 @@ guava = { module = "com.google.guava:guava", version = "30.1.1-jre" } commons-io = { module = "commons-io:commons-io", version.ref = "commons_io" } apache-commons = { module = "org.apache.commons:commons-compress", version = "1.20" } apache-commons-lang = { module = "org.apache.commons:commons-lang3", version = "3.11" } -slf4j-api = { module = "org.slf4j:slf4j-api", version = "1.7.30" } +slf4j-api = { module = "org.slf4j:slf4j-api", version.ref = "slf4j" } log4j-api = { module = "org.apache.logging.log4j:log4j-api", version.ref = "log4j" } log4j-core = { module = "org.apache.logging.log4j:log4j-core", version.ref = "log4j" } log4j-impl = { module = "org.apache.logging.log4j:log4j-slf4j-impl", version.ref = "log4j" } @@ -89,7 +89,7 @@ otel-bom = {module = "io.opentelemetry:opentelemetry-bom", version = "1.14.0"} otel-semconv = {module = "io.opentelemetry:opentelemetry-semconv", version = "1.14.0-alpha"} otel-sdk = {module = "io.opentelemetry:opentelemetry-sdk-metrics", version = "1.14.0"} otel-sdk-testing = {module = "io.opentelemetry:opentelemetry-sdk-metrics-testing", version = "1.13.0-alpha"} -micrometer-statsd = {module = "io.micrometer:micrometer-registry-statsd", version = "1.9.2"} +micrometer-statsd = {module = "io.micrometer:micrometer-registry-statsd", version = "1.9.3"} quartz-scheduler = {module="org.quartz-scheduler:quartz", version = "2.3.2"} @@ -100,18 +100,18 @@ javax-inject = { module = "javax.inject:javax.inject", version = "1" } javax-transaction = { module = "javax.transaction:javax.transaction-api", version = "1.3" } micronaut-bom = { module = "io.micronaut:micronaut-bom", version.ref = "micronaut" } micronaut-data-processor = { module = "io.micronaut.data:micronaut-data-processor", version = "3.7.2" } -micronaut-flyway = { module = "io.micronaut.flyway:micronaut-flyway", version = "5.4.0" } +micronaut-flyway = { module = "io.micronaut.flyway:micronaut-flyway", version = "5.4.1" } micronaut-inject = { module = "io.micronaut:micronaut-inject" } micronaut-http-client = { module = "io.micronaut:micronaut-http-client" } micronaut-http-server-netty = { module = "io.micronaut:micronaut-http-server-netty" } -micronaut-inject-java = { module = "io.micronaut:micronaut-inject-java" } +micronaut-inject-java = { module = "io.micronaut:micronaut-inject-java", version.ref = "micronaut" } micronaut-jaxrs-processor = { module = "io.micronaut.jaxrs:micronaut-jaxrs-processor", version = "3.4.0" } micronaut-jaxrs-server = { module = "io.micronaut.jaxrs:micronaut-jaxrs-server", version = "3.4.0" } micronaut-jdbc-hikari = { module = "io.micronaut.sql:micronaut-jdbc-hikari" } micronaut-jooq = { module = "io.micronaut.sql:micronaut-jooq" } micronaut-management = { module = "io.micronaut:micronaut-management" } micronaut-runtime = { module = "io.micronaut:micronaut-runtime" } -micronaut-security = { module = "io.micronaut.security:micronaut-security", version = "3.6.3" } +micronaut-security = { module = "io.micronaut.security:micronaut-security", version = "3.7.0" } micronaut-test-core = { module = "io.micronaut.test:micronaut-test-core", version.ref = "micronaut-test" } micronaut-test-junit5 = { module = "io.micronaut.test:micronaut-test-junit5", version.ref = "micronaut-test" } micronaut-validation = { module = "io.micronaut:micronaut-validation" } @@ -123,6 +123,7 @@ log4j = ["log4j-api", "log4j-core", "log4j-impl", "log4j-web"] slf4j = ["jul-to-slf4j", "jcl-over-slf4j", "log4j-over-slf4j"] junit = ["junit-jupiter-api", "junit-jupiter-params", "mockito-junit-jupiter"] micronaut = ["javax-inject", "javax-transaction", "micronaut-http-server-netty", "micronaut-http-client", "micronaut-inject", "micronaut-validation", "micronaut-runtime", "micronaut-management", "micronaut-security", "micronaut-jaxrs-server", "micronaut-flyway", "micronaut-jdbc-hikari", "micronaut-jooq"] +micronaut-annotation = ["javax-inject", "micronaut-inject-java"] micronaut-annotation-processor = ["micronaut-inject-java", "micronaut-management", "micronaut-validation", "micronaut-data-processor", "micronaut-jaxrs-processor"] micronaut-test = ["micronaut-test-core", "micronaut-test-junit5", "h2-database"] micronaut-test-annotation-processor = ["micronaut-inject-java"] diff --git a/docker-compose.yaml b/docker-compose.yaml index d57f7be175404..2e88f4a637a8a 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -101,10 +101,13 @@ services: - ACTIVITY_MAX_DELAY_BETWEEN_ATTEMPTS_SECONDS=${ACTIVITY_MAX_DELAY_BETWEEN_ATTEMPTS_SECONDS} - WORKFLOW_FAILURE_RESTART_DELAY_SECONDS=${WORKFLOW_FAILURE_RESTART_DELAY_SECONDS} - USE_STREAM_CAPABLE_STATE=${USE_STREAM_CAPABLE_STATE} + - MICRONAUT_ENVIRONMENTS=${WORKERS_MICRONAUT_ENVIRONMENTS} volumes: - /var/run/docker.sock:/var/run/docker.sock - workspace:${WORKSPACE_ROOT} - ${LOCAL_ROOT}:${LOCAL_ROOT} + ports: + - 9000:9000 server: image: airbyte/server:${VERSION} logging: *default-logging