Revert "Revert "source-postgres: connect with adaptiveFetch=true" (#38365)"

This reverts commit 0db33b8.

Author: Marius Posta

airbyte-integrations/connectors/source-postgres/metadata.yaml

@@ -9,7 +9,7 @@ data:
   connectorSubtype: database
   connectorType: source
   definitionId: decd338e-5647-4c0b-adf4-da0e75f5a750
-  dockerImageTag: 3.4.7
+  dockerImageTag: 3.4.6
   dockerRepository: airbyte/source-postgres
   documentationUrl: https://docs.airbyte.com/integrations/sources/postgres
   githubIssueLabel: source-postgres

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

@@ -39,6 +39,11 @@
 import static io.airbyte.integrations.source.postgres.xmin.XminCtidUtils.reclassifyCategorisedCtidStreams;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toSet;
+import static org.postgresql.PGProperty.ADAPTIVE_FETCH;
+import static org.postgresql.PGProperty.CURRENT_SCHEMA;
+import static org.postgresql.PGProperty.DEFAULT_ROW_FETCH_SIZE;
+import static org.postgresql.PGProperty.MAX_RESULT_BUFFER;
+import static org.postgresql.PGProperty.PREPARE_THRESHOLD;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
@@ -52,7 +57,6 @@
 import io.airbyte.cdk.db.jdbc.JdbcDatabase;
 import io.airbyte.cdk.db.jdbc.JdbcUtils;
 import io.airbyte.cdk.db.jdbc.StreamingJdbcDatabase;
-import io.airbyte.cdk.db.jdbc.streaming.AdaptiveStreamingQueryConfig;
 import io.airbyte.cdk.integrations.base.AirbyteTraceMessageUtility;
 import io.airbyte.cdk.integrations.base.IntegrationRunner;
 import io.airbyte.cdk.integrations.base.Source;
@@ -122,6 +126,7 @@
 import java.util.stream.Stream;
 import javax.sql.DataSource;
 import org.apache.commons.lang3.StringUtils;
+import org.postgresql.PGProperty;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -146,6 +151,14 @@ public class PostgresSource extends AbstractJdbcSource<PostgresType> implements
   public static final String SSL_MODE_DISABLE = "disable";
   public static final String SSL_MODE_REQUIRE = "require";
 
+  public static final Map<PGProperty, String> JDBC_CONNECTION_PARAMS = ImmutableMap.of(
+      // Initialize parameters with prepareThreshold=0 to mitigate pgbouncer errors
+      // https://github.com/airbytehq/airbyte/issues/24796
+      PREPARE_THRESHOLD, "0",
+      DEFAULT_ROW_FETCH_SIZE, "1",
+      ADAPTIVE_FETCH, "true",
+      MAX_RESULT_BUFFER, "10percent");
+
   private List<String> schemas;
 
   private Set<AirbyteStreamNameNamespacePair> publicizedTablesInCdc;
@@ -157,7 +170,7 @@ public static Source sshWrappedSource(PostgresSource source) {
   }
 
   PostgresSource() {
-    super(DRIVER_CLASS, AdaptiveStreamingQueryConfig::new, new PostgresSourceOperations());
+    super(DRIVER_CLASS, PostgresStreamingQueryConfig::new, new PostgresSourceOperations());
     this.stateEmissionFrequency = INTERMEDIATE_STATE_EMISSION_FREQUENCY;
   }
 
@@ -176,9 +189,9 @@ public ConnectorSpecification spec() throws Exception {
   @Override
   public JsonNode toDatabaseConfig(final JsonNode config) {
     final List<String> additionalParameters = new ArrayList<>();
-    // Initialize parameters with prepareThreshold=0 to mitigate pgbouncer errors
-    // https://github.com/airbytehq/airbyte/issues/24796
-    additionalParameters.add("prepareThreshold=0");
+    for (var e : JDBC_CONNECTION_PARAMS.entrySet()) {
+      additionalParameters.add(e.getKey().getName() + EQUALS + e.getValue());
+    }
 
     final String encodedDatabaseName = URLEncoder.encode(config.get(JdbcUtils.DATABASE_KEY).asText(), StandardCharsets.UTF_8);
 
@@ -188,7 +201,7 @@ public JsonNode toDatabaseConfig(final JsonNode config) {
         encodedDatabaseName));
 
     if (config.get(JdbcUtils.JDBC_URL_PARAMS_KEY) != null && !config.get(JdbcUtils.JDBC_URL_PARAMS_KEY).asText().isEmpty()) {
-      jdbcUrl.append(config.get(JdbcUtils.JDBC_URL_PARAMS_KEY).asText()).append(AMPERSAND);
+      additionalParameters.add(config.get(JdbcUtils.JDBC_URL_PARAMS_KEY).asText());
     }
 
     final Map<String, String> sslParameters = parseSSLConfig(config);
@@ -206,12 +219,10 @@ public JsonNode toDatabaseConfig(final JsonNode config) {
     }
 
     if (schemas != null && !schemas.isEmpty()) {
-      additionalParameters.add("currentSchema=" + String.join(",", schemas));
+      additionalParameters.add(CURRENT_SCHEMA.getName() + EQUALS + String.join(",", schemas));
     }
-
-    additionalParameters.forEach(x -> jdbcUrl.append(x).append("&"));
-
-    jdbcUrl.append(toJDBCQueryParams(sslParameters));
+    additionalParameters.addAll(toJDBCQueryParams(sslParameters));
+    jdbcUrl.append(String.join(AMPERSAND, additionalParameters));
     LOGGER.debug("jdbc url: {}", jdbcUrl);
     final ImmutableMap.Builder<Object, Object> configBuilder = ImmutableMap.builder()
         .put(JdbcUtils.USERNAME_KEY, config.get(JdbcUtils.USERNAME_KEY).asText())
@@ -225,8 +236,9 @@ public JsonNode toDatabaseConfig(final JsonNode config) {
     return Jsons.jsonNode(configBuilder.build());
   }
 
-  public String toJDBCQueryParams(final Map<String, String> sslParams) {
-    return Objects.isNull(sslParams) ? ""
+  public List<String> toJDBCQueryParams(final Map<String, String> sslParams) {
+    return Objects.isNull(sslParams)
+        ? List.of()
         : sslParams.entrySet()
             .stream()
             .map((entry) -> {
@@ -243,7 +255,7 @@ public String toJDBCQueryParams(final Map<String, String> sslParams) {
               }
             })
             .filter(s -> Objects.nonNull(s) && !s.isEmpty())
-            .collect(Collectors.joining(JdbcUtils.AMPERSAND));
+            .toList();
   }
 
   @Override

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

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.integrations.source.postgres;
+
+import io.airbyte.cdk.db.jdbc.streaming.JdbcStreamingQueryConfig;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import org.jetbrains.annotations.NotNull;
+
+public class PostgresStreamingQueryConfig implements JdbcStreamingQueryConfig {
+
+  @Override
+  public void initialize(final Connection connection, final @NotNull Statement preparedStatement) throws SQLException {
+    connection.setAutoCommit(false);
+    // Nothing else to do, adaptive streaming is enabled via JDBC connection parameters.
+  }
+
+  @Override
+  public void accept(ResultSet resultSet, Object o) {}
+
+}

airbyte-integrations/connectors/source-postgres/src/test/java/io/airbyte/integrations/source/postgres/PostgresSourceTest.java

@@ -815,10 +815,21 @@ private ConfiguredAirbyteStream createTableWithInvalidCursorType(final Database
   @Test
   void testJdbcUrlWithEscapedDatabaseName() {
     final JsonNode jdbcConfig = source().toDatabaseConfig(buildConfigEscapingNeeded());
-    assertEquals(EXPECTED_JDBC_ESCAPED_URL, jdbcConfig.get(JdbcUtils.JDBC_URL_KEY).asText());
+    assertEquals("jdbc:postgresql://localhost:1111/db%2Ffoo?" + EXPECTED_DEFAULT_PARAMS,
+        jdbcConfig.get(JdbcUtils.JDBC_URL_KEY).asText());
   }
 
-  private static final String EXPECTED_JDBC_ESCAPED_URL = "jdbc:postgresql://localhost:1111/db%2Ffoo?prepareThreshold=0&";
+  @Test
+  void testJdbcUrlWithSchemas() {
+    final JsonNode sourceConfig = buildConfigEscapingNeeded();
+    ((ObjectNode) sourceConfig).set("schemas", Jsons.arrayNode().add("bar").add("baz"));
+    final JsonNode jdbcConfig = source().toDatabaseConfig(sourceConfig);
+    assertEquals("jdbc:postgresql://localhost:1111/db%2Ffoo?" + EXPECTED_DEFAULT_PARAMS + "&currentSchema=bar,baz",
+        jdbcConfig.get(JdbcUtils.JDBC_URL_KEY).asText());
+  }
+
+  private static final String EXPECTED_DEFAULT_PARAMS =
+      "prepareThreshold=0&defaultRowFetchSize=1&adaptiveFetch=true&maxResultBuffer=10percent";
 
   private JsonNode buildConfigEscapingNeeded() {
     return Jsons.jsonNode(ImmutableMap.of(