Use StateMessageHelper in source

airbyte-integrations/connectors/source-paypal-transaction/source_paypal_transaction/source.py

@@ -7,14 +7,14 @@
 import time
 from abc import ABC
 from datetime import datetime, timedelta
-from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union, Dict
+from typing import Any, Callable, Dict, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
 
 import requests
 from airbyte_cdk.sources import AbstractSource
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.http import HttpStream
-from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 from airbyte_cdk.sources.streams.http.auth import HttpAuthenticator, Oauth2Authenticator
+from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 from dateutil.parser import isoparse
 
 
@@ -301,7 +301,7 @@ def request_params(
             "page_size": self.page_size,
             "page": page_number,
         }
-    
+
     @transformer.registerCustomTransform
     def transform_function(original_value: Any, field_schema: Dict[str, Any]) -> Any:
         if isinstance(original_value, str) and field_schema["type"] == "number":

airbyte-integrations/connectors/source-postgres/src/main/java/io/airbyte/integrations/source/postgres/PostgresSource.java

@@ -40,7 +40,6 @@
 import io.airbyte.protocol.models.AirbyteStreamState;
 import io.airbyte.protocol.models.CommonField;
 import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
-import io.airbyte.protocol.models.ConfiguredAirbyteStream;
 import io.airbyte.protocol.models.SyncMode;
 import java.sql.Connection;
 import java.sql.JDBCType;
@@ -50,7 +49,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

airbyte-integrations/connectors/source-relational-db/build.gradle

@@ -11,6 +11,7 @@ dependencies {
     implementation project(':airbyte-integrations:bases:base-java')
     implementation project(':airbyte-protocol:protocol-models')
     implementation project(':airbyte-json-validation')
+    implementation project(':airbyte-config:config-models')
 
     implementation 'org.apache.commons:commons-lang3:3.11'
 

airbyte-integrations/connectors/source-relational-db/src/main/java/io/airbyte/integrations/source/relationaldb/AbstractDbSource.java

@@ -13,14 +13,15 @@
 import io.airbyte.commons.type.Types;
 import io.airbyte.commons.util.AutoCloseableIterator;
 import io.airbyte.commons.util.AutoCloseableIterators;
+import io.airbyte.config.StateWrapper;
+import io.airbyte.config.helpers.StateMessageHelper;
 import io.airbyte.db.AbstractDatabase;
 import io.airbyte.db.IncrementalUtils;
 import io.airbyte.db.jdbc.JdbcDatabase;
 import io.airbyte.integrations.BaseConnector;
 import io.airbyte.integrations.base.AirbyteStreamNameNamespacePair;
 import io.airbyte.integrations.base.Source;
 import io.airbyte.integrations.source.relationaldb.models.DbState;
-import io.airbyte.integrations.source.relationaldb.state.AirbyteStateMessageListTypeReference;
 import io.airbyte.integrations.source.relationaldb.state.StateManager;
 import io.airbyte.integrations.source.relationaldb.state.StateManagerFactory;
 import io.airbyte.protocol.models.AirbyteCatalog;
@@ -521,16 +522,18 @@ 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) {
-    if (initialStateJson == null) {
-      return generateEmptyInitialState(config);
-    } else {
-      try {
-        return Jsons.object(initialStateJson, new AirbyteStateMessageListTypeReference());
-      } catch (final IllegalArgumentException e) {
-        LOGGER.warn("Defaulting to legacy state object...");
-        return List.of(new AirbyteStateMessage().withType(AirbyteStateType.LEGACY).withData(initialStateJson));
+    final Optional<StateWrapper> typedState = StateMessageHelper.getTypedState(initialStateJson);
+    return typedState.map((state) -> {
+      switch (state.getStateType()) {
+        case GLOBAL:
+          return List.of(state.getGlobal());
+        case STREAM:
+          return state.getStateMessages();
+        case LEGACY:
+        default:
+          return List.of(new AirbyteStateMessage().withType(AirbyteStateType.LEGACY).withData(state.getLegacyState()));
       }
-    }
+    }).orElse(generateEmptyInitialState(config));
   }
 
   /**

airbyte-integrations/connectors/source-relational-db/src/main/java/io/airbyte/integrations/source/relationaldb/state/CursorManager.java

@@ -92,7 +92,7 @@ protected Map<AirbyteStreamNameNamespacePair, CursorInfo> createCursorInfoMap(
     final Map<AirbyteStreamNameNamespacePair, CursorInfo> localMap = new HashMap<>();
     final Map<AirbyteStreamNameNamespacePair, S> pairToState = streamSupplier.get()
         .stream()
-        .collect(Collectors.toMap(namespacePairFunction,Function.identity()));
+        .collect(Collectors.toMap(namespacePairFunction, Function.identity()));
     final Map<AirbyteStreamNameNamespacePair, ConfiguredAirbyteStream> pairToConfiguredAirbyteStream = catalog.getStreams().stream()
         .collect(Collectors.toMap(AirbyteStreamNameNamespacePair::fromConfiguredAirbyteSteam, Function.identity()));
 

airbyte-integrations/connectors/source-relational-db/src/main/java/io/airbyte/integrations/source/relationaldb/state/GlobalStateManager.java

@@ -94,7 +94,8 @@ private CdcState extractCdcState(final AirbyteStateMessage airbyteStateMessage)
     if (airbyteStateMessage.getType() == AirbyteStateType.GLOBAL) {
       return Jsons.object(airbyteStateMessage.getGlobal().getSharedState(), CdcState.class);
     } else {
-      return Jsons.object(airbyteStateMessage.getData(), DbState.class).getCdcState();
+      final DbState legacyState = Jsons.object(airbyteStateMessage.getData(), DbState.class);
+      return legacyState != null ? legacyState.getCdcState() : null;
     }
   }
 

airbyte-integrations/connectors/source-relational-db/src/main/java/io/airbyte/integrations/source/relationaldb/state/StateGeneratorUtils.java

@@ -148,7 +148,8 @@ public static Optional<DbStreamState> extractState(final AirbyteStreamState stat
    * Tests whether the provided {@link StreamDescriptor} is valid. A valid descriptor is defined as
    * one that has a non-{@code null} name.
    *
-   * See https://github.com/airbytehq/airbyte/blob/e63458fabb067978beb5eaa74d2bc130919b419f/docs/understanding-airbyte/airbyte-protocol.md
+   * See
+   * https://github.com/airbytehq/airbyte/blob/e63458fabb067978beb5eaa74d2bc130919b419f/docs/understanding-airbyte/airbyte-protocol.md
    * for more details
    *
    * @param streamDescriptor A {@link StreamDescriptor} to be validated.
@@ -183,20 +184,6 @@ public static AirbyteStateMessage convertLegacyStateToGlobalState(final AirbyteS
     return new AirbyteStateMessage().withType(AirbyteStateType.GLOBAL).withGlobal(globalState);
   }
 
-  /**
-   * Converts a {@link AirbyteStateType#GLOBAL} state message into a list of
-   * {@link AirbyteStateType#STREAM} messages.
-   *
-   * @param airbyteStateMessage A {@link AirbyteStateType#GLOBAL} state message.
-   * @return A list {@link AirbyteStateType#STREAM} state messages.
-   */
-  public static List<AirbyteStateMessage> convertGlobalStateToStreamState(final AirbyteStateMessage airbyteStateMessage) {
-    return airbyteStateMessage.getGlobal().getStreamStates().stream()
-        .map(s -> new AirbyteStateMessage().withType(AirbyteStateType.STREAM)
-            .withStream(new AirbyteStreamState().withStreamDescriptor(s.getStreamDescriptor()).withStreamState(s.getStreamState())))
-        .collect(Collectors.toList());
-  }
-
   /**
    * Converts a {@link AirbyteStateType#LEGACY} state message into a list of
    * {@link AirbyteStateType#STREAM} messages.

airbyte-integrations/connectors/source-relational-db/src/test/java/io/airbyte/integrations/source/relationaldb/AbstractDbSourceTest.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.integrations.source.relationaldb;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+
+import com.fasterxml.jackson.databind.JsonNode;
+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.util.List;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Test suite for the {@link AbstractDbSource} class.
+ */
+public class AbstractDbSourceTest {
+
+  @Test
+  void testDeserializationOfLegacyState() throws IOException {
+    final AbstractDbSource dbSource = spy(AbstractDbSource.class);
+    final JsonNode config = mock(JsonNode.class);
+
+    final String legacyStateJson = MoreResources.readResource("states/legacy.json");
+    final JsonNode legacyState = Jsons.deserialize(legacyStateJson);
+
+    final List<AirbyteStateMessage> result = dbSource.deserializeInitialState(legacyState, config);
+    assertEquals(1, result.size());
+    assertEquals(AirbyteStateType.LEGACY, result.get(0).getType());
+  }
+
+  @Test
+  void testDeserializationOfGlobalState() throws IOException {
+    final AbstractDbSource dbSource = spy(AbstractDbSource.class);
+    final JsonNode config = mock(JsonNode.class);
+
+    final String globalStateJson = MoreResources.readResource("states/global.json");
+    final JsonNode globalState = Jsons.deserialize(globalStateJson);
+
+    final List<AirbyteStateMessage> result = dbSource.deserializeInitialState(globalState, config);
+    assertEquals(1, result.size());
+    assertEquals(AirbyteStateType.GLOBAL, result.get(0).getType());
+  }
+
+  @Test
+  void testDeserializationOfStreamState() throws IOException {
+    final AbstractDbSource dbSource = spy(AbstractDbSource.class);
+    final JsonNode config = mock(JsonNode.class);
+
+    final String streamStateJson = MoreResources.readResource("states/per_stream.json");
+    final JsonNode streamState = Jsons.deserialize(streamStateJson);
+
+    final List<AirbyteStateMessage> result = dbSource.deserializeInitialState(streamState, config);
+    assertEquals(2, result.size());
+    assertEquals(AirbyteStateType.STREAM, result.get(0).getType());
+  }
+
+  @Test
+  void testDeserializationOfNullState() throws IOException {
+    final AbstractDbSource dbSource = spy(AbstractDbSource.class);
+    final JsonNode config = mock(JsonNode.class);
+
+    final List<AirbyteStateMessage> result = dbSource.deserializeInitialState(null, config);
+    assertEquals(1, result.size());
+    assertEquals(dbSource.getSupportedStateType(config), result.get(0).getType());
+  }
+
+}

airbyte-integrations/connectors/source-relational-db/src/test/java/io/airbyte/integrations/source/relationaldb/state/GlobalStateManagerTest.java

@@ -30,6 +30,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import org.junit.jupiter.api.Test;
 
@@ -202,4 +203,16 @@ void testToState() {
     assertEquals(expected, actualFirstEmission);
   }
 
+  @Test
+  void testToStateWithNoState() {
+    final ConfiguredAirbyteCatalog catalog = new ConfiguredAirbyteCatalog();
+    final StateManager stateManager =
+        new GlobalStateManager(new AirbyteStateMessage(), catalog);
+
+    final AirbyteStateMessage airbyteStateMessage = stateManager.toState(Optional.empty());
+    assertNotNull(airbyteStateMessage);
+    assertEquals(AirbyteStateType.GLOBAL, airbyteStateMessage.getType());
+    assertEquals(0, airbyteStateMessage.getGlobal().getStreamStates().size());
+  }
+
 }

airbyte-integrations/connectors/source-relational-db/src/test/resources/states/global.json

@@ -0,0 +1,49 @@
+[
+  {
+    "type": "GLOBAL",
+    "global": {
+      "shared_state": {
+        "state": {
+          "foo": "bar",
+          "baz": 5
+        }
+      },
+      "stream_states": [
+        {
+          "stream_descriptor": {
+            "name": "bicycles",
+            "namespace": "public"
+          },
+          "stream_state": {
+            "stream_name": "bicycles",
+            "stream_namespace": "public",
+            "cursor_field": ["generation"]
+          }
+        },
+        {
+          "stream_descriptor": {
+            "name": "cars",
+            "namespace": "public"
+          },
+          "stream_state": {
+            "stream_name": "cars",
+            "stream_namespace": "public",
+            "cursor_field": ["year"],
+            "cursor": "a"
+          }
+        },
+        {
+          "stream_descriptor": {
+            "name": "stationary_bicycles",
+            "namespace": "public"
+          },
+          "stream_state": {
+            "stream_name": "stationary_bicycles",
+            "stream_namespace": "public",
+            "cursor_field": []
+          }
+        }
+      ]
+    }
+  }
+]

airbyte-integrations/connectors/source-relational-db/src/test/resources/states/legacy.json

@@ -0,0 +1,17 @@
+{
+  "cdc": false,
+  "streams": [
+    {
+      "cursor": "4",
+      "stream_name": "cars",
+      "cursor_field": ["id"],
+      "stream_namespace": "public"
+    },
+    {
+      "cursor": "1",
+      "stream_name": "us_states",
+      "cursor_field": ["id"],
+      "stream_namespace": "public"
+    }
+  ]
+}

airbyte-integrations/connectors/source-relational-db/src/test/resources/states/per_stream.json

@@ -0,0 +1,32 @@
+[
+  {
+    "type": "STREAM",
+    "stream": {
+      "stream_descriptor": {
+        "name": "id_and_name",
+        "namespace": "public"
+      },
+      "stream_state": {
+        "stream_name": "id_and_name",
+        "stream_namespace": "public",
+        "cursor_field": ["id"],
+        "cursor": "5"
+      }
+    }
+  },
+  {
+    "type": "STREAM",
+    "stream": {
+      "stream_descriptor": {
+        "name": "other",
+        "namespace": "public"
+      },
+      "stream_state": {
+        "stream_name": "other",
+        "stream_namespace": "public",
+        "cursor_field": ["id"],
+        "cursor": "2"
+      }
+    }
+  }
+]