diff --git a/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/operators/TableSink.java b/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/operators/TableSink.java index e5e4b54af..6c50b057a 100644 --- a/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/operators/TableSink.java +++ b/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/operators/TableSink.java @@ -19,6 +19,7 @@ package org.apache.wayang.basic.operators; import org.apache.wayang.basic.data.Record; +import org.apache.wayang.basic.util.SqlIdentifierUtils; import org.apache.wayang.core.plan.wayangplan.UnarySink; import org.apache.wayang.core.types.DataSetType; @@ -51,6 +52,8 @@ public TableSink(Properties props, String mode, String tableName, String... colu public TableSink(Properties props, String mode, String tableName, String[] columnNames, DataSetType type) { super(type); + SqlIdentifierUtils.validateTableName(tableName); + SqlIdentifierUtils.validateColumnNames(columnNames); this.tableName = tableName; this.columnNames = columnNames == null ? null : java.util.Arrays.copyOf(columnNames, columnNames.length); this.props = new Properties(); @@ -78,6 +81,7 @@ public String getTableName() { } protected void setColumnNames(String[] columnNames) { + SqlIdentifierUtils.validateColumnNames(columnNames); this.columnNames = columnNames == null ? null : java.util.Arrays.copyOf(columnNames, columnNames.length); } diff --git a/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/SqlIdentifierUtils.java b/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/SqlIdentifierUtils.java new file mode 100644 index 000000000..4403a5931 --- /dev/null +++ b/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/SqlIdentifierUtils.java @@ -0,0 +1,207 @@ +/* + * 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.wayang.basic.util; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.SQLException; +import java.util.regex.Pattern; + +/** + * Utility for validating, escaping, and quoting SQL table and column identifiers. + * Prevents SQL injection vulnerabilities (CWE-89) in DDL and DML operations. + */ +public final class SqlIdentifierUtils { + + /** + * Pattern matching a single SQL identifier segment: must start with a letter or underscore, + * followed by letters, digits, underscores, or dollar signs. + */ + private static final Pattern IDENTIFIER_PATTERN = Pattern.compile("^[a-zA-Z_][a-zA-Z0-9_$]*$"); + + /** + * Pattern matching valid table names: one or more identifier segments separated by dots + * (e.g., "table", "schema.table", or "catalog.schema.table"). + */ + private static final Pattern TABLE_NAME_PATTERN = + Pattern.compile("^[a-zA-Z_][a-zA-Z0-9_$]*(?:\\.[a-zA-Z_][a-zA-Z0-9_$]*)*$"); + + private SqlIdentifierUtils() { + // Utility class + } + + /** + * Validates a single SQL identifier (e.g., column name or schema segment). + * + * @param identifier the identifier to validate + * @throws IllegalArgumentException if the identifier is null, empty, or contains invalid characters + */ + public static void validateIdentifier(String identifier) { + if (identifier == null || identifier.trim().isEmpty()) { + throw new IllegalArgumentException("SQL identifier must not be null or blank."); + } + if (!IDENTIFIER_PATTERN.matcher(identifier).matches()) { + throw new IllegalArgumentException(String.format( + "Invalid SQL identifier: '%s'. Identifiers must start with a letter or underscore " + + "and contain only alphanumeric characters, underscores, or dollar signs.", + identifier)); + } + } + + /** + * Validates a SQL table name, which may be single-part or multipart (e.g., schema.table). + * + * @param tableName the table name to validate + * @throws IllegalArgumentException if the table name is null, empty, or contains invalid characters + */ + public static void validateTableName(String tableName) { + if (tableName == null || tableName.trim().isEmpty()) { + throw new IllegalArgumentException("SQL table name must not be null or blank."); + } + if (!TABLE_NAME_PATTERN.matcher(tableName).matches()) { + throw new IllegalArgumentException(String.format( + "Invalid SQL table name: '%s'. Table names must be valid identifiers, optionally " + + "separated by dots (e.g., 'table', 'schema.table', 'catalog.schema.table').", + tableName)); + } + } + + /** + * Validates an array of column names. + * + * @param columnNames array of column names to validate, may be null or empty + * @throws IllegalArgumentException if any column name is invalid + */ + public static void validateColumnNames(String[] columnNames) { + if (columnNames == null) { + return; + } + for (String columnName : columnNames) { + validateIdentifier(columnName); + } + } + + /** + * Validates, escapes, and quotes a table or column identifier using the database connection's + * metadata. If the connection is null or fails to provide metadata, standard ANSI double-quotes + * are used. + * + *

For multipart identifiers (e.g., {@code catalog.schema.table}), each segment is quoted + * individually (e.g., {@code "catalog"."schema"."table"}).

+ * + * @param identifier the table or column identifier to escape and quote + * @param connection active JDBC connection, may be null + * @return the escaped and quoted identifier + */ + public static String escapeAndQuoteIdentifier(String identifier, Connection connection) { + validateTableName(identifier); + + String quote = "\""; + String closingQuote = "\""; + + if (connection != null) { + try { + DatabaseMetaData meta = connection.getMetaData(); + if (meta != null) { + String metaQuote = meta.getIdentifierQuoteString(); + if (metaQuote != null && !metaQuote.trim().isEmpty()) { + quote = metaQuote.trim(); + closingQuote = quote.equals("[") ? "]" : quote; + } + } + } catch (SQLException ignored) { + // Fall back to default ANSI quoting + } + } + + return escapeAndQuoteIdentifier(identifier, quote, closingQuote, false, false); + } + + /** + * Validates, escapes, and quotes an identifier for a specific {@link DatabaseProduct}. + * + * @param identifier the identifier to escape and quote + * @param product the database product + * @return the escaped and quoted identifier + */ + public static String escapeAndQuoteIdentifier(String identifier, DatabaseProduct product) { + validateTableName(identifier); + + String quote = (product == DatabaseProduct.MYSQL) ? "`" + : (product == DatabaseProduct.MSSQL) ? "[" : "\""; + String closingQuote = (product == DatabaseProduct.MSSQL) ? "]" : quote; + return escapeAndQuoteIdentifier(identifier, quote, closingQuote, false, false); + } + + /** + * Validates, escapes, and quotes an identifier using the specified opening and closing quote strings. + * + * @param identifier the identifier to escape and quote + * @param quote the opening quote string + * @param closingQuote the closing quote string + * @return the escaped and quoted identifier + */ + public static String escapeAndQuoteIdentifier(String identifier, String quote, String closingQuote) { + validateTableName(identifier); + return escapeAndQuoteIdentifier(identifier, quote, closingQuote, false, false); + } + + /** + * Validates, escapes, and quotes an identifier with optional case folding. + * + * @param identifier the identifier to escape and quote + * @param quote the opening quote string + * @param closingQuote the closing quote string + * @param toUpper whether to fold unquoted identifiers to uppercase + * @param toLower whether to fold unquoted identifiers to lowercase + * @return the escaped and quoted identifier + */ + public static String escapeAndQuoteIdentifier( + String identifier, String quote, String closingQuote, boolean toUpper, boolean toLower) { + validateTableName(identifier); + + String q = (quote == null || quote.trim().isEmpty()) ? "\"" : quote; + String cq = (closingQuote == null || closingQuote.trim().isEmpty()) ? q : closingQuote; + + String[] parts = identifier.split("\\."); + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < parts.length; i++) { + if (i > 0) { + sb.append("."); + } + String part = parts[i]; + if (toUpper) { + part = part.toUpperCase(); + } else if (toLower) { + part = part.toLowerCase(); + } + + // Escape any occurrence of closingQuote within the identifier part by doubling it + String escapedPart; + if ("]".equals(cq)) { + escapedPart = part.replace("]", "]]"); + } else { + escapedPart = part.replace(cq, cq + cq); + } + + sb.append(q).append(escapedPart).append(cq); + } + return sb.toString(); + } +} diff --git a/wayang-commons/wayang-basic/src/test/java/org/apache/wayang/basic/util/SqlIdentifierUtilsTest.java b/wayang-commons/wayang-basic/src/test/java/org/apache/wayang/basic/util/SqlIdentifierUtilsTest.java new file mode 100644 index 000000000..d04bef38c --- /dev/null +++ b/wayang-commons/wayang-basic/src/test/java/org/apache/wayang/basic/util/SqlIdentifierUtilsTest.java @@ -0,0 +1,171 @@ +/* + * 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.wayang.basic.util; + +import org.junit.jupiter.api.Test; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.SQLException; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link SqlIdentifierUtils}. + */ +public class SqlIdentifierUtilsTest { + + @Test + void testValidateIdentifierValid() { + assertDoesNotThrow(() -> SqlIdentifierUtils.validateIdentifier("users")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateIdentifier("_temp")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateIdentifier("col_name")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateIdentifier("table$1")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateIdentifier("camelCaseColumn")); + } + + @Test + void testValidateIdentifierInvalid() { + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier(null)); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier(" ")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("1table")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("col; DROP TABLE users;")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("col\"name")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("col'name")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("col name")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("col.name")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("col-name")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateIdentifier("col/*comment*/")); + } + + @Test + void testValidateTableNameValid() { + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("users")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("_temp")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("schema.users")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("catalog.schema.users")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("cat.sch.tbl_1$")); + } + + @Test + void testValidateTableNameInvalid() { + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName(null)); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName(" ")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName(".users")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users.")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users..accounts")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users; DROP TABLE users; --")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users\" OR 1=1")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users' OR 1=1")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users name")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("123users")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users--comment")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users/*comment*/")); + assertThrows(IllegalArgumentException.class, () -> SqlIdentifierUtils.validateTableName("users(id INT)")); + } + + @Test + void testValidateColumnNames() { + assertDoesNotThrow(() -> SqlIdentifierUtils.validateColumnNames(null)); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateColumnNames(new String[0])); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateColumnNames(new String[]{"id", "name", "created_at"})); + + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateColumnNames(new String[]{"id", "name; DROP TABLE users;"})); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateColumnNames(new String[]{"id", null})); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateColumnNames(new String[]{"id", ""})); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateColumnNames(new String[]{"id", "bad col"})); + } + + @Test + void testEscapeAndQuoteIdentifierNullConnection() { + assertEquals("\"users\"", SqlIdentifierUtils.escapeAndQuoteIdentifier("users", (Connection) null)); + assertEquals("\"schema\".\"users\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("schema.users", (Connection) null)); + assertEquals("\"catalog\".\"schema\".\"users\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("catalog.schema.users", (Connection) null)); + } + + @Test + void testEscapeAndQuoteIdentifierWithMockConnection() throws SQLException { + Connection conn = mock(Connection.class); + DatabaseMetaData meta = mock(DatabaseMetaData.class); + when(conn.getMetaData()).thenReturn(meta); + + // Standard quoting + when(meta.getIdentifierQuoteString()).thenReturn("\""); + when(meta.storesUpperCaseIdentifiers()).thenReturn(false); + when(meta.storesLowerCaseIdentifiers()).thenReturn(false); + assertEquals("\"my_table\"", SqlIdentifierUtils.escapeAndQuoteIdentifier("my_table", conn)); + assertEquals("\"my_schema\".\"my_table\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my_schema.my_table", conn)); + + // MySQL backtick quoting + when(meta.getIdentifierQuoteString()).thenReturn("`"); + assertEquals("`my_table`", SqlIdentifierUtils.escapeAndQuoteIdentifier("my_table", conn)); + assertEquals("`my_schema`.`my_table`", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my_schema.my_table", conn)); + + // MSSQL bracket quoting + when(meta.getIdentifierQuoteString()).thenReturn("["); + assertEquals("[my_table]", SqlIdentifierUtils.escapeAndQuoteIdentifier("my_table", conn)); + assertEquals("[my_schema].[my_table]", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my_schema.my_table", conn)); + } + + @Test + void testEscapeAndQuoteIdentifierWithCaseFolding() { + // Explicit case folding + assertEquals("\"MY_TABLE\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my_table", "\"", "\"", true, false)); + assertEquals("\"my_table\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("MY_TABLE", "\"", "\"", false, true)); + assertEquals("\"my_table\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my_table", "\"", "\"", false, false)); + } + + @Test + void testEscapeAndQuoteIdentifierWithDatabaseProduct() { + assertEquals("`users`", SqlIdentifierUtils.escapeAndQuoteIdentifier("users", DatabaseProduct.MYSQL)); + assertEquals("`db`.`users`", SqlIdentifierUtils.escapeAndQuoteIdentifier("db.users", DatabaseProduct.MYSQL)); + + assertEquals("[users]", SqlIdentifierUtils.escapeAndQuoteIdentifier("users", DatabaseProduct.MSSQL)); + assertEquals("[dbo].[users]", SqlIdentifierUtils.escapeAndQuoteIdentifier("dbo.users", DatabaseProduct.MSSQL)); + + assertEquals("\"users\"", SqlIdentifierUtils.escapeAndQuoteIdentifier("users", DatabaseProduct.POSTGRESQL)); + assertEquals("\"users\"", SqlIdentifierUtils.escapeAndQuoteIdentifier("users", DatabaseProduct.ORACLE)); + } + + @Test + void testEscapeAndQuoteIdentifierRejectsInvalid() { + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.escapeAndQuoteIdentifier("users; DROP TABLE test;", (Connection) null)); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.escapeAndQuoteIdentifier("users\" OR 1=1", DatabaseProduct.MYSQL)); + } +} diff --git a/wayang-platforms/wayang-java/src/main/java/org/apache/wayang/java/operators/JavaTableSink.java b/wayang-platforms/wayang-java/src/main/java/org/apache/wayang/java/operators/JavaTableSink.java index 5d5086897..dee6bf976 100644 --- a/wayang-platforms/wayang-java/src/main/java/org/apache/wayang/java/operators/JavaTableSink.java +++ b/wayang-platforms/wayang-java/src/main/java/org/apache/wayang/java/operators/JavaTableSink.java @@ -21,6 +21,7 @@ import org.apache.wayang.basic.data.Record; import org.apache.wayang.basic.operators.TableSink; import org.apache.wayang.basic.util.DatabaseProduct; +import org.apache.wayang.basic.util.SqlIdentifierUtils; import org.apache.wayang.basic.util.SqlTypeUtils; import org.apache.wayang.core.optimizer.OptimizationContext; import org.apache.wayang.core.plan.wayangplan.ExecutionOperator; @@ -158,23 +159,20 @@ public Tuple, Collection> eval try { Class.forName(writeProps.getProperty("driver")); - String quote = (product == DatabaseProduct.MYSQL) ? "`" - : (product == DatabaseProduct.MSSQL) ? "[" : "\""; - String closingQuote = (product == DatabaseProduct.MSSQL) ? "]" : quote; - try (Connection conn = DriverManager.getConnection(writeProps.getProperty("url"), writeProps)) { conn.setAutoCommit(false); try (Statement stmt = conn.createStatement()) { + String quotedTable = SqlIdentifierUtils.escapeAndQuoteIdentifier(this.getTableName(), conn); if (this.getMode().equals("overwrite")) { - stmt.execute("DROP TABLE IF EXISTS " + quote + this.getTableName() + closingQuote); + stmt.execute("DROP TABLE IF EXISTS " + quotedTable); } StringBuilder sb = new StringBuilder(); - sb.append("CREATE TABLE IF NOT EXISTS ").append(quote).append(this.getTableName()) - .append(closingQuote).append(" ("); + sb.append("CREATE TABLE IF NOT EXISTS ").append(quotedTable).append(" ("); String separator = ""; for (int i = 0; i < finalColumnNames.length; i++) { - sb.append(separator).append(quote).append(finalColumnNames[i]).append(closingQuote).append(" ") + String quotedCol = SqlIdentifierUtils.escapeAndQuoteIdentifier(finalColumnNames[i], conn); + sb.append(separator).append(quotedCol).append(" ") .append(finalSqlTypes[i]); separator = ", "; } @@ -182,11 +180,11 @@ public Tuple, Collection> eval stmt.execute(sb.toString()); sb = new StringBuilder(); - sb.append("INSERT INTO ").append(quote).append(this.getTableName()).append(closingQuote) - .append(" ("); + sb.append("INSERT INTO ").append(quotedTable).append(" ("); separator = ""; for (int i = 0; i < finalColumnNames.length; i++) { - sb.append(separator).append(quote).append(finalColumnNames[i]).append(closingQuote); + String quotedCol = SqlIdentifierUtils.escapeAndQuoteIdentifier(finalColumnNames[i], conn); + sb.append(separator).append(quotedCol); separator = ", "; } sb.append(") VALUES ("); diff --git a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java index dd5792f52..e6309a8db 100644 --- a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java +++ b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java @@ -291,7 +291,7 @@ private static long executeSinkStage(final ExecutionStage stage, final Optimizat try (Statement stmt = jdbcExecutor.connection.createStatement()) { // Handle overwrite: drop existing table first if ("overwrite".equals(sinkOp.getMode())) { - stmt.execute("DROP TABLE IF EXISTS " + sinkOp.getTableName()); + stmt.execute("DROP TABLE IF EXISTS " + sinkOp.getQuotedTableName(jdbcExecutor.connection)); } // Execute the composed query: CREATE TABLE x AS SELECT ... or INSERT INTO x // SELECT ... diff --git a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/operators/JdbcTableSinkOperator.java b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/operators/JdbcTableSinkOperator.java index f58c9d3c1..b1f2f4c6f 100644 --- a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/operators/JdbcTableSinkOperator.java +++ b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/operators/JdbcTableSinkOperator.java @@ -20,6 +20,7 @@ import org.apache.wayang.basic.data.Record; import org.apache.wayang.basic.operators.TableSink; +import org.apache.wayang.basic.util.SqlIdentifierUtils; import org.apache.wayang.core.api.Configuration; import org.apache.wayang.core.optimizer.costs.LoadProfileEstimator; import org.apache.wayang.core.platform.ChannelDescriptor; @@ -47,13 +48,24 @@ public JdbcTableSinkOperator(TableSink that) { super(that); } + /** + * Returns the escaped and quoted table name appropriate for the given database connection. + * + * @param connection database connection, or {@code null} for standard ANSI quoting + * @return escaped and quoted table name + */ + public String getQuotedTableName(Connection connection) { + return SqlIdentifierUtils.escapeAndQuoteIdentifier(this.getTableName(), connection); + } + @Override public String createSqlClause(Connection connection, FunctionCompiler compiler) { String mode = this.getMode(); + String targetTable = this.getQuotedTableName(connection); if ("overwrite".equals(mode)) { - return "CREATE TABLE " + this.getTableName() + " AS"; + return "CREATE TABLE " + targetTable + " AS"; } - return "INSERT INTO " + this.getTableName(); + return "INSERT INTO " + targetTable; } /** diff --git a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcTableSinkExecutorTest.java b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcTableSinkExecutorTest.java index 483bd4f81..63d72055d 100644 --- a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcTableSinkExecutorTest.java +++ b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcTableSinkExecutorTest.java @@ -46,6 +46,7 @@ import java.util.stream.Collectors; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -282,7 +283,7 @@ void testOverwriteClauseGeneration() { JdbcTableSinkOperator sinkOp = new HsqldbTableSinkOperator( "my_table", new String[]{"col1"}); sinkOp.setMode("overwrite"); - assertEquals("CREATE TABLE my_table AS (", sinkOp.createSqlClause(null, null)); + assertEquals("CREATE TABLE \"my_table\" AS (", sinkOp.createSqlClause(null, null)); } @Test @@ -290,6 +291,43 @@ void testAppendClauseGeneration() { JdbcTableSinkOperator sinkOp = new HsqldbTableSinkOperator( "my_table", new String[]{"col1"}); sinkOp.setMode("append"); - assertEquals("INSERT INTO my_table", sinkOp.createSqlClause(null, null)); -} + assertEquals("INSERT INTO \"my_table\"", sinkOp.createSqlClause(null, null)); + } + + @Test + void testMultipartTableNameQuotedClauseGeneration() { + JdbcTableSinkOperator sinkOp = new HsqldbTableSinkOperator( + "public.my_table", new String[]{"col1"}); + sinkOp.setMode("overwrite"); + assertEquals("CREATE TABLE \"public\".\"my_table\" AS (", sinkOp.createSqlClause(null, null)); + } + + @Test + void testInvalidTableNameThrowsException() { + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("target; DROP TABLE users; --", new String[]{"col1"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("target\"table", new String[]{"col1"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("target table", new String[]{"col1"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("123table", new String[]{"col1"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("", new String[]{"col1"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator(null, new String[]{"col1"})); + } + + @Test + void testInvalidColumnNamesThrowsException() { + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("my_table", new String[]{"col; DROP TABLE users; --"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("my_table", new String[]{"col 1"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("my_table", new String[]{"col\"name"})); + assertThrows(IllegalArgumentException.class, () -> + new HsqldbTableSinkOperator("my_table", new String[]{""})); + } } + diff --git a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/test/HsqldbTableSinkOperator.java b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/test/HsqldbTableSinkOperator.java index d1369a88a..fd810176f 100644 --- a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/test/HsqldbTableSinkOperator.java +++ b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/test/HsqldbTableSinkOperator.java @@ -43,10 +43,11 @@ public HsqldbPlatform getPlatform() { @Override public String createSqlClause(Connection connection, FunctionCompiler compiler) { String mode = this.getMode(); + String targetTable = this.getQuotedTableName(connection); if ("overwrite".equals(mode)) { - return "CREATE TABLE " + this.getTableName() + " AS ("; + return "CREATE TABLE " + targetTable + " AS ("; } - return "INSERT INTO " + this.getTableName(); + return "INSERT INTO " + targetTable; } @Override diff --git a/wayang-platforms/wayang-spark/src/main/java/org/apache/wayang/spark/operators/SparkTableSink.java b/wayang-platforms/wayang-spark/src/main/java/org/apache/wayang/spark/operators/SparkTableSink.java index 466cca381..fca16d302 100644 --- a/wayang-platforms/wayang-spark/src/main/java/org/apache/wayang/spark/operators/SparkTableSink.java +++ b/wayang-platforms/wayang-spark/src/main/java/org/apache/wayang/spark/operators/SparkTableSink.java @@ -31,6 +31,7 @@ import org.apache.wayang.basic.data.Record; import org.apache.wayang.basic.operators.TableSink; import org.apache.wayang.basic.util.DatabaseProduct; +import org.apache.wayang.basic.util.SqlIdentifierUtils; import org.apache.wayang.basic.util.SqlTypeUtils; import org.apache.wayang.core.api.exception.WayangException; import org.apache.wayang.core.optimizer.OptimizationContext; @@ -124,15 +125,8 @@ public Tuple, Collection> eval writeProps.setProperty("batchSize", "250000"); } - String targetTable = this.getTableName(); - if (targetTable != null && !targetTable.startsWith("(") && !targetTable.startsWith("\"") - && !targetTable.startsWith("`")) { - DatabaseProduct product = SqlTypeUtils.detectProduct(this.getProperties().getProperty("url")); - String quote = (product == DatabaseProduct.MYSQL) ? "`" - : (product == DatabaseProduct.MSSQL) ? "[" : "\""; - String closingQuote = (product == DatabaseProduct.MSSQL) ? "]" : quote; - targetTable = quote + targetTable + closingQuote; - } + DatabaseProduct product = SqlTypeUtils.detectProduct(this.getProperties().getProperty("url")); + String targetTable = SqlIdentifierUtils.escapeAndQuoteIdentifier(this.getTableName(), product); df.write() .mode(this.mode)