Fix build (#14225)

* Fix build * Fix test
airbytehq · Jun 28, 2022 · 09798a1 · 09798a1
1 parent d0b9de1
commit 09798a1
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/...tional-db/src/main/java/io/airbyte/integrations/source/relationaldb/AbstractDbSource.java b/...tional-db/src/main/java/io/airbyte/integrations/source/relationaldb/AbstractDbSource.java
@@ -7,6 +7,8 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
+import io.airbyte.commons.features.EnvVariableFeatureFlags;
+import io.airbyte.commons.features.FeatureFlags;
 import io.airbyte.commons.functional.CheckedConsumer;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.lang.Exceptions;
@@ -66,6 +68,8 @@ public abstract class AbstractDbSource<DataType, Database extends AbstractDataba
     BaseConnector implements Source, AutoCloseable {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(AbstractDbSource.class);
+  // TODO: Remove when the flag is not use anymore
+  private final FeatureFlags featureFlags = new EnvVariableFeatureFlags();
 
   @Override
   public AirbyteConnectionStatus check(final JsonNode config) throws Exception {
@@ -522,7 +526,7 @@ private Database createDatabaseInternal(final JsonNode sourceConfig) throws Exce
    * @return The deserialized object representation of the state.
    */
   protected List<AirbyteStateMessage> deserializeInitialState(final JsonNode initialStateJson, final JsonNode config) {
-    final Optional<StateWrapper> typedState = StateMessageHelper.getTypedState(initialStateJson);
+    final Optional<StateWrapper> typedState = StateMessageHelper.getTypedState(initialStateJson, featureFlags.useStreamCapableState());
     return typedState.map((state) -> {
       switch (state.getStateType()) {
         case GLOBAL:

diff --git a/...al-db/src/test/java/io/airbyte/integrations/source/relationaldb/AbstractDbSourceTest.java b/...al-db/src/test/java/io/airbyte/integrations/source/relationaldb/AbstractDbSourceTest.java
@@ -9,12 +9,15 @@
 import static org.mockito.Mockito.spy;
 
 import com.fasterxml.jackson.databind.JsonNode;
+import io.airbyte.commons.features.EnvVariableFeatureFlags;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.resources.MoreResources;
 import io.airbyte.protocol.models.AirbyteStateMessage;
 import io.airbyte.protocol.models.AirbyteStateMessage.AirbyteStateType;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.util.List;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 
 /**
@@ -37,6 +40,7 @@ void testDeserializationOfLegacyState() throws IOException {
 
   @Test
   void testDeserializationOfGlobalState() throws IOException {
+    setEnv(EnvVariableFeatureFlags.USE_STREAM_CAPABLE_STATE, "true");
     final AbstractDbSource dbSource = spy(AbstractDbSource.class);
     final JsonNode config = mock(JsonNode.class);
 
@@ -50,6 +54,7 @@ void testDeserializationOfGlobalState() throws IOException {
 
   @Test
   void testDeserializationOfStreamState() throws IOException {
+    setEnv(EnvVariableFeatureFlags.USE_STREAM_CAPABLE_STATE, "true");
     final AbstractDbSource dbSource = spy(AbstractDbSource.class);
     final JsonNode config = mock(JsonNode.class);
 
@@ -71,4 +76,16 @@ void testDeserializationOfNullState() throws IOException {
     assertEquals(dbSource.getSupportedStateType(config), result.get(0).getType());
   }
 
+  public static void setEnv(final String key, final String value) {
+    try {
+      final Map<String, String> env = System.getenv();
+      final Class<?> cl = env.getClass();
+      final Field field = cl.getDeclaredField("m");
+      field.setAccessible(true);
+      final Map<String, String> writableEnv = (Map<String, String>) field.get(env);
+      writableEnv.put(key, value);
+    } catch (final Exception e) {
+      throw new IllegalStateException("Failed to set environment variable", e);
+    }
+  }
 }