From a99e12d904c9e4d3d42cfdc083e52ac1c690afda Mon Sep 17 00:00:00 2001 From: yuanoOo Date: Thu, 7 May 2026 17:36:57 +0800 Subject: [PATCH] feat: optimize Oracle mode write with multi-row MERGE --- .../OceanBaseConnectionProvider.java | 6 +- .../flink/dialect/OceanBaseDialect.java | 22 ++ .../flink/dialect/OceanBaseOracleDialect.java | 35 +++- .../flink/sink/OceanBaseRecordFlusher.java | 55 +++++ .../flink/OceanBaseOracleConnectorITCase.java | 197 ++++++++++++++++-- .../dialect/OceanBaseOracleDialectTest.java | 98 ++++++++- 6 files changed, 380 insertions(+), 33 deletions(-) diff --git a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/connection/OceanBaseConnectionProvider.java b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/connection/OceanBaseConnectionProvider.java index 89e8beaa..1f4fe2b5 100644 --- a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/connection/OceanBaseConnectionProvider.java +++ b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/connection/OceanBaseConnectionProvider.java @@ -109,7 +109,6 @@ protected void init() { private Properties initializeDefaultJdbcProperties(String jdbcUrl) { Properties defaultJdbcProperties = new Properties(); defaultJdbcProperties.setProperty("useSSL", "false"); - defaultJdbcProperties.setProperty("rewriteBatchedStatements", "true"); defaultJdbcProperties.setProperty("initialTimeout", "2"); defaultJdbcProperties.setProperty("autoReconnect", "true"); defaultJdbcProperties.setProperty("maxReconnects", "3"); @@ -121,6 +120,11 @@ private Properties initializeDefaultJdbcProperties(String jdbcUrl) { defaultJdbcProperties.setProperty("characterEncoding", "UTF-8"); defaultJdbcProperties.setProperty("characterSetResults", "UTF-8"); + if (dialect instanceof OceanBaseMySQLDialect) { + defaultJdbcProperties.setProperty("allowMultiQueries", "true"); + defaultJdbcProperties.setProperty("rewriteBatchedStatements", "true"); + } + // Avoid overwriting user's custom jdbc properties. List jdbcUrlProperties = defaultJdbcProperties.keySet().stream() diff --git a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseDialect.java b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseDialect.java index 6cfe0601..93c868af 100644 --- a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseDialect.java +++ b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseDialect.java @@ -76,6 +76,28 @@ String getUpsertStatement( @Nonnull List uniqueKeyFields, @Nullable SerializableFunction placeholderFunc); + /** + * Gets the upsert statement for multiple rows + * + * @param schemaName schema name + * @param tableName table name + * @param fieldNames field names list + * @param uniqueKeyFields unique key field names list + * @param rowCount number of rows to upsert in a single statement + * @param placeholderFunc function used to get placeholder for the fields + * @return the statement string + */ + default String getUpsertStatement( + @Nonnull String schemaName, + @Nonnull String tableName, + @Nonnull List fieldNames, + @Nonnull List uniqueKeyFields, + int rowCount, + @Nullable SerializableFunction placeholderFunc) { + return getUpsertStatement( + schemaName, tableName, fieldNames, uniqueKeyFields, placeholderFunc); + } + /** * Gets the insert statement * diff --git a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialect.java b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialect.java index d72a2486..57a0ea00 100644 --- a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialect.java +++ b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialect.java @@ -26,6 +26,7 @@ import java.util.List; import java.util.Objects; import java.util.stream.Collectors; +import java.util.stream.IntStream; public class OceanBaseOracleDialect implements OceanBaseDialect { @@ -53,11 +54,37 @@ public String getUpsertStatement( @Nonnull List fieldNames, @Nonnull List uniqueKeyFields, @Nullable SerializableFunction placeholderFunc) { - String sourceFields = + return getUpsertStatement( + schemaName, tableName, fieldNames, uniqueKeyFields, 1, placeholderFunc); + } + + @Override + public String getUpsertStatement( + @Nonnull String schemaName, + @Nonnull String tableName, + @Nonnull List fieldNames, + @Nonnull List uniqueKeyFields, + int rowCount, + @Nullable SerializableFunction placeholderFunc) { + String selectFields = fieldNames.stream() .map(f -> getPlaceholder(f, placeholderFunc) + " AS " + quoteIdentifier(f)) .collect(Collectors.joining(", ")); + String usingClause; + if (rowCount == 1) { + usingClause = "SELECT " + selectFields + " FROM DUAL"; + } else { + usingClause = + "SELECT " + + selectFields + + " FROM DUAL" + + IntStream.range(0, rowCount - 1) + .mapToObj( + i -> " UNION ALL SELECT " + selectFields + " FROM DUAL") + .collect(Collectors.joining()); + } + String onClause = uniqueKeyFields.stream() .map(f -> "t." + quoteIdentifier(f) + "=s." + quoteIdentifier(f)) @@ -80,9 +107,9 @@ public String getUpsertStatement( return "MERGE INTO " + getFullTableName(schemaName, tableName) + " t " - + " USING (SELECT " - + sourceFields - + " FROM DUAL) s " + + " USING (" + + usingClause + + ") s " + " ON (" + onClause + ") " diff --git a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/sink/OceanBaseRecordFlusher.java b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/sink/OceanBaseRecordFlusher.java index 0cfb8635..af16e636 100644 --- a/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/sink/OceanBaseRecordFlusher.java +++ b/flink-connector-oceanbase/src/main/java/com/oceanbase/connector/flink/sink/OceanBaseRecordFlusher.java @@ -23,6 +23,7 @@ import com.oceanbase.connector.flink.connection.OceanBaseVersion; import com.oceanbase.connector.flink.dialect.OceanBaseDialect; import com.oceanbase.connector.flink.dialect.OceanBaseMySQLDialect; +import com.oceanbase.connector.flink.dialect.OceanBaseOracleDialect; import com.oceanbase.connector.flink.table.DataChangeRecord; import com.oceanbase.connector.flink.table.SchemaChangeRecord; import com.oceanbase.connector.flink.table.TableId; @@ -33,6 +34,8 @@ import com.oceanbase.partition.calculator.model.TableEntry; import com.oceanbase.partition.calculator.model.TableEntryKey; +import org.apache.flink.util.function.SerializableFunction; + import org.apache.commons.collections.CollectionUtils; import org.apache.commons.collections.MapUtils; import org.slf4j.Logger; @@ -119,6 +122,14 @@ public synchronized void flush(List records) throws Exception tableInfo.getPlaceholderFunc()), tableInfo.getFieldNames(), upsertBatch); + } else if (dialect instanceof OceanBaseOracleDialect) { + flushMultiRowUpsert( + tableId.getSchemaName(), + tableId.getTableName(), + tableInfo.getFieldNames(), + tableInfo.getKey(), + tableInfo.getPlaceholderFunc(), + upsertBatch); } else { flush( dialect.getUpsertStatement( @@ -211,6 +222,50 @@ private void flush(String sql, List statementFields, List fieldNames, + List uniqueKeyFields, + SerializableFunction placeholderFunc, + List records) + throws Exception { + Map> group = groupRecords(records); + if (group == null) { + return; + } + for (List groupRecords : group.values()) { + String sql = + dialect.getUpsertStatement( + schemaName, + tableName, + fieldNames, + uniqueKeyFields, + groupRecords.size(), + placeholderFunc); + try (Connection connection = connectionProvider.getConnection(); + PreparedStatement statement = connection.prepareStatement(sql)) { + int fieldCount = fieldNames.size(); + for (int rowIdx = 0; rowIdx < groupRecords.size(); rowIdx++) { + DataChangeRecord record = groupRecords.get(rowIdx); + for (int i = 0; i < fieldCount; i++) { + statement.setObject( + rowIdx * fieldCount + i + 1, + record.getFieldValue(fieldNames.get(i))); + } + } + statement.executeUpdate(); + } catch (SQLException e) { + throw new RuntimeException( + "Failed to execute multi-row upsert with sql: " + + sql + + ", records: " + + groupRecords, + e); + } + } + } + private Map> groupRecords(List records) { if (CollectionUtils.isEmpty(records)) { return null; diff --git a/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/OceanBaseOracleConnectorITCase.java b/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/OceanBaseOracleConnectorITCase.java index eaa606af..92e31cb3 100644 --- a/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/OceanBaseOracleConnectorITCase.java +++ b/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/OceanBaseOracleConnectorITCase.java @@ -37,9 +37,11 @@ import org.apache.flink.types.RowKind; import org.apache.commons.collections.CollectionUtils; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; +import java.sql.SQLException; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -49,6 +51,27 @@ @Disabled public class OceanBaseOracleConnectorITCase extends OceanBaseOracleTestBase { + private static final String TABLE_MULTI_ROW_MERGE = "multi_row_merge_test"; + private static final String TABLE_A = "table_a"; + private static final String TABLE_B = "TABLE_b"; + private static final String TABLE_C = "TABLE_C"; + + @AfterEach + public void cleanup() { + dropTableIfExists(TABLE_MULTI_ROW_MERGE); + dropTableIfExists(TABLE_A); + dropTableIfExists(TABLE_B); + dropTableIfExists(TABLE_C); + } + + private void dropTableIfExists(String tableName) { + try { + dropTables(tableName); + } catch (SQLException e) { + // ignore - table may not exist + } + } + @Test public void testSink() throws Exception { StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); @@ -111,8 +134,7 @@ public void testMultipleTableSink() throws Exception { StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); env.setParallelism(1); - OceanBaseConnectorOptions connectorOptions = - new OceanBaseConnectorOptions(getBaseOptions()); + OceanBaseConnectorOptions connectorOptions = new OceanBaseConnectorOptions(getOptions()); OceanBaseSink sink = new OceanBaseSink<>( connectorOptions, @@ -123,9 +145,9 @@ public void testMultipleTableSink() throws Exception { OceanBaseDialect dialect = new OceanBaseOracleDialect(connectorOptions); String schemaName = getSchemaName(); - String tableA = "table_a"; - String tableB = "TABLE_b"; - String tableC = "TABLE_C"; + String tableA = TABLE_A; + String tableB = TABLE_B; + String tableC = TABLE_C; String tableFullNameA = dialect.getFullTableName(schemaName, tableA); String tableFullNameB = dialect.getFullTableName(schemaName, tableB); @@ -200,7 +222,7 @@ public void testMultipleTableSink() throws Exception { RowKind.INSERT, 4, StringData.fromString("4")))); env.fromCollection(dataSet).sinkTo(sink); - env.execute().wait(); + env.execute(); assertEqualsInAnyOrder(queryTable(tableFullNameA), Collections.singletonList("1,1")); assertEqualsInAnyOrder(queryTable(tableFullNameB), Arrays.asList("2,2", "3,3")); @@ -228,7 +250,7 @@ public void testMultipleTableSink() throws Exception { RowKind.DELETE, 3, StringData.fromString("3")))); env.fromCollection(dataSet).sinkTo(sink); - env.execute().wait(); + env.execute(); assertEqualsInAnyOrder(queryTable(tableFullNameA), Collections.singletonList("1,2")); assertEqualsInAnyOrder(queryTable(tableFullNameB), Collections.singletonList("2,3")); @@ -247,30 +269,163 @@ public void testMultipleTableSink() throws Exception { SchemaChangeRecord.Type.TRUNCATE, String.format("TRUNCATE TABLE %s", tableFullNameB))); env.fromCollection(dataSet).sinkTo(sink); - env.execute().wait(); + env.execute(); assertTrue(CollectionUtils.isEmpty(queryTable(tableFullNameA))); assertTrue(CollectionUtils.isEmpty(queryTable(tableFullNameB))); + } - // drop tables - dataSet = + @Test + public void testMultiRowMergeUpsert() throws Exception { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(1); + + OceanBaseConnectorOptions connectorOptions = new OceanBaseConnectorOptions(getOptions()); + OceanBaseSink sink = + new OceanBaseSink<>( + connectorOptions, + null, + new OceanBaseTestDataSerializationSchema(), + DataChangeRecord.KeyExtractor.simple(), + new OceanBaseRecordFlusher(connectorOptions)); + + OceanBaseDialect dialect = new OceanBaseOracleDialect(connectorOptions); + String schemaName = getSchemaName(); + String tableName = TABLE_MULTI_ROW_MERGE; + String tableFullName = dialect.getFullTableName(schemaName, tableName); + + ResolvedSchema tableSchema = + new ResolvedSchema( + Arrays.asList( + Column.physical("id", DataTypes.INT().notNull()), + Column.physical("name", DataTypes.VARCHAR(50).notNull()), + Column.physical("val", DataTypes.INT().notNull())), + Collections.emptyList(), + UniqueConstraint.primaryKey("pk", Collections.singletonList("id"))); + + // Create table and insert 5 rows + List insertDataSet = Arrays.asList( new OceanBaseTestData( schemaName, - tableA, - SchemaChangeRecord.Type.DROP, - String.format("DROP TABLE %s ", tableFullNameA)), + tableName, + SchemaChangeRecord.Type.CREATE, + String.format( + "CREATE TABLE %s (id NUMBER PRIMARY KEY, name VARCHAR2(50), val NUMBER)", + tableFullName)), new OceanBaseTestData( schemaName, - tableB, - SchemaChangeRecord.Type.CREATE, - String.format("DROP TABLE %s ", tableFullNameB)), + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.INSERT, 1, StringData.fromString("alice"), 10)), new OceanBaseTestData( schemaName, - tableC, - SchemaChangeRecord.Type.CREATE, - String.format("DROP TABLE %s ", tableFullNameC))); - env.fromCollection(dataSet).sinkTo(sink); - env.execute().wait(); + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.INSERT, 2, StringData.fromString("bob"), 20)), + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.INSERT, 3, StringData.fromString("carol"), 30)), + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.INSERT, 4, StringData.fromString("dave"), 40)), + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.INSERT, 5, StringData.fromString("eve"), 50))); + + env.fromCollection(insertDataSet).sinkTo(sink); + env.execute(); + + waitingAndAssertTableCount(tableName, 5); + List actual = queryTable(tableName); + assertEqualsInAnyOrder( + Arrays.asList("1,alice,10", "2,bob,20", "3,carol,30", "4,dave,40", "5,eve,50"), + actual); + + // Upsert: update id=1,2 and insert id=6,7 (multi-row MERGE) + List upsertDataSet = + Arrays.asList( + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.UPDATE_AFTER, + 1, + StringData.fromString("alice_v2"), + 111)), + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.UPDATE_AFTER, + 2, + StringData.fromString("bob_v2"), + 222)), + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.INSERT, 6, StringData.fromString("frank"), 60)), + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.INSERT, 7, StringData.fromString("grace"), 70))); + + env.fromCollection(upsertDataSet).sinkTo(sink); + env.execute(); + + waitingAndAssertTableCount(tableName, 7); + actual = queryTable(tableName); + assertEqualsInAnyOrder( + Arrays.asList( + "1,alice_v2,111", + "2,bob_v2,222", + "3,carol,30", + "4,dave,40", + "5,eve,50", + "6,frank,60", + "7,grace,70"), + actual); + + // Delete rows + List deleteDataSet = + Arrays.asList( + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.DELETE, 1, StringData.fromString("alice_v2"), 111)), + new OceanBaseTestData( + schemaName, + tableName, + tableSchema, + GenericRowData.ofKind( + RowKind.DELETE, 7, StringData.fromString("grace"), 70))); + + env.fromCollection(deleteDataSet).sinkTo(sink); + env.execute(); + + waitingAndAssertTableCount(tableName, 5); + actual = queryTable(tableName); + assertEqualsInAnyOrder( + Arrays.asList("2,bob_v2,222", "3,carol,30", "4,dave,40", "5,eve,50", "6,frank,60"), + actual); } } diff --git a/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialectTest.java b/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialectTest.java index d65fe085..fa4604f1 100644 --- a/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialectTest.java +++ b/flink-connector-oceanbase/src/test/java/com/oceanbase/connector/flink/dialect/OceanBaseOracleDialectTest.java @@ -17,19 +17,103 @@ import com.oceanbase.connector.flink.OceanBaseConnectorOptions; +import org.apache.flink.util.function.SerializableFunction; + import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.testcontainers.shaded.com.google.common.collect.Maps; -public class OceanBaseOracleDialectTest { +import java.util.stream.Collectors; +import java.util.stream.Stream; - @Test - public void testQuoteIdentifier() { - OceanBaseConnectorOptions options = new OceanBaseConnectorOptions(Maps.newHashMap()); - Assertions.assertTrue(options.getTableOracleTenantCaseInsensitive()); - OceanBaseOracleDialect oracleDialect = new OceanBaseOracleDialect(options); +class OceanBaseOracleDialectTest { + private final OceanBaseOracleDialect dialect = + new OceanBaseOracleDialect(new OceanBaseConnectorOptions(Maps.newHashMap())); + + @Test + void testQuoteIdentifier() { + Assertions.assertTrue( + new OceanBaseConnectorOptions(Maps.newHashMap()) + .getTableOracleTenantCaseInsensitive()); String identifier = "name"; - Assertions.assertEquals(identifier, oracleDialect.quoteIdentifier(identifier)); + Assertions.assertEquals(identifier, dialect.quoteIdentifier(identifier)); + } + + @Test + void getUpsertStatementSingleRow() { + String upsertStatement = + dialect.getUpsertStatement( + "sche1", + "tb1", + Stream.of("id", "name").collect(Collectors.toList()), + Stream.of("id").collect(Collectors.toList()), + 1, + (SerializableFunction) s -> "?"); + Assertions.assertEquals( + "MERGE INTO sche1.tb1 t USING (SELECT ? AS id, ? AS name FROM DUAL) s" + + " ON (t.id=s.id)" + + " WHEN MATCHED THEN UPDATE SET t.name=s.name" + + " WHEN NOT MATCHED THEN INSERT (id, name) VALUES (s.id, s.name)", + upsertStatement); + } + + @Test + void getUpsertStatementMultiRow() { + String upsertStatement = + dialect.getUpsertStatement( + "sche1", + "tb1", + Stream.of("id", "name").collect(Collectors.toList()), + Stream.of("id").collect(Collectors.toList()), + 3, + (SerializableFunction) s -> "?"); + Assertions.assertEquals( + "MERGE INTO sche1.tb1 t USING (SELECT ? AS id, ? AS name FROM DUAL" + + " UNION ALL SELECT ? AS id, ? AS name FROM DUAL" + + " UNION ALL SELECT ? AS id, ? AS name FROM DUAL) s" + + " ON (t.id=s.id)" + + " WHEN MATCHED THEN UPDATE SET t.name=s.name" + + " WHEN NOT MATCHED THEN INSERT (id, name) VALUES (s.id, s.name)", + upsertStatement); + } + + @Test + void getUpsertStatementMultiRowCompositeKey() { + String upsertStatement = + dialect.getUpsertStatement( + "sche1", + "tb1", + Stream.of("id", "name", "age").collect(Collectors.toList()), + Stream.of("id", "name").collect(Collectors.toList()), + 2, + (SerializableFunction) s -> "?"); + Assertions.assertEquals( + "MERGE INTO sche1.tb1 t USING (SELECT ? AS id, ? AS name, ? AS age FROM DUAL" + + " UNION ALL SELECT ? AS id, ? AS name, ? AS age FROM DUAL) s" + + " ON (t.id=s.id and t.name=s.name)" + + " WHEN MATCHED THEN UPDATE SET t.age=s.age" + + " WHEN NOT MATCHED THEN INSERT (id, name, age) VALUES (s.id, s.name, s.age)", + upsertStatement); + } + + @Test + void getUpsertStatementDefaultDelegatesToMultiRow() { + String singleRow = + dialect.getUpsertStatement( + "sche1", + "tb1", + Stream.of("id", "name").collect(Collectors.toList()), + Stream.of("id").collect(Collectors.toList()), + (SerializableFunction) s -> "?"); + String multiRow = + dialect.getUpsertStatement( + "sche1", + "tb1", + Stream.of("id", "name").collect(Collectors.toList()), + Stream.of("id").collect(Collectors.toList()), + 1, + (SerializableFunction) s -> "?"); + Assertions.assertEquals(singleRow, multiRow); } }