From a205d4633365417ef7e10ca861bf5ab00cc1dbfe Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Sat, 25 Sep 2021 13:27:17 -0700 Subject: [PATCH 01/23] Add migration to create latest state table --- .../V0_29_21_001__Store_last_sync_state.java | 41 +++++++++++++++++++ .../configs_database/schema_dump.txt | 7 ++++ 2 files changed, 48 insertions(+) create mode 100644 airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java new file mode 100644 index 000000000000..36f46edc79d1 --- /dev/null +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java @@ -0,0 +1,41 @@ +package io.airbyte.db.instance.configs.migrations; + +import org.flywaydb.core.api.migration.BaseJavaMigration; +import org.flywaydb.core.api.migration.Context; +import org.jooq.DSLContext; +import org.jooq.impl.DSL; +import org.jooq.impl.SQLDataType; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Create a new table to store the latest job state for each standard sync. + * Ideally the latest state should be copied to the new table in the migration. + * However, + */ +public class V0_29_21_001__Store_last_sync_state extends BaseJavaMigration { + + private static final Logger LOGGER = LoggerFactory.getLogger(V0_29_21_001__Store_last_sync_state.class); + + private static final String TABLE_NAME = "latest_sync_state"; + private static final String COLUMN_SYNC_ID = "sync_id"; + private static final String COLUMN_STATE = "state"; + private static final String COLUMN_CREATED_AT = "created_at"; + private static final String COLUMN_UPDATED_AT = "updated_at"; + + @Override + public void migrate(Context context) throws Exception { + DSLContext ctx = DSL.using(context.getConnection()); + ctx.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";"); + ctx.createTable(TABLE_NAME) + .column(COLUMN_SYNC_ID, SQLDataType.UUID.nullable(false)) + .column(COLUMN_STATE, SQLDataType.JSONB.nullable(false)) + .column(COLUMN_CREATED_AT, SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())) + .column(COLUMN_UPDATED_AT, SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())) + .execute(); + ctx.createUniqueIndexIfNotExists(String.format("%s_sync_id_idx", TABLE_NAME)) + .on(TABLE_NAME, COLUMN_SYNC_ID) + .execute(); + } + +} diff --git a/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt b/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt index 6c9e11d400e4..65caffd15822 100644 --- a/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt +++ b/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt @@ -22,6 +22,12 @@ create table "public"."airbyte_configs_migrations"( constraint "airbyte_configs_migrations_pk" primary key ("installed_rank") ); +create table "public"."latest_sync_state"( + "sync_id" uuid not null, + "state" jsonb not null, + "created_at" timestamptz(35) not null default null, + "updated_at" timestamptz(35) not null default null +); create index "airbyte_configs_id_idx" on "public"."airbyte_configs"("config_id" asc); create unique index "airbyte_configs_pkey" on "public"."airbyte_configs"("id" asc); create unique index "airbyte_configs_type_id_idx" on "public"."airbyte_configs"( @@ -30,3 +36,4 @@ create unique index "airbyte_configs_type_id_idx" on "public"."airbyte_configs"( ); create unique index "airbyte_configs_migrations_pk" on "public"."airbyte_configs_migrations"("installed_rank" asc); create index "airbyte_configs_migrations_s_idx" on "public"."airbyte_configs_migrations"("success" asc); +create unique index "latest_sync_state_sync_id_idx" on "public"."latest_sync_state"("sync_id" asc); From 31704d37c27172cda0f6285446df9a2ac2bc645c Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Sat, 25 Sep 2021 13:53:17 -0700 Subject: [PATCH 02/23] Log migration name --- ...9_15_001__Add_temporalWorkflowId_col_to_Attempts.java | 9 ++++++--- airbyte-db/lib/src/main/resources/migration_template.txt | 4 ++++ 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/migrations/V0_29_15_001__Add_temporalWorkflowId_col_to_Attempts.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/migrations/V0_29_15_001__Add_temporalWorkflowId_col_to_Attempts.java index be9cc2b2cb4d..d188ba6fe662 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/migrations/V0_29_15_001__Add_temporalWorkflowId_col_to_Attempts.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/migrations/V0_29_15_001__Add_temporalWorkflowId_col_to_Attempts.java @@ -9,14 +9,17 @@ import org.jooq.DSLContext; import org.jooq.impl.DSL; import org.jooq.impl.SQLDataType; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; public class V0_29_15_001__Add_temporalWorkflowId_col_to_Attempts extends BaseJavaMigration { + private static final Logger LOGGER = LoggerFactory.getLogger(V0_29_15_001__Add_temporalWorkflowId_col_to_Attempts.class); + @Override public void migrate(final Context context) throws Exception { - // Warning: please do not use any jOOQ generated code to write a migration. - // As database schema changes, the generated jOOQ code can be deprecated. So - // old migration may not compile if there is any generated code. + LOGGER.info("Running migration: {}", this.getClass().getSimpleName()); + final DSLContext ctx = DSL.using(context.getConnection()); ctx.alterTable("attempts") .addColumnIfNotExists(DSL.field("temporal_workflow_id", SQLDataType.VARCHAR(256).nullable(true))) diff --git a/airbyte-db/lib/src/main/resources/migration_template.txt b/airbyte-db/lib/src/main/resources/migration_template.txt index 4398aba69439..6206ab5f14a9 100644 --- a/airbyte-db/lib/src/main/resources/migration_template.txt +++ b/airbyte-db/lib/src/main/resources/migration_template.txt @@ -4,12 +4,16 @@ import org.flywaydb.core.api.migration.BaseJavaMigration; import org.flywaydb.core.api.migration.Context; import org.jooq.DSLContext; import org.jooq.impl.DSL; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; // TODO: update migration description in the class name public class V__ extends BaseJavaMigration { @Override public void migrate(Context context) throws Exception { + LOGGER.info("Running migration: {}", this.getClass().getSimpleName()); + // Warning: please do not use any jOOQ generated code to write a migration. // As database schema changes, the generated jOOQ code can be deprecated. So // old migration may not compile if there is any generated code. From ff0f13387a15496def82e9f23dc6767a64f36db3 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 27 Sep 2021 08:55:25 -0700 Subject: [PATCH 03/23] Expose db variables to airbyte-db --- docker-compose.yaml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docker-compose.yaml b/docker-compose.yaml index 60e67d5fe392..0b6c19d75962 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -26,6 +26,12 @@ services: environment: - POSTGRES_USER=${DATABASE_USER} - POSTGRES_PASSWORD=${DATABASE_PASSWORD} + - DATABASE_USER=${DATABASE_USER} + - DATABASE_PASSWORD=${DATABASE_PASSWORD} + - DATABASE_URL=${DATABASE_URL} + - CONFIG_DATABASE_USER=${CONFIG_DATABASE_USER:-} + - CONFIG_DATABASE_PASSWORD=${CONFIG_DATABASE_PASSWORD:-} + - CONFIG_DATABASE_URL=${CONFIG_DATABASE_URL:-} volumes: - db:/var/lib/postgresql/data scheduler: From 025b2ed6dcc9b98a816a8b23106a28427ad14239 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 27 Sep 2021 08:55:42 -0700 Subject: [PATCH 04/23] Implement migration --- airbyte-db/lib/build.gradle | 1 + .../V0_29_21_001__Store_last_sync_state.java | 164 ++++++++-- .../development/MigrationDevHelper.java | 4 +- .../configs_database/schema_dump.txt | 4 +- .../src/main/resources/migration_template.txt | 6 +- ...29_21_001__Store_last_sync_state_test.java | 300 ++++++++++++++++++ .../persistence/DefaultJobPersistence.java | 3 + .../scheduler/persistence/JobPersistence.java | 3 +- 8 files changed, 460 insertions(+), 25 deletions(-) create mode 100644 airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java diff --git a/airbyte-db/lib/build.gradle b/airbyte-db/lib/build.gradle index fe8cd2c8c90e..d149d0e8f90a 100644 --- a/airbyte-db/lib/build.gradle +++ b/airbyte-db/lib/build.gradle @@ -10,6 +10,7 @@ dependencies { implementation project(':airbyte-protocol:models') implementation project(':airbyte-json-validation') + implementation project(':airbyte-config:models') implementation "org.flywaydb:flyway-core:7.14.0" implementation "org.testcontainers:postgresql:1.15.3" diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java index 36f46edc79d1..782d5d7f6848 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java @@ -1,41 +1,169 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + package io.airbyte.db.instance.configs.migrations; +import com.fasterxml.jackson.databind.JsonNode; +import com.google.common.annotations.VisibleForTesting; +import io.airbyte.commons.json.Jsons; +import io.airbyte.config.Configs; +import io.airbyte.config.EnvConfigs; +import io.airbyte.db.Database; +import io.airbyte.db.instance.jobs.JobsDatabaseInstance; +import java.io.IOException; +import java.sql.SQLException; +import java.time.OffsetDateTime; +import java.util.Collections; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import java.util.stream.Collectors; import org.flywaydb.core.api.migration.BaseJavaMigration; import org.flywaydb.core.api.migration.Context; import org.jooq.DSLContext; +import org.jooq.Field; +import org.jooq.JSONB; +import org.jooq.Record2; +import org.jooq.Table; import org.jooq.impl.DSL; import org.jooq.impl.SQLDataType; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** - * Create a new table to store the latest job state for each standard sync. - * Ideally the latest state should be copied to the new table in the migration. - * However, + * Create a new table to store the latest job state for each standard sync. Issue: */ public class V0_29_21_001__Store_last_sync_state extends BaseJavaMigration { + private static final String MIGRATION_NAME = "Configs db migration 0.29.21.001"; private static final Logger LOGGER = LoggerFactory.getLogger(V0_29_21_001__Store_last_sync_state.class); - private static final String TABLE_NAME = "latest_sync_state"; - private static final String COLUMN_SYNC_ID = "sync_id"; - private static final String COLUMN_STATE = "state"; - private static final String COLUMN_CREATED_AT = "created_at"; - private static final String COLUMN_UPDATED_AT = "updated_at"; + // sync state table + static final Table SYNC_STATE_TABLE = DSL.table("sync_state"); + static final Field COLUMN_SYNC_ID = DSL.field("sync_id", SQLDataType.UUID.nullable(false)); + static final Field COLUMN_STATE = DSL.field("state", SQLDataType.JSONB.nullable(false)); + static final Field COLUMN_CREATED_AT = DSL.field("created_at", + SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())); + static final Field COLUMN_UPDATED_AT = DSL.field("updated_at", + SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())); + + private final Configs configs; + + public V0_29_21_001__Store_last_sync_state() { + this.configs = new EnvConfigs(); + } + + @VisibleForTesting + V0_29_21_001__Store_last_sync_state(Configs configs) { + this.configs = configs; + } @Override - public void migrate(Context context) throws Exception { - DSLContext ctx = DSL.using(context.getConnection()); - ctx.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";"); - ctx.createTable(TABLE_NAME) - .column(COLUMN_SYNC_ID, SQLDataType.UUID.nullable(false)) - .column(COLUMN_STATE, SQLDataType.JSONB.nullable(false)) - .column(COLUMN_CREATED_AT, SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())) - .column(COLUMN_UPDATED_AT, SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())) + public void migrate(final Context context) throws Exception { + LOGGER.info("Running migration: {}", this.getClass().getSimpleName()); + final DSLContext ctx = DSL.using(context.getConnection()); + + createTable(ctx); + + final Optional jobsDatabase = getJobsDatabase(configs); + if (jobsDatabase.isPresent()) { + copyData(ctx, getSyncToStateMap(jobsDatabase.get()), OffsetDateTime.now()); + } + } + + @VisibleForTesting + static void createTable(final DSLContext ctx) { + ctx.createTableIfNotExists(SYNC_STATE_TABLE) + .column(COLUMN_SYNC_ID) + .column(COLUMN_STATE) + .column(COLUMN_CREATED_AT) + .column(COLUMN_UPDATED_AT) .execute(); - ctx.createUniqueIndexIfNotExists(String.format("%s_sync_id_idx", TABLE_NAME)) - .on(TABLE_NAME, COLUMN_SYNC_ID) + ctx.createUniqueIndexIfNotExists(String.format("%s_sync_id_idx", SYNC_STATE_TABLE)) + .on(SYNC_STATE_TABLE, Collections.singleton(COLUMN_SYNC_ID)) .execute(); } + @VisibleForTesting + static void copyData(final DSLContext ctx, final Map syncToStateMap, final OffsetDateTime timestamp) { + for (Map.Entry entry : syncToStateMap.entrySet()) { + ctx.insertInto(SYNC_STATE_TABLE) + .set(COLUMN_SYNC_ID, UUID.fromString(entry.getKey())) + .set(COLUMN_STATE, JSONB.valueOf(Jsons.serialize(entry.getValue()))) + .set(COLUMN_CREATED_AT, timestamp) + .set(COLUMN_UPDATED_AT, timestamp) + // This migration is idempotent. If the record for a sync_id already exists, + // it means that the migration has already been run before. Abort insertion. + .onDuplicateKeyIgnore() + .execute(); + } + } + + /** + * This migration requires a connection to the job database, which may be a separate database from + * the config database. However, the job database only exists in production, not in development or + * test. We use the job database environment variables to determine how to connect to the job + * database. This approach is not 100% reliable. However, it is better than doing half of the + * migration here (creating the table), and the rest of the work during server start up (copying the + * data from the job database). + */ + @VisibleForTesting + static Optional getJobsDatabase(final Configs configs) { + try { + // If the environment variables exist, it means the migration is run in production. + // Connect to the official job database. + final Database jobsDatabase = new JobsDatabaseInstance( + configs.getDatabaseUser(), + configs.getDatabasePassword(), + configs.getDatabaseUrl()) + .getInitialized(); + LOGGER.info("[{}] Connected to jobs database: {}", MIGRATION_NAME, configs.getDatabaseUrl()); + return Optional.of(jobsDatabase); + } catch (final IllegalArgumentException e) { + // If the environment variables do not exist, it means the migration is run in development. + // Connect to a mock job database, because we don't need to copy any data in test. + LOGGER.info("[{}] This is the dev environment; there is no jobs database", MIGRATION_NAME); + return Optional.empty(); + } catch (final IOException e) { + throw new RuntimeException("Cannot connect to jobs database", e); + } + } + + /** + * @return a map from sync id to last job attempt state. + */ + @VisibleForTesting + static Map getSyncToStateMap(final Database jobsDatabase) throws SQLException { + final Table jobsTable = DSL.table("jobs"); + final Field jobIdField = DSL.field("jobs.id", SQLDataType.BIGINT); + final Field syncIdField = DSL.field("jobs.scope", SQLDataType.VARCHAR); + + final Table attemptsTable = DSL.table("attempts"); + final Field attemptJobIdField = DSL.field("attempts.job_id", SQLDataType.BIGINT); + final Field attemptNumberField = DSL.field("attempts.attempt_number", SQLDataType.INTEGER); + + // output schema: JobOutput.yaml + // sync schema: StandardSyncOutput.yaml + // state schema: State.yaml + final Field attemptStateField = DSL.field("attempts.output -> 'sync' -> 'state'", SQLDataType.JSONB); + + return jobsDatabase.query(ctx -> ctx + .select(syncIdField, attemptStateField) + .from(attemptsTable) + .innerJoin(jobsTable) + .on(jobIdField.eq(attemptJobIdField)) + .where(DSL.row(attemptJobIdField, attemptNumberField).in( + // for each job id, find the last attempt with a state + DSL.select(attemptJobIdField, DSL.max(attemptNumberField)) + .from(attemptsTable) + .where(attemptStateField.isNotNull()) + .groupBy(attemptJobIdField))) + .fetch() + .stream() + .collect(Collectors.toMap( + Record2::value1, + r -> Jsons.deserialize(r.value2().data())))); + } + } diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/development/MigrationDevHelper.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/development/MigrationDevHelper.java index 3b75359695be..076d446d8729 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/development/MigrationDevHelper.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/development/MigrationDevHelper.java @@ -65,8 +65,8 @@ public static void createNextMigrationFile(final String dbIdentifier, final Flyw final String template = MoreResources.readResource("migration_template.txt"); final String newMigration = template.replace("", dbIdentifier) - .replace("", versionId) - .replace("", description) + .replaceAll("", versionId) + .replaceAll("", description) .strip(); final String fileName = String.format("V%s__%s.java", versionId, description); diff --git a/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt b/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt index 65caffd15822..8bc8528d0993 100644 --- a/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt +++ b/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt @@ -22,7 +22,7 @@ create table "public"."airbyte_configs_migrations"( constraint "airbyte_configs_migrations_pk" primary key ("installed_rank") ); -create table "public"."latest_sync_state"( +create table "public"."sync_state"( "sync_id" uuid not null, "state" jsonb not null, "created_at" timestamptz(35) not null default null, @@ -36,4 +36,4 @@ create unique index "airbyte_configs_type_id_idx" on "public"."airbyte_configs"( ); create unique index "airbyte_configs_migrations_pk" on "public"."airbyte_configs_migrations"("installed_rank" asc); create index "airbyte_configs_migrations_s_idx" on "public"."airbyte_configs_migrations"("success" asc); -create unique index "latest_sync_state_sync_id_idx" on "public"."latest_sync_state"("sync_id" asc); +create unique index "sync_state_sync_id_idx" on "public"."sync_state"("sync_id" asc); diff --git a/airbyte-db/lib/src/main/resources/migration_template.txt b/airbyte-db/lib/src/main/resources/migration_template.txt index 6206ab5f14a9..074c2d97b9be 100644 --- a/airbyte-db/lib/src/main/resources/migration_template.txt +++ b/airbyte-db/lib/src/main/resources/migration_template.txt @@ -10,14 +10,16 @@ import org.slf4j.LoggerFactory; // TODO: update migration description in the class name public class V__ extends BaseJavaMigration { + private static final Logger LOGGER = LoggerFactory.getLogger(V__.class); + @Override - public void migrate(Context context) throws Exception { + public void migrate(final Context context) throws Exception { LOGGER.info("Running migration: {}", this.getClass().getSimpleName()); // Warning: please do not use any jOOQ generated code to write a migration. // As database schema changes, the generated jOOQ code can be deprecated. So // old migration may not compile if there is any generated code. - DSLContext ctx = DSL.using(context.getConnection()); + final DSLContext ctx = DSL.using(context.getConnection()); } } diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java new file mode 100644 index 000000000000..9e60f7fa5a20 --- /dev/null +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java @@ -0,0 +1,300 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.db.instance.configs.migrations; + +import static org.jooq.impl.DSL.field; +import static org.jooq.impl.DSL.table; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.airbyte.commons.jackson.MoreMappers; +import io.airbyte.commons.json.Jsons; +import io.airbyte.config.Configs; +import io.airbyte.config.EnvConfigs; +import io.airbyte.db.Database; +import io.airbyte.db.instance.configs.AbstractConfigsDatabaseTest; +import io.airbyte.db.instance.jobs.JobsDatabaseInstance; +import java.sql.Connection; +import java.sql.SQLException; +import java.time.OffsetDateTime; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.flywaydb.core.api.configuration.Configuration; +import org.flywaydb.core.api.migration.Context; +import org.jooq.DSLContext; +import org.jooq.Field; +import org.jooq.InsertSetMoreStep; +import org.jooq.JSONB; +import org.jooq.Table; +import org.jooq.Typed; +import org.jooq.impl.SQLDataType; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.MethodOrderer; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestMethodOrder; + +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) +class V0_29_21_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTest { + + private static final ObjectMapper OBJECT_MAPPER = MoreMappers.initMapper(); + + private static final Table JOBS_TABLE = table("jobs"); + private static final Field JOB_ID_FIELD = field("id", SQLDataType.BIGINT); + private static final Field JOB_SCOPE_FIELD = field("scope", SQLDataType.VARCHAR); + + private static final Table ATTEMPTS_TABLE = table("attempts"); + private static final Field ATTEMPT_ID_FIELD = field("id", SQLDataType.BIGINT); + private static final Field ATTEMPT_JOB_ID_FIELD = field("job_id", SQLDataType.BIGINT); + private static final Field ATTEMPT_NUMBER_FIELD = field("attempt_number", SQLDataType.INTEGER); + private static final Field ATTEMPT_OUTPUT_FIELD = field("output", SQLDataType.JSONB); + + private static final UUID SYNC_1_ID = UUID.randomUUID(); + private static final UUID SYNC_2_ID = UUID.randomUUID(); + private static final UUID SYNC_3_ID = UUID.randomUUID(); + private static final JsonNode SYNC_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }"); + private static final JsonNode SYNC_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }"); + + private static Database jobDatabase; + + @BeforeAll + public static void setupJobDatabase() throws Exception { + jobDatabase = new JobsDatabaseInstance( + container.getUsername(), + container.getPassword(), + container.getJdbcUrl()) + .getAndInitialize(); + } + + @Test + @Order(10) + public void testGetJobsDatabase() { + // when there is no database environment variable, the return value is empty + assertTrue(V0_29_21_001__Store_last_sync_state.getJobsDatabase(new EnvConfigs()).isEmpty()); + + // when there is database environment variable, return the database + final Configs configs = mock(Configs.class); + when(configs.getDatabaseUser()).thenReturn(container.getUsername()); + when(configs.getDatabasePassword()).thenReturn(container.getPassword()); + when(configs.getDatabaseUrl()).thenReturn(container.getJdbcUrl()); + + assertTrue(V0_29_21_001__Store_last_sync_state.getJobsDatabase(configs).isPresent()); + } + + @Test + @Order(20) + public void testGetSyncToStateMap() throws Exception { + jobDatabase.query(ctx -> { + // Create three jobs for three standard syncs. + // The first job has no attempt. + createJob(ctx, SYNC_1_ID); + + // The second job has one attempt. + long job2 = createJob(ctx, SYNC_2_ID); + createAttempt(ctx, job2, 1, createAttemptOutput(SYNC_2_STATE)); + + // The third job has multiple attempts. The third attempt has the latest state. + long job3 = createJob(ctx, SYNC_3_ID); + JsonNode attempt31State = Jsons.deserialize("{ \"state\": { \"cursor\": 31 } }"); + createAttempt(ctx, job3, 1, createAttemptOutput(attempt31State)); + createAttempt(ctx, job3, 2, null); + createAttempt(ctx, job3, 3, createAttemptOutput(SYNC_3_STATE)); + createAttempt(ctx, job3, 4, null); + createAttempt(ctx, job3, 5, null); + + final Map syncToStateMap = V0_29_21_001__Store_last_sync_state.getSyncToStateMap(jobDatabase); + assertEquals(2, syncToStateMap.size()); + assertEquals(SYNC_2_STATE, syncToStateMap.get(SYNC_2_ID.toString())); + assertEquals(SYNC_3_STATE, syncToStateMap.get(SYNC_3_ID.toString())); + + return null; + }); + } + + @Test + @Order(30) + public void testCreateTable() throws SQLException { + database.query(ctx -> { + checkTable(ctx, false); + V0_29_21_001__Store_last_sync_state.createTable(ctx); + checkTable(ctx, true); + return null; + }); + } + + /** + * Test the unique index on the sync_id column. + */ + @Test + @Order(40) + public void testUniqueSyncIdIndex() throws SQLException { + final UUID syncId = UUID.randomUUID(); + + database.query(ctx -> { + final InsertSetMoreStep insertRecord = ctx.insertInto(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE) + .set(V0_29_21_001__Store_last_sync_state.COLUMN_SYNC_ID, syncId) + .set(V0_29_21_001__Store_last_sync_state.COLUMN_STATE, JSONB.valueOf("{}")); + + assertDoesNotThrow(insertRecord::execute); + // insert the record with the same sync_id will violate the unique index + assertThrows(org.jooq.exception.DataAccessException.class, insertRecord::execute); + return null; + }); + } + + @Test + @Order(50) + public void testCopyData() throws SQLException { + final JsonNode sync2NewState = Jsons.deserialize("{ \"state\": { \"cursor\": 3 } }"); + + final Map syncStateMap1 = Map.of(SYNC_1_ID.toString(), SYNC_2_STATE, SYNC_2_ID.toString(), SYNC_2_STATE); + final Map syncStateMap2 = Map.of(SYNC_2_ID.toString(), sync2NewState); + + final OffsetDateTime timestamp = OffsetDateTime.now(); + + database.query(ctx -> { + V0_29_21_001__Store_last_sync_state.copyData(ctx, syncStateMap1, timestamp); + checkSyncStates(ctx, syncStateMap1, timestamp); + + // call the copyData method again with different data will not affect existing records + V0_29_21_001__Store_last_sync_state.copyData(ctx, syncStateMap2, OffsetDateTime.now()); + // the states remain the same as those in syncStateMap1 + checkSyncStates(ctx, syncStateMap1, timestamp); + + return null; + }); + } + + /** + * Clear the table and test the migration end-to-end. + */ + @Test + @Order(60) + public void testMigration() throws Exception { + database.query(ctx -> ctx.dropTableIfExists(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE).execute()); + + final Configs configs = mock(Configs.class); + when(configs.getDatabaseUser()).thenReturn(container.getUsername()); + when(configs.getDatabasePassword()).thenReturn(container.getPassword()); + when(configs.getDatabaseUrl()).thenReturn(container.getJdbcUrl()); + + final var migration = new V0_29_21_001__Store_last_sync_state(configs); + // this context is a flyway class; only the getConnection method is needed to run the migration + final Context context = new Context() { + + @Override + public Configuration getConfiguration() { + return null; + } + + @Override + public Connection getConnection() { + try { + return database.getDataSource().getConnection(); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + + }; + migration.migrate(context); + database.query(ctx -> { + checkSyncStates(ctx, Map.of(SYNC_2_ID.toString(), SYNC_2_STATE, SYNC_3_ID.toString(), SYNC_3_STATE), null); + return null; + }); + } + + /** + * Create a job record whose scope equals to the passed in standard sync id, and return the job id. + */ + private static long createJob(DSLContext ctx, UUID standardSyncId) { + final int insertCount = ctx.insertInto(JOBS_TABLE, JOB_SCOPE_FIELD) + .values(standardSyncId.toString()) + .execute(); + assertEquals(1, insertCount); + + return ctx.select(JOB_ID_FIELD) + .from(JOBS_TABLE) + .where(JOB_SCOPE_FIELD.eq(standardSyncId.toString())) + .fetchOne() + .get(JOB_ID_FIELD); + } + + private static void createAttempt(DSLContext ctx, long jobId, int attemptNumber, JsonNode attemptOutput) { + final int insertCount = ctx.insertInto(ATTEMPTS_TABLE, ATTEMPT_JOB_ID_FIELD, ATTEMPT_NUMBER_FIELD, ATTEMPT_OUTPUT_FIELD) + .values(jobId, attemptNumber, JSONB.valueOf(Jsons.serialize(attemptOutput))) + .execute(); + assertEquals(1, insertCount); + + ctx.select(ATTEMPT_ID_FIELD) + .from(ATTEMPTS_TABLE) + .where(ATTEMPT_JOB_ID_FIELD.eq(jobId), ATTEMPT_NUMBER_FIELD.eq(attemptNumber)) + .fetchOne() + .get(ATTEMPT_ID_FIELD); + } + + /** + * Create an JobOutput object whose output type is StandardSyncOutput. + * + * @param state The state object within a StandardSyncOutput. + */ + private static JsonNode createAttemptOutput(JsonNode state) { + final ObjectNode standardSyncOutput = OBJECT_MAPPER.createObjectNode() + .set("state", state); + return OBJECT_MAPPER.createObjectNode() + .put("output_type", "sync") + .set("sync", standardSyncOutput); + } + + private static void checkTable(final DSLContext ctx, final boolean tableExists) { + final List> tables = ctx.meta().getTables(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE.getName()); + assertEquals(tableExists, tables.size() > 0); + + if (!tableExists) { + assertTrue(tables.isEmpty()); + } else { + System.out.println(Arrays.stream(tables.get(0).fields()).map(Typed::getDataType).collect(Collectors.toSet())); + final Set actualFields = Arrays.stream(tables.get(0).fields()).map(Field::getName).collect(Collectors.toSet()); + final Set expectedFields = Set.of( + V0_29_21_001__Store_last_sync_state.COLUMN_SYNC_ID.getName(), + V0_29_21_001__Store_last_sync_state.COLUMN_STATE.getName(), + V0_29_21_001__Store_last_sync_state.COLUMN_CREATED_AT.getName(), + V0_29_21_001__Store_last_sync_state.COLUMN_UPDATED_AT.getName()); + assertEquals(expectedFields, actualFields); + } + } + + private static void checkSyncStates(final DSLContext ctx, + final Map expectedSyncStates, + @Nullable OffsetDateTime expectedTimestamp) { + for (final Map.Entry entry : expectedSyncStates.entrySet()) { + final var record = ctx + .select(V0_29_21_001__Store_last_sync_state.COLUMN_STATE, + V0_29_21_001__Store_last_sync_state.COLUMN_CREATED_AT, + V0_29_21_001__Store_last_sync_state.COLUMN_UPDATED_AT) + .from(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE) + .where(V0_29_21_001__Store_last_sync_state.COLUMN_SYNC_ID.eq(UUID.fromString(entry.getKey()))) + .fetchOne(); + assertEquals(entry.getValue(), Jsons.deserialize(record.value1().data())); + if (expectedTimestamp != null) { + assertEquals(expectedTimestamp, record.value2()); + assertEquals(expectedTimestamp, record.value3()); + } + } + } + +} diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java index 13c9379b9768..8050a75bbc1a 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java @@ -125,6 +125,9 @@ public DefaultJobPersistence(final Database database) { this(database, Instant::now, 30, 500, 10); } + /** + * @param scope This is the primary id of a standard sync (StandardSync#connectionId). + */ @Override public Optional enqueueJob(final String scope, final JobConfig jobConfig) throws IOException { LOGGER.info("enqueuing pending job for scope: {}", scope); diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java index a059e9dc47d5..825936ef4b41 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java @@ -36,7 +36,8 @@ public interface JobPersistence { /** * Enqueue a new job. Its initial status will be pending. * - * @param scope key that will be used to determine if two jobs should not be run at the same time. + * @param scope key that will be used to determine if two jobs should not be run at the same time; + * it is the primary id of the standard sync (StandardSync#connectionId) * @param jobConfig configuration for the job * @return job id * @throws IOException exception due to interaction with persistence From ddd05a6dac6c72cff6f1839e2d78f471b936a6eb Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Fri, 15 Oct 2021 12:50:51 -0700 Subject: [PATCH 05/23] Fix migration test --- .../config/persistence/ConfigRepository.java | 3 + .../DatabaseConfigPersistence.java | 4 +- .../FileSystemConfigPersistence.java | 3 +- .../configs/ConfigsDatabaseTestProvider.java | 58 ++++++++ .../jobs/JobsDatabaseTestProvider.java | 40 ++++++ .../instance/test/TestDatabaseProvider.java | 14 ++ .../instance/test/TestDatabaseProviders.java | 84 ++++++++++++ .../db/instance/AbstractDatabaseTest.java | 7 +- .../configs/AbstractConfigsDatabaseTest.java | 21 +-- .../jobs/AbstractJobsDatabaseTest.java | 5 +- .../toys/AbstractToysDatabaseTest.java | 17 --- .../toys/ToysDatabaseMigratorTest.java | 10 +- .../airbyte/scheduler/app/SchedulerApp.java | 3 +- airbyte-scheduler/persistence/build.gradle | 1 + .../persistence/DefaultJobPersistence.java | 99 +++++++------- .../DefaultJobPersistenceTest.java | 38 +++--- .../java/io/airbyte/server/ServerApp.java | 5 +- .../server/handlers/ArchiveHandlerTest.java | 14 +- .../migration/DatabaseArchiverTest.java | 16 ++- .../server/migration/RunMigrationTest.java | 129 ++++++++++++------ .../server/migration/StubAirbyteDB.java | 45 ------ .../temporal/TemporalAttemptExecution.java | 8 +- .../TemporalAttemptExecutionTest.java | 40 ++---- 23 files changed, 425 insertions(+), 239 deletions(-) create mode 100644 airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java create mode 100644 airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/JobsDatabaseTestProvider.java create mode 100644 airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java create mode 100644 airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProviders.java delete mode 100644 airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/AbstractToysDatabaseTest.java delete mode 100644 airbyte-server/src/test/java/io/airbyte/server/migration/StubAirbyteDB.java diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java index c03d7b44fca6..0825eac538a0 100644 --- a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java +++ b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java @@ -5,6 +5,7 @@ package io.airbyte.config.persistence; import com.fasterxml.jackson.databind.JsonNode; +import com.google.api.client.util.Preconditions; import io.airbyte.commons.docker.DockerUtils; import io.airbyte.commons.json.Jsons; import io.airbyte.commons.lang.Exceptions; @@ -42,6 +43,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +@SuppressWarnings("OptionalUsedAsFieldOrParameterType") public class ConfigRepository { private static final Logger LOGGER = LoggerFactory.getLogger(ConfigRepository.class); @@ -453,6 +455,7 @@ public void replaceAllConfigsDeserializing(final Map> c public void replaceAllConfigs(final Map> configs, final boolean dryRun) throws IOException { if (longLivedSecretPersistence.isPresent()) { + Preconditions.checkNotNull(specFetcherFn); final var augmentedMap = new HashMap<>(configs); // get all source defs so that we can use their specs when storing secrets. diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/DatabaseConfigPersistence.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/DatabaseConfigPersistence.java index b4efefc427a1..19035ed42295 100644 --- a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/DatabaseConfigPersistence.java +++ b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/DatabaseConfigPersistence.java @@ -61,7 +61,7 @@ public DatabaseConfigPersistence(final Database database) { * If this is a migration deployment from an old version that relies on file system config * persistence, copy the existing configs from local files. */ - public void migrateFileConfigs(final Configs serverConfigs) throws IOException { + public DatabaseConfigPersistence migrateFileConfigs(final Configs serverConfigs) throws IOException { database.transaction(ctx -> { final boolean isInitialized = ctx.fetchExists(AIRBYTE_CONFIGS); if (isInitialized) { @@ -77,6 +77,8 @@ public void migrateFileConfigs(final Configs serverConfigs) throws IOException { return null; }); + + return this; } @Override diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/FileSystemConfigPersistence.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/FileSystemConfigPersistence.java index 9a1c9befbf29..8681788cb6f8 100644 --- a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/FileSystemConfigPersistence.java +++ b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/FileSystemConfigPersistence.java @@ -11,7 +11,6 @@ import io.airbyte.config.AirbyteConfig; import io.airbyte.validation.json.JsonValidationException; import java.io.IOException; -import java.io.UnsupportedEncodingException; import java.nio.file.Files; import java.nio.file.Path; import java.util.ArrayList; @@ -186,7 +185,7 @@ public void replaceAllConfigs(final Map> configs, final @Override public void loadData(final ConfigPersistence seedPersistence) throws IOException { - throw new UnsupportedEncodingException("This method is not supported in this implementation"); + // this method is not supported in this implementation, but needed in tests; do nothing } private T getConfigInternal(final AirbyteConfig configType, final String configId, final Class clazz) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java new file mode 100644 index 000000000000..243007746900 --- /dev/null +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java @@ -0,0 +1,58 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.db.instance.configs; + +import static org.jooq.impl.DSL.field; +import static org.jooq.impl.DSL.table; + +import io.airbyte.db.Database; +import io.airbyte.db.ExceptionWrappingDatabase; +import io.airbyte.db.instance.DatabaseMigrator; +import io.airbyte.db.instance.test.TestDatabaseProvider; +import java.io.IOException; +import java.time.OffsetDateTime; +import java.util.UUID; +import org.jooq.JSONB; + +public class ConfigsDatabaseTestProvider implements TestDatabaseProvider { + + private final String user; + private final String password; + private final String jdbcUrl; + + public ConfigsDatabaseTestProvider(String user, String password, String jdbcUrl) { + this.user = user; + this.password = password; + this.jdbcUrl = jdbcUrl; + } + + @Override + public Database create(final boolean runMigration) throws IOException { + final Database database = new ConfigsDatabaseInstance(user, password, jdbcUrl) + .getAndInitialize(); + + if (runMigration) { + final DatabaseMigrator migrator = new ConfigsDatabaseMigrator( + database, + ConfigsDatabaseTestProvider.class.getSimpleName()); + migrator.createBaseline(); + migrator.migrate(); + } + + // The configs database is considered ready only if there are some seed records. + // So we need to create at least one record here. + OffsetDateTime timestamp = OffsetDateTime.now(); + new ExceptionWrappingDatabase(database).transaction(ctx -> ctx.insertInto(table("airbyte_configs")) + .set(field("config_id"), UUID.randomUUID().toString()) + .set(field("config_type"), "STANDARD_SOURCE_DEFINITION") + .set(field("config_blob"), JSONB.valueOf("{}")) + .set(field("created_at"), timestamp) + .set(field("updated_at"), timestamp) + .execute()); + + return database; + } + +} diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/JobsDatabaseTestProvider.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/JobsDatabaseTestProvider.java new file mode 100644 index 000000000000..133b9c290807 --- /dev/null +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/jobs/JobsDatabaseTestProvider.java @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.db.instance.jobs; + +import io.airbyte.db.Database; +import io.airbyte.db.instance.DatabaseMigrator; +import io.airbyte.db.instance.test.TestDatabaseProvider; +import java.io.IOException; + +public class JobsDatabaseTestProvider implements TestDatabaseProvider { + + private final String user; + private final String password; + private final String jdbcUrl; + + public JobsDatabaseTestProvider(String user, String password, String jdbcUrl) { + this.user = user; + this.password = password; + this.jdbcUrl = jdbcUrl; + } + + @Override + public Database create(final boolean runMigration) throws IOException { + final Database jobsDatabase = new JobsDatabaseInstance(user, password, jdbcUrl) + .getAndInitialize(); + + if (runMigration) { + final DatabaseMigrator migrator = new JobsDatabaseMigrator( + jobsDatabase, + JobsDatabaseTestProvider.class.getSimpleName()); + migrator.createBaseline(); + migrator.migrate(); + } + + return jobsDatabase; + } + +} diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java new file mode 100644 index 000000000000..c623bc1243f6 --- /dev/null +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java @@ -0,0 +1,14 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.db.instance.test; + +import io.airbyte.db.Database; +import java.io.IOException; + +public interface TestDatabaseProvider { + + Database create(final boolean runMigration) throws IOException; + +} diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProviders.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProviders.java new file mode 100644 index 000000000000..7bb34b7dc588 --- /dev/null +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProviders.java @@ -0,0 +1,84 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.db.instance.test; + +import com.google.api.client.util.Preconditions; +import io.airbyte.config.Configs; +import io.airbyte.db.Database; +import io.airbyte.db.instance.configs.ConfigsDatabaseTestProvider; +import io.airbyte.db.instance.jobs.JobsDatabaseTestProvider; +import java.io.IOException; +import java.util.Optional; +import org.testcontainers.containers.PostgreSQLContainer; + +/** + * Use this class to create mock databases in unit tests. This class takes care of database + * initialization and migration. + */ +@SuppressWarnings("OptionalUsedAsFieldOrParameterType") +public class TestDatabaseProviders { + + private final Optional configs; + private final Optional> container; + private boolean runMigration = true; + + public TestDatabaseProviders(final Configs configs) { + this.configs = Optional.of(configs); + this.container = Optional.empty(); + } + + public TestDatabaseProviders(final PostgreSQLContainer container) { + this.configs = Optional.empty(); + this.container = Optional.of(container); + } + + /** + * When creating mock databases in unit tests, migration should be run by default. Call this method + * to turn migration off, which is needed when unit testing migration code. + */ + public TestDatabaseProviders turnOffMigration() { + this.runMigration = false; + return this; + } + + public Database createNewConfigsDatabase() throws IOException { + Preconditions.checkArgument(configs.isPresent() || container.isPresent()); + if (configs.isPresent()) { + final Configs c = configs.get(); + return new ConfigsDatabaseTestProvider( + c.getConfigDatabaseUser(), + c.getConfigDatabasePassword(), + c.getConfigDatabaseUrl()) + .create(runMigration); + } else { + final PostgreSQLContainer c = container.get(); + return new ConfigsDatabaseTestProvider( + c.getUsername(), + c.getPassword(), + c.getJdbcUrl()) + .create(runMigration); + } + } + + public Database createNewJobsDatabase() throws IOException { + Preconditions.checkArgument(configs.isPresent() || container.isPresent()); + if (configs.isPresent()) { + final Configs c = configs.get(); + return new JobsDatabaseTestProvider( + c.getDatabaseUser(), + c.getDatabasePassword(), + c.getDatabaseUrl()) + .create(runMigration); + } else { + final PostgreSQLContainer c = container.get(); + return new JobsDatabaseTestProvider( + c.getUsername(), + c.getPassword(), + c.getJdbcUrl()) + .create(runMigration); + } + } + +} diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/AbstractDatabaseTest.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/AbstractDatabaseTest.java index 1fae9c7f320c..de040b91c344 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/AbstractDatabaseTest.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/AbstractDatabaseTest.java @@ -34,7 +34,7 @@ public static void dbDown() { @BeforeEach public void setup() throws Exception { - database = getAndInitializeDatabase(container.getUsername(), container.getPassword(), container.getJdbcUrl()); + database = getDatabase(); } @AfterEach @@ -44,8 +44,9 @@ void tearDown() throws Exception { /** * Create an initialized database. The downstream implementation should do it by calling - * {@link DatabaseInstance#getAndInitialize}. + * {@link DatabaseInstance#getAndInitialize} or {@link DatabaseInstance#getInitialized}, and + * {@link DatabaseMigrator#migrate} if necessary. */ - public abstract Database getAndInitializeDatabase(String username, String password, String connectionString) throws IOException; + public abstract Database getDatabase() throws IOException; } diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/AbstractConfigsDatabaseTest.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/AbstractConfigsDatabaseTest.java index bd7b1695145d..edb6a454c66a 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/AbstractConfigsDatabaseTest.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/AbstractConfigsDatabaseTest.java @@ -8,11 +8,10 @@ import static org.jooq.impl.DSL.table; import io.airbyte.db.Database; -import io.airbyte.db.ExceptionWrappingDatabase; import io.airbyte.db.instance.AbstractDatabaseTest; +import io.airbyte.db.instance.test.TestDatabaseProviders; import java.io.IOException; import java.time.OffsetDateTime; -import java.util.UUID; import org.jooq.Field; import org.jooq.JSONB; import org.jooq.Record; @@ -27,21 +26,9 @@ public abstract class AbstractConfigsDatabaseTest extends AbstractDatabaseTest { public static final Field CREATED_AT = field("created_at", OffsetDateTime.class); public static final Field UPDATED_AT = field("updated_at", OffsetDateTime.class); - public Database getAndInitializeDatabase(final String username, final String password, final String connectionString) throws IOException { - final Database database = - new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - - // The configs database is considered ready only if there are some seed records. - // So we need to create at least one record here. - final OffsetDateTime timestamp = OffsetDateTime.now(); - new ExceptionWrappingDatabase(database).transaction(ctx -> ctx.insertInto(AIRBYTE_CONFIGS) - .set(CONFIG_ID, UUID.randomUUID().toString()) - .set(CONFIG_TYPE, "STANDARD_SOURCE_DEFINITION") - .set(CONFIG_BLOB, JSONB.valueOf("{}")) - .set(CREATED_AT, timestamp) - .set(UPDATED_AT, timestamp) - .execute()); - return database; + @Override + public Database getDatabase() throws IOException { + return new TestDatabaseProviders(container).turnOffMigration().createNewConfigsDatabase(); } } diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/jobs/AbstractJobsDatabaseTest.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/jobs/AbstractJobsDatabaseTest.java index a7c49167e73e..b14fc53b8a63 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/jobs/AbstractJobsDatabaseTest.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/jobs/AbstractJobsDatabaseTest.java @@ -6,13 +6,14 @@ import io.airbyte.db.Database; import io.airbyte.db.instance.AbstractDatabaseTest; +import io.airbyte.db.instance.test.TestDatabaseProviders; import java.io.IOException; public abstract class AbstractJobsDatabaseTest extends AbstractDatabaseTest { @Override - public Database getAndInitializeDatabase(final String username, final String password, final String connectionString) throws IOException { - return new JobsDatabaseInstance(username, password, connectionString).getAndInitialize(); + public Database getDatabase() throws IOException { + return new TestDatabaseProviders(container).turnOffMigration().createNewJobsDatabase(); } } diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/AbstractToysDatabaseTest.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/AbstractToysDatabaseTest.java deleted file mode 100644 index cccf8919c4e8..000000000000 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/AbstractToysDatabaseTest.java +++ /dev/null @@ -1,17 +0,0 @@ -/* - * Copyright (c) 2021 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.db.instance.toys; - -import io.airbyte.db.Database; -import io.airbyte.db.instance.AbstractDatabaseTest; -import java.io.IOException; - -public abstract class AbstractToysDatabaseTest extends AbstractDatabaseTest { - - public Database getAndInitializeDatabase(final String username, final String password, final String connectionString) throws IOException { - return new ToysDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - } - -} diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/ToysDatabaseMigratorTest.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/ToysDatabaseMigratorTest.java index 8733a37a716c..4c432e844262 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/ToysDatabaseMigratorTest.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/toys/ToysDatabaseMigratorTest.java @@ -7,14 +7,22 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import io.airbyte.commons.resources.MoreResources; +import io.airbyte.db.Database; +import io.airbyte.db.instance.AbstractDatabaseTest; import io.airbyte.db.instance.DatabaseMigrator; +import java.io.IOException; import org.junit.jupiter.api.Test; -class ToysDatabaseMigratorTest extends AbstractToysDatabaseTest { +class ToysDatabaseMigratorTest extends AbstractDatabaseTest { private static final String PRE_MIGRATION_SCHEMA_DUMP = "toys_database/pre_migration_schema.txt"; private static final String POST_MIGRATION_SCHEMA_DUMP = "toys_database/schema_dump.txt"; + @Override + public Database getDatabase() throws IOException { + return new ToysDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + } + @Test public void testMigration() throws Exception { final DatabaseMigrator migrator = new ToysDatabaseMigrator(database, ToysDatabaseMigratorTest.class.getSimpleName()); diff --git a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java index 835440ec30de..64a59a4667dd 100644 --- a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java +++ b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java @@ -204,7 +204,6 @@ public static void main(final String[] args) throws IOException, InterruptedExce configs.getDatabaseUrl()) .getInitialized(); - final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); final Database configDatabase = new ConfigsDatabaseInstance( configs.getConfigDatabaseUser(), configs.getConfigDatabasePassword(), @@ -215,6 +214,8 @@ public static void main(final String[] args) throws IOException, InterruptedExce final Optional ephemeralSecretPersistence = SecretPersistence.getEphemeral(configs); final SecretsHydrator secretsHydrator = SecretPersistence.getSecretsHydrator(configs); final ConfigRepository configRepository = new ConfigRepository(configPersistence, secretsHydrator, secretPersistence, ephemeralSecretPersistence); + + final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); final JobCleaner jobCleaner = new JobCleaner( configs.getWorkspaceRetentionConfig(), workspaceRoot, diff --git a/airbyte-scheduler/persistence/build.gradle b/airbyte-scheduler/persistence/build.gradle index fa9abd9512a3..73ae03582082 100644 --- a/airbyte-scheduler/persistence/build.gradle +++ b/airbyte-scheduler/persistence/build.gradle @@ -8,6 +8,7 @@ dependencies { implementation project(':airbyte-config:models') implementation project(':airbyte-config:persistence') implementation project(':airbyte-db:lib') + implementation project(':airbyte-db:jooq') implementation project(':airbyte-json-validation') implementation project(':airbyte-notification') implementation project(':airbyte-oauth') diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java index 8050a75bbc1a..f8cbef5752f4 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java @@ -18,10 +18,10 @@ import io.airbyte.config.JobConfig; import io.airbyte.config.JobConfig.ConfigType; import io.airbyte.config.JobOutput; -import io.airbyte.config.StandardSyncOutput; import io.airbyte.config.State; import io.airbyte.db.Database; import io.airbyte.db.ExceptionWrappingDatabase; +import io.airbyte.db.instance.configs.jooq.tables.SyncState; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; import io.airbyte.scheduler.models.Attempt; import io.airbyte.scheduler.models.AttemptStatus; @@ -51,10 +51,12 @@ import org.jooq.DSLContext; import org.jooq.Field; import org.jooq.InsertValuesStepN; +import org.jooq.JSONB; import org.jooq.JSONFormat; import org.jooq.JSONFormat.RecordFormat; import org.jooq.Named; import org.jooq.Record; +import org.jooq.Record1; import org.jooq.Result; import org.jooq.Sequence; import org.jooq.Table; @@ -105,24 +107,30 @@ public class DefaultJobPersistence implements JobPersistence { public static final String ORDER_BY_JOB_TIME_ATTEMPT_TIME = "ORDER BY jobs.created_at DESC, jobs.id DESC, attempts.created_at ASC, attempts.id ASC "; - private final ExceptionWrappingDatabase database; + private final ExceptionWrappingDatabase jobDatabase; + private final ExceptionWrappingDatabase configDatabase; private final Supplier timeSupplier; @VisibleForTesting - DefaultJobPersistence(final Database database, + DefaultJobPersistence(final Database jobDatabase, + final Database configDatabase, final Supplier timeSupplier, final int minimumAgeInDays, final int excessiveNumberOfJobs, final int minimumRecencyCount) { - this.database = new ExceptionWrappingDatabase(database); + this.jobDatabase = new ExceptionWrappingDatabase(jobDatabase); + this.configDatabase = new ExceptionWrappingDatabase(configDatabase); this.timeSupplier = timeSupplier; JOB_HISTORY_MINIMUM_AGE_IN_DAYS = minimumAgeInDays; JOB_HISTORY_EXCESSIVE_NUMBER_OF_JOBS = excessiveNumberOfJobs; JOB_HISTORY_MINIMUM_RECENCY = minimumRecencyCount; } - public DefaultJobPersistence(final Database database) { - this(database, Instant::now, 30, 500, 10); + /** + * @param configDatabase The config database is only needed for reading and writing sync state. + */ + public DefaultJobPersistence(final Database jobDatabase, final Database configDatabase) { + this(jobDatabase, configDatabase, Instant::now, 30, 500, 10); } /** @@ -140,7 +148,7 @@ public Optional enqueueJob(final String scope, final JobConfig jobConfig) JobStatus.TERMINAL_STATUSES.stream().map(Sqls::toSqlName).map(Names::singleQuote).collect(Collectors.joining(","))) : ""; - return database.query( + return jobDatabase.query( ctx -> ctx.fetch( "INSERT INTO jobs(config_type, scope, created_at, updated_at, status, config) " + "SELECT CAST(? AS JOB_CONFIG_TYPE), ?, ?, ?, CAST(? AS JOB_STATUS), CAST(? as JSONB) " + @@ -160,7 +168,7 @@ public Optional enqueueJob(final String scope, final JobConfig jobConfig) @Override public void resetJob(final long jobId) throws IOException { final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - database.query(ctx -> { + jobDatabase.query(ctx -> { updateJobStatusIfNotInTerminalState(ctx, jobId, JobStatus.PENDING, now, new IllegalStateException(String.format("Attempt to reset a job that is in a terminal state. job id: %s", jobId))); return null; @@ -170,7 +178,7 @@ public void resetJob(final long jobId) throws IOException { @Override public void cancelJob(final long jobId) throws IOException { final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - database.query(ctx -> { + jobDatabase.query(ctx -> { updateJobStatusIfNotInTerminalState(ctx, jobId, JobStatus.CANCELLED, now); return null; }); @@ -179,7 +187,7 @@ public void cancelJob(final long jobId) throws IOException { @Override public void failJob(final long jobId) throws IOException { final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - database.query(ctx -> { + jobDatabase.query(ctx -> { updateJobStatusIfNotInTerminalState(ctx, jobId, JobStatus.FAILED, now); return null; }); @@ -215,7 +223,7 @@ private void updateJobStatus(final DSLContext ctx, final long jobId, final JobSt public int createAttempt(final long jobId, final Path logPath) throws IOException { final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - return database.transaction(ctx -> { + return jobDatabase.transaction(ctx -> { final Job job = getJob(ctx, jobId); if (job.isJobInTerminalState()) { final var errMsg = @@ -253,7 +261,7 @@ public int createAttempt(final long jobId, final Path logPath) throws IOExceptio @Override public void failAttempt(final long jobId, final int attemptNumber) throws IOException { final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - database.transaction(ctx -> { + jobDatabase.transaction(ctx -> { // do not overwrite terminal states. updateJobStatusIfNotInTerminalState(ctx, jobId, JobStatus.INCOMPLETE, now); @@ -271,7 +279,7 @@ public void failAttempt(final long jobId, final int attemptNumber) throws IOExce @Override public void succeedAttempt(final long jobId, final int attemptNumber) throws IOException { final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - database.transaction(ctx -> { + jobDatabase.transaction(ctx -> { // override any other terminal statuses if we are now succeeded. updateJobStatus(ctx, jobId, JobStatus.SUCCEEDED, now); @@ -288,7 +296,7 @@ public void succeedAttempt(final long jobId, final int attemptNumber) throws IOE @Override public void setAttemptTemporalWorkflowId(final long jobId, final int attemptNumber, final String temporalWorkflowId) throws IOException { - database.query(ctx -> ctx.execute( + jobDatabase.query(ctx -> ctx.execute( " UPDATE attempts SET temporal_workflow_id = ? WHERE job_id = ? AND attempt_number = ?", temporalWorkflowId, jobId, @@ -297,7 +305,7 @@ public void setAttemptTemporalWorkflowId(final long jobId, final int attemptNumb @Override public Optional getAttemptTemporalWorkflowId(final long jobId, final int attemptNumber) throws IOException { - final var result = database.query(ctx -> ctx.fetch( + final var result = jobDatabase.query(ctx -> ctx.fetch( " SELECT temporal_workflow_id from attempts WHERE job_id = ? AND attempt_number = ?", jobId, attemptNumber)).stream().findFirst(); @@ -313,7 +321,7 @@ public Optional getAttemptTemporalWorkflowId(final long jobId, final int public void writeOutput(final long jobId, final int attemptNumber, final T output) throws IOException { final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - database.query( + jobDatabase.query( ctx -> ctx.execute( "UPDATE attempts SET output = CAST(? as JSONB), updated_at = ? WHERE job_id = ? AND attempt_number = ?", Jsons.serialize(output), @@ -324,7 +332,7 @@ public void writeOutput(final long jobId, final int attemptNumber, final T o @Override public Job getJob(final long jobId) throws IOException { - return database.query(ctx -> getJob(ctx, jobId)); + return jobDatabase.query(ctx -> getJob(ctx, jobId)); } private Job getJob(final DSLContext ctx, final long jobId) { @@ -342,7 +350,7 @@ public List listJobs(final ConfigType configType, final String configId, fi @Override public List listJobs(final Set configTypes, final String configId, final int pagesize, final int offset) throws IOException { - return database.query(ctx -> getJobsFromResult(ctx.fetch( + return jobDatabase.query(ctx -> getJobsFromResult(ctx.fetch( BASE_JOB_SELECT_AND_JOIN + "WHERE " + "CAST(config_type AS VARCHAR) in " + Sqls.toSqlInFragment(configTypes) + " " + "AND scope = ? " + @@ -358,7 +366,7 @@ public List listJobsWithStatus(final JobStatus status) throws IOException { @Override public List listJobsWithStatus(final Set configTypes, final JobStatus status) throws IOException { - return database.query(ctx -> getJobsFromResult(ctx + return jobDatabase.query(ctx -> getJobsFromResult(ctx .fetch(BASE_JOB_SELECT_AND_JOIN + "WHERE " + "CAST(config_type AS VARCHAR) IN " + Sqls.toSqlInFragment(configTypes) + " AND " + "CAST(jobs.status AS VARCHAR) = ? " + @@ -373,7 +381,7 @@ public List listJobsWithStatus(final ConfigType configType, final JobStatus @Override public Optional getLastReplicationJob(final UUID connectionId) throws IOException { - return database.query(ctx -> ctx + return jobDatabase.query(ctx -> ctx .fetch(BASE_JOB_SELECT_AND_JOIN + "WHERE " + "CAST(jobs.config_type AS VARCHAR) in " + Sqls.toSqlInFragment(Job.REPLICATION_TYPES) + " AND " + "scope = ? AND " + @@ -388,20 +396,17 @@ public Optional getLastReplicationJob(final UUID connectionId) throws IOExc @Override public Optional getCurrentState(final UUID connectionId) throws IOException { - return database.query(ctx -> ctx - .fetch(BASE_JOB_SELECT_AND_JOIN + "WHERE " + - "CAST(jobs.config_type AS VARCHAR) in " + Sqls.toSqlInFragment(Job.REPLICATION_TYPES) + " AND " + - "scope = ? AND " + - "output->'sync'->'state' IS NOT NULL " + - "ORDER BY attempts.created_at DESC LIMIT 1", - connectionId.toString()) - .stream() - .findFirst() - .flatMap(r -> getJobOptional(ctx, r.get("job_id", Long.class))) - .flatMap(Job::getLastAttemptWithOutput) - .flatMap(Attempt::getOutput) - .map(JobOutput::getSync) - .map(StandardSyncOutput::getState)); + return configDatabase.query(ctx -> { + final Record1 record = ctx.select(SyncState.SYNC_STATE.STATE) + .from(SyncState.SYNC_STATE) + .where(SyncState.SYNC_STATE.SYNC_ID.eq(connectionId)) + .fetchAny(); + if (record == null) { + return Optional.empty(); + } + + return Optional.of(Jsons.deserialize(record.value1().data(), State.class)); + }); } @Override @@ -410,7 +415,7 @@ public Optional getNextJob() throws IOException { // 1. get oldest, pending job // 2. job is excluded if another job of the same scope is already running // 3. job is excluded if another job of the same scope is already incomplete - return database.query(ctx -> ctx + return jobDatabase.query(ctx -> ctx .fetch(BASE_JOB_SELECT_AND_JOIN + "WHERE " + "CAST(jobs.status AS VARCHAR) = 'pending' AND " + "jobs.scope NOT IN ( SELECT scope FROM jobs WHERE status = 'running' OR status = 'incomplete' ) " + @@ -423,7 +428,7 @@ public Optional getNextJob() throws IOException { @Override public List listJobs(final ConfigType configType, final Instant attemptEndedAtTimestamp) throws IOException { final LocalDateTime timeConvertedIntoLocalDateTime = LocalDateTime.ofInstant(attemptEndedAtTimestamp, ZoneOffset.UTC); - return database.query(ctx -> getJobsFromResult(ctx + return jobDatabase.query(ctx -> getJobsFromResult(ctx .fetch(BASE_JOB_SELECT_AND_JOIN + "WHERE " + "CAST(config_type AS VARCHAR) = ? AND " + " attempts.ended_at > ? ORDER BY jobs.created_at ASC, attempts.created_at ASC", Sqls.toSqlName(configType), @@ -476,7 +481,7 @@ private static long getEpoch(final Record record, final String fieldName) { @Override public Optional getVersion() throws IOException { - final Result result = database.query(ctx -> ctx.select() + final Result result = jobDatabase.query(ctx -> ctx.select() .from(AIRBYTE_METADATA_TABLE) .where(DSL.field(METADATA_KEY_COL).eq(AirbyteVersion.AIRBYTE_VERSION_KEY_NAME)) .fetch()); @@ -485,7 +490,7 @@ public Optional getVersion() throws IOException { @Override public void setVersion(final String airbyteVersion) throws IOException { - database.query(ctx -> ctx.execute(String.format( + jobDatabase.query(ctx -> ctx.execute(String.format( "INSERT INTO %s(%s, %s) VALUES('%s', '%s'), ('%s_init_db', '%s') ON CONFLICT (%s) DO UPDATE SET %s = '%s'", AIRBYTE_METADATA_TABLE, METADATA_KEY_COL, @@ -501,7 +506,7 @@ public void setVersion(final String airbyteVersion) throws IOException { @Override public Optional getDeployment() throws IOException { - final Result result = database.query(ctx -> ctx.select() + final Result result = jobDatabase.query(ctx -> ctx.select() .from(AIRBYTE_METADATA_TABLE) .where(DSL.field(METADATA_KEY_COL).eq(DEPLOYMENT_ID_KEY)) .fetch()); @@ -511,7 +516,7 @@ public Optional getDeployment() throws IOException { @Override public void setDeployment(final UUID deployment) throws IOException { // if an existing deployment id already exists, on conflict, return it so we can log it. - final UUID committedDeploymentId = database.query(ctx -> ctx.fetch(String.format( + final UUID committedDeploymentId = jobDatabase.query(ctx -> ctx.fetch(String.format( "INSERT INTO %s(%s, %s) VALUES('%s', '%s') ON CONFLICT (%s) DO NOTHING RETURNING (SELECT %s FROM %s WHERE %s='%s') as existing_deployment_id", AIRBYTE_METADATA_TABLE, METADATA_KEY_COL, @@ -580,7 +585,7 @@ private Map> exportDatabase(final String sc */ private List listTables(final String schema) throws IOException { if (schema != null) { - return database.query(context -> context.meta().getSchemas(schema).stream() + return jobDatabase.query(context -> context.meta().getSchemas(schema).stream() .flatMap(s -> context.meta(s).getTables().stream()) .map(Named::getName) .filter(table -> JobsDatabaseSchema.getTableNames().contains(table.toLowerCase())) @@ -601,7 +606,7 @@ public void purgeJobHistory(final LocalDateTime asOfDate) { final String JOB_HISTORY_PURGE_SQL = MoreResources.readResource("job_history_purge.sql"); // interval '?' days cannot use a ? bind, so we're using %d instead. final String sql = String.format(JOB_HISTORY_PURGE_SQL, (JOB_HISTORY_MINIMUM_AGE_IN_DAYS - 1)); - final Integer rows = database.query(ctx -> ctx.execute(sql, + final Integer rows = jobDatabase.query(ctx -> ctx.execute(sql, asOfDate.format(DateTimeFormatter.ofPattern("YYYY-MM-dd")), JOB_HISTORY_EXCESSIVE_NUMBER_OF_JOBS, JOB_HISTORY_MINIMUM_RECENCY)); @@ -612,7 +617,7 @@ public void purgeJobHistory(final LocalDateTime asOfDate) { private List listAllTables(final String schema) throws IOException { if (schema != null) { - return database.query(context -> context.meta().getSchemas(schema).stream() + return jobDatabase.query(context -> context.meta().getSchemas(schema).stream() .flatMap(s -> context.meta(s).getTables().stream()) .map(Named::getName) .collect(Collectors.toList())); @@ -622,7 +627,7 @@ private List listAllTables(final String schema) throws IOException { } private List listSchemas() throws IOException { - return database.query(context -> context.meta().getSchemas().stream() + return jobDatabase.query(context -> context.meta().getSchemas().stream() .map(Named::getName) .filter(c -> !SYSTEM_SCHEMA.contains(c)) .collect(Collectors.toList())); @@ -631,7 +636,7 @@ private List listSchemas() throws IOException { private Stream exportTable(final String schema, final String tableName) throws IOException { final Table tableSql = getTable(schema, tableName); - try (final Stream records = database.query(ctx -> ctx.select(DSL.asterisk()).from(tableSql).fetchStream())) { + try (final Stream records = jobDatabase.query(ctx -> ctx.select(DSL.asterisk()).from(tableSql).fetchStream())) { return records.map(record -> { final Set jsonFieldNames = Arrays.stream(record.fields()) .filter(f -> f.getDataType().getTypeName().equals("jsonb")) @@ -658,7 +663,7 @@ private void importDatabase(final String airbyteVersion, throws IOException { if (!data.isEmpty()) { createSchema(BACKUP_SCHEMA); - database.transaction(ctx -> { + jobDatabase.transaction(ctx -> { for (final JobsDatabaseSchema tableType : data.keySet()) { if (!incrementalImport) { truncateTable(ctx, targetSchema, tableType.name(), BACKUP_SCHEMA); @@ -673,7 +678,7 @@ private void importDatabase(final String airbyteVersion, } private void createSchema(final String schema) throws IOException { - database.query(ctx -> ctx.createSchemaIfNotExists(schema).execute()); + jobDatabase.query(ctx -> ctx.createSchemaIfNotExists(schema).execute()); } /** diff --git a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java index 84c3379f8e7b..48978670f816 100644 --- a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java +++ b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java @@ -28,10 +28,8 @@ import io.airbyte.config.StandardSyncOutput; import io.airbyte.config.State; import io.airbyte.db.Database; -import io.airbyte.db.instance.DatabaseMigrator; -import io.airbyte.db.instance.jobs.JobsDatabaseInstance; -import io.airbyte.db.instance.jobs.JobsDatabaseMigrator; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; +import io.airbyte.db.instance.test.TestDatabaseProviders; import io.airbyte.scheduler.models.Attempt; import io.airbyte.scheduler.models.AttemptStatus; import io.airbyte.scheduler.models.Job; @@ -90,7 +88,8 @@ class DefaultJobPersistenceTest { .withSync(new JobSyncConfig()); private static PostgreSQLContainer container; - private Database database; + private Database jobDatabase; + private Database configDatabase; private Supplier timeSupplier; private JobPersistence jobPersistence; @@ -158,33 +157,31 @@ private static Job createJob( @SuppressWarnings("unchecked") @BeforeEach public void setup() throws Exception { - database = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + final TestDatabaseProviders databaseProviders = new TestDatabaseProviders(container); + jobDatabase = databaseProviders.createNewJobsDatabase(); + configDatabase = databaseProviders.createNewConfigsDatabase(); resetDb(); - final DatabaseMigrator jobDbMigrator = new JobsDatabaseMigrator(database, "test"); - jobDbMigrator.createBaseline(); - jobDbMigrator.migrate(); - timeSupplier = mock(Supplier.class); when(timeSupplier.get()).thenReturn(NOW); - jobPersistence = new DefaultJobPersistence(database, timeSupplier, 30, 500, 10); + jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase, timeSupplier, 30, 500, 10); } @AfterEach void tearDown() throws Exception { - database.close(); + jobDatabase.close(); } private void resetDb() throws SQLException { // todo (cgardens) - truncate whole db. - database.query(ctx -> ctx.execute("TRUNCATE TABLE jobs")); - database.query(ctx -> ctx.execute("TRUNCATE TABLE attempts")); - database.query(ctx -> ctx.execute("TRUNCATE TABLE airbyte_metadata")); + jobDatabase.query(ctx -> ctx.execute("TRUNCATE TABLE jobs")); + jobDatabase.query(ctx -> ctx.execute("TRUNCATE TABLE attempts")); + jobDatabase.query(ctx -> ctx.execute("TRUNCATE TABLE airbyte_metadata")); } private Result getJobRecord(final long jobId) throws SQLException { - return database.query(ctx -> ctx.fetch(DefaultJobPersistence.BASE_JOB_SELECT_AND_JOIN + "WHERE jobs.id = ?", jobId)); + return jobDatabase.query(ctx -> ctx.fetch(DefaultJobPersistence.BASE_JOB_SELECT_AND_JOIN + "WHERE jobs.id = ?", jobId)); } @Test @@ -336,7 +333,7 @@ void testListJobsWithTimestamp() throws IOException { now.plusSeconds(14), now.plusSeconds(15), now.plusSeconds(16)); - jobPersistence = new DefaultJobPersistence(database, timeSupplier, 30, 500, 10); + jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase, timeSupplier, 30, 500, 10); final long syncJobId = jobPersistence.enqueueJob(SCOPE, SYNC_JOB_CONFIG).orElseThrow(); final int syncJobAttemptNumber0 = jobPersistence.createAttempt(syncJobId, LOG_PATH); jobPersistence.failAttempt(syncJobId, syncJobAttemptNumber0); @@ -439,7 +436,7 @@ void testSuccessfulGet() throws IOException, SQLException { final var defaultWorkflowId = jobPersistence.getAttemptTemporalWorkflowId(jobId, attemptNumber); assertTrue(defaultWorkflowId.isEmpty()); - database.query(ctx -> ctx.execute( + jobDatabase.query(ctx -> ctx.execute( "UPDATE attempts SET temporal_workflow_id = '56a81f3a-006c-42d7-bce2-29d675d08ea4' WHERE job_id = ? AND attempt_number =?", jobId, attemptNumber)); final var workflowId = jobPersistence.getAttemptTemporalWorkflowId(jobId, attemptNumber).get(); @@ -1181,7 +1178,7 @@ class PurgeJobHistory { private Job persistJobForJobHistoryTesting(final String scope, final JobConfig jobConfig, final JobStatus status, final LocalDateTime runDate) throws IOException, SQLException { final String when = runDate.toString(); - final Optional id = database.query( + final Optional id = jobDatabase.query( ctx -> ctx.fetch( "INSERT INTO jobs(config_type, scope, created_at, updated_at, status, config) " + "SELECT CAST(? AS JOB_CONFIG_TYPE), ?, ?, ?, CAST(? AS JOB_STATUS), CAST(? as JSONB) " + @@ -1210,7 +1207,7 @@ private void persistAttemptForJobHistoryTesting(final Job job, final String logP + " \"sync\": {\n" + " \"output_catalog\": {" + "}}}"; - final Integer attemptNumber = database.query(ctx -> ctx.fetch( + final Integer attemptNumber = jobDatabase.query(ctx -> ctx.fetch( "INSERT INTO attempts(job_id, attempt_number, log_path, status, created_at, updated_at, output) " + "VALUES(?, ?, ?, CAST(? AS ATTEMPT_STATUS), ?, ?, CAST(? as JSONB)) RETURNING attempt_number", job.getId(), @@ -1295,7 +1292,8 @@ void testPurgeJobHistory(final int numJobs, final String DECOY_SCOPE = UUID.randomUUID().toString(); // Reconfigure constants to test various combinations of tuning knobs and make sure all work. - final DefaultJobPersistence jobPersistence = new DefaultJobPersistence(database, timeSupplier, ageCutoff, tooManyJobs, recencyCutoff); + final DefaultJobPersistence jobPersistence = + new DefaultJobPersistence(jobDatabase, configDatabase, timeSupplier, ageCutoff, tooManyJobs, recencyCutoff); final LocalDateTime fakeNow = LocalDateTime.of(2021, 6, 20, 0, 0); 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 6d03bd2c040e..242c4267714b 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java @@ -168,8 +168,7 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con configs.getConfigDatabasePassword(), configs.getConfigDatabaseUrl()) .getAndInitialize(); - final DatabaseConfigPersistence configPersistence = new DatabaseConfigPersistence(configDatabase); - configPersistence.migrateFileConfigs(configs); + final DatabaseConfigPersistence configPersistence = new DatabaseConfigPersistence(configDatabase).migrateFileConfigs(configs); final SecretsHydrator secretsHydrator = SecretPersistence.getSecretsHydrator(configs); final Optional secretPersistence = SecretPersistence.getLongLived(configs); @@ -184,7 +183,7 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con configs.getDatabasePassword(), configs.getDatabaseUrl()) .getAndInitialize(); - final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); + final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); createDeploymentIfNoneExists(jobPersistence); 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 2aa2546a191d..e387da384786 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 @@ -68,7 +68,8 @@ public class ArchiveHandlerTest { private static final String VERSION = "0.6.8"; private static PostgreSQLContainer container; - private Database database; + private Database jobDatabase; + private Database configDatabase; private JobPersistence jobPersistence; private DatabaseConfigPersistence configPersistence; private ConfigPersistence seedPersistence; @@ -102,11 +103,11 @@ public static void dbDown() { @BeforeEach public void setup() throws Exception { - database = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - jobPersistence = new DefaultJobPersistence(database); - database = new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + jobDatabase = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + configDatabase = new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); seedPersistence = YamlSeedConfigPersistence.getDefault(); - configPersistence = new DatabaseConfigPersistence(database); + configPersistence = new DatabaseConfigPersistence(jobDatabase); configPersistence.replaceAllConfigs(Collections.emptyMap(), false); configPersistence.loadData(seedPersistence); configRepository = new ConfigRepository(configPersistence, new NoOpSecretsHydrator(), Optional.empty(), Optional.empty()); @@ -131,7 +132,8 @@ public void setup() throws Exception { @AfterEach void tearDown() throws Exception { - database.close(); + jobDatabase.close(); + configDatabase.close(); } /** diff --git a/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java b/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java index 3bef88fc15c6..864fb7da8795 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java @@ -9,6 +9,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import io.airbyte.db.Database; +import io.airbyte.db.instance.configs.ConfigsDatabaseInstance; import io.airbyte.db.instance.jobs.JobsDatabaseInstance; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; import io.airbyte.scheduler.persistence.DefaultJobPersistence; @@ -28,7 +29,8 @@ public class DatabaseArchiverTest { private static final String TEMP_PREFIX = "testDatabaseArchive"; private PostgreSQLContainer container; - private Database database; + private Database jobDatabase; + private Database configDatabase; private DatabaseArchiver databaseArchiver; @BeforeEach @@ -39,21 +41,23 @@ void setUp() throws IOException { .withPassword("docker"); container.start(); - database = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - final JobPersistence persistence = new DefaultJobPersistence(database); + jobDatabase = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + configDatabase = new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + final JobPersistence persistence = new DefaultJobPersistence(jobDatabase, configDatabase); databaseArchiver = new DatabaseArchiver(persistence); } @AfterEach void tearDown() throws Exception { - database.close(); + jobDatabase.close(); + configDatabase.close(); container.close(); } @Test void testUnknownTableExport() throws Exception { // Create a table that is not declared in JobsDatabaseSchema - database.query(ctx -> { + jobDatabase.query(ctx -> { ctx.fetch("CREATE TABLE id_and_name(id INTEGER, name VARCHAR(200), updated_at DATE);"); ctx.fetch( "INSERT INTO id_and_name (id, name, updated_at) VALUES (1,'picard', '2004-10-19'), (2, 'crusher', '2005-10-19'), (3, 'vash', '2006-10-19');"); @@ -74,7 +78,7 @@ void testUnknownTableExport() throws Exception { @Test void testDatabaseExportImport() throws Exception { - database.query(ctx -> { + jobDatabase.query(ctx -> { ctx.fetch( "INSERT INTO jobs(id, scope, config_type, config, status, created_at, started_at, updated_at) VALUES " + "(1,'get_spec_scope', 'get_spec', '{ \"type\" : \"getSpec\" }', 'succeeded', '2004-10-19', null, '2004-10-19'), " diff --git a/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java b/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java index 8fe1365079c1..97a4e010fa3a 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java @@ -11,9 +11,12 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.fail; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import com.google.common.io.Resources; import io.airbyte.commons.io.Archives; +import io.airbyte.commons.json.Jsons; +import io.airbyte.config.Configs; import io.airbyte.config.DestinationConnection; import io.airbyte.config.OperatorNormalization.Option; import io.airbyte.config.SourceConnection; @@ -24,14 +27,17 @@ import io.airbyte.config.StandardSyncOperation.OperatorType; import io.airbyte.config.StandardWorkspace; import io.airbyte.config.persistence.ConfigNotFoundException; +import io.airbyte.config.persistence.ConfigPersistence; import io.airbyte.config.persistence.ConfigRepository; -import io.airbyte.config.persistence.FileSystemConfigPersistence; +import io.airbyte.config.persistence.DatabaseConfigPersistence; import io.airbyte.config.persistence.YamlSeedConfigPersistence; import io.airbyte.config.persistence.split_secrets.MemorySecretPersistence; import io.airbyte.config.persistence.split_secrets.NoOpSecretsHydrator; import io.airbyte.config.persistence.split_secrets.SecretPersistence; import io.airbyte.db.Database; +import io.airbyte.db.instance.test.TestDatabaseProviders; import io.airbyte.migrate.Migrations; +import io.airbyte.protocol.models.ConnectorSpecification; import io.airbyte.scheduler.persistence.DefaultJobPersistence; import io.airbyte.scheduler.persistence.JobPersistence; import io.airbyte.server.RunMigration; @@ -49,10 +55,12 @@ import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; import org.jetbrains.annotations.NotNull; +import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; +import org.testcontainers.containers.PostgreSQLContainer; public class RunMigrationTest { @@ -60,9 +68,35 @@ public class RunMigrationTest { private static final String TARGET_VERSION = Migrations.MIGRATIONS.get(Migrations.MIGRATIONS.size() - 1).getVersion(); private static final String DEPRECATED_SOURCE_DEFINITION_NOT_BEING_USED = "d2147be5-fa36-4936-977e-f031affa5895"; private static final String DEPRECATED_SOURCE_DEFINITION_BEING_USED = "4eb22946-2a79-4d20-a3e6-effd234613c3"; + private static final ConnectorSpecification MOCK_CONNECTOR_SPEC = new ConnectorSpecification() + .withConnectionSpecification(Jsons.deserialize("{}")); + + private static PostgreSQLContainer container; + private static Database jobDatabase; + private static Database configDatabase; + private List resourceToBeCleanedUp; private SecretPersistence secretPersistence; + @BeforeAll + public static void prepareContainer() throws IOException { + container = new PostgreSQLContainer<>("postgres:13-alpine") + .withDatabaseName("airbyte") + .withUsername("docker") + .withPassword("docker"); + container.start(); + final TestDatabaseProviders databaseProviders = new TestDatabaseProviders(container); + jobDatabase = databaseProviders.createNewJobsDatabase(); + configDatabase = databaseProviders.createNewConfigsDatabase(); + } + + @AfterAll + public static void closeContainer() { + if (container != null) { + container.close(); + } + } + @BeforeEach public void setup() { resourceToBeCleanedUp = new ArrayList<>(); @@ -84,35 +118,48 @@ public void cleanup() throws IOException { @SuppressWarnings("UnstableApiUsage") @Test - @Disabled - // TODO(#5857): Make migration tests compatible with writing new migrations. public void testRunMigration() throws Exception { - try (final StubAirbyteDB stubAirbyteDB = new StubAirbyteDB()) { - final File file = Path - .of(Resources.getResource("migration/03a4c904-c91d-447f-ab59-27a43b52c2fd.gz").toURI()) - .toFile(); + final Path configRoot = Files.createTempDirectory(Path.of("/tmp"), "dummy_data"); + final ConfigPersistence configPersistence = getConfigPersistence(configRoot); - final Path dummyDataSource = Path.of(Resources.getResource("migration/dummy_data").toURI()); + final File file = Path + .of(Resources.getResource("migration/03a4c904-c91d-447f-ab59-27a43b52c2fd.gz").toURI()) + .toFile(); + final JobPersistence jobPersistence = getJobPersistence(file, INITIAL_VERSION); + assertPreMigrationConfigs(configPersistence, jobPersistence); - final Path configRoot = Files.createTempDirectory(Path.of("/tmp"), "dummy_data"); - FileUtils.copyDirectory(dummyDataSource.toFile(), configRoot.toFile()); - resourceToBeCleanedUp.add(configRoot.toFile()); - final JobPersistence jobPersistence = getJobPersistence(stubAirbyteDB.getDatabase(), file, INITIAL_VERSION); - assertPreMigrationConfigs(configRoot, jobPersistence); + runMigration(configPersistence, jobPersistence); - runMigration(jobPersistence, configRoot); + assertDatabaseVersion(jobPersistence, TARGET_VERSION); + assertPostMigrationConfigs(configPersistence); + FileUtils.deleteDirectory(configRoot.toFile()); + } - assertDatabaseVersion(jobPersistence, TARGET_VERSION); - assertPostMigrationConfigs(configRoot); - FileUtils.deleteDirectory(configRoot.toFile()); - } + @SuppressWarnings("UnstableApiUsage") + private ConfigPersistence getConfigPersistence(final Path configRoot) throws Exception { + final Path dummyDataSource = Path.of(Resources.getResource("migration/dummy_data").toURI()); + + FileUtils.copyDirectory(dummyDataSource.toFile(), configRoot.toFile()); + resourceToBeCleanedUp.add(configRoot.toFile()); + final Configs configs = mock(Configs.class); + when(configs.getConfigRoot()).thenReturn(configRoot); + // The database provider creates a config database that is ready to use, which contains a mock + // config + // entry. However, the migrateFileConfigs method will only copy the file configs if the database is + // empty. So we need to purge the database first. + configDatabase.transaction(ctx -> ctx.execute("TRUNCATE TABLE airbyte_configs;")); + return new DatabaseConfigPersistence(configDatabase) + .migrateFileConfigs(configs); } - private void assertPreMigrationConfigs(final Path configRoot, final JobPersistence jobPersistence) throws Exception { + private void assertPreMigrationConfigs(final ConfigPersistence configPersistence, final JobPersistence jobPersistence) throws Exception { assertDatabaseVersion(jobPersistence, INITIAL_VERSION); - final ConfigRepository configRepository = - new ConfigRepository(FileSystemConfigPersistence.createWithValidation(configRoot), new NoOpSecretsHydrator(), Optional.of(secretPersistence), - Optional.of(secretPersistence)); + final ConfigRepository configRepository = new ConfigRepository( + configPersistence, + new NoOpSecretsHydrator(), + Optional.of(secretPersistence), + Optional.of(secretPersistence)); + configRepository.setSpecFetcher(s -> MOCK_CONNECTOR_SPEC); final Map sourceDefinitionsBeforeMigration = configRepository.listStandardSourceDefinitions().stream() .collect(Collectors.toMap(c -> c.getSourceDefinitionId().toString(), c -> c)); assertTrue(sourceDefinitionsBeforeMigration.containsKey(DEPRECATED_SOURCE_DEFINITION_NOT_BEING_USED)); @@ -125,10 +172,13 @@ private void assertDatabaseVersion(final JobPersistence jobPersistence, final St assertEquals(versionFromDb.get(), version); } - private void assertPostMigrationConfigs(final Path importRoot) throws Exception { - final ConfigRepository configRepository = - new ConfigRepository(FileSystemConfigPersistence.createWithValidation(importRoot), new NoOpSecretsHydrator(), Optional.of(secretPersistence), - Optional.of(secretPersistence)); + private void assertPostMigrationConfigs(final ConfigPersistence configPersistence) throws Exception { + final ConfigRepository configRepository = new ConfigRepository( + configPersistence, + new NoOpSecretsHydrator(), + Optional.of(secretPersistence), + Optional.of(secretPersistence)); + configRepository.setSpecFetcher(s -> MOCK_CONNECTOR_SPEC); final UUID workspaceId = configRepository.listStandardWorkspaces(true).get(0).getWorkspaceId(); // originally the default workspace started with a hardcoded id. the migration in version 0.29.0 // took that id and randomized it. we want to check that the id is now NOT that hardcoded id and @@ -147,9 +197,9 @@ private void assertSourceDefinitions(final ConfigRepository configRepository) th final Map sourceDefinitions = configRepository.listStandardSourceDefinitions() .stream() .collect(Collectors.toMap(c -> c.getSourceDefinitionId().toString(), c -> c)); - assertTrue(sourceDefinitions.size() >= 59); - // the definition is not present in latest seeds so it should be deleted - assertFalse(sourceDefinitions.containsKey(DEPRECATED_SOURCE_DEFINITION_NOT_BEING_USED)); + assertTrue(sourceDefinitions.size() >= 98); + // the definition is not present in latest seeds but we no longer purge unused deprecated definition + assertTrue(sourceDefinitions.containsKey(DEPRECATED_SOURCE_DEFINITION_NOT_BEING_USED)); // the definition is not present in latest seeds but it was being used as a connection so it should // not be deleted assertTrue(sourceDefinitions.containsKey(DEPRECATED_SOURCE_DEFINITION_BEING_USED)); @@ -294,23 +344,26 @@ private void assertDestinations(final ConfigRepository configRepository, final U } } - private void runMigration(final JobPersistence jobPersistence, final Path configRoot) throws Exception { + private void runMigration(final ConfigPersistence configPersistence, final JobPersistence jobPersistence) throws Exception { + final ConfigRepository configRepository = new ConfigRepository( + configPersistence, + new NoOpSecretsHydrator(), + Optional.of(secretPersistence), + Optional.of(secretPersistence)); + configRepository.setSpecFetcher(s -> MOCK_CONNECTOR_SPEC); try (final RunMigration runMigration = new RunMigration( jobPersistence, - new ConfigRepository(FileSystemConfigPersistence.createWithValidation(configRoot), new NoOpSecretsHydrator(), Optional.of(secretPersistence), - Optional.of(secretPersistence)), + configRepository, TARGET_VERSION, YamlSeedConfigPersistence.getDefault(), - mock(SpecFetcher.class) // this test was disabled/broken when this fetcher mock was added. apologies if you have to fix this - // in the future. - )) { + mock(SpecFetcher.class))) { runMigration.run(); } } @SuppressWarnings("SameParameterValue") - private JobPersistence getJobPersistence(final Database database, final File file, final String version) throws IOException { - final DefaultJobPersistence jobPersistence = new DefaultJobPersistence(database); + private JobPersistence getJobPersistence(final File file, final String version) throws IOException { + final DefaultJobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); final Path tempFolder = Files.createTempDirectory(Path.of("/tmp"), "db_init"); resourceToBeCleanedUp.add(tempFolder.toFile()); diff --git a/airbyte-server/src/test/java/io/airbyte/server/migration/StubAirbyteDB.java b/airbyte-server/src/test/java/io/airbyte/server/migration/StubAirbyteDB.java deleted file mode 100644 index c35427637efc..000000000000 --- a/airbyte-server/src/test/java/io/airbyte/server/migration/StubAirbyteDB.java +++ /dev/null @@ -1,45 +0,0 @@ -/* - * Copyright (c) 2021 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.server.migration; - -import io.airbyte.commons.resources.MoreResources; -import io.airbyte.db.Database; -import io.airbyte.db.instance.jobs.JobsDatabaseInstance; -import java.io.IOException; -import org.testcontainers.containers.PostgreSQLContainer; - -public class StubAirbyteDB implements AutoCloseable { - - private final PostgreSQLContainer container; - private final Database database; - - public Database getDatabase() { - return database; - } - - public StubAirbyteDB() throws IOException { - container = - new PostgreSQLContainer<>("postgres:13-alpine") - .withDatabaseName("airbyte") - .withUsername("docker") - .withPassword("docker"); - container.start(); - - final String jobsDatabaseSchema = MoreResources.readResource("migration/schema.sql"); - database = new JobsDatabaseInstance( - container.getUsername(), - container.getPassword(), - container.getJdbcUrl(), - jobsDatabaseSchema) - .getAndInitialize(); - } - - @Override - public void close() throws Exception { - database.close(); - container.close(); - } - -} diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java index d3ec246d936c..3580852ee6dd 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java @@ -10,6 +10,7 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.helpers.LogClientSingleton; import io.airbyte.db.Database; +import io.airbyte.db.instance.configs.ConfigsDatabaseInstance; import io.airbyte.db.instance.jobs.JobsDatabaseInstance; import io.airbyte.scheduler.models.JobRunConfig; import io.airbyte.scheduler.persistence.DefaultJobPersistence; @@ -131,7 +132,12 @@ private void saveWorkflowIdForCancellation() throws IOException { configs.getDatabasePassword(), configs.getDatabaseUrl()) .getInitialized(); - final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); + final Database configDatabase = new ConfigsDatabaseInstance( + configs.getConfigDatabaseUser(), + configs.getConfigDatabasePassword(), + configs.getConfigDatabaseUrl()) + .getInitialized(); + final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); final String workflowId = workflowIdProvider.get(); jobPersistence.setAttemptTemporalWorkflowId(Long.parseLong(jobRunConfig.getJobId()), jobRunConfig.getAttemptId().intValue(), workflowId); } diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java index f2d0f3a70ac1..0eb7cd78540e 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java @@ -14,20 +14,15 @@ import io.airbyte.commons.functional.CheckedSupplier; import io.airbyte.config.Configs; -import io.airbyte.db.Database; -import io.airbyte.db.instance.DatabaseMigrator; -import io.airbyte.db.instance.jobs.JobsDatabaseInstance; -import io.airbyte.db.instance.jobs.JobsDatabaseMigrator; +import io.airbyte.db.instance.test.TestDatabaseProviders; import io.airbyte.scheduler.models.JobRunConfig; import io.airbyte.workers.Worker; import io.temporal.internal.common.CheckedExceptionWrapper; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; -import java.sql.SQLException; import java.util.function.Consumer; import org.junit.jupiter.api.AfterAll; -import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -42,9 +37,8 @@ class TemporalAttemptExecutionTest { private static final String SOURCE_USERNAME = "sourceusername"; private static final String SOURCE_PASSWORD = "hunter2"; - private static PostgreSQLContainer container; + private static PostgreSQLContainer container; private static Configs configs; - private static Database database; private Path jobRoot; @@ -54,8 +48,8 @@ class TemporalAttemptExecutionTest { private TemporalAttemptExecution attemptExecution; @BeforeAll - static void setUpAll() throws IOException { - container = new PostgreSQLContainer("postgres:13-alpine") + static void setUpAll() { + container = new PostgreSQLContainer<>("postgres:13-alpine") .withUsername(SOURCE_USERNAME) .withPassword(SOURCE_PASSWORD); container.start(); @@ -63,23 +57,18 @@ static void setUpAll() throws IOException { when(configs.getDatabaseUrl()).thenReturn(container.getJdbcUrl()); when(configs.getDatabaseUser()).thenReturn(SOURCE_USERNAME); when(configs.getDatabasePassword()).thenReturn(SOURCE_PASSWORD); - - // create the initial schema - database = new JobsDatabaseInstance( - configs.getDatabaseUser(), - configs.getDatabasePassword(), - configs.getDatabaseUrl()) - .getAndInitialize(); - - // make sure schema is up-to-date - final DatabaseMigrator jobDbMigrator = new JobsDatabaseMigrator(database, "test"); - jobDbMigrator.createBaseline(); - jobDbMigrator.migrate(); + when(configs.getConfigDatabaseUrl()).thenReturn(container.getJdbcUrl()); + when(configs.getConfigDatabaseUser()).thenReturn(SOURCE_USERNAME); + when(configs.getConfigDatabasePassword()).thenReturn(SOURCE_PASSWORD); } @SuppressWarnings("unchecked") @BeforeEach void setup() throws IOException { + final TestDatabaseProviders databaseProviders = new TestDatabaseProviders(container); + databaseProviders.createNewJobsDatabase(); + databaseProviders.createNewConfigsDatabase(); + final Path workspaceRoot = Files.createTempDirectory(Path.of("/tmp"), "temporal_attempt_execution_test"); jobRoot = workspaceRoot.resolve(JOB_ID).resolve(String.valueOf(ATTEMPT_ID)); @@ -96,13 +85,6 @@ void setup() throws IOException { configs); } - @AfterEach - void tearDown() throws SQLException { - database.query(ctx -> ctx.execute("TRUNCATE TABLE jobs")); - database.query(ctx -> ctx.execute("TRUNCATE TABLE attempts")); - database.query(ctx -> ctx.execute("TRUNCATE TABLE airbyte_metadata")); - } - @AfterAll static void tearDownAll() { container.close(); From 56f6c1334e5ced872168ccdc0bb56bb8d9d30a4e Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Tue, 19 Oct 2021 11:45:13 -0700 Subject: [PATCH 06/23] temp --- .../persistence/DefaultJobPersistence.java | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java index f8cbef5752f4..297e76ed59d0 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java @@ -193,8 +193,7 @@ public void failJob(final long jobId) throws IOException { }); } - private void updateJobStatusIfNotInTerminalState( - final DSLContext ctx, + private void updateJobStatusIfNotInTerminalState(final DSLContext ctx, final long jobId, final JobStatus newStatus, final LocalDateTime now, @@ -226,16 +225,16 @@ public int createAttempt(final long jobId, final Path logPath) throws IOExceptio return jobDatabase.transaction(ctx -> { final Job job = getJob(ctx, jobId); if (job.isJobInTerminalState()) { - final var errMsg = - String.format("Cannot create an attempt for a job id: %s that is in a terminal state: %s for connection id: %s", job.getId(), - job.getStatus(), job.getScope()); + final var errMsg = String.format( + "Cannot create an attempt for a job id: %s that is in a terminal state: %s for connection id: %s", + job.getId(), job.getStatus(), job.getScope()); throw new IllegalStateException(errMsg); } if (job.hasRunningAttempt()) { - final var errMsg = - String.format("Cannot create an attempt for a job id: %s that has a running attempt: %s for connection id: %s", job.getId(), - job.getStatus(), job.getScope()); + final var errMsg = String.format( + "Cannot create an attempt for a job id: %s that has a running attempt: %s for connection id: %s", + job.getId(), job.getStatus(), job.getScope()); throw new IllegalStateException(errMsg); } @@ -799,4 +798,6 @@ private static Table getTable(final String schema, final String tableNam return DSL.table(String.format("%s.%s", schema, tableName)); } + // TODO: add a method to update attempt state in job database, and another in config database + } From 7b655233289355c5bd2d1485cb2942d76e3a03e9 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Wed, 20 Oct 2021 09:30:55 -0700 Subject: [PATCH 07/23] Rebase on master --- .../V0_29_21_001__Store_last_sync_state.java | 8 +- ...29_21_001__Store_last_sync_state_test.java | 18 +++-- .../persistence/DefaultJobPersistence.java | 79 ++++++++++++++++--- .../DefaultJobPersistenceTest.java | 13 ++- 4 files changed, 92 insertions(+), 26 deletions(-) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java index 782d5d7f6848..9ab05a77b19a 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java @@ -32,7 +32,9 @@ import org.slf4j.LoggerFactory; /** - * Create a new table to store the latest job state for each standard sync. Issue: + * Create a new table to store the latest job state for each standard sync. + *
  • Column sync_id: the connectionId in StandardSync
  • Column state: a json node representing + * a State object
    • */ public class V0_29_21_001__Store_last_sync_state extends BaseJavaMigration { @@ -55,7 +57,7 @@ public V0_29_21_001__Store_last_sync_state() { } @VisibleForTesting - V0_29_21_001__Store_last_sync_state(Configs configs) { + V0_29_21_001__Store_last_sync_state(final Configs configs) { this.configs = configs; } @@ -87,7 +89,7 @@ static void createTable(final DSLContext ctx) { @VisibleForTesting static void copyData(final DSLContext ctx, final Map syncToStateMap, final OffsetDateTime timestamp) { - for (Map.Entry entry : syncToStateMap.entrySet()) { + for (final Map.Entry entry : syncToStateMap.entrySet()) { ctx.insertInto(SYNC_STATE_TABLE) .set(COLUMN_SYNC_ID, UUID.fromString(entry.getKey())) .set(COLUMN_STATE, JSONB.valueOf(Jsons.serialize(entry.getValue()))) diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java index 9e60f7fa5a20..b54ba3e87be6 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java @@ -66,6 +66,8 @@ class V0_29_21_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTe private static final UUID SYNC_1_ID = UUID.randomUUID(); private static final UUID SYNC_2_ID = UUID.randomUUID(); private static final UUID SYNC_3_ID = UUID.randomUUID(); + // these are State objects, see State.yaml for its schema; + // we cannot construct the POJO directly because State is defined in an downstream module private static final JsonNode SYNC_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }"); private static final JsonNode SYNC_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }"); @@ -104,12 +106,12 @@ public void testGetSyncToStateMap() throws Exception { createJob(ctx, SYNC_1_ID); // The second job has one attempt. - long job2 = createJob(ctx, SYNC_2_ID); + final long job2 = createJob(ctx, SYNC_2_ID); createAttempt(ctx, job2, 1, createAttemptOutput(SYNC_2_STATE)); // The third job has multiple attempts. The third attempt has the latest state. - long job3 = createJob(ctx, SYNC_3_ID); - JsonNode attempt31State = Jsons.deserialize("{ \"state\": { \"cursor\": 31 } }"); + final long job3 = createJob(ctx, SYNC_3_ID); + final JsonNode attempt31State = Jsons.deserialize("{ \"state\": { \"cursor\": 31 } }"); createAttempt(ctx, job3, 1, createAttemptOutput(attempt31State)); createAttempt(ctx, job3, 2, null); createAttempt(ctx, job3, 3, createAttemptOutput(SYNC_3_STATE)); @@ -205,7 +207,7 @@ public Configuration getConfiguration() { public Connection getConnection() { try { return database.getDataSource().getConnection(); - } catch (SQLException e) { + } catch (final SQLException e) { throw new RuntimeException(e); } } @@ -221,7 +223,7 @@ public Connection getConnection() { /** * Create a job record whose scope equals to the passed in standard sync id, and return the job id. */ - private static long createJob(DSLContext ctx, UUID standardSyncId) { + private static long createJob(final DSLContext ctx, final UUID standardSyncId) { final int insertCount = ctx.insertInto(JOBS_TABLE, JOB_SCOPE_FIELD) .values(standardSyncId.toString()) .execute(); @@ -234,7 +236,7 @@ private static long createJob(DSLContext ctx, UUID standardSyncId) { .get(JOB_ID_FIELD); } - private static void createAttempt(DSLContext ctx, long jobId, int attemptNumber, JsonNode attemptOutput) { + private static void createAttempt(final DSLContext ctx, final long jobId, final int attemptNumber, final JsonNode attemptOutput) { final int insertCount = ctx.insertInto(ATTEMPTS_TABLE, ATTEMPT_JOB_ID_FIELD, ATTEMPT_NUMBER_FIELD, ATTEMPT_OUTPUT_FIELD) .values(jobId, attemptNumber, JSONB.valueOf(Jsons.serialize(attemptOutput))) .execute(); @@ -252,7 +254,7 @@ private static void createAttempt(DSLContext ctx, long jobId, int attemptNumber, * * @param state The state object within a StandardSyncOutput. */ - private static JsonNode createAttemptOutput(JsonNode state) { + private static JsonNode createAttemptOutput(final JsonNode state) { final ObjectNode standardSyncOutput = OBJECT_MAPPER.createObjectNode() .set("state", state); return OBJECT_MAPPER.createObjectNode() @@ -280,7 +282,7 @@ private static void checkTable(final DSLContext ctx, final boolean tableExists) private static void checkSyncStates(final DSLContext ctx, final Map expectedSyncStates, - @Nullable OffsetDateTime expectedTimestamp) { + @Nullable final OffsetDateTime expectedTimestamp) { for (final Map.Entry entry : expectedSyncStates.entrySet()) { final var record = ctx .select(V0_29_21_001__Store_last_sync_state.COLUMN_STATE, diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java index 297e76ed59d0..50dd07ccd110 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java @@ -4,6 +4,10 @@ package io.airbyte.scheduler.persistence; +import static io.airbyte.db.instance.configs.jooq.Tables.SYNC_STATE; +import static io.airbyte.db.instance.jobs.jooq.Tables.ATTEMPTS; +import static io.airbyte.db.instance.jobs.jooq.Tables.JOBS; + import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.JsonNodeType; import com.fasterxml.jackson.databind.node.ObjectNode; @@ -21,7 +25,6 @@ import io.airbyte.config.State; import io.airbyte.db.Database; import io.airbyte.db.ExceptionWrappingDatabase; -import io.airbyte.db.instance.configs.jooq.tables.SyncState; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; import io.airbyte.scheduler.models.Attempt; import io.airbyte.scheduler.models.AttemptStatus; @@ -32,6 +35,7 @@ import java.nio.file.Path; import java.time.Instant; import java.time.LocalDateTime; +import java.time.OffsetDateTime; import java.time.ZoneOffset; import java.time.ZonedDateTime; import java.time.format.DateTimeFormatter; @@ -318,15 +322,66 @@ public Optional getAttemptTemporalWorkflowId(final long jobId, final int @Override public void writeOutput(final long jobId, final int attemptNumber, final T output) throws IOException { - final LocalDateTime now = LocalDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); + final OffsetDateTime now = OffsetDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); + writeOutputToAttemptTable(jobId, attemptNumber, output, now); + writeOutputToSyncStateTable(jobId, output, now); + } - jobDatabase.query( - ctx -> ctx.execute( - "UPDATE attempts SET output = CAST(? as JSONB), updated_at = ? WHERE job_id = ? AND attempt_number = ?", - Jsons.serialize(output), - now, - jobId, - attemptNumber)); + private void writeOutputToAttemptTable(final long jobId, + final int attemptNumber, + final T output, + final OffsetDateTime now) + throws IOException { + jobDatabase.transaction( + ctx -> ctx.update(ATTEMPTS) + .set(ATTEMPTS.OUTPUT, JSONB.valueOf(Jsons.serialize(output))) + .set(ATTEMPTS.UPDATED_AT, now) + .where(ATTEMPTS.JOB_ID.eq(jobId), ATTEMPTS.ATTEMPT_NUMBER.eq(attemptNumber)) + .execute()); + } + + private void writeOutputToSyncStateTable(final long jobId, final T output, final OffsetDateTime now) throws IOException { + if (!(output instanceof JobOutput)) { + return; + } + final JobOutput jobOutput = (JobOutput) output; + if (jobOutput.getSync() == null) { + return; + } + + final Record1 jobConnectionId = jobDatabase.query(ctx -> ctx + .select(JOBS.SCOPE) + .from(JOBS) + .where(JOBS.ID.eq(jobId)) + .fetchAny()); + final State syncState = jobOutput.getSync().getState(); + + if (jobConnectionId == null) { + LOGGER.error("No job can be found for id {}", jobId); + return; + } + + final UUID connectionId = UUID.fromString(jobConnectionId.value1()); + configDatabase.transaction( + ctx -> { + final boolean hasExistingRecord = ctx.fetchExists(SYNC_STATE, SYNC_STATE.SYNC_ID.eq(connectionId)); + if (hasExistingRecord) { + LOGGER.info("Updating connection {} state", connectionId); + return ctx.update(SYNC_STATE) + .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(syncState))) + .set(SYNC_STATE.UPDATED_AT, now) + .where(SYNC_STATE.SYNC_ID.eq(connectionId)) + .execute(); + } else { + LOGGER.info("Inserting new state for connection {}", connectionId); + return ctx.insertInto(SYNC_STATE) + .set(SYNC_STATE.SYNC_ID, UUID.fromString(jobConnectionId.value1())) + .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(syncState))) + .set(SYNC_STATE.CREATED_AT, now) + .set(SYNC_STATE.UPDATED_AT, now) + .execute(); + } + }); } @Override @@ -396,9 +451,9 @@ public Optional getLastReplicationJob(final UUID connectionId) throws IOExc @Override public Optional getCurrentState(final UUID connectionId) throws IOException { return configDatabase.query(ctx -> { - final Record1 record = ctx.select(SyncState.SYNC_STATE.STATE) - .from(SyncState.SYNC_STATE) - .where(SyncState.SYNC_STATE.SYNC_ID.eq(connectionId)) + final Record1 record = ctx.select(SYNC_STATE.STATE) + .from(SYNC_STATE) + .where(SYNC_STATE.SYNC_ID.eq(connectionId)) .fetchAny(); if (record == null) { return Optional.empty(); diff --git a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java index 48978670f816..08c950c1cea5 100644 --- a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java +++ b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java @@ -4,6 +4,10 @@ package io.airbyte.scheduler.persistence; +import static io.airbyte.db.instance.configs.jooq.Tables.SYNC_STATE; +import static io.airbyte.db.instance.jobs.jooq.Tables.AIRBYTE_METADATA; +import static io.airbyte.db.instance.jobs.jooq.Tables.ATTEMPTS; +import static io.airbyte.db.instance.jobs.jooq.Tables.JOBS; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; @@ -175,9 +179,12 @@ void tearDown() throws Exception { private void resetDb() throws SQLException { // todo (cgardens) - truncate whole db. - jobDatabase.query(ctx -> ctx.execute("TRUNCATE TABLE jobs")); - jobDatabase.query(ctx -> ctx.execute("TRUNCATE TABLE attempts")); - jobDatabase.query(ctx -> ctx.execute("TRUNCATE TABLE airbyte_metadata")); + jobDatabase.query(ctx -> ctx.truncateTable(JOBS).execute()); + jobDatabase.query(ctx -> ctx.truncateTable(ATTEMPTS).execute()); + jobDatabase.query(ctx -> ctx.truncateTable(AIRBYTE_METADATA).execute()); + // the airbyte_configs table cannot be truncated because the config database + // is considered ready for consumption only when there are records in this table + configDatabase.query(ctx -> ctx.truncateTable(SYNC_STATE).execute()); } private Result getJobRecord(final long jobId) throws SQLException { From bba689a2636d922fd129c4317d64eb3a3e2c9eb7 Mon Sep 17 00:00:00 2001 From: Charles Date: Sun, 24 Oct 2021 09:18:51 -0700 Subject: [PATCH 08/23] Save state in temporal (#7253) --- .../persistence/ConfigPersistence2.java | 83 +++++++++++++ .../persistence/ConfigPersistence2Test.java | 7 ++ .../airbyte/scheduler/app/JobScheduler.java | 4 +- .../airbyte/scheduler/app/SchedulerApp.java | 24 +++- .../worker_run/TemporalWorkerRunFactory.java | 12 +- .../persistence/DefaultJobCreator.java | 7 +- .../persistence/DefaultJobPersistence.java | 75 +----------- .../scheduler/persistence/JobPersistence.java | 14 --- .../DefaultJobPersistenceTest.java | 112 +----------------- .../server/ConfigurationApiFactory.java | 5 + .../java/io/airbyte/server/ServerApp.java | 8 +- .../java/io/airbyte/server/ServerFactory.java | 4 + .../airbyte/server/apis/ConfigurationApi.java | 3 + .../server/handlers/SchedulerHandler.java | 10 +- .../server/handlers/ArchiveHandlerTest.java | 2 +- .../server/handlers/SchedulerHandlerTest.java | 8 +- .../migration/DatabaseArchiverTest.java | 2 +- .../server/migration/RunMigrationTest.java | 2 +- .../java/io/airbyte/workers/WorkerApp.java | 24 +++- .../workers/temporal/SyncWorkflow.java | 75 +++++++++--- .../temporal/TemporalAttemptExecution.java | 2 +- .../workers/temporal/TemporalClient.java | 5 +- 22 files changed, 248 insertions(+), 240 deletions(-) create mode 100644 airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java create mode 100644 airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java new file mode 100644 index 000000000000..c15bdafd72d5 --- /dev/null +++ b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java @@ -0,0 +1,83 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.config.persistence; + +import static io.airbyte.db.instance.configs.jooq.Tables.SYNC_STATE; + +import com.google.common.annotations.VisibleForTesting; +import io.airbyte.commons.json.Jsons; +import io.airbyte.config.State; +import io.airbyte.db.Database; +import io.airbyte.db.ExceptionWrappingDatabase; +import java.io.IOException; +import java.time.Instant; +import java.time.OffsetDateTime; +import java.time.ZoneOffset; +import java.util.Optional; +import java.util.UUID; +import java.util.function.Supplier; +import org.jooq.JSONB; +import org.jooq.Record1; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class ConfigPersistence2 { + + private static final Logger LOGGER = LoggerFactory.getLogger(ConfigPersistence2.class); + + private final ExceptionWrappingDatabase configDatabase; + private final Supplier timeSupplier; + + public ConfigPersistence2(final Database configDatabase) { + this(configDatabase, Instant::now); + } + + @VisibleForTesting + ConfigPersistence2(final Database configDatabase, final Supplier timeSupplier) { + this.configDatabase = new ExceptionWrappingDatabase(configDatabase); + + this.timeSupplier = timeSupplier; + } + + public Optional getCurrentState(final UUID connectionId) throws IOException { + return configDatabase.query(ctx -> { + final Record1 record = ctx.select(SYNC_STATE.STATE) + .from(SYNC_STATE) + .where(SYNC_STATE.SYNC_ID.eq(connectionId)) + .fetchAny(); + if (record == null) { + return Optional.empty(); + } + + return Optional.of(Jsons.deserialize(record.value1().data(), State.class)); + }); + } + + public void updateSyncState(final UUID connectionId, final State state) throws IOException { + final OffsetDateTime now = OffsetDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); + + configDatabase.transaction( + ctx -> { + final boolean hasExistingRecord = ctx.fetchExists(SYNC_STATE, SYNC_STATE.SYNC_ID.eq(connectionId)); + if (hasExistingRecord) { + LOGGER.info("Updating connection {} state", connectionId); + return ctx.update(SYNC_STATE) + .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(state))) + .set(SYNC_STATE.UPDATED_AT, now) + .where(SYNC_STATE.SYNC_ID.eq(connectionId)) + .execute(); + } else { + LOGGER.info("Inserting new state for connection {}", connectionId); + return ctx.insertInto(SYNC_STATE) + .set(SYNC_STATE.SYNC_ID, connectionId) + .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(state))) + .set(SYNC_STATE.CREATED_AT, now) + .set(SYNC_STATE.UPDATED_AT, now) + .execute(); + } + }); + } + +} diff --git a/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java b/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java new file mode 100644 index 000000000000..fca39e71cf21 --- /dev/null +++ b/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java @@ -0,0 +1,7 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.config.persistence; + +class ConfigPersistence2Test {} diff --git a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java index ddd943bb8c36..d6004a797ef6 100644 --- a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java +++ b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java @@ -9,6 +9,7 @@ import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSync.Status; import io.airbyte.config.persistence.ConfigNotFoundException; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.scheduler.models.Job; import io.airbyte.scheduler.persistence.DefaultJobCreator; @@ -47,6 +48,7 @@ public class JobScheduler implements Runnable { } public JobScheduler(final JobPersistence jobPersistence, + final ConfigPersistence2 configPersistence2, final ConfigRepository configRepository, final TrackingClient trackingClient) { this( @@ -54,7 +56,7 @@ public JobScheduler(final JobPersistence jobPersistence, configRepository, new ScheduleJobPredicate(Instant::now), new DefaultSyncJobFactory( - new DefaultJobCreator(jobPersistence), + new DefaultJobCreator(jobPersistence, configPersistence2), configRepository, new OAuthConfigSupplier(configRepository, false, trackingClient))); } diff --git a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java index 64a59a4667dd..0793e16a165d 100644 --- a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java +++ b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java @@ -18,6 +18,7 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.helpers.LogClientSingleton; import io.airbyte.config.persistence.ConfigPersistence; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.config.persistence.DatabaseConfigPersistence; import io.airbyte.config.persistence.split_secrets.SecretPersistence; @@ -75,6 +76,7 @@ public class SchedulerApp { private final Path workspaceRoot; private final JobPersistence jobPersistence; private final ConfigRepository configRepository; + private final ConfigPersistence2 configPersistence2; private final JobCleaner jobCleaner; private final JobNotifier jobNotifier; private final TemporalClient temporalClient; @@ -85,6 +87,7 @@ public class SchedulerApp { public SchedulerApp(final Path workspaceRoot, final JobPersistence jobPersistence, final ConfigRepository configRepository, + final ConfigPersistence2 configPersistence2, final JobCleaner jobCleaner, final JobNotifier jobNotifier, final TemporalClient temporalClient, @@ -94,6 +97,7 @@ public SchedulerApp(final Path workspaceRoot, this.workspaceRoot = workspaceRoot; this.jobPersistence = jobPersistence; this.configRepository = configRepository; + this.configPersistence2 = configPersistence2; this.jobCleaner = jobCleaner; this.jobNotifier = jobNotifier; this.temporalClient = temporalClient; @@ -110,7 +114,7 @@ public void start() throws IOException { final TemporalWorkerRunFactory temporalWorkerRunFactory = new TemporalWorkerRunFactory(temporalClient, workspaceRoot, airbyteVersionOrWarnings); final JobRetrier jobRetrier = new JobRetrier(jobPersistence, Instant::now, jobNotifier, maxSyncJobAttempts); final TrackingClient trackingClient = TrackingClientSingleton.get(); - final JobScheduler jobScheduler = new JobScheduler(jobPersistence, configRepository, trackingClient); + final JobScheduler jobScheduler = new JobScheduler(jobPersistence, configPersistence2, configRepository, trackingClient); final JobSubmitter jobSubmitter = new JobSubmitter( workerThreadPool, jobPersistence, @@ -209,18 +213,19 @@ public static void main(final String[] args) throws IOException, InterruptedExce configs.getConfigDatabasePassword(), configs.getConfigDatabaseUrl()) .getInitialized(); + final ConfigPersistence2 configPersistence2 = new ConfigPersistence2(configDatabase); final ConfigPersistence configPersistence = new DatabaseConfigPersistence(configDatabase).withValidation(); final Optional secretPersistence = SecretPersistence.getLongLived(configs); final Optional ephemeralSecretPersistence = SecretPersistence.getEphemeral(configs); final SecretsHydrator secretsHydrator = SecretPersistence.getSecretsHydrator(configs); final ConfigRepository configRepository = new ConfigRepository(configPersistence, secretsHydrator, secretPersistence, ephemeralSecretPersistence); - final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); + final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); final JobCleaner jobCleaner = new JobCleaner( configs.getWorkspaceRetentionConfig(), workspaceRoot, jobPersistence); - AirbyteVersion.assertIsCompatible(configs.getAirbyteVersion(), jobPersistence.getVersion().get()); + AirbyteVersion.assertIsCompatible(configs.getAirbyteVersion(), jobPersistence.getVersion().orElseThrow()); TrackingClientSingleton.initialize( configs.getTrackingStrategy(), @@ -239,8 +244,17 @@ public static void main(final String[] args) throws IOException, InterruptedExce MetricSingleton.initializeMonitoringServiceDaemon("8082", mdc, configs.getPublishMetrics()); LOGGER.info("Launching scheduler..."); - new SchedulerApp(workspaceRoot, jobPersistence, configRepository, jobCleaner, jobNotifier, temporalClient, - Integer.parseInt(configs.getSubmitterNumThreads()), configs.getMaxSyncJobAttempts(), configs.getAirbyteVersionOrWarning()) + new SchedulerApp( + workspaceRoot, + jobPersistence, + configRepository, + configPersistence2, + jobCleaner, + jobNotifier, + temporalClient, + Integer.parseInt(configs.getSubmitterNumThreads()), + configs.getMaxSyncJobAttempts(), + configs.getAirbyteVersionOrWarning()) .start(); } diff --git a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactory.java b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactory.java index 579e7b5e9266..65aed379df31 100644 --- a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactory.java +++ b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactory.java @@ -20,6 +20,7 @@ import io.airbyte.workers.temporal.TemporalJobType; import io.airbyte.workers.temporal.TemporalResponse; import java.nio.file.Path; +import java.util.UUID; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,9 +45,14 @@ public WorkerRun create(final Job job) { public CheckedSupplier, Exception> createSupplier(final Job job, final int attemptId) { final TemporalJobType temporalJobType = toTemporalJobType(job.getConfigType()); + final UUID connectionId = UUID.fromString(job.getScope()); return switch (job.getConfigType()) { case SYNC -> () -> { - final TemporalResponse output = temporalClient.submitSync(job.getId(), attemptId, job.getConfig().getSync()); + final TemporalResponse output = temporalClient.submitSync( + job.getId(), + attemptId, + job.getConfig().getSync(), + connectionId); return toOutputAndStatus(output); }; case RESET_CONNECTION -> () -> { @@ -63,7 +69,7 @@ public CheckedSupplier, Exception> createSupplier(fin .withOperationSequence(resetConnection.getOperationSequence()) .withResourceRequirements(resetConnection.getResourceRequirements()); - final TemporalResponse output = temporalClient.submitSync(job.getId(), attemptId, config); + final TemporalResponse output = temporalClient.submitSync(job.getId(), attemptId, config, connectionId); return toOutputAndStatus(output); }; default -> throw new IllegalArgumentException("Does not support job type: " + temporalJobType); @@ -84,7 +90,7 @@ private OutputAndStatus toOutputAndStatus(final TemporalResponse createSyncJob(final SourceConnection source, .withState(null) .withResourceRequirements(standardSync.getResourceRequirements()); - jobPersistence.getCurrentState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState); + configPersistence2.getCurrentState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState); final JobConfig jobConfig = new JobConfig() .withConfigType(ConfigType.SYNC) diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java index 50dd07ccd110..4313844715fd 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java @@ -4,9 +4,7 @@ package io.airbyte.scheduler.persistence; -import static io.airbyte.db.instance.configs.jooq.Tables.SYNC_STATE; import static io.airbyte.db.instance.jobs.jooq.Tables.ATTEMPTS; -import static io.airbyte.db.instance.jobs.jooq.Tables.JOBS; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.JsonNodeType; @@ -22,7 +20,6 @@ import io.airbyte.config.JobConfig; import io.airbyte.config.JobConfig.ConfigType; import io.airbyte.config.JobOutput; -import io.airbyte.config.State; import io.airbyte.db.Database; import io.airbyte.db.ExceptionWrappingDatabase; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; @@ -60,7 +57,6 @@ import org.jooq.JSONFormat.RecordFormat; import org.jooq.Named; import org.jooq.Record; -import org.jooq.Record1; import org.jooq.Result; import org.jooq.Sequence; import org.jooq.Table; @@ -112,29 +108,23 @@ public class DefaultJobPersistence implements JobPersistence { "ORDER BY jobs.created_at DESC, jobs.id DESC, attempts.created_at ASC, attempts.id ASC "; private final ExceptionWrappingDatabase jobDatabase; - private final ExceptionWrappingDatabase configDatabase; private final Supplier timeSupplier; @VisibleForTesting DefaultJobPersistence(final Database jobDatabase, - final Database configDatabase, final Supplier timeSupplier, final int minimumAgeInDays, final int excessiveNumberOfJobs, final int minimumRecencyCount) { this.jobDatabase = new ExceptionWrappingDatabase(jobDatabase); - this.configDatabase = new ExceptionWrappingDatabase(configDatabase); this.timeSupplier = timeSupplier; JOB_HISTORY_MINIMUM_AGE_IN_DAYS = minimumAgeInDays; JOB_HISTORY_EXCESSIVE_NUMBER_OF_JOBS = excessiveNumberOfJobs; JOB_HISTORY_MINIMUM_RECENCY = minimumRecencyCount; } - /** - * @param configDatabase The config database is only needed for reading and writing sync state. - */ - public DefaultJobPersistence(final Database jobDatabase, final Database configDatabase) { - this(jobDatabase, configDatabase, Instant::now, 30, 500, 10); + public DefaultJobPersistence(final Database jobDatabase) { + this(jobDatabase, Instant::now, 30, 500, 10); } /** @@ -324,7 +314,7 @@ public Optional getAttemptTemporalWorkflowId(final long jobId, final int public void writeOutput(final long jobId, final int attemptNumber, final T output) throws IOException { final OffsetDateTime now = OffsetDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); writeOutputToAttemptTable(jobId, attemptNumber, output, now); - writeOutputToSyncStateTable(jobId, output, now); + // writeOutputToSyncStateTable(jobId, output, now); } private void writeOutputToAttemptTable(final long jobId, @@ -340,50 +330,6 @@ private void writeOutputToAttemptTable(final long jobId, .execute()); } - private void writeOutputToSyncStateTable(final long jobId, final T output, final OffsetDateTime now) throws IOException { - if (!(output instanceof JobOutput)) { - return; - } - final JobOutput jobOutput = (JobOutput) output; - if (jobOutput.getSync() == null) { - return; - } - - final Record1 jobConnectionId = jobDatabase.query(ctx -> ctx - .select(JOBS.SCOPE) - .from(JOBS) - .where(JOBS.ID.eq(jobId)) - .fetchAny()); - final State syncState = jobOutput.getSync().getState(); - - if (jobConnectionId == null) { - LOGGER.error("No job can be found for id {}", jobId); - return; - } - - final UUID connectionId = UUID.fromString(jobConnectionId.value1()); - configDatabase.transaction( - ctx -> { - final boolean hasExistingRecord = ctx.fetchExists(SYNC_STATE, SYNC_STATE.SYNC_ID.eq(connectionId)); - if (hasExistingRecord) { - LOGGER.info("Updating connection {} state", connectionId); - return ctx.update(SYNC_STATE) - .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(syncState))) - .set(SYNC_STATE.UPDATED_AT, now) - .where(SYNC_STATE.SYNC_ID.eq(connectionId)) - .execute(); - } else { - LOGGER.info("Inserting new state for connection {}", connectionId); - return ctx.insertInto(SYNC_STATE) - .set(SYNC_STATE.SYNC_ID, UUID.fromString(jobConnectionId.value1())) - .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(syncState))) - .set(SYNC_STATE.CREATED_AT, now) - .set(SYNC_STATE.UPDATED_AT, now) - .execute(); - } - }); - } - @Override public Job getJob(final long jobId) throws IOException { return jobDatabase.query(ctx -> getJob(ctx, jobId)); @@ -448,21 +394,6 @@ public Optional getLastReplicationJob(final UUID connectionId) throws IOExc .flatMap(r -> getJobOptional(ctx, r.get("job_id", Long.class)))); } - @Override - public Optional getCurrentState(final UUID connectionId) throws IOException { - return configDatabase.query(ctx -> { - final Record1 record = ctx.select(SYNC_STATE.STATE) - .from(SYNC_STATE) - .where(SYNC_STATE.SYNC_ID.eq(connectionId)) - .fetchAny(); - if (record == null) { - return Optional.empty(); - } - - return Optional.of(Jsons.deserialize(record.value1().data(), State.class)); - }); - } - @Override public Optional getNextJob() throws IOException { // rules: diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java index 825936ef4b41..26260753e566 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java @@ -7,7 +7,6 @@ import com.fasterxml.jackson.databind.JsonNode; import io.airbyte.config.JobConfig; import io.airbyte.config.JobConfig.ConfigType; -import io.airbyte.config.State; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; import io.airbyte.scheduler.models.Job; import io.airbyte.scheduler.models.JobStatus; @@ -147,19 +146,6 @@ public interface JobPersistence { Optional getLastReplicationJob(UUID connectionId) throws IOException; - /** - * if a job does not succeed, we assume that it synced nothing. that is the most conservative - * assumption we can make. as long as all destinations write the final data output in a - * transactional way, this will be true. if this changes, then we may end up writing duplicate data - * with our incremental append only. this is preferable to failing to send data at all. our - * incremental append only most closely resembles a deliver at least once strategy anyway. - * - * @param connectionId - id of the connection whose state we want to fetch. - * @return the current state, if any of, the connection - * @throws IOException exception due to interaction with persistence - */ - Optional getCurrentState(UUID connectionId) throws IOException; - Optional getNextJob() throws IOException; /// ARCHIVE diff --git a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java index 08c950c1cea5..43836c83336f 100644 --- a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java +++ b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java @@ -18,7 +18,6 @@ import static org.mockito.Mockito.when; import com.fasterxml.jackson.databind.JsonNode; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.collect.Sets; import io.airbyte.commons.json.Jsons; @@ -27,10 +26,7 @@ import io.airbyte.config.JobConfig.ConfigType; import io.airbyte.config.JobGetSpecConfig; import io.airbyte.config.JobOutput; -import io.airbyte.config.JobOutput.OutputType; import io.airbyte.config.JobSyncConfig; -import io.airbyte.config.StandardSyncOutput; -import io.airbyte.config.State; import io.airbyte.db.Database; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; import io.airbyte.db.instance.test.TestDatabaseProviders; @@ -163,13 +159,12 @@ private static Job createJob( public void setup() throws Exception { final TestDatabaseProviders databaseProviders = new TestDatabaseProviders(container); jobDatabase = databaseProviders.createNewJobsDatabase(); - configDatabase = databaseProviders.createNewConfigsDatabase(); resetDb(); timeSupplier = mock(Supplier.class); when(timeSupplier.get()).thenReturn(NOW); - jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase, timeSupplier, 30, 500, 10); + jobPersistence = new DefaultJobPersistence(jobDatabase, timeSupplier, 30, 500, 10); } @AfterEach @@ -340,7 +335,7 @@ void testListJobsWithTimestamp() throws IOException { now.plusSeconds(14), now.plusSeconds(15), now.plusSeconds(16)); - jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase, timeSupplier, 30, 500, 10); + jobPersistence = new DefaultJobPersistence(jobDatabase, timeSupplier, 30, 500, 10); final long syncJobId = jobPersistence.enqueueJob(SCOPE, SYNC_JOB_CONFIG).orElseThrow(); final int syncJobAttemptNumber0 = jobPersistence.createAttempt(syncJobId, LOG_PATH); jobPersistence.failAttempt(syncJobId, syncJobAttemptNumber0); @@ -728,107 +723,6 @@ public void testGetLastSyncJobForConnectionId() throws IOException { } - @Nested - @DisplayName("When getting current state") - class GetCurrentState { - - @Test - @DisplayName("Should take latest state (regardless of attempt status)") - void testGetCurrentStateWithMultipleAttempts() throws IOException { - // just have the time monotonically increase. - when(timeSupplier.get()).thenAnswer(a -> Instant.now()); - - // create a job - final long jobId = jobPersistence.enqueueJob(SCOPE, SYNC_JOB_CONFIG).orElseThrow(); - - // fail the first attempt - jobPersistence.failAttempt(jobId, jobPersistence.createAttempt(jobId, LOG_PATH)); - assertTrue(jobPersistence.getCurrentState(UUID.fromString(SCOPE)).isEmpty()); - - // succeed the second attempt - final State state = new State().withState(Jsons.jsonNode(ImmutableMap.of("checkpoint", 4))); - final JobOutput jobOutput = new JobOutput().withOutputType(OutputType.SYNC).withSync(new StandardSyncOutput().withState(state)); - final int attemptId = jobPersistence.createAttempt(jobId, LOG_PATH); - jobPersistence.writeOutput(jobId, attemptId, jobOutput); - jobPersistence.succeedAttempt(jobId, attemptId); - - // verify we get that state back. - assertEquals(Optional.of(state), jobPersistence.getCurrentState(UUID.fromString(SCOPE))); - - // create a second job - final long jobId2 = jobPersistence.enqueueJob(SCOPE, SYNC_JOB_CONFIG).orElseThrow(); - - // check that when we have a failed attempt with no state, we still use old state. - jobPersistence.failAttempt(jobId2, jobPersistence.createAttempt(jobId2, LOG_PATH)); - assertEquals(Optional.of(state), jobPersistence.getCurrentState(UUID.fromString(SCOPE))); - - // check that when we have a failed attempt, if we have a state we use it. - final State state2 = new State().withState(Jsons.jsonNode(ImmutableMap.of("checkpoint", 5))); - final JobOutput jobOutput2 = new JobOutput().withSync(new StandardSyncOutput().withState(state2)); - final int attemptId2 = jobPersistence.createAttempt(jobId2, LOG_PATH); - jobPersistence.writeOutput(jobId2, attemptId2, jobOutput2); - jobPersistence.failAttempt(jobId2, attemptId2); - assertEquals(Optional.of(state2), jobPersistence.getCurrentState(UUID.fromString(SCOPE))); - } - - @Test - @DisplayName("Should not have state if the latest attempt did not succeeded and have state otherwise") - public void testGetCurrentStateForConnectionIdNoState() throws IOException { - // no state when the connection has never had a job. - assertEquals(Optional.empty(), jobPersistence.getCurrentState(CONNECTION_ID)); - - final long jobId = jobPersistence.enqueueJob(SCOPE, SYNC_JOB_CONFIG).orElseThrow(); - final int attemptNumber = jobPersistence.createAttempt(jobId, LOG_PATH); - - // no state when connection has a job but it has not completed that has not completed - assertEquals(Optional.empty(), jobPersistence.getCurrentState(CONNECTION_ID)); - - jobPersistence.failJob(jobId); - - // no state when connection has a job but it is failed. - assertEquals(Optional.empty(), jobPersistence.getCurrentState(CONNECTION_ID)); - - jobPersistence.cancelJob(jobId); - - // no state when connection has a job but it is cancelled. - assertEquals(Optional.empty(), jobPersistence.getCurrentState(CONNECTION_ID)); - - final JobOutput jobOutput1 = new JobOutput() - .withSync(new StandardSyncOutput().withState(new State().withState(Jsons.jsonNode(ImmutableMap.of("checkpoint", "1"))))); - jobPersistence.writeOutput(jobId, attemptNumber, jobOutput1); - jobPersistence.succeedAttempt(jobId, attemptNumber); - - // job 1 state, after first success. - assertEquals(Optional.of(jobOutput1.getSync().getState()), jobPersistence.getCurrentState(CONNECTION_ID)); - - when(timeSupplier.get()).thenReturn(NOW.plusSeconds(1000)); - final long jobId2 = jobPersistence.enqueueJob(SCOPE, SYNC_JOB_CONFIG).orElseThrow(); - final int attemptNumber2 = jobPersistence.createAttempt(jobId2, LOG_PATH); - - // job 1 state, second job created. - assertEquals(Optional.of(jobOutput1.getSync().getState()), jobPersistence.getCurrentState(CONNECTION_ID)); - - jobPersistence.failJob(jobId2); - - // job 1 state, second job failed. - assertEquals(Optional.of(jobOutput1.getSync().getState()), jobPersistence.getCurrentState(CONNECTION_ID)); - - jobPersistence.cancelJob(jobId2); - - // job 1 state, second job cancelled - assertEquals(Optional.of(jobOutput1.getSync().getState()), jobPersistence.getCurrentState(CONNECTION_ID)); - - final JobOutput jobOutput2 = new JobOutput() - .withSync(new StandardSyncOutput().withState(new State().withState(Jsons.jsonNode(ImmutableMap.of("checkpoint", "2"))))); - jobPersistence.writeOutput(jobId2, attemptNumber2, jobOutput2); - jobPersistence.succeedAttempt(jobId2, attemptNumber2); - - // job 2 state, after second job success. - assertEquals(Optional.of(jobOutput2.getSync().getState()), jobPersistence.getCurrentState(CONNECTION_ID)); - } - - } - @Nested @DisplayName("When getting next job") class GetNextJob { @@ -1300,7 +1194,7 @@ void testPurgeJobHistory(final int numJobs, // Reconfigure constants to test various combinations of tuning knobs and make sure all work. final DefaultJobPersistence jobPersistence = - new DefaultJobPersistence(jobDatabase, configDatabase, timeSupplier, ageCutoff, tooManyJobs, recencyCutoff); + new DefaultJobPersistence(jobDatabase, timeSupplier, ageCutoff, tooManyJobs, recencyCutoff); final LocalDateTime fakeNow = LocalDateTime.of(2021, 6, 20, 0, 0); diff --git a/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java b/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java index cb1a52c61e2a..519b3cfef094 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java @@ -8,6 +8,7 @@ import io.airbyte.commons.io.FileTtlManager; import io.airbyte.config.Configs; import io.airbyte.config.persistence.ConfigPersistence; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.db.Database; import io.airbyte.scheduler.client.CachingSynchronousSchedulerClient; @@ -23,6 +24,7 @@ public class ConfigurationApiFactory implements Factory { private static WorkflowServiceStubs temporalService; private static ConfigRepository configRepository; + private static ConfigPersistence2 configPersistence2; private static JobPersistence jobPersistence; private static ConfigPersistence seed; private static SchedulerJobClient schedulerJobClient; @@ -37,6 +39,7 @@ public class ConfigurationApiFactory implements Factory { public static void setValues( final WorkflowServiceStubs temporalService, final ConfigRepository configRepository, + final ConfigPersistence2 configPersistence2, final JobPersistence jobPersistence, final ConfigPersistence seed, final SchedulerJobClient schedulerJobClient, @@ -48,6 +51,7 @@ public static void setValues( final Database jobsDatabase, final TrackingClient trackingClient) { ConfigurationApiFactory.configRepository = configRepository; + ConfigurationApiFactory.configPersistence2 = configPersistence2; ConfigurationApiFactory.jobPersistence = jobPersistence; ConfigurationApiFactory.seed = seed; ConfigurationApiFactory.schedulerJobClient = schedulerJobClient; @@ -67,6 +71,7 @@ public ConfigurationApi provide() { return new ConfigurationApi( ConfigurationApiFactory.configRepository, + ConfigurationApiFactory.configPersistence2, ConfigurationApiFactory.jobPersistence, ConfigurationApiFactory.seed, ConfigurationApiFactory.schedulerJobClient, 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 242c4267714b..78e8f2417b79 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java @@ -16,6 +16,7 @@ import io.airbyte.config.StandardWorkspace; import io.airbyte.config.helpers.LogClientSingleton; import io.airbyte.config.persistence.ConfigPersistence; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.config.persistence.DatabaseConfigPersistence; import io.airbyte.config.persistence.YamlSeedConfigPersistence; @@ -176,6 +177,7 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con final ConfigRepository configRepository = new ConfigRepository(configPersistence.withValidation(), secretsHydrator, secretPersistence, ephemeralSecretPersistence); + final ConfigPersistence2 configPersistence2 = new ConfigPersistence2(configDatabase); LOGGER.info("Creating Scheduler persistence..."); final Database jobDatabase = new JobsDatabaseInstance( @@ -183,7 +185,7 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con configs.getDatabasePassword(), configs.getDatabaseUrl()) .getAndInitialize(); - final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); + final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); createDeploymentIfNoneExists(jobPersistence); @@ -209,7 +211,8 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con final WorkflowServiceStubs temporalService = TemporalUtils.createTemporalService(configs.getTemporalHost()); final TemporalClient temporalClient = TemporalClient.production(configs.getTemporalHost(), configs.getWorkspaceRoot()); final OAuthConfigSupplier oAuthConfigSupplier = new OAuthConfigSupplier(configRepository, false, trackingClient); - final SchedulerJobClient schedulerJobClient = new DefaultSchedulerJobClient(jobPersistence, new DefaultJobCreator(jobPersistence)); + final SchedulerJobClient schedulerJobClient = + new DefaultSchedulerJobClient(jobPersistence, new DefaultJobCreator(jobPersistence, configPersistence2)); final DefaultSynchronousSchedulerClient syncSchedulerClient = new DefaultSynchronousSchedulerClient(temporalClient, jobTracker, oAuthConfigSupplier); final SynchronousSchedulerClient bucketSpecCacheSchedulerClient = @@ -245,6 +248,7 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con cachingSchedulerClient, temporalService, configRepository, + configPersistence2, jobPersistence, seed, configDatabase, diff --git a/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java b/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java index d58f0fce32d8..93ebb0a6b05c 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java @@ -8,6 +8,7 @@ import io.airbyte.commons.io.FileTtlManager; import io.airbyte.config.Configs; import io.airbyte.config.persistence.ConfigPersistence; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.db.Database; import io.airbyte.scheduler.client.SchedulerJobClient; @@ -25,6 +26,7 @@ ServerRunnable create(SchedulerJobClient schedulerJobClient, SpecCachingSynchronousSchedulerClient cachingSchedulerClient, WorkflowServiceStubs temporalService, ConfigRepository configRepository, + ConfigPersistence2 configPersistence2, JobPersistence jobPersistence, ConfigPersistence seed, Database configsDatabase, @@ -39,6 +41,7 @@ public ServerRunnable create(final SchedulerJobClient schedulerJobClient, final SpecCachingSynchronousSchedulerClient cachingSchedulerClient, final WorkflowServiceStubs temporalService, final ConfigRepository configRepository, + final ConfigPersistence2 configPersistence2, final JobPersistence jobPersistence, final ConfigPersistence seed, final Database configsDatabase, @@ -49,6 +52,7 @@ public ServerRunnable create(final SchedulerJobClient schedulerJobClient, ConfigurationApiFactory.setValues( temporalService, configRepository, + configPersistence2, jobPersistence, seed, schedulerJobClient, diff --git a/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java b/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java index 3b2dc4f0e846..3acee99bd292 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java +++ b/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java @@ -84,6 +84,7 @@ import io.airbyte.config.Configs; import io.airbyte.config.persistence.ConfigNotFoundException; import io.airbyte.config.persistence.ConfigPersistence; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.db.Database; import io.airbyte.scheduler.client.CachingSynchronousSchedulerClient; @@ -145,6 +146,7 @@ public class ConfigurationApi implements io.airbyte.api.V1Api { private final Configs configs; public ConfigurationApi(final ConfigRepository configRepository, + final ConfigPersistence2 configPersistence2, final JobPersistence jobPersistence, final ConfigPersistence seed, final SchedulerJobClient schedulerJobClient, @@ -164,6 +166,7 @@ public ConfigurationApi(final ConfigRepository configRepository, trackingClient); schedulerHandler = new SchedulerHandler( configRepository, + configPersistence2, schedulerJobClient, synchronousSchedulerClient, jobPersistence, diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java index b9709985d165..a75beab35f06 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java @@ -37,6 +37,7 @@ import io.airbyte.config.StandardSyncOperation; import io.airbyte.config.State; import io.airbyte.config.persistence.ConfigNotFoundException; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.protocol.models.AirbyteCatalog; import io.airbyte.protocol.models.ConnectorSpecification; @@ -68,7 +69,6 @@ public class SchedulerHandler { private static final Logger LOGGER = LoggerFactory.getLogger(SchedulerHandler.class); - private static final UUID NO_WORKSPACE = UUID.fromString("00000000-0000-0000-0000-000000000000"); private final ConfigRepository configRepository; private final SchedulerJobClient schedulerJobClient; @@ -80,8 +80,10 @@ public class SchedulerHandler { private final JobNotifier jobNotifier; private final WorkflowServiceStubs temporalService; private final OAuthConfigSupplier oAuthConfigSupplier; + private final ConfigPersistence2 configPersistence2; public SchedulerHandler(final ConfigRepository configRepository, + final ConfigPersistence2 configPersistence2, final SchedulerJobClient schedulerJobClient, final SynchronousSchedulerClient synchronousSchedulerClient, final JobPersistence jobPersistence, @@ -90,6 +92,7 @@ public SchedulerHandler(final ConfigRepository configRepository, final OAuthConfigSupplier oAuthConfigSupplier) { this( configRepository, + configPersistence2, schedulerJobClient, synchronousSchedulerClient, new ConfigurationUpdate(configRepository, new SpecFetcher(synchronousSchedulerClient)), @@ -103,6 +106,7 @@ public SchedulerHandler(final ConfigRepository configRepository, @VisibleForTesting SchedulerHandler(final ConfigRepository configRepository, + final ConfigPersistence2 configPersistence2, final SchedulerJobClient schedulerJobClient, final SynchronousSchedulerClient synchronousSchedulerClient, final ConfigurationUpdate configurationUpdate, @@ -113,6 +117,7 @@ public SchedulerHandler(final ConfigRepository configRepository, final WorkflowServiceStubs temporalService, final OAuthConfigSupplier oAuthConfigSupplier) { this.configRepository = configRepository; + this.configPersistence2 = configPersistence2; this.schedulerJobClient = schedulerJobClient; this.synchronousSchedulerClient = synchronousSchedulerClient; this.configurationUpdate = configurationUpdate; @@ -354,7 +359,7 @@ public JobInfoRead resetConnection(final ConnectionIdRequestBody connectionIdReq } public ConnectionState getState(final ConnectionIdRequestBody connectionIdRequestBody) throws IOException { - final Optional currentState = jobPersistence.getCurrentState(connectionIdRequestBody.getConnectionId()); + final Optional currentState = configPersistence2.getCurrentState(connectionIdRequestBody.getConnectionId()); LOGGER.info("currentState server: {}", currentState); final ConnectionState connectionState = new ConnectionState() @@ -365,6 +370,7 @@ public ConnectionState getState(final ConnectionIdRequestBody connectionIdReques return connectionState; } + // todo (cgardens) - this method needs a test. public JobInfoRead cancelJob(final JobIdRequestBody jobIdRequestBody) throws IOException { final long jobId = jobIdRequestBody.getId(); 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 e387da384786..5784c58bbbfd 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 @@ -105,7 +105,7 @@ public static void dbDown() { public void setup() throws Exception { jobDatabase = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); configDatabase = new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); + jobPersistence = new DefaultJobPersistence(jobDatabase); seedPersistence = YamlSeedConfigPersistence.getDefault(); configPersistence = new DatabaseConfigPersistence(jobDatabase); configPersistence.replaceAllConfigs(Collections.emptyMap(), false); diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java index 85c1b0b76ba5..ccfcf7c6e82f 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java @@ -49,6 +49,7 @@ import io.airbyte.config.StandardSyncOperation.OperatorType; import io.airbyte.config.State; import io.airbyte.config.persistence.ConfigNotFoundException; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.protocol.models.AirbyteCatalog; import io.airbyte.protocol.models.CatalogHelpers; @@ -117,6 +118,7 @@ class SchedulerHandlerTest { private SchedulerHandler schedulerHandler; private ConfigRepository configRepository; + private ConfigPersistence2 configPersistence2; private Job completedJob; private SchedulerJobClient schedulerJobClient; private SynchronousSchedulerClient synchronousSchedulerClient; @@ -140,11 +142,13 @@ void setup() { schedulerJobClient = spy(SchedulerJobClient.class); synchronousSchedulerClient = mock(SynchronousSchedulerClient.class); configRepository = mock(ConfigRepository.class); + configPersistence2 = mock(ConfigPersistence2.class); jobPersistence = mock(JobPersistence.class); final JobNotifier jobNotifier = mock(JobNotifier.class); schedulerHandler = new SchedulerHandler( configRepository, + configPersistence2, schedulerJobClient, synchronousSchedulerClient, configurationUpdate, @@ -557,7 +561,7 @@ void testResetConnection() throws JsonValidationException, IOException, ConfigNo void testGetCurrentState() throws IOException { final UUID connectionId = UUID.randomUUID(); final State state = new State().withState(Jsons.jsonNode(ImmutableMap.of("checkpoint", 1))); - when(jobPersistence.getCurrentState(connectionId)).thenReturn(Optional.of(state)); + when(configPersistence2.getCurrentState(connectionId)).thenReturn(Optional.of(state)); final ConnectionState connectionState = schedulerHandler.getState(new ConnectionIdRequestBody().connectionId(connectionId)); assertEquals(new ConnectionState().connectionId(connectionId).state(state.getState()), connectionState); @@ -566,7 +570,7 @@ void testGetCurrentState() throws IOException { @Test void testGetCurrentStateEmpty() throws IOException { final UUID connectionId = UUID.randomUUID(); - when(jobPersistence.getCurrentState(connectionId)).thenReturn(Optional.empty()); + when(configPersistence2.getCurrentState(connectionId)).thenReturn(Optional.empty()); final ConnectionState connectionState = schedulerHandler.getState(new ConnectionIdRequestBody().connectionId(connectionId)); assertEquals(new ConnectionState().connectionId(connectionId), connectionState); diff --git a/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java b/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java index 864fb7da8795..f9f4d060301a 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java @@ -43,7 +43,7 @@ void setUp() throws IOException { jobDatabase = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); configDatabase = new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - final JobPersistence persistence = new DefaultJobPersistence(jobDatabase, configDatabase); + final JobPersistence persistence = new DefaultJobPersistence(jobDatabase); databaseArchiver = new DatabaseArchiver(persistence); } diff --git a/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java b/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java index 97a4e010fa3a..2eb6aa7dc951 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java @@ -363,7 +363,7 @@ private void runMigration(final ConfigPersistence configPersistence, final JobPe @SuppressWarnings("SameParameterValue") private JobPersistence getJobPersistence(final File file, final String version) throws IOException { - final DefaultJobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); + final DefaultJobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); final Path tempFolder = Files.createTempDirectory(Path.of("/tmp"), "db_init"); resourceToBeCleanedUp.add(tempFolder.toFile()); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java index c7cf3c100ec5..307eddd162ba 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java @@ -9,8 +9,11 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.MaxWorkersConfig; import io.airbyte.config.helpers.LogClientSingleton; +import io.airbyte.config.persistence.ConfigPersistence2; 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.instance.configs.ConfigsDatabaseInstance; import io.airbyte.workers.process.DockerProcessFactory; import io.airbyte.workers.process.KubeProcessFactory; import io.airbyte.workers.process.ProcessFactory; @@ -50,7 +53,7 @@ public class WorkerApp { private final WorkflowServiceStubs temporalService; private final MaxWorkersConfig maxWorkers; private final WorkerEnvironment workerEnvironment; - private final String airbyteVersion; + private final ConfigPersistence2 configPersistence2; public WorkerApp(final Path workspaceRoot, final ProcessFactory processFactory, @@ -58,14 +61,14 @@ public WorkerApp(final Path workspaceRoot, final WorkflowServiceStubs temporalService, final MaxWorkersConfig maxWorkers, final WorkerEnvironment workerEnvironment, - final String airbyteVersion) { + final ConfigPersistence2 configPersistence2) { this.workspaceRoot = workspaceRoot; this.processFactory = processFactory; this.secretsHydrator = secretsHydrator; this.temporalService = temporalService; this.maxWorkers = maxWorkers; this.workerEnvironment = workerEnvironment; - this.airbyteVersion = airbyteVersion; + this.configPersistence2 = configPersistence2; } public void start() { @@ -101,8 +104,9 @@ public void start() { syncWorker.registerWorkflowImplementationTypes(SyncWorkflow.WorkflowImpl.class); syncWorker.registerActivitiesImplementations( new SyncWorkflow.ReplicationActivityImpl(processFactory, secretsHydrator, workspaceRoot), - new SyncWorkflow.NormalizationActivityImpl(processFactory, secretsHydrator, workspaceRoot, workerEnvironment, airbyteVersion), - new SyncWorkflow.DbtTransformationActivityImpl(processFactory, secretsHydrator, workspaceRoot, airbyteVersion)); + new SyncWorkflow.NormalizationActivityImpl(processFactory, secretsHydrator, workspaceRoot, workerEnvironment), + new SyncWorkflow.DbtTransformationActivityImpl(processFactory, secretsHydrator, workspaceRoot), + new SyncWorkflow.PersistStateActivityImpl(workspaceRoot, configPersistence2)); factory.start(); } @@ -146,6 +150,14 @@ public static void main(final String[] args) throws IOException, InterruptedExce final WorkflowServiceStubs temporalService = TemporalUtils.createTemporalService(temporalHost); + // todo (cgardens) - make sure appropriate env variables are passed to this container. + final Database configDatabase = new ConfigsDatabaseInstance( + configs.getConfigDatabaseUser(), + configs.getConfigDatabasePassword(), + configs.getConfigDatabaseUrl()) + .getInitialized(); + final ConfigPersistence2 configPersistence2 = new ConfigPersistence2(configDatabase); + new WorkerApp( workspaceRoot, processFactory, @@ -153,7 +165,7 @@ public static void main(final String[] args) throws IOException, InterruptedExce temporalService, configs.getMaxWorkers(), configs.getWorkerEnvironment(), - configs.getAirbyteVersion()).start(); + configPersistence2).start(); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java index 496c6d83c758..b3725e978e96 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java @@ -20,6 +20,8 @@ import io.airbyte.config.StandardSyncOperation.OperatorType; import io.airbyte.config.StandardSyncOutput; import io.airbyte.config.StandardSyncSummary; +import io.airbyte.config.State; +import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.split_secrets.SecretsHydrator; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; @@ -43,11 +45,14 @@ import io.temporal.activity.ActivityInterface; import io.temporal.activity.ActivityMethod; import io.temporal.activity.ActivityOptions; +import io.temporal.common.RetryOptions; import io.temporal.workflow.Workflow; import io.temporal.workflow.WorkflowInterface; import io.temporal.workflow.WorkflowMethod; +import java.io.IOException; import java.nio.file.Path; import java.time.Duration; +import java.util.UUID; import java.util.function.Supplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -59,7 +64,8 @@ public interface SyncWorkflow { StandardSyncOutput run(JobRunConfig jobRunConfig, IntegrationLauncherConfig sourceLauncherConfig, IntegrationLauncherConfig destinationLauncherConfig, - StandardSyncInput syncInput); + StandardSyncInput syncInput, + UUID connectionId); class WorkflowImpl implements SyncWorkflow { @@ -73,15 +79,23 @@ class WorkflowImpl implements SyncWorkflow { .setRetryOptions(TemporalUtils.NO_RETRY) .build(); + private static final ActivityOptions persistOptions = options.toBuilder() + .setRetryOptions(RetryOptions.newBuilder() + .setMaximumAttempts(10) + .build()) + .build(); + private final ReplicationActivity replicationActivity = Workflow.newActivityStub(ReplicationActivity.class, options); private final NormalizationActivity normalizationActivity = Workflow.newActivityStub(NormalizationActivity.class, options); private final DbtTransformationActivity dbtTransformationActivity = Workflow.newActivityStub(DbtTransformationActivity.class, options); + private final PersistStateActivity persistActivity = Workflow.newActivityStub(PersistStateActivity.class, persistOptions); @Override public StandardSyncOutput run(final JobRunConfig jobRunConfig, final IntegrationLauncherConfig sourceLauncherConfig, final IntegrationLauncherConfig destinationLauncherConfig, - final StandardSyncInput syncInput) { + final StandardSyncInput syncInput, + final UUID connectionId) { final StandardSyncOutput run = replicationActivity.replicate(jobRunConfig, sourceLauncherConfig, destinationLauncherConfig, syncInput); if (syncInput.getOperationSequence() != null && !syncInput.getOperationSequence().isEmpty()) { @@ -104,6 +118,7 @@ public StandardSyncOutput run(final JobRunConfig jobRunConfig, LOGGER.error(message); throw new IllegalArgumentException(message); } + persistActivity.persist(connectionId, run); } } @@ -255,14 +270,12 @@ class NormalizationActivityImpl implements NormalizationActivity { private final Path workspaceRoot; private final AirbyteConfigValidator validator; private final WorkerEnvironment workerEnvironment; - private final String airbyteVersion; public NormalizationActivityImpl(final ProcessFactory processFactory, final SecretsHydrator secretsHydrator, final Path workspaceRoot, - final WorkerEnvironment workerEnvironment, - final String airbyteVersion) { - this(processFactory, secretsHydrator, workspaceRoot, new AirbyteConfigValidator(), workerEnvironment, airbyteVersion); + final WorkerEnvironment workerEnvironment) { + this(processFactory, secretsHydrator, workspaceRoot, new AirbyteConfigValidator(), workerEnvironment); } @VisibleForTesting @@ -270,14 +283,12 @@ public NormalizationActivityImpl(final ProcessFactory processFactory, final SecretsHydrator secretsHydrator, final Path workspaceRoot, final AirbyteConfigValidator validator, - final WorkerEnvironment workerEnvironment, - final String airbyteVersion) { + final WorkerEnvironment workerEnvironment) { this.processFactory = processFactory; this.secretsHydrator = secretsHydrator; this.workspaceRoot = workspaceRoot; this.validator = validator; this.workerEnvironment = workerEnvironment; - this.airbyteVersion = airbyteVersion; } @Override @@ -335,27 +346,23 @@ class DbtTransformationActivityImpl implements DbtTransformationActivity { private final SecretsHydrator secretsHydrator; private final Path workspaceRoot; private final AirbyteConfigValidator validator; - private final String airbyteVersion; public DbtTransformationActivityImpl( final ProcessFactory processFactory, final SecretsHydrator secretsHydrator, - final Path workspaceRoot, - final String airbyteVersion) { - this(processFactory, secretsHydrator, workspaceRoot, new AirbyteConfigValidator(), airbyteVersion); + final Path workspaceRoot) { + this(processFactory, secretsHydrator, workspaceRoot, new AirbyteConfigValidator()); } @VisibleForTesting DbtTransformationActivityImpl(final ProcessFactory processFactory, final SecretsHydrator secretsHydrator, final Path workspaceRoot, - final AirbyteConfigValidator validator, - final String airbyteVersion) { + final AirbyteConfigValidator validator) { this.processFactory = processFactory; this.secretsHydrator = secretsHydrator; this.workspaceRoot = workspaceRoot; this.validator = validator; - this.airbyteVersion = airbyteVersion; } @Override @@ -397,4 +404,40 @@ private CheckedSupplier, Exception> getWorkerFact } + @ActivityInterface + interface PersistStateActivity { + + @ActivityMethod + boolean persist(final UUID connectionId, final StandardSyncOutput syncOutput); + + } + + class PersistStateActivityImpl implements PersistStateActivity { + + private static final Logger LOGGER = LoggerFactory.getLogger(PersistStateActivityImpl.class); + private final Path workspaceRoot; + private final ConfigPersistence2 configPersistence2; + + public PersistStateActivityImpl(final Path workspaceRoot, final ConfigPersistence2 configPersistence2) { + this.workspaceRoot = workspaceRoot; + this.configPersistence2 = configPersistence2; + } + + @Override + public boolean persist(final UUID connectionId, final StandardSyncOutput syncOutput) { + final State state = syncOutput.getState(); + if (state != null) { + try { + configPersistence2.updateSyncState(connectionId, state); + } catch (final IOException e) { + throw new RuntimeException(e); + } + return true; + } else { + return false; + } + } + + } + } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java index 3580852ee6dd..4c1dcd9380d8 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java @@ -137,7 +137,7 @@ private void saveWorkflowIdForCancellation() throws IOException { configs.getConfigDatabasePassword(), configs.getConfigDatabaseUrl()) .getInitialized(); - final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase, configDatabase); + final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); final String workflowId = workflowIdProvider.get(); jobPersistence.setAttemptTemporalWorkflowId(Long.parseLong(jobRunConfig.getJobId()), jobRunConfig.getAttemptId().intValue(), workflowId); } 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 94510d187799..d3139315a766 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 @@ -78,7 +78,7 @@ public TemporalResponse submitDiscoverSchema(final UUID jobId, f () -> getWorkflowStub(DiscoverCatalogWorkflow.class, TemporalJobType.DISCOVER_SCHEMA).run(jobRunConfig, launcherConfig, input)); } - public TemporalResponse submitSync(final long jobId, final int attempt, final JobSyncConfig config) { + public TemporalResponse submitSync(final long jobId, final int attempt, final JobSyncConfig config, final UUID connectionId) { final JobRunConfig jobRunConfig = TemporalUtils.createJobRunConfig(jobId, attempt); final IntegrationLauncherConfig sourceLauncherConfig = new IntegrationLauncherConfig() @@ -107,7 +107,8 @@ public TemporalResponse submitSync(final long jobId, final i jobRunConfig, sourceLauncherConfig, destinationLauncherConfig, - input)); + input, + connectionId)); } private T getWorkflowStub(final Class workflowClass, final TemporalJobType jobType) { From 145765c1be2ff73e88e98bfbff0a38c72f209420 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 04:39:18 -0700 Subject: [PATCH 09/23] Copy state to airbyte_configs table --- ... V0_30_22_001__Store_last_sync_state.java} | 74 +++++----- ...0_22_001__Store_last_sync_state_test.java} | 138 ++++++------------ 2 files changed, 83 insertions(+), 129 deletions(-) rename airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/{V0_29_21_001__Store_last_sync_state.java => V0_30_22_001__Store_last_sync_state.java} (73%) rename airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/{V0_29_21_001__Store_last_sync_state_test.java => V0_30_22_001__Store_last_sync_state_test.java} (61%) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java similarity index 73% rename from airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java rename to airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java index 9ab05a77b19a..abf2563bf21e 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java @@ -5,7 +5,9 @@ package io.airbyte.db.instance.configs.migrations; import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; +import io.airbyte.commons.jackson.MoreMappers; import io.airbyte.commons.json.Jsons; import io.airbyte.config.Configs; import io.airbyte.config.EnvConfigs; @@ -14,7 +16,6 @@ import java.io.IOException; import java.sql.SQLException; import java.time.OffsetDateTime; -import java.util.Collections; import java.util.Map; import java.util.Optional; import java.util.UUID; @@ -36,28 +37,30 @@ *
  • Column sync_id: the connectionId in StandardSync
  • Column state: a json node representing * a State object
    • */ -public class V0_29_21_001__Store_last_sync_state extends BaseJavaMigration { - - private static final String MIGRATION_NAME = "Configs db migration 0.29.21.001"; - private static final Logger LOGGER = LoggerFactory.getLogger(V0_29_21_001__Store_last_sync_state.class); - - // sync state table - static final Table SYNC_STATE_TABLE = DSL.table("sync_state"); - static final Field COLUMN_SYNC_ID = DSL.field("sync_id", SQLDataType.UUID.nullable(false)); - static final Field COLUMN_STATE = DSL.field("state", SQLDataType.JSONB.nullable(false)); - static final Field COLUMN_CREATED_AT = DSL.field("created_at", - SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())); - static final Field COLUMN_UPDATED_AT = DSL.field("updated_at", - SQLDataType.OFFSETDATETIME.nullable(false).defaultValue(DSL.currentOffsetDateTime())); +public class V0_30_22_001__Store_last_sync_state extends BaseJavaMigration { + + private static final ObjectMapper MAPPER = MoreMappers.initMapper(); + private static final String MIGRATION_NAME = "Configs db migration 0.30.22.001"; + private static final Logger LOGGER = LoggerFactory.getLogger(V0_30_22_001__Store_last_sync_state.class); + + // airbyte configs table + // (we cannot use the jooq generated code here to avoid circular dependency) + static final String STANDARD_SYNC_STATE = "STANDARD_SYNC_STATE"; + static final Table TABLE_AIRBYTE_CONFIGS = DSL.table("airbyte_configs"); + static final Field COLUMN_CONFIG_TYPE = DSL.field("config_type", SQLDataType.VARCHAR(60).nullable(false)); + static final Field COLUMN_CONFIG_ID = DSL.field("config_id", SQLDataType.VARCHAR(36).nullable(false)); + static final Field COLUMN_CONFIG_BLOB = DSL.field("config_blob", SQLDataType.JSONB.nullable(false)); + static final Field COLUMN_CREATED_AT = DSL.field("created_at", SQLDataType.TIMESTAMPWITHTIMEZONE); + static final Field COLUMN_UPDATED_AT = DSL.field("updated_at", SQLDataType.TIMESTAMPWITHTIMEZONE); private final Configs configs; - public V0_29_21_001__Store_last_sync_state() { + public V0_30_22_001__Store_last_sync_state() { this.configs = new EnvConfigs(); } @VisibleForTesting - V0_29_21_001__Store_last_sync_state(final Configs configs) { + V0_30_22_001__Store_last_sync_state(final Configs configs) { this.configs = configs; } @@ -66,33 +69,19 @@ public void migrate(final Context context) throws Exception { LOGGER.info("Running migration: {}", this.getClass().getSimpleName()); final DSLContext ctx = DSL.using(context.getConnection()); - createTable(ctx); - final Optional jobsDatabase = getJobsDatabase(configs); if (jobsDatabase.isPresent()) { - copyData(ctx, getSyncToStateMap(jobsDatabase.get()), OffsetDateTime.now()); + copyData(ctx, getStandardSyncStates(jobsDatabase.get()), OffsetDateTime.now()); } } - @VisibleForTesting - static void createTable(final DSLContext ctx) { - ctx.createTableIfNotExists(SYNC_STATE_TABLE) - .column(COLUMN_SYNC_ID) - .column(COLUMN_STATE) - .column(COLUMN_CREATED_AT) - .column(COLUMN_UPDATED_AT) - .execute(); - ctx.createUniqueIndexIfNotExists(String.format("%s_sync_id_idx", SYNC_STATE_TABLE)) - .on(SYNC_STATE_TABLE, Collections.singleton(COLUMN_SYNC_ID)) - .execute(); - } - @VisibleForTesting static void copyData(final DSLContext ctx, final Map syncToStateMap, final OffsetDateTime timestamp) { for (final Map.Entry entry : syncToStateMap.entrySet()) { - ctx.insertInto(SYNC_STATE_TABLE) - .set(COLUMN_SYNC_ID, UUID.fromString(entry.getKey())) - .set(COLUMN_STATE, JSONB.valueOf(Jsons.serialize(entry.getValue()))) + ctx.insertInto(TABLE_AIRBYTE_CONFIGS) + .set(COLUMN_CONFIG_TYPE, STANDARD_SYNC_STATE) + .set(COLUMN_CONFIG_ID, UUID.fromString(entry.getKey()).toString()) + .set(COLUMN_CONFIG_BLOB, JSONB.valueOf(Jsons.serialize(entry.getValue()))) .set(COLUMN_CREATED_AT, timestamp) .set(COLUMN_UPDATED_AT, timestamp) // This migration is idempotent. If the record for a sync_id already exists, @@ -133,10 +122,10 @@ static Optional getJobsDatabase(final Configs configs) { } /** - * @return a map from sync id to last job attempt state. + * @return a map from connection id (UUID) to connection state (StandardSyncState). */ @VisibleForTesting - static Map getSyncToStateMap(final Database jobsDatabase) throws SQLException { + static Map getStandardSyncStates(final Database jobsDatabase) throws SQLException { final Table jobsTable = DSL.table("jobs"); final Field jobIdField = DSL.field("jobs.id", SQLDataType.BIGINT); final Field syncIdField = DSL.field("jobs.scope", SQLDataType.VARCHAR); @@ -147,7 +136,7 @@ static Map getSyncToStateMap(final Database jobsDatabase) thro // output schema: JobOutput.yaml // sync schema: StandardSyncOutput.yaml - // state schema: State.yaml + // state schema: State.yaml, e.g. { "state": { "cursor": 1000 } } final Field attemptStateField = DSL.field("attempts.output -> 'sync' -> 'state'", SQLDataType.JSONB); return jobsDatabase.query(ctx -> ctx @@ -165,7 +154,14 @@ static Map getSyncToStateMap(final Database jobsDatabase) thro .stream() .collect(Collectors.toMap( Record2::value1, - r -> Jsons.deserialize(r.value2().data())))); + r -> getStandardSyncState(r.value1(), Jsons.deserialize(r.value2().data()))))); + } + + @VisibleForTesting + static JsonNode getStandardSyncState(final String connectionId, final JsonNode state) { + return MAPPER.createObjectNode() + .put("connectionId", connectionId) + .set("state", state); } } diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java similarity index 61% rename from airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java rename to airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java index b54ba3e87be6..7934e84f6dc0 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_29_21_001__Store_last_sync_state_test.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java @@ -4,11 +4,16 @@ package io.airbyte.db.instance.configs.migrations; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_CONFIG_BLOB; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_CONFIG_ID; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_CONFIG_TYPE; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_CREATED_AT; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_UPDATED_AT; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.STANDARD_SYNC_STATE; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.TABLE_AIRBYTE_CONFIGS; import static org.jooq.impl.DSL.field; import static org.jooq.impl.DSL.table; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -26,21 +31,15 @@ import java.sql.Connection; import java.sql.SQLException; import java.time.OffsetDateTime; -import java.util.Arrays; -import java.util.List; import java.util.Map; -import java.util.Set; import java.util.UUID; -import java.util.stream.Collectors; import javax.annotation.Nullable; import org.flywaydb.core.api.configuration.Configuration; import org.flywaydb.core.api.migration.Context; import org.jooq.DSLContext; import org.jooq.Field; -import org.jooq.InsertSetMoreStep; import org.jooq.JSONB; import org.jooq.Table; -import org.jooq.Typed; import org.jooq.impl.SQLDataType; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.MethodOrderer; @@ -49,7 +48,7 @@ import org.junit.jupiter.api.TestMethodOrder; @TestMethodOrder(MethodOrderer.OrderAnnotation.class) -class V0_29_21_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTest { +class V0_30_22_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTest { private static final ObjectMapper OBJECT_MAPPER = MoreMappers.initMapper(); @@ -63,13 +62,19 @@ class V0_29_21_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTe private static final Field ATTEMPT_NUMBER_FIELD = field("attempt_number", SQLDataType.INTEGER); private static final Field ATTEMPT_OUTPUT_FIELD = field("output", SQLDataType.JSONB); - private static final UUID SYNC_1_ID = UUID.randomUUID(); - private static final UUID SYNC_2_ID = UUID.randomUUID(); - private static final UUID SYNC_3_ID = UUID.randomUUID(); + private static final UUID CONNECTION_1_ID = UUID.randomUUID(); + private static final UUID CONNECTION_2_ID = UUID.randomUUID(); + private static final UUID CONNECTION_3_ID = UUID.randomUUID(); // these are State objects, see State.yaml for its schema; // we cannot construct the POJO directly because State is defined in an downstream module - private static final JsonNode SYNC_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }"); - private static final JsonNode SYNC_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }"); + private static final JsonNode CONNECTION_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }"); + private static final JsonNode CONNECTION_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }"); + + private static final JsonNode STD_CONNECTION_STATE_2 = V0_30_22_001__Store_last_sync_state.getStandardSyncState(CONNECTION_2_ID.toString(), CONNECTION_2_STATE); + private static final JsonNode STD_CONNECTION_STATE_3 = V0_30_22_001__Store_last_sync_state.getStandardSyncState(CONNECTION_3_ID.toString(), CONNECTION_3_STATE); + private static final Map CONNECTION_STATE_MAP = Map.of( + CONNECTION_2_ID.toString(), STD_CONNECTION_STATE_2, + CONNECTION_3_ID.toString(), STD_CONNECTION_STATE_3); private static Database jobDatabase; @@ -86,7 +91,7 @@ public static void setupJobDatabase() throws Exception { @Order(10) public void testGetJobsDatabase() { // when there is no database environment variable, the return value is empty - assertTrue(V0_29_21_001__Store_last_sync_state.getJobsDatabase(new EnvConfigs()).isEmpty()); + assertTrue(V0_30_22_001__Store_last_sync_state.getJobsDatabase(new EnvConfigs()).isEmpty()); // when there is database environment variable, return the database final Configs configs = mock(Configs.class); @@ -94,7 +99,7 @@ public void testGetJobsDatabase() { when(configs.getDatabasePassword()).thenReturn(container.getPassword()); when(configs.getDatabaseUrl()).thenReturn(container.getJdbcUrl()); - assertTrue(V0_29_21_001__Store_last_sync_state.getJobsDatabase(configs).isPresent()); + assertTrue(V0_30_22_001__Store_last_sync_state.getJobsDatabase(configs).isPresent()); } @Test @@ -103,25 +108,25 @@ public void testGetSyncToStateMap() throws Exception { jobDatabase.query(ctx -> { // Create three jobs for three standard syncs. // The first job has no attempt. - createJob(ctx, SYNC_1_ID); + createJob(ctx, CONNECTION_1_ID); // The second job has one attempt. - final long job2 = createJob(ctx, SYNC_2_ID); - createAttempt(ctx, job2, 1, createAttemptOutput(SYNC_2_STATE)); + final long job2 = createJob(ctx, CONNECTION_2_ID); + createAttempt(ctx, job2, 1, createAttemptOutput(CONNECTION_2_STATE)); // The third job has multiple attempts. The third attempt has the latest state. - final long job3 = createJob(ctx, SYNC_3_ID); + final long job3 = createJob(ctx, CONNECTION_3_ID); final JsonNode attempt31State = Jsons.deserialize("{ \"state\": { \"cursor\": 31 } }"); createAttempt(ctx, job3, 1, createAttemptOutput(attempt31State)); createAttempt(ctx, job3, 2, null); - createAttempt(ctx, job3, 3, createAttemptOutput(SYNC_3_STATE)); + createAttempt(ctx, job3, 3, createAttemptOutput(CONNECTION_3_STATE)); createAttempt(ctx, job3, 4, null); createAttempt(ctx, job3, 5, null); - final Map syncToStateMap = V0_29_21_001__Store_last_sync_state.getSyncToStateMap(jobDatabase); + final Map syncToStateMap = V0_30_22_001__Store_last_sync_state.getStandardSyncStates(jobDatabase); assertEquals(2, syncToStateMap.size()); - assertEquals(SYNC_2_STATE, syncToStateMap.get(SYNC_2_ID.toString())); - assertEquals(SYNC_3_STATE, syncToStateMap.get(SYNC_3_ID.toString())); + assertEquals(STD_CONNECTION_STATE_2, syncToStateMap.get(CONNECTION_2_ID.toString())); + assertEquals(STD_CONNECTION_STATE_3, syncToStateMap.get(CONNECTION_3_ID.toString())); return null; }); @@ -129,53 +134,21 @@ public void testGetSyncToStateMap() throws Exception { @Test @Order(30) - public void testCreateTable() throws SQLException { - database.query(ctx -> { - checkTable(ctx, false); - V0_29_21_001__Store_last_sync_state.createTable(ctx); - checkTable(ctx, true); - return null; - }); - } - - /** - * Test the unique index on the sync_id column. - */ - @Test - @Order(40) - public void testUniqueSyncIdIndex() throws SQLException { - final UUID syncId = UUID.randomUUID(); - - database.query(ctx -> { - final InsertSetMoreStep insertRecord = ctx.insertInto(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE) - .set(V0_29_21_001__Store_last_sync_state.COLUMN_SYNC_ID, syncId) - .set(V0_29_21_001__Store_last_sync_state.COLUMN_STATE, JSONB.valueOf("{}")); - - assertDoesNotThrow(insertRecord::execute); - // insert the record with the same sync_id will violate the unique index - assertThrows(org.jooq.exception.DataAccessException.class, insertRecord::execute); - return null; - }); - } - - @Test - @Order(50) public void testCopyData() throws SQLException { - final JsonNode sync2NewState = Jsons.deserialize("{ \"state\": { \"cursor\": 3 } }"); - - final Map syncStateMap1 = Map.of(SYNC_1_ID.toString(), SYNC_2_STATE, SYNC_2_ID.toString(), SYNC_2_STATE); - final Map syncStateMap2 = Map.of(SYNC_2_ID.toString(), sync2NewState); + final Map newConnectionStateMap = Map.of( + CONNECTION_2_ID.toString(), + Jsons.deserialize("{ \"connectionId\": \"invalid\", \"state\": { \"state\": { \"cursor\": 3 } } }")); final OffsetDateTime timestamp = OffsetDateTime.now(); database.query(ctx -> { - V0_29_21_001__Store_last_sync_state.copyData(ctx, syncStateMap1, timestamp); - checkSyncStates(ctx, syncStateMap1, timestamp); + V0_30_22_001__Store_last_sync_state.copyData(ctx, CONNECTION_STATE_MAP, timestamp); + checkSyncStates(ctx, CONNECTION_STATE_MAP, timestamp); // call the copyData method again with different data will not affect existing records - V0_29_21_001__Store_last_sync_state.copyData(ctx, syncStateMap2, OffsetDateTime.now()); + V0_30_22_001__Store_last_sync_state.copyData(ctx, CONNECTION_STATE_MAP, OffsetDateTime.now()); // the states remain the same as those in syncStateMap1 - checkSyncStates(ctx, syncStateMap1, timestamp); + checkSyncStates(ctx, CONNECTION_STATE_MAP, timestamp); return null; }); @@ -185,16 +158,18 @@ public void testCopyData() throws SQLException { * Clear the table and test the migration end-to-end. */ @Test - @Order(60) + @Order(40) public void testMigration() throws Exception { - database.query(ctx -> ctx.dropTableIfExists(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE).execute()); + database.query(ctx -> ctx.deleteFrom(TABLE_AIRBYTE_CONFIGS) + .where(COLUMN_CONFIG_TYPE.eq(STANDARD_SYNC_STATE)) + .execute()); final Configs configs = mock(Configs.class); when(configs.getDatabaseUser()).thenReturn(container.getUsername()); when(configs.getDatabasePassword()).thenReturn(container.getPassword()); when(configs.getDatabaseUrl()).thenReturn(container.getJdbcUrl()); - final var migration = new V0_29_21_001__Store_last_sync_state(configs); + final var migration = new V0_30_22_001__Store_last_sync_state(configs); // this context is a flyway class; only the getConnection method is needed to run the migration final Context context = new Context() { @@ -215,7 +190,7 @@ public Connection getConnection() { }; migration.migrate(context); database.query(ctx -> { - checkSyncStates(ctx, Map.of(SYNC_2_ID.toString(), SYNC_2_STATE, SYNC_3_ID.toString(), SYNC_3_STATE), null); + checkSyncStates(ctx, CONNECTION_STATE_MAP, null); return null; }); } @@ -262,34 +237,17 @@ private static JsonNode createAttemptOutput(final JsonNode state) { .set("sync", standardSyncOutput); } - private static void checkTable(final DSLContext ctx, final boolean tableExists) { - final List> tables = ctx.meta().getTables(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE.getName()); - assertEquals(tableExists, tables.size() > 0); - - if (!tableExists) { - assertTrue(tables.isEmpty()); - } else { - System.out.println(Arrays.stream(tables.get(0).fields()).map(Typed::getDataType).collect(Collectors.toSet())); - final Set actualFields = Arrays.stream(tables.get(0).fields()).map(Field::getName).collect(Collectors.toSet()); - final Set expectedFields = Set.of( - V0_29_21_001__Store_last_sync_state.COLUMN_SYNC_ID.getName(), - V0_29_21_001__Store_last_sync_state.COLUMN_STATE.getName(), - V0_29_21_001__Store_last_sync_state.COLUMN_CREATED_AT.getName(), - V0_29_21_001__Store_last_sync_state.COLUMN_UPDATED_AT.getName()); - assertEquals(expectedFields, actualFields); - } - } - private static void checkSyncStates(final DSLContext ctx, final Map expectedSyncStates, @Nullable final OffsetDateTime expectedTimestamp) { for (final Map.Entry entry : expectedSyncStates.entrySet()) { final var record = ctx - .select(V0_29_21_001__Store_last_sync_state.COLUMN_STATE, - V0_29_21_001__Store_last_sync_state.COLUMN_CREATED_AT, - V0_29_21_001__Store_last_sync_state.COLUMN_UPDATED_AT) - .from(V0_29_21_001__Store_last_sync_state.SYNC_STATE_TABLE) - .where(V0_29_21_001__Store_last_sync_state.COLUMN_SYNC_ID.eq(UUID.fromString(entry.getKey()))) + .select(COLUMN_CONFIG_BLOB, + COLUMN_CREATED_AT, + COLUMN_UPDATED_AT) + .from(TABLE_AIRBYTE_CONFIGS) + .where(COLUMN_CONFIG_ID.eq(entry.getKey()), + COLUMN_CONFIG_TYPE.eq(STANDARD_SYNC_STATE)) .fetchOne(); assertEquals(entry.getValue(), Jsons.deserialize(record.value1().data())); if (expectedTimestamp != null) { From 31ee0a22eda8abbc15f6bf2175e3b148a467c6ad Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 04:39:33 -0700 Subject: [PATCH 10/23] Add standard sync state --- .../java/io/airbyte/config/ConfigSchema.java | 4 ++++ .../main/resources/types/StandardSyncState.yaml | 17 +++++++++++++++++ 2 files changed, 21 insertions(+) create mode 100644 airbyte-config/models/src/main/resources/types/StandardSyncState.yaml diff --git a/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java b/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java index 81626ced0c2e..bf5a9436b52d 100644 --- a/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java +++ b/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java @@ -46,6 +46,10 @@ public enum ConfigSchema implements AirbyteConfig { StandardSyncOperation.class, standardSyncOperation -> standardSyncOperation.getOperationId().toString(), "operationId"), + STANDARD_SYNC_STATE("StandardSyncState.yaml", + StandardSyncState.class, + standardSyncState -> standardSyncState.getConnectionId().toString(), + "connectionId"), SOURCE_OAUTH_PARAM("SourceOAuthParameter.yaml", SourceOAuthParameter.class, sourceOAuthParameter -> sourceOAuthParameter.getOauthParameterId().toString(), diff --git a/airbyte-config/models/src/main/resources/types/StandardSyncState.yaml b/airbyte-config/models/src/main/resources/types/StandardSyncState.yaml new file mode 100644 index 000000000000..964e084bc696 --- /dev/null +++ b/airbyte-config/models/src/main/resources/types/StandardSyncState.yaml @@ -0,0 +1,17 @@ +--- +"$schema": http://json-schema.org/draft-07/schema# +"$id": https://github.com/airbytehq/airbyte/blob/master/airbyte-config/models/src/main/resources/types/StandardSyncState.yaml +title: StandardSyncState +description: The current state of a connection (i.e. StandardSync). +type: object +additionalProperties: false +required: + - connectionId +properties: + connectionId: + type: string + format: uuid + description: This is a foreign key that references a connection (i.e. StandardSync). + state: + "$ref": State.yaml + description: The current (latest) connection state. From 1c97b34840f687bd73baa8494036b1d4eced51c5 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 05:07:03 -0700 Subject: [PATCH 11/23] Move state methods to config repository --- .../persistence/ConfigPersistence2.java | 83 ------------------- .../config/persistence/ConfigRepository.java | 25 ++++++ .../persistence/ConfigPersistence2Test.java | 7 -- .../airbyte/scheduler/app/JobScheduler.java | 4 +- .../airbyte/scheduler/app/SchedulerApp.java | 8 +- .../persistence/DefaultJobCreator.java | 10 +-- .../persistence/DefaultJobCreatorTest.java | 7 +- .../server/ConfigurationApiFactory.java | 5 -- .../java/io/airbyte/server/ServerApp.java | 5 +- .../java/io/airbyte/server/ServerFactory.java | 4 - .../airbyte/server/apis/ConfigurationApi.java | 3 - .../server/handlers/SchedulerHandler.java | 8 +- .../server/handlers/SchedulerHandlerTest.java | 8 +- .../java/io/airbyte/workers/WorkerApp.java | 20 +++-- .../workers/temporal/SyncWorkflow.java | 10 +-- 15 files changed, 59 insertions(+), 148 deletions(-) delete mode 100644 airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java delete mode 100644 airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java deleted file mode 100644 index c15bdafd72d5..000000000000 --- a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigPersistence2.java +++ /dev/null @@ -1,83 +0,0 @@ -/* - * Copyright (c) 2021 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.config.persistence; - -import static io.airbyte.db.instance.configs.jooq.Tables.SYNC_STATE; - -import com.google.common.annotations.VisibleForTesting; -import io.airbyte.commons.json.Jsons; -import io.airbyte.config.State; -import io.airbyte.db.Database; -import io.airbyte.db.ExceptionWrappingDatabase; -import java.io.IOException; -import java.time.Instant; -import java.time.OffsetDateTime; -import java.time.ZoneOffset; -import java.util.Optional; -import java.util.UUID; -import java.util.function.Supplier; -import org.jooq.JSONB; -import org.jooq.Record1; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -public class ConfigPersistence2 { - - private static final Logger LOGGER = LoggerFactory.getLogger(ConfigPersistence2.class); - - private final ExceptionWrappingDatabase configDatabase; - private final Supplier timeSupplier; - - public ConfigPersistence2(final Database configDatabase) { - this(configDatabase, Instant::now); - } - - @VisibleForTesting - ConfigPersistence2(final Database configDatabase, final Supplier timeSupplier) { - this.configDatabase = new ExceptionWrappingDatabase(configDatabase); - - this.timeSupplier = timeSupplier; - } - - public Optional getCurrentState(final UUID connectionId) throws IOException { - return configDatabase.query(ctx -> { - final Record1 record = ctx.select(SYNC_STATE.STATE) - .from(SYNC_STATE) - .where(SYNC_STATE.SYNC_ID.eq(connectionId)) - .fetchAny(); - if (record == null) { - return Optional.empty(); - } - - return Optional.of(Jsons.deserialize(record.value1().data(), State.class)); - }); - } - - public void updateSyncState(final UUID connectionId, final State state) throws IOException { - final OffsetDateTime now = OffsetDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - - configDatabase.transaction( - ctx -> { - final boolean hasExistingRecord = ctx.fetchExists(SYNC_STATE, SYNC_STATE.SYNC_ID.eq(connectionId)); - if (hasExistingRecord) { - LOGGER.info("Updating connection {} state", connectionId); - return ctx.update(SYNC_STATE) - .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(state))) - .set(SYNC_STATE.UPDATED_AT, now) - .where(SYNC_STATE.SYNC_ID.eq(connectionId)) - .execute(); - } else { - LOGGER.info("Inserting new state for connection {}", connectionId); - return ctx.insertInto(SYNC_STATE) - .set(SYNC_STATE.SYNC_ID, connectionId) - .set(SYNC_STATE.STATE, JSONB.valueOf(Jsons.serialize(state))) - .set(SYNC_STATE.CREATED_AT, now) - .set(SYNC_STATE.UPDATED_AT, now) - .execute(); - } - }); - } - -} diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java index 0825eac538a0..b42e5823b44a 100644 --- a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java +++ b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java @@ -20,7 +20,9 @@ import io.airbyte.config.StandardSourceDefinition; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSyncOperation; +import io.airbyte.config.StandardSyncState; import io.airbyte.config.StandardWorkspace; +import io.airbyte.config.State; import io.airbyte.config.persistence.split_secrets.SecretPersistence; import io.airbyte.config.persistence.split_secrets.SecretsHelpers; import io.airbyte.config.persistence.split_secrets.SecretsHydrator; @@ -433,6 +435,29 @@ public List listDestinationOAuthParam() throws JsonVa return persistence.listConfigs(ConfigSchema.DESTINATION_OAUTH_PARAM, DestinationOAuthParameter.class); } + public Optional getCurrentState(final UUID connectionId) throws IOException { + try { + final StandardSyncState connectionState = persistence.getConfig( + ConfigSchema.STANDARD_SYNC_STATE, + connectionId.toString(), + StandardSyncState.class); + return Optional.of(connectionState.getState()); + } catch (final ConfigNotFoundException e) { + return Optional.empty(); + } catch (final JsonValidationException e) { + throw new IllegalStateException(e); + } + } + + public void updateSyncState(final UUID connectionId, final State state) throws IOException { + final StandardSyncState connectionState = new StandardSyncState().withConnectionId(connectionId).withState(state); + try { + persistence.writeConfig(ConfigSchema.STANDARD_SYNC_STATE, connectionId.toString(), connectionState); + } catch (final JsonValidationException e) { + throw new IllegalStateException(e); + } + } + /** * Converts between a dumpConfig() output and a replaceAllConfigs() input, by deserializing the * string/jsonnode into the AirbyteConfig, Stream diff --git a/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java b/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java deleted file mode 100644 index fca39e71cf21..000000000000 --- a/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigPersistence2Test.java +++ /dev/null @@ -1,7 +0,0 @@ -/* - * Copyright (c) 2021 Airbyte, Inc., all rights reserved. - */ - -package io.airbyte.config.persistence; - -class ConfigPersistence2Test {} diff --git a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java index d6004a797ef6..7e2db6e7b900 100644 --- a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java +++ b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/JobScheduler.java @@ -9,7 +9,6 @@ import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSync.Status; import io.airbyte.config.persistence.ConfigNotFoundException; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.scheduler.models.Job; import io.airbyte.scheduler.persistence.DefaultJobCreator; @@ -48,7 +47,6 @@ public class JobScheduler implements Runnable { } public JobScheduler(final JobPersistence jobPersistence, - final ConfigPersistence2 configPersistence2, final ConfigRepository configRepository, final TrackingClient trackingClient) { this( @@ -56,7 +54,7 @@ public JobScheduler(final JobPersistence jobPersistence, configRepository, new ScheduleJobPredicate(Instant::now), new DefaultSyncJobFactory( - new DefaultJobCreator(jobPersistence, configPersistence2), + new DefaultJobCreator(jobPersistence, configRepository), configRepository, new OAuthConfigSupplier(configRepository, false, trackingClient))); } diff --git a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java index 0793e16a165d..d02b1cd275c9 100644 --- a/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java +++ b/airbyte-scheduler/app/src/main/java/io/airbyte/scheduler/app/SchedulerApp.java @@ -18,7 +18,6 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.helpers.LogClientSingleton; import io.airbyte.config.persistence.ConfigPersistence; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.config.persistence.DatabaseConfigPersistence; import io.airbyte.config.persistence.split_secrets.SecretPersistence; @@ -76,7 +75,6 @@ public class SchedulerApp { private final Path workspaceRoot; private final JobPersistence jobPersistence; private final ConfigRepository configRepository; - private final ConfigPersistence2 configPersistence2; private final JobCleaner jobCleaner; private final JobNotifier jobNotifier; private final TemporalClient temporalClient; @@ -87,7 +85,6 @@ public class SchedulerApp { public SchedulerApp(final Path workspaceRoot, final JobPersistence jobPersistence, final ConfigRepository configRepository, - final ConfigPersistence2 configPersistence2, final JobCleaner jobCleaner, final JobNotifier jobNotifier, final TemporalClient temporalClient, @@ -97,7 +94,6 @@ public SchedulerApp(final Path workspaceRoot, this.workspaceRoot = workspaceRoot; this.jobPersistence = jobPersistence; this.configRepository = configRepository; - this.configPersistence2 = configPersistence2; this.jobCleaner = jobCleaner; this.jobNotifier = jobNotifier; this.temporalClient = temporalClient; @@ -114,7 +110,7 @@ public void start() throws IOException { final TemporalWorkerRunFactory temporalWorkerRunFactory = new TemporalWorkerRunFactory(temporalClient, workspaceRoot, airbyteVersionOrWarnings); final JobRetrier jobRetrier = new JobRetrier(jobPersistence, Instant::now, jobNotifier, maxSyncJobAttempts); final TrackingClient trackingClient = TrackingClientSingleton.get(); - final JobScheduler jobScheduler = new JobScheduler(jobPersistence, configPersistence2, configRepository, trackingClient); + final JobScheduler jobScheduler = new JobScheduler(jobPersistence, configRepository, trackingClient); final JobSubmitter jobSubmitter = new JobSubmitter( workerThreadPool, jobPersistence, @@ -213,7 +209,6 @@ public static void main(final String[] args) throws IOException, InterruptedExce configs.getConfigDatabasePassword(), configs.getConfigDatabaseUrl()) .getInitialized(); - final ConfigPersistence2 configPersistence2 = new ConfigPersistence2(configDatabase); final ConfigPersistence configPersistence = new DatabaseConfigPersistence(configDatabase).withValidation(); final Optional secretPersistence = SecretPersistence.getLongLived(configs); final Optional ephemeralSecretPersistence = SecretPersistence.getEphemeral(configs); @@ -248,7 +243,6 @@ public static void main(final String[] args) throws IOException, InterruptedExce workspaceRoot, jobPersistence, configRepository, - configPersistence2, jobCleaner, jobNotifier, temporalClient, diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java index 6fbd19bdf570..616064195b6a 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java @@ -12,7 +12,7 @@ import io.airbyte.config.SourceConnection; import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSyncOperation; -import io.airbyte.config.persistence.ConfigPersistence2; +import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.protocol.models.ConfiguredAirbyteCatalog; import io.airbyte.protocol.models.DestinationSyncMode; import io.airbyte.protocol.models.SyncMode; @@ -23,11 +23,11 @@ public class DefaultJobCreator implements JobCreator { private final JobPersistence jobPersistence; - private final ConfigPersistence2 configPersistence2; + private final ConfigRepository configRepository; - public DefaultJobCreator(final JobPersistence jobPersistence, final ConfigPersistence2 configPersistence2) { + public DefaultJobCreator(final JobPersistence jobPersistence, final ConfigRepository configRepository) { this.jobPersistence = jobPersistence; - this.configPersistence2 = configPersistence2; + this.configRepository = configRepository; } @Override @@ -52,7 +52,7 @@ public Optional createSyncJob(final SourceConnection source, .withState(null) .withResourceRequirements(standardSync.getResourceRequirements()); - configPersistence2.getCurrentState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState); + configRepository.getCurrentState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState); final JobConfig jobConfig = new JobConfig() .withConfigType(ConfigType.SYNC) diff --git a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobCreatorTest.java b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobCreatorTest.java index fd0acbea3704..e47d14d9d79b 100644 --- a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobCreatorTest.java +++ b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobCreatorTest.java @@ -24,6 +24,7 @@ import io.airbyte.config.StandardSync; import io.airbyte.config.StandardSyncOperation; import io.airbyte.config.StandardSyncOperation.OperatorType; +import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.protocol.models.CatalogHelpers; import io.airbyte.protocol.models.ConfiguredAirbyteCatalog; import io.airbyte.protocol.models.ConfiguredAirbyteStream; @@ -52,6 +53,7 @@ public class DefaultJobCreatorTest { private static final long JOB_ID = 12L; private JobPersistence jobPersistence; + private ConfigRepository configRepository; private JobCreator jobCreator; static { @@ -109,9 +111,10 @@ public class DefaultJobCreatorTest { } @BeforeEach - void setup() throws IOException { + void setup() { jobPersistence = mock(JobPersistence.class); - jobCreator = new DefaultJobCreator(jobPersistence); + configRepository = mock(ConfigRepository.class); + jobCreator = new DefaultJobCreator(jobPersistence, configRepository); } @Test diff --git a/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java b/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java index 519b3cfef094..cb1a52c61e2a 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ConfigurationApiFactory.java @@ -8,7 +8,6 @@ import io.airbyte.commons.io.FileTtlManager; import io.airbyte.config.Configs; import io.airbyte.config.persistence.ConfigPersistence; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.db.Database; import io.airbyte.scheduler.client.CachingSynchronousSchedulerClient; @@ -24,7 +23,6 @@ public class ConfigurationApiFactory implements Factory { private static WorkflowServiceStubs temporalService; private static ConfigRepository configRepository; - private static ConfigPersistence2 configPersistence2; private static JobPersistence jobPersistence; private static ConfigPersistence seed; private static SchedulerJobClient schedulerJobClient; @@ -39,7 +37,6 @@ public class ConfigurationApiFactory implements Factory { public static void setValues( final WorkflowServiceStubs temporalService, final ConfigRepository configRepository, - final ConfigPersistence2 configPersistence2, final JobPersistence jobPersistence, final ConfigPersistence seed, final SchedulerJobClient schedulerJobClient, @@ -51,7 +48,6 @@ public static void setValues( final Database jobsDatabase, final TrackingClient trackingClient) { ConfigurationApiFactory.configRepository = configRepository; - ConfigurationApiFactory.configPersistence2 = configPersistence2; ConfigurationApiFactory.jobPersistence = jobPersistence; ConfigurationApiFactory.seed = seed; ConfigurationApiFactory.schedulerJobClient = schedulerJobClient; @@ -71,7 +67,6 @@ public ConfigurationApi provide() { return new ConfigurationApi( ConfigurationApiFactory.configRepository, - ConfigurationApiFactory.configPersistence2, ConfigurationApiFactory.jobPersistence, ConfigurationApiFactory.seed, ConfigurationApiFactory.schedulerJobClient, 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 78e8f2417b79..f12273be2df9 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ServerApp.java @@ -16,7 +16,6 @@ import io.airbyte.config.StandardWorkspace; import io.airbyte.config.helpers.LogClientSingleton; import io.airbyte.config.persistence.ConfigPersistence; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.config.persistence.DatabaseConfigPersistence; import io.airbyte.config.persistence.YamlSeedConfigPersistence; @@ -177,7 +176,6 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con final ConfigRepository configRepository = new ConfigRepository(configPersistence.withValidation(), secretsHydrator, secretPersistence, ephemeralSecretPersistence); - final ConfigPersistence2 configPersistence2 = new ConfigPersistence2(configDatabase); LOGGER.info("Creating Scheduler persistence..."); final Database jobDatabase = new JobsDatabaseInstance( @@ -212,7 +210,7 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con final TemporalClient temporalClient = TemporalClient.production(configs.getTemporalHost(), configs.getWorkspaceRoot()); final OAuthConfigSupplier oAuthConfigSupplier = new OAuthConfigSupplier(configRepository, false, trackingClient); final SchedulerJobClient schedulerJobClient = - new DefaultSchedulerJobClient(jobPersistence, new DefaultJobCreator(jobPersistence, configPersistence2)); + new DefaultSchedulerJobClient(jobPersistence, new DefaultJobCreator(jobPersistence, configRepository)); final DefaultSynchronousSchedulerClient syncSchedulerClient = new DefaultSynchronousSchedulerClient(temporalClient, jobTracker, oAuthConfigSupplier); final SynchronousSchedulerClient bucketSpecCacheSchedulerClient = @@ -248,7 +246,6 @@ public static ServerRunnable getServer(final ServerFactory apiFactory, final Con cachingSchedulerClient, temporalService, configRepository, - configPersistence2, jobPersistence, seed, configDatabase, diff --git a/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java b/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java index 93ebb0a6b05c..d58f0fce32d8 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java +++ b/airbyte-server/src/main/java/io/airbyte/server/ServerFactory.java @@ -8,7 +8,6 @@ import io.airbyte.commons.io.FileTtlManager; import io.airbyte.config.Configs; import io.airbyte.config.persistence.ConfigPersistence; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.db.Database; import io.airbyte.scheduler.client.SchedulerJobClient; @@ -26,7 +25,6 @@ ServerRunnable create(SchedulerJobClient schedulerJobClient, SpecCachingSynchronousSchedulerClient cachingSchedulerClient, WorkflowServiceStubs temporalService, ConfigRepository configRepository, - ConfigPersistence2 configPersistence2, JobPersistence jobPersistence, ConfigPersistence seed, Database configsDatabase, @@ -41,7 +39,6 @@ public ServerRunnable create(final SchedulerJobClient schedulerJobClient, final SpecCachingSynchronousSchedulerClient cachingSchedulerClient, final WorkflowServiceStubs temporalService, final ConfigRepository configRepository, - final ConfigPersistence2 configPersistence2, final JobPersistence jobPersistence, final ConfigPersistence seed, final Database configsDatabase, @@ -52,7 +49,6 @@ public ServerRunnable create(final SchedulerJobClient schedulerJobClient, ConfigurationApiFactory.setValues( temporalService, configRepository, - configPersistence2, jobPersistence, seed, schedulerJobClient, diff --git a/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java b/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java index 3acee99bd292..3b2dc4f0e846 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java +++ b/airbyte-server/src/main/java/io/airbyte/server/apis/ConfigurationApi.java @@ -84,7 +84,6 @@ import io.airbyte.config.Configs; import io.airbyte.config.persistence.ConfigNotFoundException; import io.airbyte.config.persistence.ConfigPersistence; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.db.Database; import io.airbyte.scheduler.client.CachingSynchronousSchedulerClient; @@ -146,7 +145,6 @@ public class ConfigurationApi implements io.airbyte.api.V1Api { private final Configs configs; public ConfigurationApi(final ConfigRepository configRepository, - final ConfigPersistence2 configPersistence2, final JobPersistence jobPersistence, final ConfigPersistence seed, final SchedulerJobClient schedulerJobClient, @@ -166,7 +164,6 @@ public ConfigurationApi(final ConfigRepository configRepository, trackingClient); schedulerHandler = new SchedulerHandler( configRepository, - configPersistence2, schedulerJobClient, synchronousSchedulerClient, jobPersistence, diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java index a75beab35f06..15a5064d3d3a 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java @@ -37,7 +37,6 @@ import io.airbyte.config.StandardSyncOperation; import io.airbyte.config.State; import io.airbyte.config.persistence.ConfigNotFoundException; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.protocol.models.AirbyteCatalog; import io.airbyte.protocol.models.ConnectorSpecification; @@ -80,10 +79,8 @@ public class SchedulerHandler { private final JobNotifier jobNotifier; private final WorkflowServiceStubs temporalService; private final OAuthConfigSupplier oAuthConfigSupplier; - private final ConfigPersistence2 configPersistence2; public SchedulerHandler(final ConfigRepository configRepository, - final ConfigPersistence2 configPersistence2, final SchedulerJobClient schedulerJobClient, final SynchronousSchedulerClient synchronousSchedulerClient, final JobPersistence jobPersistence, @@ -92,7 +89,6 @@ public SchedulerHandler(final ConfigRepository configRepository, final OAuthConfigSupplier oAuthConfigSupplier) { this( configRepository, - configPersistence2, schedulerJobClient, synchronousSchedulerClient, new ConfigurationUpdate(configRepository, new SpecFetcher(synchronousSchedulerClient)), @@ -106,7 +102,6 @@ public SchedulerHandler(final ConfigRepository configRepository, @VisibleForTesting SchedulerHandler(final ConfigRepository configRepository, - final ConfigPersistence2 configPersistence2, final SchedulerJobClient schedulerJobClient, final SynchronousSchedulerClient synchronousSchedulerClient, final ConfigurationUpdate configurationUpdate, @@ -117,7 +112,6 @@ public SchedulerHandler(final ConfigRepository configRepository, final WorkflowServiceStubs temporalService, final OAuthConfigSupplier oAuthConfigSupplier) { this.configRepository = configRepository; - this.configPersistence2 = configPersistence2; this.schedulerJobClient = schedulerJobClient; this.synchronousSchedulerClient = synchronousSchedulerClient; this.configurationUpdate = configurationUpdate; @@ -359,7 +353,7 @@ public JobInfoRead resetConnection(final ConnectionIdRequestBody connectionIdReq } public ConnectionState getState(final ConnectionIdRequestBody connectionIdRequestBody) throws IOException { - final Optional currentState = configPersistence2.getCurrentState(connectionIdRequestBody.getConnectionId()); + final Optional currentState = configRepository.getCurrentState(connectionIdRequestBody.getConnectionId()); LOGGER.info("currentState server: {}", currentState); final ConnectionState connectionState = new ConnectionState() diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java index ccfcf7c6e82f..e972248310f6 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java @@ -49,7 +49,6 @@ import io.airbyte.config.StandardSyncOperation.OperatorType; import io.airbyte.config.State; import io.airbyte.config.persistence.ConfigNotFoundException; -import io.airbyte.config.persistence.ConfigPersistence2; import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.protocol.models.AirbyteCatalog; import io.airbyte.protocol.models.CatalogHelpers; @@ -118,7 +117,6 @@ class SchedulerHandlerTest { private SchedulerHandler schedulerHandler; private ConfigRepository configRepository; - private ConfigPersistence2 configPersistence2; private Job completedJob; private SchedulerJobClient schedulerJobClient; private SynchronousSchedulerClient synchronousSchedulerClient; @@ -142,13 +140,11 @@ void setup() { schedulerJobClient = spy(SchedulerJobClient.class); synchronousSchedulerClient = mock(SynchronousSchedulerClient.class); configRepository = mock(ConfigRepository.class); - configPersistence2 = mock(ConfigPersistence2.class); jobPersistence = mock(JobPersistence.class); final JobNotifier jobNotifier = mock(JobNotifier.class); schedulerHandler = new SchedulerHandler( configRepository, - configPersistence2, schedulerJobClient, synchronousSchedulerClient, configurationUpdate, @@ -561,7 +557,7 @@ void testResetConnection() throws JsonValidationException, IOException, ConfigNo void testGetCurrentState() throws IOException { final UUID connectionId = UUID.randomUUID(); final State state = new State().withState(Jsons.jsonNode(ImmutableMap.of("checkpoint", 1))); - when(configPersistence2.getCurrentState(connectionId)).thenReturn(Optional.of(state)); + when(configRepository.getCurrentState(connectionId)).thenReturn(Optional.of(state)); final ConnectionState connectionState = schedulerHandler.getState(new ConnectionIdRequestBody().connectionId(connectionId)); assertEquals(new ConnectionState().connectionId(connectionId).state(state.getState()), connectionState); @@ -570,7 +566,7 @@ void testGetCurrentState() throws IOException { @Test void testGetCurrentStateEmpty() throws IOException { final UUID connectionId = UUID.randomUUID(); - when(configPersistence2.getCurrentState(connectionId)).thenReturn(Optional.empty()); + when(configRepository.getCurrentState(connectionId)).thenReturn(Optional.empty()); final ConnectionState connectionState = schedulerHandler.getState(new ConnectionIdRequestBody().connectionId(connectionId)); assertEquals(new ConnectionState().connectionId(connectionId), connectionState); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java index 307eddd162ba..acb64f336acd 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java @@ -9,7 +9,9 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.MaxWorkersConfig; import io.airbyte.config.helpers.LogClientSingleton; -import io.airbyte.config.persistence.ConfigPersistence2; +import io.airbyte.config.persistence.ConfigPersistence; +import io.airbyte.config.persistence.ConfigRepository; +import io.airbyte.config.persistence.DatabaseConfigPersistence; import io.airbyte.config.persistence.split_secrets.SecretPersistence; import io.airbyte.config.persistence.split_secrets.SecretsHydrator; import io.airbyte.db.Database; @@ -37,6 +39,7 @@ import java.net.InetAddress; import java.nio.file.Path; import java.util.Map; +import java.util.Optional; import java.util.concurrent.Executors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -53,7 +56,7 @@ public class WorkerApp { private final WorkflowServiceStubs temporalService; private final MaxWorkersConfig maxWorkers; private final WorkerEnvironment workerEnvironment; - private final ConfigPersistence2 configPersistence2; + private final ConfigRepository configRepository; public WorkerApp(final Path workspaceRoot, final ProcessFactory processFactory, @@ -61,14 +64,14 @@ public WorkerApp(final Path workspaceRoot, final WorkflowServiceStubs temporalService, final MaxWorkersConfig maxWorkers, final WorkerEnvironment workerEnvironment, - final ConfigPersistence2 configPersistence2) { + final ConfigRepository configRepository) { this.workspaceRoot = workspaceRoot; this.processFactory = processFactory; this.secretsHydrator = secretsHydrator; this.temporalService = temporalService; this.maxWorkers = maxWorkers; this.workerEnvironment = workerEnvironment; - this.configPersistence2 = configPersistence2; + this.configRepository = configRepository; } public void start() { @@ -106,7 +109,7 @@ public void start() { new SyncWorkflow.ReplicationActivityImpl(processFactory, secretsHydrator, workspaceRoot), new SyncWorkflow.NormalizationActivityImpl(processFactory, secretsHydrator, workspaceRoot, workerEnvironment), new SyncWorkflow.DbtTransformationActivityImpl(processFactory, secretsHydrator, workspaceRoot), - new SyncWorkflow.PersistStateActivityImpl(workspaceRoot, configPersistence2)); + new SyncWorkflow.PersistStateActivityImpl(workspaceRoot, configRepository)); factory.start(); } @@ -156,7 +159,10 @@ public static void main(final String[] args) throws IOException, InterruptedExce configs.getConfigDatabasePassword(), configs.getConfigDatabaseUrl()) .getInitialized(); - final ConfigPersistence2 configPersistence2 = new ConfigPersistence2(configDatabase); + final ConfigPersistence configPersistence = new DatabaseConfigPersistence(configDatabase).withValidation(); + final Optional secretPersistence = SecretPersistence.getLongLived(configs); + final Optional ephemeralSecretPersistence = SecretPersistence.getEphemeral(configs); + final ConfigRepository configRepository = new ConfigRepository(configPersistence, secretsHydrator, secretPersistence, ephemeralSecretPersistence); new WorkerApp( workspaceRoot, @@ -165,7 +171,7 @@ public static void main(final String[] args) throws IOException, InterruptedExce temporalService, configs.getMaxWorkers(), configs.getWorkerEnvironment(), - configPersistence2).start(); + configRepository).start(); } } diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java index b3725e978e96..a19d8f020ebb 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java @@ -21,7 +21,7 @@ import io.airbyte.config.StandardSyncOutput; import io.airbyte.config.StandardSyncSummary; import io.airbyte.config.State; -import io.airbyte.config.persistence.ConfigPersistence2; +import io.airbyte.config.persistence.ConfigRepository; import io.airbyte.config.persistence.split_secrets.SecretsHydrator; import io.airbyte.scheduler.models.IntegrationLauncherConfig; import io.airbyte.scheduler.models.JobRunConfig; @@ -416,11 +416,11 @@ class PersistStateActivityImpl implements PersistStateActivity { private static final Logger LOGGER = LoggerFactory.getLogger(PersistStateActivityImpl.class); private final Path workspaceRoot; - private final ConfigPersistence2 configPersistence2; + private final ConfigRepository configRepository; - public PersistStateActivityImpl(final Path workspaceRoot, final ConfigPersistence2 configPersistence2) { + public PersistStateActivityImpl(final Path workspaceRoot, final ConfigRepository configRepository) { this.workspaceRoot = workspaceRoot; - this.configPersistence2 = configPersistence2; + this.configRepository = configRepository; } @Override @@ -428,7 +428,7 @@ public boolean persist(final UUID connectionId, final StandardSyncOutput syncOut final State state = syncOutput.getState(); if (state != null) { try { - configPersistence2.updateSyncState(connectionId, state); + configRepository.updateSyncState(connectionId, state); } catch (final IOException e) { throw new RuntimeException(e); } From 04e4fd9613fc2b1eb9470713c397524e49e02e60 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 05:20:27 -0700 Subject: [PATCH 12/23] Add unit tests --- .../config/persistence/ConfigRepository.java | 4 +- .../persistence/ConfigRepositoryTest.java | 42 +++++++++++++++++++ .../persistence/DefaultJobCreator.java | 2 +- .../server/handlers/SchedulerHandler.java | 2 +- .../server/handlers/SchedulerHandlerTest.java | 4 +- .../workers/temporal/SyncWorkflow.java | 2 +- 6 files changed, 49 insertions(+), 7 deletions(-) diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java index b42e5823b44a..af49d6e9a09b 100644 --- a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java +++ b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java @@ -435,7 +435,7 @@ public List listDestinationOAuthParam() throws JsonVa return persistence.listConfigs(ConfigSchema.DESTINATION_OAUTH_PARAM, DestinationOAuthParameter.class); } - public Optional getCurrentState(final UUID connectionId) throws IOException { + public Optional getConnectionState(final UUID connectionId) throws IOException { try { final StandardSyncState connectionState = persistence.getConfig( ConfigSchema.STANDARD_SYNC_STATE, @@ -449,7 +449,7 @@ public Optional getCurrentState(final UUID connectionId) throws IOExcepti } } - public void updateSyncState(final UUID connectionId, final State state) throws IOException { + public void updateConnectionState(final UUID connectionId, final State state) throws IOException { final StandardSyncState connectionState = new StandardSyncState().withConnectionId(connectionId).withState(state); try { persistence.writeConfig(ConfigSchema.STANDARD_SYNC_STATE, connectionId.toString(), connectionState); diff --git a/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigRepositoryTest.java b/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigRepositoryTest.java index 6818f4e13128..7580216f12e4 100644 --- a/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigRepositoryTest.java +++ b/airbyte-config/persistence/src/test/java/io/airbyte/config/persistence/ConfigRepositoryTest.java @@ -6,16 +6,23 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.reset; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import io.airbyte.commons.json.Jsons; import io.airbyte.config.ConfigSchema; +import io.airbyte.config.StandardSyncState; import io.airbyte.config.StandardWorkspace; +import io.airbyte.config.State; import io.airbyte.config.persistence.split_secrets.MemorySecretPersistence; import io.airbyte.config.persistence.split_secrets.NoOpSecretsHydrator; import io.airbyte.validation.json.JsonValidationException; import java.io.IOException; import java.util.Optional; import java.util.UUID; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -34,6 +41,11 @@ void setup() { new ConfigRepository(configPersistence, new NoOpSecretsHydrator(), Optional.of(secretPersistence), Optional.of(secretPersistence)); } + @AfterEach + void cleanUp() { + reset(configPersistence); + } + @Test void testWorkspaceWithNullTombstone() throws ConfigNotFoundException, IOException, JsonValidationException { assertReturnsWorkspace(new StandardWorkspace().withWorkspaceId(WORKSPACE_ID)); @@ -55,4 +67,34 @@ void assertReturnsWorkspace(final StandardWorkspace workspace) throws ConfigNotF assertEquals(workspace, configRepository.getStandardWorkspace(WORKSPACE_ID, true)); } + @Test + void testGetConnectionState() throws Exception { + final UUID connectionId = UUID.randomUUID(); + final State state = new State().withState(Jsons.deserialize("{ \"cursor\": 1000 }")); + final StandardSyncState connectionState = new StandardSyncState().withConnectionId(connectionId).withState(state); + + when(configPersistence.getConfig(ConfigSchema.STANDARD_SYNC_STATE, connectionId.toString(), StandardSyncState.class)) + .thenThrow(new ConfigNotFoundException(ConfigSchema.STANDARD_SYNC_STATE, connectionId)); + assertEquals(Optional.empty(), configRepository.getConnectionState(connectionId)); + + reset(configPersistence); + when(configPersistence.getConfig(ConfigSchema.STANDARD_SYNC_STATE, connectionId.toString(), StandardSyncState.class)) + .thenReturn(connectionState); + assertEquals(Optional.of(state), configRepository.getConnectionState(connectionId)); + } + + @Test + void testUpdateConnectionState() throws Exception { + final UUID connectionId = UUID.randomUUID(); + final State state1 = new State().withState(Jsons.deserialize("{ \"cursor\": 1 }")); + final StandardSyncState connectionState1 = new StandardSyncState().withConnectionId(connectionId).withState(state1); + final State state2 = new State().withState(Jsons.deserialize("{ \"cursor\": 2 }")); + final StandardSyncState connectionState2 = new StandardSyncState().withConnectionId(connectionId).withState(state2); + + configRepository.updateConnectionState(connectionId, state1); + verify(configPersistence, times(1)).writeConfig(ConfigSchema.STANDARD_SYNC_STATE, connectionId.toString(), connectionState1); + configRepository.updateConnectionState(connectionId, state2); + verify(configPersistence, times(1)).writeConfig(ConfigSchema.STANDARD_SYNC_STATE, connectionId.toString(), connectionState2); + } + } diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java index 616064195b6a..dd2194cad4dc 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobCreator.java @@ -52,7 +52,7 @@ public Optional createSyncJob(final SourceConnection source, .withState(null) .withResourceRequirements(standardSync.getResourceRequirements()); - configRepository.getCurrentState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState); + configRepository.getConnectionState(standardSync.getConnectionId()).ifPresent(jobSyncConfig::withState); final JobConfig jobConfig = new JobConfig() .withConfigType(ConfigType.SYNC) diff --git a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java index 15a5064d3d3a..1f025d29672a 100644 --- a/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java +++ b/airbyte-server/src/main/java/io/airbyte/server/handlers/SchedulerHandler.java @@ -353,7 +353,7 @@ public JobInfoRead resetConnection(final ConnectionIdRequestBody connectionIdReq } public ConnectionState getState(final ConnectionIdRequestBody connectionIdRequestBody) throws IOException { - final Optional currentState = configRepository.getCurrentState(connectionIdRequestBody.getConnectionId()); + final Optional currentState = configRepository.getConnectionState(connectionIdRequestBody.getConnectionId()); LOGGER.info("currentState server: {}", currentState); final ConnectionState connectionState = new ConnectionState() diff --git a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java index e972248310f6..24681797ab91 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/handlers/SchedulerHandlerTest.java @@ -557,7 +557,7 @@ void testResetConnection() throws JsonValidationException, IOException, ConfigNo void testGetCurrentState() throws IOException { final UUID connectionId = UUID.randomUUID(); final State state = new State().withState(Jsons.jsonNode(ImmutableMap.of("checkpoint", 1))); - when(configRepository.getCurrentState(connectionId)).thenReturn(Optional.of(state)); + when(configRepository.getConnectionState(connectionId)).thenReturn(Optional.of(state)); final ConnectionState connectionState = schedulerHandler.getState(new ConnectionIdRequestBody().connectionId(connectionId)); assertEquals(new ConnectionState().connectionId(connectionId).state(state.getState()), connectionState); @@ -566,7 +566,7 @@ void testGetCurrentState() throws IOException { @Test void testGetCurrentStateEmpty() throws IOException { final UUID connectionId = UUID.randomUUID(); - when(configRepository.getCurrentState(connectionId)).thenReturn(Optional.empty()); + when(configRepository.getConnectionState(connectionId)).thenReturn(Optional.empty()); final ConnectionState connectionState = schedulerHandler.getState(new ConnectionIdRequestBody().connectionId(connectionId)); assertEquals(new ConnectionState().connectionId(connectionId), connectionState); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java index a19d8f020ebb..6b894c55ed88 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java @@ -428,7 +428,7 @@ public boolean persist(final UUID connectionId, final StandardSyncOutput syncOut final State state = syncOutput.getState(); if (state != null) { try { - configRepository.updateSyncState(connectionId, state); + configRepository.updateConnectionState(connectionId, state); } catch (final IOException e) { throw new RuntimeException(e); } From b7788dc3d44a5dadb3bedb615c4c4e5d2dd24e89 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 06:00:10 -0700 Subject: [PATCH 13/23] Fix unit tests --- .../V0_30_22_001__Store_last_sync_state.java | 1 - .../main/resources/configs_database/schema_dump.txt | 7 ------- .../V0_30_22_001__Store_last_sync_state_test.java | 5 +++-- .../app/worker_run/TemporalWorkerRunFactoryTest.java | 11 +++++++---- .../persistence/DefaultJobPersistenceTest.java | 4 ---- .../io/airbyte/workers/temporal/SyncWorkflowTest.java | 9 +++++++-- .../airbyte/workers/temporal/TemporalClientTest.java | 5 +++-- 7 files changed, 20 insertions(+), 22 deletions(-) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java index abf2563bf21e..fdac8d487e92 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java @@ -113,7 +113,6 @@ static Optional getJobsDatabase(final Configs configs) { return Optional.of(jobsDatabase); } catch (final IllegalArgumentException e) { // If the environment variables do not exist, it means the migration is run in development. - // Connect to a mock job database, because we don't need to copy any data in test. LOGGER.info("[{}] This is the dev environment; there is no jobs database", MIGRATION_NAME); return Optional.empty(); } catch (final IOException e) { diff --git a/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt b/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt index 8bc8528d0993..6c9e11d400e4 100644 --- a/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt +++ b/airbyte-db/lib/src/main/resources/configs_database/schema_dump.txt @@ -22,12 +22,6 @@ create table "public"."airbyte_configs_migrations"( constraint "airbyte_configs_migrations_pk" primary key ("installed_rank") ); -create table "public"."sync_state"( - "sync_id" uuid not null, - "state" jsonb not null, - "created_at" timestamptz(35) not null default null, - "updated_at" timestamptz(35) not null default null -); create index "airbyte_configs_id_idx" on "public"."airbyte_configs"("config_id" asc); create unique index "airbyte_configs_pkey" on "public"."airbyte_configs"("id" asc); create unique index "airbyte_configs_type_id_idx" on "public"."airbyte_configs"( @@ -36,4 +30,3 @@ create unique index "airbyte_configs_type_id_idx" on "public"."airbyte_configs"( ); create unique index "airbyte_configs_migrations_pk" on "public"."airbyte_configs_migrations"("installed_rank" asc); create index "airbyte_configs_migrations_s_idx" on "public"."airbyte_configs_migrations"("success" asc); -create unique index "sync_state_sync_id_idx" on "public"."sync_state"("sync_id" asc); diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java index 7934e84f6dc0..9cc34694d210 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java @@ -11,6 +11,7 @@ import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_UPDATED_AT; import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.STANDARD_SYNC_STATE; import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.TABLE_AIRBYTE_CONFIGS; +import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.getStandardSyncState; import static org.jooq.impl.DSL.field; import static org.jooq.impl.DSL.table; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -70,8 +71,8 @@ class V0_30_22_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTe private static final JsonNode CONNECTION_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }"); private static final JsonNode CONNECTION_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }"); - private static final JsonNode STD_CONNECTION_STATE_2 = V0_30_22_001__Store_last_sync_state.getStandardSyncState(CONNECTION_2_ID.toString(), CONNECTION_2_STATE); - private static final JsonNode STD_CONNECTION_STATE_3 = V0_30_22_001__Store_last_sync_state.getStandardSyncState(CONNECTION_3_ID.toString(), CONNECTION_3_STATE); + private static final JsonNode STD_CONNECTION_STATE_2 = getStandardSyncState(CONNECTION_2_ID.toString(), CONNECTION_2_STATE); + private static final JsonNode STD_CONNECTION_STATE_3 = getStandardSyncState(CONNECTION_3_ID.toString(), CONNECTION_3_STATE); private static final Map CONNECTION_STATE_MAP = Map.of( CONNECTION_2_ID.toString(), STD_CONNECTION_STATE_2, CONNECTION_3_ID.toString(), STD_CONNECTION_STATE_3); diff --git a/airbyte-scheduler/app/src/test/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactoryTest.java b/airbyte-scheduler/app/src/test/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactoryTest.java index c3edfffc72f0..2e305afc7ec8 100644 --- a/airbyte-scheduler/app/src/test/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactoryTest.java +++ b/airbyte-scheduler/app/src/test/java/io/airbyte/scheduler/app/worker_run/TemporalWorkerRunFactoryTest.java @@ -27,12 +27,14 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.List; +import java.util.UUID; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; class TemporalWorkerRunFactoryTest { + private static final UUID CONNECTION_ID = UUID.randomUUID(); private static final long JOB_ID = 10L; private static final int ATTEMPT_ID = 20; @@ -50,6 +52,7 @@ void setup() throws IOException { job = mock(Job.class, RETURNS_DEEP_STUBS); when(job.getId()).thenReturn(JOB_ID); when(job.getAttemptsCount()).thenReturn(ATTEMPT_ID); + when(job.getScope()).thenReturn(CONNECTION_ID.toString()); } @SuppressWarnings("unchecked") @@ -57,11 +60,11 @@ void setup() throws IOException { void testSync() throws Exception { when(job.getConfigType()).thenReturn(ConfigType.SYNC); final TemporalResponse mockResponse = mock(TemporalResponse.class); - when(temporalClient.submitSync(JOB_ID, ATTEMPT_ID, job.getConfig().getSync())).thenReturn(mockResponse); + when(temporalClient.submitSync(JOB_ID, ATTEMPT_ID, job.getConfig().getSync(), CONNECTION_ID)).thenReturn(mockResponse); final WorkerRun workerRun = workerRunFactory.create(job); workerRun.call(); - verify(temporalClient).submitSync(JOB_ID, ATTEMPT_ID, job.getConfig().getSync()); + verify(temporalClient).submitSync(JOB_ID, ATTEMPT_ID, job.getConfig().getSync(), CONNECTION_ID); assertEquals(jobRoot, workerRun.getJobRoot()); } @@ -83,13 +86,13 @@ void testResetConnection() throws Exception { when(job.getConfigType()).thenReturn(ConfigType.RESET_CONNECTION); when(job.getConfig().getResetConnection()).thenReturn(resetConfig); final TemporalResponse mockResponse = mock(TemporalResponse.class); - when(temporalClient.submitSync(JOB_ID, ATTEMPT_ID, syncConfig)).thenReturn(mockResponse); + when(temporalClient.submitSync(JOB_ID, ATTEMPT_ID, syncConfig, CONNECTION_ID)).thenReturn(mockResponse); final WorkerRun workerRun = workerRunFactory.create(job); workerRun.call(); final ArgumentCaptor argument = ArgumentCaptor.forClass(JobSyncConfig.class); - verify(temporalClient).submitSync(eq(JOB_ID), eq(ATTEMPT_ID), argument.capture()); + verify(temporalClient).submitSync(eq(JOB_ID), eq(ATTEMPT_ID), argument.capture(), eq(CONNECTION_ID)); assertEquals(syncConfig, argument.getValue()); assertEquals(jobRoot, workerRun.getJobRoot()); } diff --git a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java index 43836c83336f..cb886bb5db55 100644 --- a/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java +++ b/airbyte-scheduler/persistence/src/test/java/io/airbyte/scheduler/persistence/DefaultJobPersistenceTest.java @@ -4,7 +4,6 @@ package io.airbyte.scheduler.persistence; -import static io.airbyte.db.instance.configs.jooq.Tables.SYNC_STATE; import static io.airbyte.db.instance.jobs.jooq.Tables.AIRBYTE_METADATA; import static io.airbyte.db.instance.jobs.jooq.Tables.ATTEMPTS; import static io.airbyte.db.instance.jobs.jooq.Tables.JOBS; @@ -177,9 +176,6 @@ private void resetDb() throws SQLException { jobDatabase.query(ctx -> ctx.truncateTable(JOBS).execute()); jobDatabase.query(ctx -> ctx.truncateTable(ATTEMPTS).execute()); jobDatabase.query(ctx -> ctx.truncateTable(AIRBYTE_METADATA).execute()); - // the airbyte_configs table cannot be truncated because the config database - // is considered ready for consumption only when there are records in this table - configDatabase.query(ctx -> ctx.truncateTable(SYNC_STATE).execute()); } private Result getJobRecord(final long jobId) throws SQLException { diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java index 66a2b867d3bb..5a3b01c49c93 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java @@ -26,6 +26,7 @@ import io.airbyte.workers.temporal.SyncWorkflow.DbtTransformationActivityImpl; import io.airbyte.workers.temporal.SyncWorkflow.NormalizationActivity; import io.airbyte.workers.temporal.SyncWorkflow.NormalizationActivityImpl; +import io.airbyte.workers.temporal.SyncWorkflow.PersistStateActivityImpl; import io.airbyte.workers.temporal.SyncWorkflow.ReplicationActivity; import io.airbyte.workers.temporal.SyncWorkflow.ReplicationActivityImpl; import io.temporal.api.common.v1.WorkflowExecution; @@ -52,6 +53,7 @@ class SyncWorkflowTest { private ReplicationActivityImpl replicationActivity; private NormalizationActivityImpl normalizationActivity; private DbtTransformationActivityImpl dbtTransformationActivity; + private PersistStateActivityImpl persistStateActivity; // AIRBYTE CONFIGURATION private static final long JOB_ID = 11L; @@ -70,6 +72,7 @@ class SyncWorkflowTest { .withAttemptId((long) ATTEMPT_ID) .withDockerImage(IMAGE_NAME2); + private StandardSync sync; private StandardSyncInput syncInput; private NormalizationInput normalizationInput; private OperatorDbtInput operatorDbtInput; @@ -85,6 +88,7 @@ public void setUp() { client = testEnv.getWorkflowClient(); final ImmutablePair syncPair = TestConfigHelpers.createSyncConfig(); + sync = syncPair.getKey(); syncInput = syncPair.getValue(); replicationSuccessOutput = new StandardSyncOutput().withOutputCatalog(syncInput.getCatalog()); @@ -99,15 +103,16 @@ public void setUp() { replicationActivity = mock(ReplicationActivityImpl.class); normalizationActivity = mock(NormalizationActivityImpl.class); dbtTransformationActivity = mock(DbtTransformationActivityImpl.class); + persistStateActivity = mock(PersistStateActivityImpl.class); } // bundle up all of the temporal worker setup / execution into one method. private StandardSyncOutput execute() { - worker.registerActivitiesImplementations(replicationActivity, normalizationActivity, dbtTransformationActivity); + worker.registerActivitiesImplementations(replicationActivity, normalizationActivity, dbtTransformationActivity, persistStateActivity); testEnv.start(); final SyncWorkflow workflow = client.newWorkflowStub(SyncWorkflow.class, WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build()); - return workflow.run(JOB_RUN_CONFIG, SOURCE_LAUNCHER_CONFIG, DESTINATION_LAUNCHER_CONFIG, syncInput); + return workflow.run(JOB_RUN_CONFIG, SOURCE_LAUNCHER_CONFIG, DESTINATION_LAUNCHER_CONFIG, syncInput, sync.getConnectionId()); } @Test 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 213649b422d4..7b91e484ecca 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 @@ -38,6 +38,7 @@ class TemporalClientTest { + private static final UUID CONNECTION_ID = UUID.randomUUID(); private static final UUID JOB_UUID = UUID.randomUUID(); private static final long JOB_ID = 11L; private static final int ATTEMPT_ID = 21; @@ -176,8 +177,8 @@ void testSubmitSync() { .withAttemptId((long) ATTEMPT_ID) .withDockerImage(IMAGE_NAME2); - temporalClient.submitSync(JOB_ID, ATTEMPT_ID, syncConfig); - discoverCatalogWorkflow.run(JOB_RUN_CONFIG, LAUNCHER_CONFIG, destinationLauncherConfig, input); + 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)); } From 6698c88cdff3e1ae9987ae9c4208a394e16f90ff Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 06:39:10 -0700 Subject: [PATCH 14/23] Register standard sync state in migration --- .../java/io/airbyte/migrate/Migrations.java | 3 +- .../migrate/migrations/MigrationV0_30_0.java | 64 +++++++++++++++++++ .../airbyte_config/StandardSyncState.yaml | 17 +++++ .../airbyte_config/State.yaml | 14 ++++ .../migrations/MigrationV0_30_0Test.java | 22 +++++++ 5 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 airbyte-migration/src/main/java/io/airbyte/migrate/migrations/MigrationV0_30_0.java create mode 100644 airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/StandardSyncState.yaml create mode 100644 airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/State.yaml create mode 100644 airbyte-migration/src/test/java/io/airbyte/migrate/migrations/MigrationV0_30_0Test.java diff --git a/airbyte-migration/src/main/java/io/airbyte/migrate/Migrations.java b/airbyte-migration/src/main/java/io/airbyte/migrate/Migrations.java index d7629b278d11..1725ec346b1b 100644 --- a/airbyte-migration/src/main/java/io/airbyte/migrate/Migrations.java +++ b/airbyte-migration/src/main/java/io/airbyte/migrate/Migrations.java @@ -17,6 +17,7 @@ import io.airbyte.migrate.migrations.MigrationV0_27_0; import io.airbyte.migrate.migrations.MigrationV0_28_0; import io.airbyte.migrate.migrations.MigrationV0_29_0; +import io.airbyte.migrate.migrations.MigrationV0_30_0; import io.airbyte.migrate.migrations.NoOpMigration; import java.util.List; @@ -39,7 +40,7 @@ public class Migrations { private static final Migration MIGRATION_V_0_27_0 = new MigrationV0_27_0(MIGRATION_V_0_26_0); public static final Migration MIGRATION_V_0_28_0 = new MigrationV0_28_0(MIGRATION_V_0_27_0); public static final Migration MIGRATION_V_0_29_0 = new MigrationV0_29_0(MIGRATION_V_0_28_0); - public static final Migration MIGRATION_V_0_30_0 = new NoOpMigration(MIGRATION_V_0_29_0, "0.30.0-alpha"); + public static final Migration MIGRATION_V_0_30_0 = new MigrationV0_30_0(MIGRATION_V_0_29_0); // all migrations must be added to the list in the order that they should be applied. public static final List MIGRATIONS = ImmutableList.of( diff --git a/airbyte-migration/src/main/java/io/airbyte/migrate/migrations/MigrationV0_30_0.java b/airbyte-migration/src/main/java/io/airbyte/migrate/migrations/MigrationV0_30_0.java new file mode 100644 index 000000000000..3ed383e32b35 --- /dev/null +++ b/airbyte-migration/src/main/java/io/airbyte/migrate/migrations/MigrationV0_30_0.java @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.migrate.migrations; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.common.annotations.VisibleForTesting; +import io.airbyte.migrate.Migration; +import io.airbyte.migrate.MigrationUtils; +import io.airbyte.migrate.ResourceId; +import io.airbyte.migrate.ResourceType; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import java.util.function.Consumer; +import java.util.stream.Stream; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * This migration adds the StandardSyncState in the migration resource. + */ +public class MigrationV0_30_0 extends BaseMigration implements Migration { + + private static final Logger LOGGER = LoggerFactory.getLogger(MigrationV0_30_0.class); + private static final String MIGRATION_VERSION = "0.30.0-alpha"; + + private static final Path RESOURCE_PATH = Path.of("migrations/migrationV0_30_0/airbyte_config"); + + @VisibleForTesting + protected static final ResourceId STANDARD_SYNC_STATE_RESOURCE_ID = ResourceId + .fromConstantCase(ResourceType.CONFIG, "STANDARD_SYNC_STATE"); + + private final Migration previousMigration; + + public MigrationV0_30_0(final Migration previousMigration) { + super(previousMigration); + this.previousMigration = previousMigration; + } + + @Override + public String getVersion() { + return MIGRATION_VERSION; + } + + @Override + public Map getOutputSchema() { + final Map outputSchema = new HashMap<>(previousMigration.getOutputSchema()); + outputSchema.put(STANDARD_SYNC_STATE_RESOURCE_ID, + MigrationUtils.getSchemaFromResourcePath(RESOURCE_PATH, STANDARD_SYNC_STATE_RESOURCE_ID)); + return outputSchema; + } + + // no op migration. + @Override + public void migrate(final Map> inputData, final Map> outputData) { + for (final Map.Entry> entry : inputData.entrySet()) { + final Consumer recordConsumer = outputData.get(entry.getKey()); + entry.getValue().forEach(recordConsumer); + } + } + +} diff --git a/airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/StandardSyncState.yaml b/airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/StandardSyncState.yaml new file mode 100644 index 000000000000..964e084bc696 --- /dev/null +++ b/airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/StandardSyncState.yaml @@ -0,0 +1,17 @@ +--- +"$schema": http://json-schema.org/draft-07/schema# +"$id": https://github.com/airbytehq/airbyte/blob/master/airbyte-config/models/src/main/resources/types/StandardSyncState.yaml +title: StandardSyncState +description: The current state of a connection (i.e. StandardSync). +type: object +additionalProperties: false +required: + - connectionId +properties: + connectionId: + type: string + format: uuid + description: This is a foreign key that references a connection (i.e. StandardSync). + state: + "$ref": State.yaml + description: The current (latest) connection state. diff --git a/airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/State.yaml b/airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/State.yaml new file mode 100644 index 000000000000..666941a7e3bb --- /dev/null +++ b/airbyte-migration/src/main/resources/migrations/migrationV0_30_0/airbyte_config/State.yaml @@ -0,0 +1,14 @@ +--- +"$schema": http://json-schema.org/draft-07/schema# +"$id": https://github.com/airbytehq/airbyte/blob/master/airbyte-config/models/src/main/resources/types/State.yaml +title: State +description: information output by the connection. +type: object +required: + - state +additionalProperties: false +properties: + state: + description: Integration specific blob. Must be a valid JSON string. + type: object + existingJavaType: com.fasterxml.jackson.databind.JsonNode diff --git a/airbyte-migration/src/test/java/io/airbyte/migrate/migrations/MigrationV0_30_0Test.java b/airbyte-migration/src/test/java/io/airbyte/migrate/migrations/MigrationV0_30_0Test.java new file mode 100644 index 000000000000..688a7d0df410 --- /dev/null +++ b/airbyte-migration/src/test/java/io/airbyte/migrate/migrations/MigrationV0_30_0Test.java @@ -0,0 +1,22 @@ +/* + * Copyright (c) 2021 Airbyte, Inc., all rights reserved. + */ + +package io.airbyte.migrate.migrations; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.airbyte.migrate.Migrations; +import org.junit.jupiter.api.Test; + +class MigrationV0_30_0Test { + + @Test + void testMigration() { + // standard sync state does not exist before migration v0.30.0 + assertFalse(Migrations.MIGRATION_V_0_29_0.getOutputSchema().containsKey(MigrationV0_30_0.STANDARD_SYNC_STATE_RESOURCE_ID)); + assertTrue(Migrations.MIGRATION_V_0_30_0.getOutputSchema().containsKey(MigrationV0_30_0.STANDARD_SYNC_STATE_RESOURCE_ID)); + } + +} From f25bef9d37022ba23862decbc30e66306ffa6564 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 06:48:31 -0700 Subject: [PATCH 15/23] Add comment --- .../models/src/main/java/io/airbyte/config/ConfigSchema.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java b/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java index bf5a9436b52d..2a11f656ebe8 100644 --- a/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java +++ b/airbyte-config/models/src/main/java/io/airbyte/config/ConfigSchema.java @@ -37,7 +37,7 @@ public enum ConfigSchema implements AirbyteConfig { destinationConnection -> destinationConnection.getDestinationId().toString(), "destinationId"), - // sync + // sync (i.e. connection) STANDARD_SYNC("StandardSync.yaml", StandardSync.class, standardSync -> standardSync.getConnectionId().toString(), From b57195843042022f237decfdcf59569ec0110c86 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 07:14:19 -0700 Subject: [PATCH 16/23] Use config model instead of json node --- .../V0_30_22_001__Store_last_sync_state.java | 33 ++++---- ...30_22_001__Store_last_sync_state_test.java | 77 +++++++++---------- 2 files changed, 52 insertions(+), 58 deletions(-) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java index fdac8d487e92..a5bc6f0408e0 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java @@ -4,20 +4,22 @@ package io.airbyte.db.instance.configs.migrations; -import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.annotations.VisibleForTesting; import io.airbyte.commons.jackson.MoreMappers; import io.airbyte.commons.json.Jsons; +import io.airbyte.config.ConfigSchema; import io.airbyte.config.Configs; import io.airbyte.config.EnvConfigs; +import io.airbyte.config.StandardSyncState; +import io.airbyte.config.State; import io.airbyte.db.Database; import io.airbyte.db.instance.jobs.JobsDatabaseInstance; import java.io.IOException; import java.sql.SQLException; import java.time.OffsetDateTime; -import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.UUID; import java.util.stream.Collectors; import org.flywaydb.core.api.migration.BaseJavaMigration; @@ -25,7 +27,6 @@ import org.jooq.DSLContext; import org.jooq.Field; import org.jooq.JSONB; -import org.jooq.Record2; import org.jooq.Table; import org.jooq.impl.DSL; import org.jooq.impl.SQLDataType; @@ -45,7 +46,6 @@ public class V0_30_22_001__Store_last_sync_state extends BaseJavaMigration { // airbyte configs table // (we cannot use the jooq generated code here to avoid circular dependency) - static final String STANDARD_SYNC_STATE = "STANDARD_SYNC_STATE"; static final Table TABLE_AIRBYTE_CONFIGS = DSL.table("airbyte_configs"); static final Field COLUMN_CONFIG_TYPE = DSL.field("config_type", SQLDataType.VARCHAR(60).nullable(false)); static final Field COLUMN_CONFIG_ID = DSL.field("config_id", SQLDataType.VARCHAR(36).nullable(false)); @@ -76,12 +76,12 @@ public void migrate(final Context context) throws Exception { } @VisibleForTesting - static void copyData(final DSLContext ctx, final Map syncToStateMap, final OffsetDateTime timestamp) { - for (final Map.Entry entry : syncToStateMap.entrySet()) { + static void copyData(final DSLContext ctx, final Set standardSyncStates, final OffsetDateTime timestamp) { + for (final StandardSyncState standardSyncState : standardSyncStates) { ctx.insertInto(TABLE_AIRBYTE_CONFIGS) - .set(COLUMN_CONFIG_TYPE, STANDARD_SYNC_STATE) - .set(COLUMN_CONFIG_ID, UUID.fromString(entry.getKey()).toString()) - .set(COLUMN_CONFIG_BLOB, JSONB.valueOf(Jsons.serialize(entry.getValue()))) + .set(COLUMN_CONFIG_TYPE, ConfigSchema.STANDARD_SYNC_STATE.name()) + .set(COLUMN_CONFIG_ID, standardSyncState.getConnectionId().toString()) + .set(COLUMN_CONFIG_BLOB, JSONB.valueOf(Jsons.serialize(standardSyncState))) .set(COLUMN_CREATED_AT, timestamp) .set(COLUMN_UPDATED_AT, timestamp) // This migration is idempotent. If the record for a sync_id already exists, @@ -121,10 +121,10 @@ static Optional getJobsDatabase(final Configs configs) { } /** - * @return a map from connection id (UUID) to connection state (StandardSyncState). + * @return a set of StandardSyncStates from the latest attempt for each connection. */ @VisibleForTesting - static Map getStandardSyncStates(final Database jobsDatabase) throws SQLException { + static Set getStandardSyncStates(final Database jobsDatabase) throws SQLException { final Table jobsTable = DSL.table("jobs"); final Field jobIdField = DSL.field("jobs.id", SQLDataType.BIGINT); final Field syncIdField = DSL.field("jobs.scope", SQLDataType.VARCHAR); @@ -151,16 +151,13 @@ static Map getStandardSyncStates(final Database jobsDatabase) .groupBy(attemptJobIdField))) .fetch() .stream() - .collect(Collectors.toMap( - Record2::value1, - r -> getStandardSyncState(r.value1(), Jsons.deserialize(r.value2().data()))))); + .map(r -> getStandardSyncState(UUID.fromString(r.value1()), Jsons.deserialize(r.value2().data(), State.class)))) + .collect(Collectors.toSet()); } @VisibleForTesting - static JsonNode getStandardSyncState(final String connectionId, final JsonNode state) { - return MAPPER.createObjectNode() - .put("connectionId", connectionId) - .set("state", state); + static StandardSyncState getStandardSyncState(final UUID connectionId, final State state) { + return new StandardSyncState().withConnectionId(connectionId).withState(state); } } diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java index 9cc34694d210..c9758cb38834 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java @@ -9,7 +9,6 @@ import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_CONFIG_TYPE; import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_CREATED_AT; import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.COLUMN_UPDATED_AT; -import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.STANDARD_SYNC_STATE; import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.TABLE_AIRBYTE_CONFIGS; import static io.airbyte.db.instance.configs.migrations.V0_30_22_001__Store_last_sync_state.getStandardSyncState; import static org.jooq.impl.DSL.field; @@ -19,20 +18,25 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.node.ObjectNode; import io.airbyte.commons.jackson.MoreMappers; import io.airbyte.commons.json.Jsons; +import io.airbyte.config.ConfigSchema; import io.airbyte.config.Configs; import io.airbyte.config.EnvConfigs; +import io.airbyte.config.JobOutput; +import io.airbyte.config.JobOutput.OutputType; +import io.airbyte.config.StandardSyncOutput; +import io.airbyte.config.StandardSyncState; +import io.airbyte.config.State; import io.airbyte.db.Database; import io.airbyte.db.instance.configs.AbstractConfigsDatabaseTest; import io.airbyte.db.instance.jobs.JobsDatabaseInstance; import java.sql.Connection; import java.sql.SQLException; import java.time.OffsetDateTime; -import java.util.Map; +import java.util.Collections; +import java.util.Set; import java.util.UUID; import javax.annotation.Nullable; import org.flywaydb.core.api.configuration.Configuration; @@ -66,16 +70,13 @@ class V0_30_22_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTe private static final UUID CONNECTION_1_ID = UUID.randomUUID(); private static final UUID CONNECTION_2_ID = UUID.randomUUID(); private static final UUID CONNECTION_3_ID = UUID.randomUUID(); - // these are State objects, see State.yaml for its schema; - // we cannot construct the POJO directly because State is defined in an downstream module - private static final JsonNode CONNECTION_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }"); - private static final JsonNode CONNECTION_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }"); - private static final JsonNode STD_CONNECTION_STATE_2 = getStandardSyncState(CONNECTION_2_ID.toString(), CONNECTION_2_STATE); - private static final JsonNode STD_CONNECTION_STATE_3 = getStandardSyncState(CONNECTION_3_ID.toString(), CONNECTION_3_STATE); - private static final Map CONNECTION_STATE_MAP = Map.of( - CONNECTION_2_ID.toString(), STD_CONNECTION_STATE_2, - CONNECTION_3_ID.toString(), STD_CONNECTION_STATE_3); + private static final State CONNECTION_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }", State.class); + private static final State CONNECTION_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }", State.class); + + private static final StandardSyncState STD_CONNECTION_STATE_2 = getStandardSyncState(CONNECTION_2_ID, CONNECTION_2_STATE); + private static final StandardSyncState STD_CONNECTION_STATE_3 = getStandardSyncState(CONNECTION_3_ID, CONNECTION_3_STATE); + private static final Set STD_CONNECTION_STATES = Set.of(STD_CONNECTION_STATE_2, STD_CONNECTION_STATE_3); private static Database jobDatabase; @@ -117,17 +118,14 @@ public void testGetSyncToStateMap() throws Exception { // The third job has multiple attempts. The third attempt has the latest state. final long job3 = createJob(ctx, CONNECTION_3_ID); - final JsonNode attempt31State = Jsons.deserialize("{ \"state\": { \"cursor\": 31 } }"); + final State attempt31State = new State().withState(Jsons.deserialize("{\"cursor\": 31 }")); createAttempt(ctx, job3, 1, createAttemptOutput(attempt31State)); createAttempt(ctx, job3, 2, null); createAttempt(ctx, job3, 3, createAttemptOutput(CONNECTION_3_STATE)); createAttempt(ctx, job3, 4, null); createAttempt(ctx, job3, 5, null); - final Map syncToStateMap = V0_30_22_001__Store_last_sync_state.getStandardSyncStates(jobDatabase); - assertEquals(2, syncToStateMap.size()); - assertEquals(STD_CONNECTION_STATE_2, syncToStateMap.get(CONNECTION_2_ID.toString())); - assertEquals(STD_CONNECTION_STATE_3, syncToStateMap.get(CONNECTION_3_ID.toString())); + assertEquals(STD_CONNECTION_STATES, V0_30_22_001__Store_last_sync_state.getStandardSyncStates(jobDatabase)); return null; }); @@ -136,20 +134,22 @@ public void testGetSyncToStateMap() throws Exception { @Test @Order(30) public void testCopyData() throws SQLException { - final Map newConnectionStateMap = Map.of( - CONNECTION_2_ID.toString(), - Jsons.deserialize("{ \"connectionId\": \"invalid\", \"state\": { \"state\": { \"cursor\": 3 } } }")); + + final Set newConnectionStates = Collections.singleton( + new StandardSyncState() + .withConnectionId(CONNECTION_2_ID) + .withState(new State().withState(Jsons.deserialize("{ \"cursor\": 3 }")))); final OffsetDateTime timestamp = OffsetDateTime.now(); database.query(ctx -> { - V0_30_22_001__Store_last_sync_state.copyData(ctx, CONNECTION_STATE_MAP, timestamp); - checkSyncStates(ctx, CONNECTION_STATE_MAP, timestamp); + V0_30_22_001__Store_last_sync_state.copyData(ctx, STD_CONNECTION_STATES, timestamp); + checkSyncStates(ctx, STD_CONNECTION_STATES, timestamp); // call the copyData method again with different data will not affect existing records - V0_30_22_001__Store_last_sync_state.copyData(ctx, CONNECTION_STATE_MAP, OffsetDateTime.now()); - // the states remain the same as those in syncStateMap1 - checkSyncStates(ctx, CONNECTION_STATE_MAP, timestamp); + V0_30_22_001__Store_last_sync_state.copyData(ctx, newConnectionStates, OffsetDateTime.now()); + // the states remain the same as those in STD_CONNECTION_STATES + checkSyncStates(ctx, STD_CONNECTION_STATES, timestamp); return null; }); @@ -162,7 +162,7 @@ public void testCopyData() throws SQLException { @Order(40) public void testMigration() throws Exception { database.query(ctx -> ctx.deleteFrom(TABLE_AIRBYTE_CONFIGS) - .where(COLUMN_CONFIG_TYPE.eq(STANDARD_SYNC_STATE)) + .where(COLUMN_CONFIG_TYPE.eq(ConfigSchema.STANDARD_SYNC_STATE.name())) .execute()); final Configs configs = mock(Configs.class); @@ -191,7 +191,7 @@ public Connection getConnection() { }; migration.migrate(context); database.query(ctx -> { - checkSyncStates(ctx, CONNECTION_STATE_MAP, null); + checkSyncStates(ctx, STD_CONNECTION_STATES, null); return null; }); } @@ -212,7 +212,7 @@ private static long createJob(final DSLContext ctx, final UUID standardSyncId) { .get(JOB_ID_FIELD); } - private static void createAttempt(final DSLContext ctx, final long jobId, final int attemptNumber, final JsonNode attemptOutput) { + private static void createAttempt(final DSLContext ctx, final long jobId, final int attemptNumber, final JobOutput attemptOutput) { final int insertCount = ctx.insertInto(ATTEMPTS_TABLE, ATTEMPT_JOB_ID_FIELD, ATTEMPT_NUMBER_FIELD, ATTEMPT_OUTPUT_FIELD) .values(jobId, attemptNumber, JSONB.valueOf(Jsons.serialize(attemptOutput))) .execute(); @@ -230,27 +230,24 @@ private static void createAttempt(final DSLContext ctx, final long jobId, final * * @param state The state object within a StandardSyncOutput. */ - private static JsonNode createAttemptOutput(final JsonNode state) { - final ObjectNode standardSyncOutput = OBJECT_MAPPER.createObjectNode() - .set("state", state); - return OBJECT_MAPPER.createObjectNode() - .put("output_type", "sync") - .set("sync", standardSyncOutput); + private static JobOutput createAttemptOutput(final State state) { + final StandardSyncOutput standardSyncOutput = new StandardSyncOutput().withState(state); + return new JobOutput().withOutputType(OutputType.SYNC).withSync(standardSyncOutput); } private static void checkSyncStates(final DSLContext ctx, - final Map expectedSyncStates, + final Set standardSyncStates, @Nullable final OffsetDateTime expectedTimestamp) { - for (final Map.Entry entry : expectedSyncStates.entrySet()) { + for (final StandardSyncState standardSyncState : standardSyncStates) { final var record = ctx .select(COLUMN_CONFIG_BLOB, COLUMN_CREATED_AT, COLUMN_UPDATED_AT) .from(TABLE_AIRBYTE_CONFIGS) - .where(COLUMN_CONFIG_ID.eq(entry.getKey()), - COLUMN_CONFIG_TYPE.eq(STANDARD_SYNC_STATE)) + .where(COLUMN_CONFIG_ID.eq(standardSyncState.getConnectionId().toString()), + COLUMN_CONFIG_TYPE.eq(ConfigSchema.STANDARD_SYNC_STATE.name())) .fetchOne(); - assertEquals(entry.getValue(), Jsons.deserialize(record.value1().data())); + assertEquals(standardSyncState, Jsons.deserialize(record.value1().data(), StandardSyncState.class)); if (expectedTimestamp != null) { assertEquals(expectedTimestamp, record.value2()); assertEquals(expectedTimestamp, record.value3()); From a1c950cf1f07ff606257e0ba23e36c16044c6540 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 07:42:08 -0700 Subject: [PATCH 17/23] Add comments --- .../db/instance/configs/ConfigsDatabaseTestProvider.java | 7 ++++--- .../migrations/V0_30_22_001__Store_last_sync_state.java | 6 +++--- .../airbyte/db/instance/test/TestDatabaseProvider.java | 9 +++++++++ .../scheduler/persistence/DefaultJobPersistence.java | 2 -- .../io/airbyte/scheduler/persistence/JobPersistence.java | 5 +++++ .../io/airbyte/server/handlers/ArchiveHandlerTest.java | 8 ++++---- .../airbyte/server/migration/DatabaseArchiverTest.java | 8 ++++---- .../io/airbyte/server/migration/RunMigrationTest.java | 5 ++--- .../src/main/java/io/airbyte/workers/WorkerApp.java | 1 - 9 files changed, 31 insertions(+), 20 deletions(-) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java index 243007746900..99d4fbfd1c1c 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/ConfigsDatabaseTestProvider.java @@ -7,6 +7,7 @@ import static org.jooq.impl.DSL.field; import static org.jooq.impl.DSL.table; +import io.airbyte.config.ConfigSchema; import io.airbyte.db.Database; import io.airbyte.db.ExceptionWrappingDatabase; import io.airbyte.db.instance.DatabaseMigrator; @@ -22,7 +23,7 @@ public class ConfigsDatabaseTestProvider implements TestDatabaseProvider { private final String password; private final String jdbcUrl; - public ConfigsDatabaseTestProvider(String user, String password, String jdbcUrl) { + public ConfigsDatabaseTestProvider(final String user, final String password, final String jdbcUrl) { this.user = user; this.password = password; this.jdbcUrl = jdbcUrl; @@ -43,10 +44,10 @@ public Database create(final boolean runMigration) throws IOException { // The configs database is considered ready only if there are some seed records. // So we need to create at least one record here. - OffsetDateTime timestamp = OffsetDateTime.now(); + final OffsetDateTime timestamp = OffsetDateTime.now(); new ExceptionWrappingDatabase(database).transaction(ctx -> ctx.insertInto(table("airbyte_configs")) .set(field("config_id"), UUID.randomUUID().toString()) - .set(field("config_type"), "STANDARD_SOURCE_DEFINITION") + .set(field("config_type"), ConfigSchema.STATE.name()) .set(field("config_blob"), JSONB.valueOf("{}")) .set(field("created_at"), timestamp) .set(field("updated_at"), timestamp) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java index a5bc6f0408e0..b4ea92ab8c26 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java @@ -34,9 +34,7 @@ import org.slf4j.LoggerFactory; /** - * Create a new table to store the latest job state for each standard sync. - *
  • Column sync_id: the connectionId in StandardSync
  • Column state: a json node representing - * a State object
    • + * Copy the latest job state for each standard sync to the config database. */ public class V0_30_22_001__Store_last_sync_state extends BaseJavaMigration { @@ -77,6 +75,8 @@ public void migrate(final Context context) throws Exception { @VisibleForTesting static void copyData(final DSLContext ctx, final Set standardSyncStates, final OffsetDateTime timestamp) { + LOGGER.info("[{}] Number of connection states to copy: {}", MIGRATION_NAME, standardSyncStates.size()); + for (final StandardSyncState standardSyncState : standardSyncStates) { ctx.insertInto(TABLE_AIRBYTE_CONFIGS) .set(COLUMN_CONFIG_TYPE, ConfigSchema.STANDARD_SYNC_STATE.name()) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java index c623bc1243f6..a5f0b9abe474 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/test/TestDatabaseProvider.java @@ -7,8 +7,17 @@ import io.airbyte.db.Database; import java.io.IOException; +/** + * Create mock database in unit tests. The implementation will be responsible for: 1) constructing + * and preparing the database, and 2) running the Flyway migration. + */ public interface TestDatabaseProvider { + /** + * @param runMigration Whether the mock database should run Flyway migration before it is used in + * unit test. Usually this parameter should be false only when the migration itself is being + * tested. + */ Database create(final boolean runMigration) throws IOException; } diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java index 4313844715fd..6fd4f22cb3d7 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java @@ -784,6 +784,4 @@ private static Table getTable(final String schema, final String tableNam return DSL.table(String.format("%s.%s", schema, tableName)); } - // TODO: add a method to update attempt state in job database, and another in config database - } diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java index 26260753e566..d784c1f3f51c 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/JobPersistence.java @@ -117,6 +117,11 @@ public interface JobPersistence { */ Optional getAttemptTemporalWorkflowId(long jobId, int attemptNumber) throws IOException; + /** + * When the output is a StandardSyncOutput, caller of this method should persiste + * StandardSyncOutput#state in the configs database by calling + * ConfigRepository#updateConnectionState, which takes care of persisting the connection state. + */ void writeOutput(long jobId, int attemptNumber, T output) throws IOException; /** 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 5784c58bbbfd..971778031dae 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 @@ -31,8 +31,7 @@ import io.airbyte.config.persistence.YamlSeedConfigPersistence; import io.airbyte.config.persistence.split_secrets.NoOpSecretsHydrator; import io.airbyte.db.Database; -import io.airbyte.db.instance.configs.ConfigsDatabaseInstance; -import io.airbyte.db.instance.jobs.JobsDatabaseInstance; +import io.airbyte.db.instance.test.TestDatabaseProviders; import io.airbyte.protocol.models.ConnectorSpecification; import io.airbyte.scheduler.persistence.DefaultJobPersistence; import io.airbyte.scheduler.persistence.JobPersistence; @@ -103,8 +102,9 @@ public static void dbDown() { @BeforeEach public void setup() throws Exception { - jobDatabase = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - configDatabase = new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + final TestDatabaseProviders databaseProviders = new TestDatabaseProviders(container); + jobDatabase = databaseProviders.createNewJobsDatabase(); + configDatabase = databaseProviders.createNewConfigsDatabase(); jobPersistence = new DefaultJobPersistence(jobDatabase); seedPersistence = YamlSeedConfigPersistence.getDefault(); configPersistence = new DatabaseConfigPersistence(jobDatabase); diff --git a/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java b/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java index f9f4d060301a..d67f98fd73f9 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/migration/DatabaseArchiverTest.java @@ -9,9 +9,8 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import io.airbyte.db.Database; -import io.airbyte.db.instance.configs.ConfigsDatabaseInstance; -import io.airbyte.db.instance.jobs.JobsDatabaseInstance; import io.airbyte.db.instance.jobs.JobsDatabaseSchema; +import io.airbyte.db.instance.test.TestDatabaseProviders; import io.airbyte.scheduler.persistence.DefaultJobPersistence; import io.airbyte.scheduler.persistence.JobPersistence; import java.io.IOException; @@ -41,8 +40,9 @@ void setUp() throws IOException { .withPassword("docker"); container.start(); - jobDatabase = new JobsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); - configDatabase = new ConfigsDatabaseInstance(container.getUsername(), container.getPassword(), container.getJdbcUrl()).getAndInitialize(); + final TestDatabaseProviders databaseProviders = new TestDatabaseProviders(container); + jobDatabase = databaseProviders.createNewJobsDatabase(); + configDatabase = databaseProviders.createNewConfigsDatabase(); final JobPersistence persistence = new DefaultJobPersistence(jobDatabase); databaseArchiver = new DatabaseArchiver(persistence); } diff --git a/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java b/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java index 2eb6aa7dc951..79129e554937 100644 --- a/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java +++ b/airbyte-server/src/test/java/io/airbyte/server/migration/RunMigrationTest.java @@ -144,9 +144,8 @@ private ConfigPersistence getConfigPersistence(final Path configRoot) throws Exc final Configs configs = mock(Configs.class); when(configs.getConfigRoot()).thenReturn(configRoot); // The database provider creates a config database that is ready to use, which contains a mock - // config - // entry. However, the migrateFileConfigs method will only copy the file configs if the database is - // empty. So we need to purge the database first. + // config entry. However, the migrateFileConfigs method will only copy the file configs if the + // database is empty. So we need to purge the database first. configDatabase.transaction(ctx -> ctx.execute("TRUNCATE TABLE airbyte_configs;")); return new DatabaseConfigPersistence(configDatabase) .migrateFileConfigs(configs); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java index acb64f336acd..8ad5318df6e8 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/WorkerApp.java @@ -153,7 +153,6 @@ public static void main(final String[] args) throws IOException, InterruptedExce final WorkflowServiceStubs temporalService = TemporalUtils.createTemporalService(temporalHost); - // todo (cgardens) - make sure appropriate env variables are passed to this container. final Database configDatabase = new ConfigsDatabaseInstance( configs.getConfigDatabaseUser(), configs.getConfigDatabasePassword(), From 9e0dc1e7e40765ae2f6ce59f1892be4135a965b7 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 08:14:32 -0700 Subject: [PATCH 18/23] Remove unnecessary method --- .../scheduler/persistence/DefaultJobPersistence.java | 9 --------- 1 file changed, 9 deletions(-) diff --git a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java index 6fd4f22cb3d7..8feceddbda9a 100644 --- a/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java +++ b/airbyte-scheduler/persistence/src/main/java/io/airbyte/scheduler/persistence/DefaultJobPersistence.java @@ -313,15 +313,6 @@ public Optional getAttemptTemporalWorkflowId(final long jobId, final int @Override public void writeOutput(final long jobId, final int attemptNumber, final T output) throws IOException { final OffsetDateTime now = OffsetDateTime.ofInstant(timeSupplier.get(), ZoneOffset.UTC); - writeOutputToAttemptTable(jobId, attemptNumber, output, now); - // writeOutputToSyncStateTable(jobId, output, now); - } - - private void writeOutputToAttemptTable(final long jobId, - final int attemptNumber, - final T output, - final OffsetDateTime now) - throws IOException { jobDatabase.transaction( ctx -> ctx.update(ATTEMPTS) .set(ATTEMPTS.OUTPUT, JSONB.valueOf(Jsons.serialize(output))) From aa20d11f138e14b7ac49914c3b8388c74845e5fd Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 10:23:01 -0700 Subject: [PATCH 19/23] Fix migration query --- .../V0_30_22_001__Store_last_sync_state.java | 25 +++--- ...30_22_001__Store_last_sync_state_test.java | 80 +++++++++++++------ 2 files changed, 67 insertions(+), 38 deletions(-) diff --git a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java index b4ea92ab8c26..79567bb34c5c 100644 --- a/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java +++ b/airbyte-db/lib/src/main/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state.java @@ -126,29 +126,28 @@ static Optional getJobsDatabase(final Configs configs) { @VisibleForTesting static Set getStandardSyncStates(final Database jobsDatabase) throws SQLException { final Table jobsTable = DSL.table("jobs"); - final Field jobIdField = DSL.field("jobs.id", SQLDataType.BIGINT); - final Field syncIdField = DSL.field("jobs.scope", SQLDataType.VARCHAR); + final Field jobId = DSL.field("jobs.id", SQLDataType.BIGINT); + final Field connectionId = DSL.field("jobs.scope", SQLDataType.VARCHAR); final Table attemptsTable = DSL.table("attempts"); - final Field attemptJobIdField = DSL.field("attempts.job_id", SQLDataType.BIGINT); - final Field attemptNumberField = DSL.field("attempts.attempt_number", SQLDataType.INTEGER); + final Field attemptJobId = DSL.field("attempts.job_id", SQLDataType.BIGINT); + final Field attemptNumber = DSL.field("attempts.attempt_number", SQLDataType.INTEGER); + final Field attemptCreatedAt = DSL.field("attempts.created_at", SQLDataType.TIMESTAMPWITHTIMEZONE); // output schema: JobOutput.yaml // sync schema: StandardSyncOutput.yaml // state schema: State.yaml, e.g. { "state": { "cursor": 1000 } } - final Field attemptStateField = DSL.field("attempts.output -> 'sync' -> 'state'", SQLDataType.JSONB); + final Field attemptState = DSL.field("attempts.output -> 'sync' -> 'state'", SQLDataType.JSONB); return jobsDatabase.query(ctx -> ctx - .select(syncIdField, attemptStateField) + .select(connectionId, attemptState) + .distinctOn(connectionId) .from(attemptsTable) .innerJoin(jobsTable) - .on(jobIdField.eq(attemptJobIdField)) - .where(DSL.row(attemptJobIdField, attemptNumberField).in( - // for each job id, find the last attempt with a state - DSL.select(attemptJobIdField, DSL.max(attemptNumberField)) - .from(attemptsTable) - .where(attemptStateField.isNotNull()) - .groupBy(attemptJobIdField))) + .on(jobId.eq(attemptJobId)) + .where(attemptState.isNotNull()) + // this query assumes that an attempt with larger created_at field is always a newer attempt + .orderBy(connectionId, attemptCreatedAt.desc()) .fetch() .stream() .map(r -> getStandardSyncState(UUID.fromString(r.value1()), Jsons.deserialize(r.value2().data(), State.class)))) diff --git a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java index c9758cb38834..e7fbd7e371a7 100644 --- a/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java +++ b/airbyte-db/lib/src/test/java/io/airbyte/db/instance/configs/migrations/V0_30_22_001__Store_last_sync_state_test.java @@ -14,6 +14,7 @@ import static org.jooq.impl.DSL.field; import static org.jooq.impl.DSL.table; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -56,16 +57,19 @@ class V0_30_22_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTest { private static final ObjectMapper OBJECT_MAPPER = MoreMappers.initMapper(); + private static final OffsetDateTime TIMESTAMP = OffsetDateTime.now(); private static final Table JOBS_TABLE = table("jobs"); private static final Field JOB_ID_FIELD = field("id", SQLDataType.BIGINT); private static final Field JOB_SCOPE_FIELD = field("scope", SQLDataType.VARCHAR); + private static final Field JOB_CREATED_AT_FIELD = field("created_at", SQLDataType.TIMESTAMPWITHTIMEZONE); private static final Table ATTEMPTS_TABLE = table("attempts"); private static final Field ATTEMPT_ID_FIELD = field("id", SQLDataType.BIGINT); private static final Field ATTEMPT_JOB_ID_FIELD = field("job_id", SQLDataType.BIGINT); private static final Field ATTEMPT_NUMBER_FIELD = field("attempt_number", SQLDataType.INTEGER); private static final Field ATTEMPT_OUTPUT_FIELD = field("output", SQLDataType.JSONB); + private static final Field ATTEMPT_CREATED_AT_FIELD = field("created_at", SQLDataType.TIMESTAMPWITHTIMEZONE); private static final UUID CONNECTION_1_ID = UUID.randomUUID(); private static final UUID CONNECTION_2_ID = UUID.randomUUID(); @@ -73,6 +77,7 @@ class V0_30_22_001__Store_last_sync_state_test extends AbstractConfigsDatabaseTe private static final State CONNECTION_2_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 2222 } }", State.class); private static final State CONNECTION_3_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": 3333 } }", State.class); + private static final State CONNECTION_OLD_STATE = Jsons.deserialize("{ \"state\": { \"cursor\": -1 } }", State.class); private static final StandardSyncState STD_CONNECTION_STATE_2 = getStandardSyncState(CONNECTION_2_ID, CONNECTION_2_STATE); private static final StandardSyncState STD_CONNECTION_STATE_3 = getStandardSyncState(CONNECTION_3_ID, CONNECTION_3_STATE); @@ -106,24 +111,34 @@ public void testGetJobsDatabase() { @Test @Order(20) - public void testGetSyncToStateMap() throws Exception { + public void testGetStandardSyncStates() throws Exception { jobDatabase.query(ctx -> { - // Create three jobs for three standard syncs. - // The first job has no attempt. - createJob(ctx, CONNECTION_1_ID); - - // The second job has one attempt. - final long job2 = createJob(ctx, CONNECTION_2_ID); - createAttempt(ctx, job2, 1, createAttemptOutput(CONNECTION_2_STATE)); - - // The third job has multiple attempts. The third attempt has the latest state. - final long job3 = createJob(ctx, CONNECTION_3_ID); - final State attempt31State = new State().withState(Jsons.deserialize("{\"cursor\": 31 }")); - createAttempt(ctx, job3, 1, createAttemptOutput(attempt31State)); - createAttempt(ctx, job3, 2, null); - createAttempt(ctx, job3, 3, createAttemptOutput(CONNECTION_3_STATE)); - createAttempt(ctx, job3, 4, null); - createAttempt(ctx, job3, 5, null); + // Connection 1 has 1 job, no attempt. + // This is to test that connection without no state is not returned. + createJob(ctx, CONNECTION_1_ID, 30); + + // Connection 2 has two jobs, each has one attempt. + // This is to test that only the state from the latest job is returned. + final long job21 = createJob(ctx, CONNECTION_2_ID, 10); + final long job22 = createJob(ctx, CONNECTION_2_ID, 20); + assertNotEquals(job21, job22); + createAttempt(ctx, job21, 1, createAttemptOutput(CONNECTION_OLD_STATE), 11); + createAttempt(ctx, job22, 1, createAttemptOutput(CONNECTION_2_STATE), 21); + + // Connection 3 has two jobs. + // The first job has multiple attempts. Its third attempt has the latest state. + // The second job has two attempts with no state. + // This is to test that only the state from the latest attempt is returned. + final long job31 = createJob(ctx, CONNECTION_3_ID, 5); + final long job32 = createJob(ctx, CONNECTION_3_ID, 15); + assertNotEquals(job31, job32); + createAttempt(ctx, job31, 1, createAttemptOutput(CONNECTION_OLD_STATE), 6); + createAttempt(ctx, job31, 2, null, 7); + createAttempt(ctx, job31, 3, createAttemptOutput(CONNECTION_3_STATE), 8); + createAttempt(ctx, job31, 4, null, 9); + createAttempt(ctx, job31, 5, null, 10); + createAttempt(ctx, job32, 1, null, 20); + createAttempt(ctx, job32, 2, null, 25); assertEquals(STD_CONNECTION_STATES, V0_30_22_001__Store_last_sync_state.getStandardSyncStates(jobDatabase)); @@ -197,24 +212,39 @@ public Connection getConnection() { } /** - * Create a job record whose scope equals to the passed in standard sync id, and return the job id. + * Create a job record whose scope equals to the passed in connection id, and return the job id. + * + * @param creationOffset Set the creation timestamp to {@code TIMESTAMP} + this passed in offset. */ - private static long createJob(final DSLContext ctx, final UUID standardSyncId) { - final int insertCount = ctx.insertInto(JOBS_TABLE, JOB_SCOPE_FIELD) - .values(standardSyncId.toString()) + private static long createJob(final DSLContext ctx, final UUID connectionId, final long creationOffset) { + final int insertCount = ctx.insertInto(JOBS_TABLE) + .set(JOB_SCOPE_FIELD, connectionId.toString()) + .set(JOB_CREATED_AT_FIELD, TIMESTAMP.plusDays(creationOffset)) .execute(); assertEquals(1, insertCount); return ctx.select(JOB_ID_FIELD) .from(JOBS_TABLE) - .where(JOB_SCOPE_FIELD.eq(standardSyncId.toString())) + .where(JOB_SCOPE_FIELD.eq(connectionId.toString())) + .orderBy(JOB_CREATED_AT_FIELD.desc()) + .limit(1) .fetchOne() .get(JOB_ID_FIELD); } - private static void createAttempt(final DSLContext ctx, final long jobId, final int attemptNumber, final JobOutput attemptOutput) { - final int insertCount = ctx.insertInto(ATTEMPTS_TABLE, ATTEMPT_JOB_ID_FIELD, ATTEMPT_NUMBER_FIELD, ATTEMPT_OUTPUT_FIELD) - .values(jobId, attemptNumber, JSONB.valueOf(Jsons.serialize(attemptOutput))) + /** + * @param creationOffset Set the creation timestamp to {@code TIMESTAMP} + this passed in offset. + */ + private static void createAttempt(final DSLContext ctx, + final long jobId, + final int attemptNumber, + final JobOutput attemptOutput, + final long creationOffset) { + final int insertCount = ctx.insertInto(ATTEMPTS_TABLE) + .set(ATTEMPT_JOB_ID_FIELD, jobId) + .set(ATTEMPT_NUMBER_FIELD, attemptNumber) + .set(ATTEMPT_OUTPUT_FIELD, JSONB.valueOf(Jsons.serialize(attemptOutput))) + .set(ATTEMPT_CREATED_AT_FIELD, TIMESTAMP.plusDays(creationOffset)) .execute(); assertEquals(1, insertCount); From 1f24699d721ea371d881b07a12c7d80a1e4361c2 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 10:25:48 -0700 Subject: [PATCH 20/23] Remove unused config database --- .../airbyte/workers/temporal/TemporalAttemptExecution.java | 6 ------ .../workers/temporal/TemporalAttemptExecutionTest.java | 4 ---- 2 files changed, 10 deletions(-) diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java index 4c1dcd9380d8..d3ec246d936c 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/TemporalAttemptExecution.java @@ -10,7 +10,6 @@ import io.airbyte.config.EnvConfigs; import io.airbyte.config.helpers.LogClientSingleton; import io.airbyte.db.Database; -import io.airbyte.db.instance.configs.ConfigsDatabaseInstance; import io.airbyte.db.instance.jobs.JobsDatabaseInstance; import io.airbyte.scheduler.models.JobRunConfig; import io.airbyte.scheduler.persistence.DefaultJobPersistence; @@ -132,11 +131,6 @@ private void saveWorkflowIdForCancellation() throws IOException { configs.getDatabasePassword(), configs.getDatabaseUrl()) .getInitialized(); - final Database configDatabase = new ConfigsDatabaseInstance( - configs.getConfigDatabaseUser(), - configs.getConfigDatabasePassword(), - configs.getConfigDatabaseUrl()) - .getInitialized(); final JobPersistence jobPersistence = new DefaultJobPersistence(jobDatabase); final String workflowId = workflowIdProvider.get(); jobPersistence.setAttemptTemporalWorkflowId(Long.parseLong(jobRunConfig.getJobId()), jobRunConfig.getAttemptId().intValue(), workflowId); diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java index 0eb7cd78540e..de7e5827a2e8 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/TemporalAttemptExecutionTest.java @@ -57,9 +57,6 @@ static void setUpAll() { when(configs.getDatabaseUrl()).thenReturn(container.getJdbcUrl()); when(configs.getDatabaseUser()).thenReturn(SOURCE_USERNAME); when(configs.getDatabasePassword()).thenReturn(SOURCE_PASSWORD); - when(configs.getConfigDatabaseUrl()).thenReturn(container.getJdbcUrl()); - when(configs.getConfigDatabaseUser()).thenReturn(SOURCE_USERNAME); - when(configs.getConfigDatabasePassword()).thenReturn(SOURCE_PASSWORD); } @SuppressWarnings("unchecked") @@ -67,7 +64,6 @@ static void setUpAll() { void setup() throws IOException { final TestDatabaseProviders databaseProviders = new TestDatabaseProviders(container); databaseProviders.createNewJobsDatabase(); - databaseProviders.createNewConfigsDatabase(); final Path workspaceRoot = Files.createTempDirectory(Path.of("/tmp"), "temporal_attempt_execution_test"); jobRoot = workspaceRoot.resolve(JOB_ID).resolve(String.valueOf(ATTEMPT_ID)); From 1ed500656c172552f5792b5a855f98ecba2e0033 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 15:02:13 -0700 Subject: [PATCH 21/23] Move persist statement and log the call --- .../java/io/airbyte/config/persistence/ConfigRepository.java | 1 + .../src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java index af49d6e9a09b..d7d47242df73 100644 --- a/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java +++ b/airbyte-config/persistence/src/main/java/io/airbyte/config/persistence/ConfigRepository.java @@ -450,6 +450,7 @@ public Optional getConnectionState(final UUID connectionId) throws IOExce } public void updateConnectionState(final UUID connectionId, final State state) throws IOException { + LOGGER.info("Updating connection {} state: {}", connectionId, state); final StandardSyncState connectionState = new StandardSyncState().withConnectionId(connectionId).withState(state); try { persistence.writeConfig(ConfigSchema.STANDARD_SYNC_STATE, connectionId.toString(), connectionState); diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java index 6b894c55ed88..7e74557d696a 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java @@ -97,6 +97,7 @@ public StandardSyncOutput run(final JobRunConfig jobRunConfig, final StandardSyncInput syncInput, final UUID connectionId) { final StandardSyncOutput run = replicationActivity.replicate(jobRunConfig, sourceLauncherConfig, destinationLauncherConfig, syncInput); + persistActivity.persist(connectionId, run); if (syncInput.getOperationSequence() != null && !syncInput.getOperationSequence().isEmpty()) { for (final StandardSyncOperation standardSyncOperation : syncInput.getOperationSequence()) { @@ -118,7 +119,6 @@ public StandardSyncOutput run(final JobRunConfig jobRunConfig, LOGGER.error(message); throw new IllegalArgumentException(message); } - persistActivity.persist(connectionId, run); } } From 4467d7d633d62fc28b82832282a2803f71277c84 Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 15:31:05 -0700 Subject: [PATCH 22/23] Update dev doc --- docs/contributing-to-airbyte/developing-locally.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/contributing-to-airbyte/developing-locally.md b/docs/contributing-to-airbyte/developing-locally.md index aaccabd08b24..bfaa11d6283c 100644 --- a/docs/contributing-to-airbyte/developing-locally.md +++ b/docs/contributing-to-airbyte/developing-locally.md @@ -89,7 +89,7 @@ In `dev` mode, all data will be persisted in `/tmp/dev_root`. To run acceptance \(end-to-end\) tests, you must have the Airbyte running locally. ```bash -SUB_BUILD=PLATFORM ./gradlew build +SUB_BUILD=PLATFORM ./gradlew clean build VERSION=dev docker-compose up SUB_BUILD=PLATFORM ./gradlew :airbyte-tests:acceptanceTests ``` @@ -164,7 +164,7 @@ Sometimes you'll want to reset the data in your local environment. One common ca * Rebuild the project ```bash - SUB_BUILD=PLATFORM ./gradlew build + SUB_BUILD=PLATFORM ./gradlew clean build VERSION=dev docker-compose up -V ``` From ee3b251ad6e2bc706a94076a59a7ff827b8be05c Mon Sep 17 00:00:00 2001 From: Liren Tu Date: Mon, 25 Oct 2021 15:48:31 -0700 Subject: [PATCH 23/23] Add unit tests for sync workflow --- .../airbyte/workers/temporal/SyncWorkflow.java | 3 +++ .../workers/temporal/SyncWorkflowTest.java | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java index 7e74557d696a..bd298ecf7411 100644 --- a/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java +++ b/airbyte-workers/src/main/java/io/airbyte/workers/temporal/SyncWorkflow.java @@ -97,6 +97,9 @@ public StandardSyncOutput run(final JobRunConfig jobRunConfig, final StandardSyncInput syncInput, final UUID connectionId) { final StandardSyncOutput run = replicationActivity.replicate(jobRunConfig, sourceLauncherConfig, destinationLauncherConfig, syncInput); + // the state is persisted immediately after the replication succeeded, because the + // state is a checkpoint of the raw data that has been copied to the destination; + // normalization & dbt does not depend on it persistActivity.persist(connectionId, run); if (syncInput.getOperationSequence() != null && !syncInput.getOperationSequence().isEmpty()) { diff --git a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java index 5a3b01c49c93..a8799ef430b5 100644 --- a/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java +++ b/airbyte-workers/src/test/java/io/airbyte/workers/temporal/SyncWorkflowTest.java @@ -26,6 +26,7 @@ import io.airbyte.workers.temporal.SyncWorkflow.DbtTransformationActivityImpl; import io.airbyte.workers.temporal.SyncWorkflow.NormalizationActivity; import io.airbyte.workers.temporal.SyncWorkflow.NormalizationActivityImpl; +import io.airbyte.workers.temporal.SyncWorkflow.PersistStateActivity; import io.airbyte.workers.temporal.SyncWorkflow.PersistStateActivityImpl; import io.airbyte.workers.temporal.SyncWorkflow.ReplicationActivity; import io.airbyte.workers.temporal.SyncWorkflow.ReplicationActivityImpl; @@ -127,6 +128,7 @@ void testSuccess() { assertEquals(replicationSuccessOutput, actualOutput); verifyReplication(replicationActivity, syncInput); + verifyPersistState(persistStateActivity, sync, actualOutput); verifyNormalize(normalizationActivity, normalizationInput); verifyDbtTransform(dbtTransformationActivity, syncInput.getResourceRequirements(), operatorDbtInput); } @@ -142,7 +144,9 @@ void testReplicationFailure() { assertThrows(WorkflowFailedException.class, this::execute); verifyReplication(replicationActivity, syncInput); + verifyNoInteractions(persistStateActivity); verifyNoInteractions(normalizationActivity); + verifyNoInteractions(dbtTransformationActivity); } @Test @@ -161,7 +165,9 @@ void testNormalizationFailure() { assertThrows(WorkflowFailedException.class, this::execute); verifyReplication(replicationActivity, syncInput); + verifyPersistState(persistStateActivity, sync, replicationSuccessOutput); verifyNormalize(normalizationActivity, normalizationInput); + verifyNoInteractions(dbtTransformationActivity); } @Test @@ -178,7 +184,9 @@ void testCancelDuringReplication() { assertThrows(WorkflowFailedException.class, this::execute); verifyReplication(replicationActivity, syncInput); + verifyNoInteractions(persistStateActivity); verifyNoInteractions(normalizationActivity); + verifyNoInteractions(dbtTransformationActivity); } @Test @@ -200,7 +208,9 @@ void testCancelDuringNormalization() { assertThrows(WorkflowFailedException.class, this::execute); verifyReplication(replicationActivity, syncInput); + verifyPersistState(persistStateActivity, sync, replicationSuccessOutput); verifyNormalize(normalizationActivity, normalizationInput); + verifyNoInteractions(dbtTransformationActivity); } @SuppressWarnings("ResultOfMethodCallIgnored") @@ -228,6 +238,14 @@ private static void verifyReplication(final ReplicationActivity replicationActiv syncInput); } + private static void verifyPersistState(final PersistStateActivity persistStateActivity, + final StandardSync sync, + final StandardSyncOutput syncOutput) { + verify(persistStateActivity).persist( + sync.getConnectionId(), + syncOutput); + } + private static void verifyNormalize(final NormalizationActivity normalizationActivity, final NormalizationInput normalizationInput) { verify(normalizationActivity).normalize( JOB_RUN_CONFIG,