From 47448b331bd4d613bc9aa7eddc145feea2e3d844 Mon Sep 17 00:00:00 2001 From: yuqi Date: Tue, 20 Oct 2026 11:39:35 +0800 Subject: [PATCH] fix(jdbc): Escape Doris and StarRocks partition values --- .../catalog/doris/utils/DorisUtils.java | 17 ++- .../TestDorisPartitionSqlGeneration.java | 115 +++++++++++++++++ .../TestDorisTablePartitionOperations.java | 87 +++++++++++++ .../catalog/doris/utils/TestDorisUtils.java | 49 ++++++++ .../starrocks/utils/StarRocksUtils.java | 9 +- .../TestStarRocksPartitionSqlGeneration.java | 117 ++++++++++++++++++ ...TestStarRocksTablePartitionOperations.java | 87 +++++++++++++ .../starrocks/utils/TestStarRocksUtils.java | 41 ++++++ 8 files changed, 514 insertions(+), 8 deletions(-) create mode 100644 catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisPartitionSqlGeneration.java create mode 100644 catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksPartitionSqlGeneration.java diff --git a/catalogs/catalog-jdbc-doris/src/main/java/org/apache/gravitino/catalog/doris/utils/DorisUtils.java b/catalogs/catalog-jdbc-doris/src/main/java/org/apache/gravitino/catalog/doris/utils/DorisUtils.java index 80a376411ed..c0a437e9662 100644 --- a/catalogs/catalog-jdbc-doris/src/main/java/org/apache/gravitino/catalog/doris/utils/DorisUtils.java +++ b/catalogs/catalog-jdbc-doris/src/main/java/org/apache/gravitino/catalog/doris/utils/DorisUtils.java @@ -171,11 +171,18 @@ private static String generateRangePartitionValues(RangePartition rangePartition if (Literals.NULL.equals(upper) && Literals.NULL.equals(lower)) { partitionValues = "LESS THAN MAXVALUE"; } else if (Literals.NULL.equals(lower)) { - partitionValues = String.format("LESS THAN (\"%s\")", upper.value()); + partitionValues = + String.format("LESS THAN (\"%s\")", escapeSqlLiteral(String.valueOf(upper.value()), '"')); } else if (Literals.NULL.equals(upper)) { - partitionValues = String.format("[(\"%s\"), (MAXVALUE))", lower.value()); + partitionValues = + String.format( + "[(\"%s\"), (MAXVALUE))", escapeSqlLiteral(String.valueOf(lower.value()), '"')); } else { - partitionValues = String.format("[(\"%s\"), (\"%s\"))", lower.value(), upper.value()); + partitionValues = + String.format( + "[(\"%s\"), (\"%s\"))", + escapeSqlLiteral(String.valueOf(lower.value()), '"'), + escapeSqlLiteral(String.valueOf(upper.value()), '"')); } return partitionValues; } @@ -190,10 +197,10 @@ private static String generateListPartitionSqlValues(ListPartition listPartition String.format( "(%s)", Arrays.stream(part) - .map(p -> "\"" + p.value() + "\"") + .map(p -> "\"" + escapeSqlLiteral(String.valueOf(p.value()), '"') + "\"") .collect(Collectors.joining(","))); } else { - values = String.format("\"%s\"", part[0].value()); + values = String.format("\"%s\"", escapeSqlLiteral(String.valueOf(part[0].value()), '"')); } listValues.add(values); } diff --git a/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisPartitionSqlGeneration.java b/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisPartitionSqlGeneration.java new file mode 100644 index 00000000000..b8ff6a750bf --- /dev/null +++ b/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisPartitionSqlGeneration.java @@ -0,0 +1,115 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.catalog.doris.operation; + +import java.util.Collections; +import java.util.Map; +import org.apache.gravitino.catalog.doris.converter.DorisColumnDefaultValueConverter; +import org.apache.gravitino.catalog.doris.converter.DorisExceptionConverter; +import org.apache.gravitino.catalog.doris.converter.DorisTypeConverter; +import org.apache.gravitino.catalog.jdbc.JdbcColumn; +import org.apache.gravitino.catalog.jdbc.JdbcTable; +import org.apache.gravitino.catalog.jdbc.utils.ConnectionCountingDataSource; +import org.apache.gravitino.rel.expressions.NamedReference; +import org.apache.gravitino.rel.expressions.distributions.Distributions; +import org.apache.gravitino.rel.expressions.literals.Literal; +import org.apache.gravitino.rel.expressions.literals.Literals; +import org.apache.gravitino.rel.expressions.transforms.Transform; +import org.apache.gravitino.rel.expressions.transforms.Transforms; +import org.apache.gravitino.rel.indexes.Indexes; +import org.apache.gravitino.rel.partitions.ListPartition; +import org.apache.gravitino.rel.partitions.Partition; +import org.apache.gravitino.rel.partitions.Partitions; +import org.apache.gravitino.rel.partitions.RangePartition; +import org.apache.gravitino.rel.types.Types; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +class TestDorisPartitionSqlGeneration { + @ParameterizedTest + @ValueSource(strings = {"range", "list", "multi-list"}) + void testCreateAndAddPartitionEscapeValues(String kind) { + Literal value = Literals.stringLiteral("say \"hi\" C:\\temp\\; --"); + String quoted = "\"say \"\"hi\"\" C:\\\\temp\\\\; --\""; + Partition partition; + Transform transform; + String fragment; + if ("range".equals(kind)) { + RangePartition range = Partitions.range("p", value, Literals.NULL, Collections.emptyMap()); + partition = range; + transform = Transforms.range(new String[] {"col1"}, new RangePartition[] {range}); + fragment = "PARTITION `p` VALUES LESS THAN (" + quoted + ")"; + } else { + boolean multiColumn = "multi-list".equals(kind); + ListPartition list = + Partitions.list( + "p", + multiColumn ? new Literal[][] {{value, value}} : new Literal[][] {{value}}, + Collections.emptyMap()); + partition = list; + transform = + Transforms.list( + multiColumn ? new String[][] {{"col1"}, {"col2"}} : new String[][] {{"col1"}}, + new ListPartition[] {list}); + fragment = + "PARTITION `p` VALUES IN (" + + (multiColumn ? "(" + quoted + "," + quoted + ")" : quoted) + + ")"; + } + JdbcColumn[] columns = { + JdbcColumn.builder().withName("col1").withType(Types.VarCharType.of(128)).build(), + JdbcColumn.builder().withName("col2").withType(Types.VarCharType.of(128)).build() + }; + ConnectionCountingDataSource connections = new ConnectionCountingDataSource(); + DorisTableOperations operations = new DorisTableOperations(); + DorisExceptionConverter exceptions = new DorisExceptionConverter(); + DorisTypeConverter types = new DorisTypeConverter(); + operations.initialize( + connections.dataSource(), + exceptions, + types, + new DorisColumnDefaultValueConverter(), + Collections.emptyMap()); + operations.create( + "database", + "test_table", + columns, + null, + Map.of("replication_num", "1"), + new Transform[] {transform}, + Distributions.hash(1, NamedReference.field("col1")), + Indexes.EMPTY_INDEXES); + Assertions.assertTrue(connections.executedSql().get(0).contains(fragment)); + JdbcTable table = + JdbcTable.builder() + .withName("test_table") + .withDatabaseName("database") + .withColumns(columns) + .withPartitioning(new Transform[] {transform}) + .build(); + DorisTablePartitionOperations partitions = + new DorisTablePartitionOperations(connections.dataSource(), table, exceptions, types); + partitions.addPartition(partition); + Assertions.assertEquals(2, connections.executedSql().size()); + Assertions.assertEquals( + "ALTER TABLE `test_table` ADD " + fragment, connections.executedSql().get(1)); + Assertions.assertEquals(0, connections.borrowed()); + } +} diff --git a/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisTablePartitionOperations.java b/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisTablePartitionOperations.java index ffbcb98d3c6..18c91f2a51d 100644 --- a/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisTablePartitionOperations.java +++ b/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/operation/TestDorisTablePartitionOperations.java @@ -24,6 +24,9 @@ import static org.junit.Assert.assertTrue; import com.google.common.collect.ImmutableMap; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; import java.time.LocalDate; import java.util.Arrays; import java.util.Collections; @@ -56,6 +59,8 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; @Tag("gravitino-docker-test") public class TestDorisTablePartitionOperations extends TestDoris { @@ -357,4 +362,86 @@ public void testListPartition() { // drop non-existing partition assertFalse(tablePartitionOperations.dropPartition("p3")); } + + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testEscapedListPartitionValues(int partitionColumns) throws Exception { + String tableName = GravitinoITUtils.genRandomName("escaped_list_partition"); + JdbcColumn[] columns = { + JdbcColumn.builder() + .withName("col1") + .withType(Types.VarCharType.of(128)) + .withNullable(false) + .build(), + JdbcColumn.builder() + .withName("col2") + .withType(Types.VarCharType.of(128)) + .withNullable(false) + .build() + }; + String[] values = {"say \"hi\" C:\\temp\\; --", "5\" pipe C:\\next\\"}; + Literal[][] initial = + partitionColumns == 1 + ? new Literal[][] {{Literals.stringLiteral(values[0])}} + : new Literal[][] { + {Literals.stringLiteral(values[0]), Literals.stringLiteral(values[0])} + }; + ListPartition created = Partitions.list("p_created", initial, Collections.emptyMap()); + Transform transform = + Transforms.list( + partitionColumns == 1 ? new String[][] {{"col1"}} : new String[][] {{"col1"}, {"col2"}}, + new ListPartition[] {created}); + TABLE_OPERATIONS.create( + databaseName, + tableName, + columns, + "escaped partition values", + Map.of("replication_num", "1"), + new Transform[] {transform}, + Distributions.hash(1, NamedReference.field("col1")), + new Index[0]); + try { + JdbcTable table = + JdbcTable.builder() + .withName(tableName) + .withDatabaseName(databaseName) + .withColumns(columns) + .withPartitioning(new Transform[] {transform}) + .build(); + JdbcTablePartitionOperations partitions = + new DorisTablePartitionOperations( + DATA_SOURCE, table, JDBC_EXCEPTION_CONVERTER, TYPE_CONVERTER); + Literal[][] added = + partitionColumns == 1 + ? new Literal[][] {{Literals.stringLiteral(values[1])}} + : new Literal[][] { + {Literals.stringLiteral(values[1]), Literals.stringLiteral(values[1])} + }; + partitions.addPartition(Partitions.list("p_added", added, Collections.emptyMap())); + // Inserting the original values proves both DDL paths preserved the partition keys. + try (Connection connection = DATA_SOURCE.getConnection()) { + connection.setCatalog(databaseName); + for (String value : values) { + try (PreparedStatement insert = + connection.prepareStatement("INSERT INTO `" + tableName + "` VALUES (?, ?)")) { + insert.setString(1, value); + insert.setString(2, value); + insert.executeUpdate(); + } + try (PreparedStatement select = + connection.prepareStatement( + "SELECT COUNT(*) FROM `" + tableName + "` WHERE col1 = ? AND col2 = ?")) { + select.setString(1, value); + select.setString(2, value); + try (ResultSet result = select.executeQuery()) { + assertTrue(result.next()); + assertEquals(1, result.getInt(1)); + } + } + } + } + } finally { + TABLE_OPERATIONS.drop(databaseName, tableName); + } + } } diff --git a/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/utils/TestDorisUtils.java b/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/utils/TestDorisUtils.java index 4a2b37981e9..017254fa362 100644 --- a/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/utils/TestDorisUtils.java +++ b/catalogs/catalog-jdbc-doris/src/test/java/org/apache/gravitino/catalog/doris/utils/TestDorisUtils.java @@ -26,6 +26,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Optional; +import java.util.stream.Stream; import org.apache.gravitino.rel.expressions.distributions.Distribution; import org.apache.gravitino.rel.expressions.literals.Literal; import org.apache.gravitino.rel.expressions.literals.Literals; @@ -35,6 +36,9 @@ import org.apache.gravitino.rel.partitions.Partitions; import org.apache.gravitino.rel.types.Types; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; public class TestDorisUtils { @Test @@ -268,4 +272,49 @@ public void testExtractBucketNumWithWhitespace() { DorisUtils.extractDistributionInfoFromSql(createTableSqlWithAutoLeadingWhitespace); assertEquals(-1, distributionAutoLeading.number()); } + + @ParameterizedTest + @MethodSource("partitionLiteralValues") + void testPartitionLiteralEscaping(String value, String escaped) { + Literal literal = Literals.stringLiteral(value); + assertEquals( + "PARTITION `p` VALUES LESS THAN (\"" + escaped + "\")", + DorisUtils.generatePartitionSqlFragment( + Partitions.range("p", literal, Literals.NULL, Collections.emptyMap()))); + assertEquals( + "PARTITION `p` VALUES [(\"" + escaped + "\"), (MAXVALUE))", + DorisUtils.generatePartitionSqlFragment( + Partitions.range("p", Literals.NULL, literal, Collections.emptyMap()))); + assertEquals( + "PARTITION `p` VALUES [(\"" + escaped + "\"), (\"" + escaped + "\"))", + DorisUtils.generatePartitionSqlFragment( + Partitions.range("p", literal, literal, Collections.emptyMap()))); + + assertEquals( + "PARTITION `p` VALUES IN (\"" + escaped + "\",\"plain\")", + DorisUtils.generatePartitionSqlFragment( + Partitions.list( + "p", + new Literal[][] {{literal}, {Literals.stringLiteral("plain")}}, + Collections.emptyMap()))); + assertEquals( + "PARTITION `p` VALUES IN ((\"" + escaped + "\",\"plain\"),(\"plain\",\"" + escaped + "\"))", + DorisUtils.generatePartitionSqlFragment( + Partitions.list( + "p", + new Literal[][] { + {literal, Literals.stringLiteral("plain")}, + {Literals.stringLiteral("plain"), literal} + }, + Collections.emptyMap()))); + } + + private static Stream partitionLiteralValues() { + return Stream.of( + Arguments.of("say \"hi\"", "say \"\"hi\"\""), + Arguments.of("C:\\temp\\", "C:\\\\temp\\\\"), + Arguments.of("5\" pipe' ; --", "5\"\" pipe' ; --"), + Arguments.of("slash\\\"quote", "slash\\\\\"\"quote"), + Arguments.of("line\nnext\tvalue", "line\nnext\tvalue")); + } } diff --git a/catalogs/catalog-jdbc-starrocks/src/main/java/org/apache/gravitino/catalog/starrocks/utils/StarRocksUtils.java b/catalogs/catalog-jdbc-starrocks/src/main/java/org/apache/gravitino/catalog/starrocks/utils/StarRocksUtils.java index a918b26dc0c..6356e6e22ee 100644 --- a/catalogs/catalog-jdbc-starrocks/src/main/java/org/apache/gravitino/catalog/starrocks/utils/StarRocksUtils.java +++ b/catalogs/catalog-jdbc-starrocks/src/main/java/org/apache/gravitino/catalog/starrocks/utils/StarRocksUtils.java @@ -18,6 +18,8 @@ */ package org.apache.gravitino.catalog.starrocks.utils; +import static org.apache.gravitino.catalog.jdbc.utils.JdbcConnectorUtils.escapeSqlLiteral; + import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import java.sql.ResultSet; @@ -239,7 +241,8 @@ private static String generateRangePartitionValues(RangePartition rangePartition if (Literals.NULL.equals(upper)) { partitionValues = "LESS THAN MAXVALUE"; } else { - partitionValues = String.format("LESS THAN (\"%s\")", upper.value()); + partitionValues = + String.format("LESS THAN (\"%s\")", escapeSqlLiteral(String.valueOf(upper.value()), '"')); } return partitionValues; } @@ -254,10 +257,10 @@ private static String generateListPartitionSqlValues(ListPartition listPartition String.format( "(%s)", Arrays.stream(part) - .map(p -> "\"" + p.value() + "\"") + .map(p -> "\"" + escapeSqlLiteral(String.valueOf(p.value()), '"') + "\"") .collect(Collectors.joining(","))); } else { - values = String.format("\"%s\"", part[0].value()); + values = String.format("\"%s\"", escapeSqlLiteral(String.valueOf(part[0].value()), '"')); } listValues.add(values); } diff --git a/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksPartitionSqlGeneration.java b/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksPartitionSqlGeneration.java new file mode 100644 index 00000000000..a84868b60b9 --- /dev/null +++ b/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksPartitionSqlGeneration.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.catalog.starrocks.operation; + +import java.util.Collections; +import java.util.Map; +import org.apache.gravitino.catalog.jdbc.JdbcColumn; +import org.apache.gravitino.catalog.jdbc.JdbcTable; +import org.apache.gravitino.catalog.jdbc.utils.ConnectionCountingDataSource; +import org.apache.gravitino.catalog.starrocks.converter.StarRocksColumnDefaultValueConverter; +import org.apache.gravitino.catalog.starrocks.converter.StarRocksExceptionConverter; +import org.apache.gravitino.catalog.starrocks.converter.StarRocksTypeConverter; +import org.apache.gravitino.catalog.starrocks.operations.StarRocksTableOperations; +import org.apache.gravitino.catalog.starrocks.operations.StarRocksTablePartitionOperations; +import org.apache.gravitino.rel.expressions.NamedReference; +import org.apache.gravitino.rel.expressions.distributions.Distributions; +import org.apache.gravitino.rel.expressions.literals.Literal; +import org.apache.gravitino.rel.expressions.literals.Literals; +import org.apache.gravitino.rel.expressions.transforms.Transform; +import org.apache.gravitino.rel.expressions.transforms.Transforms; +import org.apache.gravitino.rel.indexes.Indexes; +import org.apache.gravitino.rel.partitions.ListPartition; +import org.apache.gravitino.rel.partitions.Partition; +import org.apache.gravitino.rel.partitions.Partitions; +import org.apache.gravitino.rel.partitions.RangePartition; +import org.apache.gravitino.rel.types.Types; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +class TestStarRocksPartitionSqlGeneration { + @ParameterizedTest + @ValueSource(strings = {"range", "list", "multi-list"}) + void testCreateAndAddPartitionEscapeValues(String kind) { + Literal value = Literals.stringLiteral("say \"hi\" C:\\temp\\; --"); + String quoted = "\"say \"\"hi\"\" C:\\\\temp\\\\; --\""; + Partition partition; + Transform transform; + String fragment; + if ("range".equals(kind)) { + RangePartition range = Partitions.range("p", value, Literals.NULL, Collections.emptyMap()); + partition = range; + transform = Transforms.range(new String[] {"col1"}, new RangePartition[] {range}); + fragment = "PARTITION `p` VALUES LESS THAN (" + quoted + ")"; + } else { + boolean multiColumn = "multi-list".equals(kind); + ListPartition list = + Partitions.list( + "p", + multiColumn ? new Literal[][] {{value, value}} : new Literal[][] {{value}}, + Collections.emptyMap()); + partition = list; + transform = + Transforms.list( + multiColumn ? new String[][] {{"col1"}, {"col2"}} : new String[][] {{"col1"}}, + new ListPartition[] {list}); + fragment = + "PARTITION `p` VALUES IN (" + + (multiColumn ? "(" + quoted + "," + quoted + ")" : quoted) + + ")"; + } + JdbcColumn[] columns = { + JdbcColumn.builder().withName("col1").withType(Types.VarCharType.of(128)).build(), + JdbcColumn.builder().withName("col2").withType(Types.VarCharType.of(128)).build() + }; + ConnectionCountingDataSource connections = new ConnectionCountingDataSource(); + StarRocksTableOperations operations = new StarRocksTableOperations(); + StarRocksExceptionConverter exceptions = new StarRocksExceptionConverter(); + StarRocksTypeConverter types = new StarRocksTypeConverter(); + operations.initialize( + connections.dataSource(), + exceptions, + types, + new StarRocksColumnDefaultValueConverter(), + Collections.emptyMap()); + operations.create( + "database", + "test_table", + columns, + null, + Map.of("replication_num", "1"), + new Transform[] {transform}, + Distributions.hash(1, NamedReference.field("col1")), + Indexes.EMPTY_INDEXES); + Assertions.assertTrue(connections.executedSql().get(0).contains(fragment)); + JdbcTable table = + JdbcTable.builder() + .withName("test_table") + .withDatabaseName("database") + .withColumns(columns) + .withPartitioning(new Transform[] {transform}) + .build(); + StarRocksTablePartitionOperations partitions = + new StarRocksTablePartitionOperations(connections.dataSource(), table, exceptions, types); + partitions.addPartition(partition); + Assertions.assertEquals(2, connections.executedSql().size()); + Assertions.assertEquals( + "ALTER TABLE `test_table` ADD " + fragment, connections.executedSql().get(1)); + Assertions.assertEquals(0, connections.borrowed()); + } +} diff --git a/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksTablePartitionOperations.java b/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksTablePartitionOperations.java index 1a722f9b982..52b0f952201 100644 --- a/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksTablePartitionOperations.java +++ b/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/operation/TestStarRocksTablePartitionOperations.java @@ -24,6 +24,9 @@ import static org.junit.Assert.assertTrue; import com.google.common.collect.ImmutableMap; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; import java.time.LocalDate; import java.util.Arrays; import java.util.Collections; @@ -57,6 +60,8 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Tag; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; @Tag("gravitino-docker-test") public class TestStarRocksTablePartitionOperations extends TestStarRocks { @@ -358,4 +363,86 @@ public void testListPartition() { // drop non-existing partition assertFalse(tablePartitionOperations.dropPartition("p3")); } + + @ParameterizedTest + @ValueSource(ints = {1, 2}) + void testEscapedListPartitionValues(int partitionColumns) throws Exception { + String tableName = GravitinoITUtils.genRandomName("escaped_list_partition"); + JdbcColumn[] columns = { + JdbcColumn.builder() + .withName("col1") + .withType(Types.VarCharType.of(128)) + .withNullable(false) + .build(), + JdbcColumn.builder() + .withName("col2") + .withType(Types.VarCharType.of(128)) + .withNullable(false) + .build() + }; + String[] values = {"say \"hi\" C:\\temp\\; --", "5\" pipe C:\\next\\"}; + Literal[][] initial = + partitionColumns == 1 + ? new Literal[][] {{Literals.stringLiteral(values[0])}} + : new Literal[][] { + {Literals.stringLiteral(values[0]), Literals.stringLiteral(values[0])} + }; + ListPartition created = Partitions.list("p_created", initial, Collections.emptyMap()); + Transform transform = + Transforms.list( + partitionColumns == 1 ? new String[][] {{"col1"}} : new String[][] {{"col1"}, {"col2"}}, + new ListPartition[] {created}); + TABLE_OPERATIONS.create( + databaseName, + tableName, + columns, + "escaped partition values", + Map.of("replication_num", "1"), + new Transform[] {transform}, + Distributions.hash(1, NamedReference.field("col1")), + new Index[0]); + try { + JdbcTable table = + JdbcTable.builder() + .withName(tableName) + .withDatabaseName(databaseName) + .withColumns(columns) + .withPartitioning(new Transform[] {transform}) + .build(); + JdbcTablePartitionOperations partitions = + new StarRocksTablePartitionOperations( + DATA_SOURCE, table, JDBC_EXCEPTION_CONVERTER, TYPE_CONVERTER); + Literal[][] added = + partitionColumns == 1 + ? new Literal[][] {{Literals.stringLiteral(values[1])}} + : new Literal[][] { + {Literals.stringLiteral(values[1]), Literals.stringLiteral(values[1])} + }; + partitions.addPartition(Partitions.list("p_added", added, Collections.emptyMap())); + // Inserting the original values proves both DDL paths preserved the partition keys. + try (Connection connection = DATA_SOURCE.getConnection()) { + connection.setCatalog(databaseName); + for (String value : values) { + try (PreparedStatement insert = + connection.prepareStatement("INSERT INTO `" + tableName + "` VALUES (?, ?)")) { + insert.setString(1, value); + insert.setString(2, value); + insert.executeUpdate(); + } + try (PreparedStatement select = + connection.prepareStatement( + "SELECT COUNT(*) FROM `" + tableName + "` WHERE col1 = ? AND col2 = ?")) { + select.setString(1, value); + select.setString(2, value); + try (ResultSet result = select.executeQuery()) { + assertTrue(result.next()); + assertEquals(1, result.getInt(1)); + } + } + } + } + } finally { + TABLE_OPERATIONS.drop(databaseName, tableName); + } + } } diff --git a/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/utils/TestStarRocksUtils.java b/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/utils/TestStarRocksUtils.java index 3eb449327ed..0c023d9014c 100644 --- a/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/utils/TestStarRocksUtils.java +++ b/catalogs/catalog-jdbc-starrocks/src/test/java/org/apache/gravitino/catalog/starrocks/utils/TestStarRocksUtils.java @@ -26,6 +26,7 @@ import java.util.HashMap; import java.util.Map; import java.util.Optional; +import java.util.stream.Stream; import org.apache.gravitino.rel.expressions.distributions.Distribution; import org.apache.gravitino.rel.expressions.distributions.Distributions; import org.apache.gravitino.rel.expressions.literals.Literal; @@ -37,6 +38,9 @@ import org.apache.gravitino.rel.types.Types; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; public class TestStarRocksUtils { @@ -251,4 +255,41 @@ public void testDistributedInfoPattern() { RuntimeException.class, () -> StarRocksUtils.extractDistributionInfoFromSql(createTableSqlWithInvalidDistribution)); } + + @ParameterizedTest + @MethodSource("partitionLiteralValues") + void testPartitionLiteralEscaping(String value, String escaped) { + Literal literal = Literals.stringLiteral(value); + assertEquals( + "PARTITION `p` VALUES LESS THAN (\"" + escaped + "\")", + StarRocksUtils.generatePartitionSqlFragment( + Partitions.range("p", literal, Literals.NULL, Collections.emptyMap()))); + + assertEquals( + "PARTITION `p` VALUES IN (\"" + escaped + "\",\"plain\")", + StarRocksUtils.generatePartitionSqlFragment( + Partitions.list( + "p", + new Literal[][] {{literal}, {Literals.stringLiteral("plain")}}, + Collections.emptyMap()))); + assertEquals( + "PARTITION `p` VALUES IN ((\"" + escaped + "\",\"plain\"),(\"plain\",\"" + escaped + "\"))", + StarRocksUtils.generatePartitionSqlFragment( + Partitions.list( + "p", + new Literal[][] { + {literal, Literals.stringLiteral("plain")}, + {Literals.stringLiteral("plain"), literal} + }, + Collections.emptyMap()))); + } + + private static Stream partitionLiteralValues() { + return Stream.of( + Arguments.of("say \"hi\"", "say \"\"hi\"\""), + Arguments.of("C:\\temp\\", "C:\\\\temp\\\\"), + Arguments.of("5\" pipe' ; --", "5\"\" pipe' ; --"), + Arguments.of("slash\\\"quote", "slash\\\\\"\"quote"), + Arguments.of("line\nnext\tvalue", "line\nnext\tvalue")); + } }