[source-mysql]: support Planetscale MySQL's per-query row limit (#40561)

Fixes airbytehq/oncall#5051 Planet MySQL has defined multiple system limits. One of our customers has been hitting the 100K per-query row limit, which causes our connector to emit system errors. This patch defines a chunk size limit for our connector. By default, it is the largest long, and if we find it is a Planetscale MySQL, we will adjust to 100K.

Author: theyueli

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

@@ -9,7 +9,7 @@ data:
  connectorSubtype: database
  connectorType: source
  definitionId: 435bb9a5-7887-4809-aa58-28c27df0d7ad
- dockerImageTag: 3.4.10
+ dockerImageTag: 3.4.11
  dockerRepository: airbyte/source-mysql
  documentationUrl: https://docs.airbyte.com/integrations/sources/mysql
  githubIssueLabel: source-mysql

airbyte-integrations/connectors/source-mysql/src/main/java/io/airbyte/integrations/source/mysql/initialsync/MySqlInitialLoadHandler.java

@@ -12,6 +12,7 @@
 import com.mysql.cj.MysqlType;
 import io.airbyte.cdk.db.jdbc.AirbyteRecordData;
 import io.airbyte.cdk.db.jdbc.JdbcDatabase;
+import io.airbyte.cdk.db.jdbc.JdbcUtils;
 import io.airbyte.cdk.integrations.debezium.DebeziumIteratorConstants;
 import io.airbyte.cdk.integrations.source.relationaldb.DbSourceDiscoverUtil;
 import io.airbyte.cdk.integrations.source.relationaldb.InitialLoadHandler;
@@ -57,6 +58,7 @@ public class MySqlInitialLoadHandler implements InitialLoadHandler<MysqlType> {
 
  private static final long QUERY_TARGET_SIZE_GB = 1_073_741_824;
  private static final long DEFAULT_CHUNK_SIZE = 1_000_000;
+ private long MAX_CHUNK_SIZE = Long.MAX_VALUE;
  final Map<AirbyteStreamNameNamespacePair, TableSizeInfo> tableSizeInfoMap;
 
  public MySqlInitialLoadHandler(final JsonNode config,
@@ -73,6 +75,16 @@ public MySqlInitialLoadHandler(final JsonNode config,
  this.initialLoadStateManager = initialLoadStateManager;
  this.streamStateForIncrementalRunSupplier = streamStateForIncrementalRunSupplier;
  this.tableSizeInfoMap = tableSizeInfoMap;
+ adjustChunkSizeLimitForMySQLVariants();
+ }
+
+ private void adjustChunkSizeLimitForMySQLVariants() {
+ // For PSDB, we need to limit the chunk size to 100k rows to avoid the query being killed by the
+ // server.
+ // Reference:
+ // https://planetscale.com/docs/reference/planetscale-system-limits
+ if (config.get(JdbcUtils.HOST_KEY).asText().toLowerCase().contains("psdb.cloud"))
+ MAX_CHUNK_SIZE = 100_000;
  }
 
  public List<AutoCloseableIterator<AirbyteMessage>> getIncrementalIterators(
@@ -122,7 +134,7 @@ public AutoCloseableIterator<AirbyteMessage> getIteratorForStream(
  .collect(Collectors.toList());
  final AutoCloseableIterator<AirbyteRecordData> queryStream =
  new MySqlInitialLoadRecordIterator(database, sourceOperations, quoteString, initialLoadStateManager, selectedDatabaseFields, pair,
- calculateChunkSize(tableSizeInfoMap.get(pair), pair), isCompositePrimaryKey(airbyteStream));
+ Long.min(calculateChunkSize(tableSizeInfoMap.get(pair), pair), MAX_CHUNK_SIZE), isCompositePrimaryKey(airbyteStream));
  final AutoCloseableIterator<AirbyteMessage> recordIterator =
  getRecordIterator(queryStream, streamName, namespace, emittedAt.toEpochMilli());
  final AutoCloseableIterator<AirbyteMessage> recordAndMessageIterator = augmentWithState(recordIterator, airbyteStream, pair);

docs/integrations/sources/mysql.md

@@ -233,6 +233,7 @@ Any database or table encoding combination of charset and collation is supported
 
 | Version | Date | Pull Request | Subject |
 |:--------|:-----------|:-----------------------------------------------------------|:------------------------------------------------------------------------------------------------------------------------------------------------|
+| 3.4.11 | 2024-06-26 | [40561](https://github.com/airbytehq/airbyte/pull/40561) | Support PlanetScale MySQL's per-query row limit. |
 | 3.4.10 | 2024-06-14 | [39349](https://github.com/airbytehq/airbyte/pull/39349) | Full refresh stream sending internal count metadata. |
 | 3.4.9 | 2024-06-11 | [39405](https://github.com/airbytehq/airbyte/pull/39405) | Adopt latest CDK. |
 | 3.4.8 | 2024-06-05 | [39144](https://github.com/airbytehq/airbyte/pull/39144) | Upgrade Debezium to 2.5.4 |