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..4f0fdd59d --- /dev/null +++ b/wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/SqlIdentifierUtils.java @@ -0,0 +1,336 @@ +/* + * 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 an unquoted table name segment, optionally containing hyphens (e.g., BigQuery project IDs). + * Disallows consecutive hyphens (to prevent SQL comments '--') and trailing/leading hyphens. + */ + private static final Pattern UNQUOTED_TABLE_SEGMENT_PATTERN = + Pattern.compile("^[a-zA-Z_][a-zA-Z0-9_$]*(?:-[a-zA-Z0-9_$]+)*$"); + + /** + * Pattern matching a quoted table name segment enclosed in backticks (`...`), double quotes ("..."), + * or square brackets ([...]). + */ + private static final Pattern QUOTED_TABLE_SEGMENT_PATTERN = + Pattern.compile("^(?:`[a-zA-Z0-9_$-]+`|\"[a-zA-Z0-9_$-]+\"|\\[[a-zA-Z0-9_$-]+\\])$"); + + private SqlIdentifierUtils() { + // Utility class + } + + /** + * Checks whether an identifier or multipart table name is already quoted. + * + * @param identifier the identifier to check + * @return {@code true} if already quoted, {@code false} otherwise + */ + public static boolean isAlreadyQuoted(String identifier) { + if (identifier == null) { + return false; + } + String trimmed = identifier.trim(); + if (trimmed.length() < 2) { + return false; + } + // Entire string wrapped in backticks (e.g. BigQuery `project.dataset.table` or `table`) + if (trimmed.startsWith("`") && trimmed.endsWith("`") + && !trimmed.substring(1, trimmed.length() - 1).contains("`")) { + return true; + } + // Entire string wrapped in double quotes (e.g. "table") + if (trimmed.startsWith("\"") && trimmed.endsWith("\"") + && !trimmed.substring(1, trimmed.length() - 1).contains("\"")) { + return true; + } + // Entire string wrapped in brackets (e.g. [table]) + if (trimmed.startsWith("[") && trimmed.endsWith("]") + && !trimmed.substring(1, trimmed.length() - 1).contains("[")) { + return true; + } + // Segment-by-segment quoted (e.g. `proj`.`ds`.`tbl` or "schema"."table" or [dbo].[users]) + String[] parts = trimmed.split("\\."); + if (parts.length > 1) { + boolean allPartsQuoted = true; + for (String part : parts) { + String p = part.trim(); + boolean partQuoted = (p.length() >= 2) && ( + (p.startsWith("`") && p.endsWith("`")) + || (p.startsWith("\"") && p.endsWith("\"")) + || (p.startsWith("[") && p.endsWith("]")) + ); + if (!partQuoted) { + allPartsQuoted = false; + break; + } + } + if (allPartsQuoted) { + return true; + } + } + return false; + } + + /** + * 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), + * dialect-specific (e.g., BigQuery project IDs with hyphens), or quoted (e.g., `project.dataset.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."); + } + String trimmed = tableName.trim(); + + // Check for disallowed characters / SQL injection tokens + if (trimmed.contains(";") || trimmed.contains("'") || trimmed.contains("/*") + || trimmed.contains("*/") || trimmed.contains("(") || trimmed.contains(")") + || trimmed.contains(" ") || trimmed.contains("\t") || trimmed.contains("\n") + || trimmed.contains("\r") || trimmed.contains("--")) { + throw new IllegalArgumentException(String.format( + "Invalid SQL table name: '%s'. Contains disallowed characters or SQL injection tokens.", + tableName)); + } + + // Check for BigQuery whole backtick-quoted table name: `project.dataset.table` or `table` + if (trimmed.startsWith("`") && trimmed.endsWith("`") && trimmed.length() >= 2) { + String inner = trimmed.substring(1, trimmed.length() - 1); + if (!inner.contains("`") && !inner.startsWith(".") && !inner.endsWith(".") && !inner.contains("..")) { + String[] segments = inner.split("\\."); + boolean allSegmentsValid = segments.length > 0; + for (String segment : segments) { + if (!UNQUOTED_TABLE_SEGMENT_PATTERN.matcher(segment).matches()) { + allSegmentsValid = false; + break; + } + } + if (allSegmentsValid) { + return; + } + } + } + + // Check for multipart or single-part table name + if (trimmed.startsWith(".") || trimmed.endsWith(".") || trimmed.contains("..")) { + throw new IllegalArgumentException(String.format( + "Invalid SQL table name: '%s'. Table name cannot start, end with, or contain consecutive dots.", + tableName)); + } + + String[] segments = trimmed.split("\\."); + for (String segment : segments) { + if (!UNQUOTED_TABLE_SEGMENT_PATTERN.matcher(segment).matches() + && !QUOTED_TABLE_SEGMENT_PATTERN.matcher(segment).matches()) { + throw new IllegalArgumentException(String.format( + "Invalid SQL table name: '%s'. Segment '%s' is not a valid identifier.", + tableName, segment)); + } + } + } + + /** + * 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 identifier is already quoted, it is returned as-is to preserve dialect-specific + * quoting and case. If unquoted, the database metadata's case-folding preferences + * ({@link DatabaseMetaData#storesUpperCaseIdentifiers()} and {@link DatabaseMetaData#storesLowerCaseIdentifiers()}) + * are applied so lookup semantics are preserved. + * + *

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); + + if (isAlreadyQuoted(identifier)) { + return identifier; + } + + String quote = "\""; + String closingQuote = "\""; + boolean toUpper = false; + boolean toLower = false; + + 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; + } + if (meta.storesUpperCaseIdentifiers()) { + toUpper = true; + } else if (meta.storesLowerCaseIdentifiers()) { + toLower = true; + } + } + } catch (SQLException ignored) { + // Fall back to default ANSI quoting + } + } + + return escapeAndQuoteIdentifier(identifier, quote, closingQuote, toUpper, toLower); + } + + /** + * Validates, escapes, and quotes an identifier for a specific {@link DatabaseProduct}. + * If the identifier is already quoted, it is returned as-is. + * + * @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); + + if (isAlreadyQuoted(identifier)) { + return 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. + * If the identifier is already quoted, it is returned as-is. + * + * @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); + + if (isAlreadyQuoted(identifier)) { + return identifier; + } + + return escapeAndQuoteIdentifier(identifier, quote, closingQuote, false, false); + } + + /** + * Validates, escapes, and quotes an identifier with optional case folding. + * If the identifier is already quoted, it is returned as-is. + * + * @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); + + if (isAlreadyQuoted(identifier)) { + return 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..ec986cc5d --- /dev/null +++ b/wayang-commons/wayang-basic/src/test/java/org/apache/wayang/basic/util/SqlIdentifierUtilsTest.java @@ -0,0 +1,276 @@ +/* + * 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +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)); + } + + @Test + void testValidateTableNameBigQueryNames() { + // Unquoted table names with hyphens (e.g. BigQuery project IDs) + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("my-project.sales.orders")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("my-project.orders")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("my-gcp-project-123.sales_dataset.order_table")); + + // Documented whole backtick-quoted BigQuery form + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("`my-project.sales.orders`")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("`my-project.orders`")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("`orders`")); + + // Per-segment backtick-quoted form + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("`my-project`.`sales`.`orders`")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("`sales`.`orders`")); + + // Quoted forms in other dialects + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("\"my_schema\".\"my_table\"")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("\"users\"")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("[dbo].[users]")); + assertDoesNotThrow(() -> SqlIdentifierUtils.validateTableName("[users]")); + + // Invalid forms and injection attempts + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateTableName("`my-project; DROP TABLE users;`")); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateTableName("`my-project\" OR 1=1`")); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateTableName("my-project--comment.sales.orders")); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateTableName("my-project..orders")); + assertThrows(IllegalArgumentException.class, + () -> SqlIdentifierUtils.validateTableName("`my-project`..`orders`")); + } + + @Test + void testIsAlreadyQuoted() { + assertTrue(SqlIdentifierUtils.isAlreadyQuoted("`my-project.sales.orders`")); + assertTrue(SqlIdentifierUtils.isAlreadyQuoted("`my-project`.`sales`.`orders`")); + assertTrue(SqlIdentifierUtils.isAlreadyQuoted("\"my_table\"")); + assertTrue(SqlIdentifierUtils.isAlreadyQuoted("\"my_schema\".\"my_table\"")); + assertTrue(SqlIdentifierUtils.isAlreadyQuoted("[users]")); + assertTrue(SqlIdentifierUtils.isAlreadyQuoted("[dbo].[users]")); + + assertFalse(SqlIdentifierUtils.isAlreadyQuoted("my-project.sales.orders")); + assertFalse(SqlIdentifierUtils.isAlreadyQuoted("target_append")); + assertFalse(SqlIdentifierUtils.isAlreadyQuoted(null)); + assertFalse(SqlIdentifierUtils.isAlreadyQuoted("")); + } + + @Test + void testEscapeAndQuoteIdentifierBigQueryNames() { + // Documented whole backtick form is preserved as-is + assertEquals("`my-project.sales.orders`", + SqlIdentifierUtils.escapeAndQuoteIdentifier("`my-project.sales.orders`", (Connection) null)); + + // Segment-by-segment backtick form is preserved as-is + assertEquals("`my-project`.`sales`.`orders`", + SqlIdentifierUtils.escapeAndQuoteIdentifier("`my-project`.`sales`.`orders`", (Connection) null)); + + // Unquoted form with MySQL/BigQuery quote character + assertEquals("`my-project`.`sales`.`orders`", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my-project.sales.orders", "`", "`")); + + // Unquoted form with ANSI quote character + assertEquals("\"my-project\".\"sales\".\"orders\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my-project.sales.orders", "\"", "\"")); + } + + @Test + void testEscapeAndQuoteIdentifierPreservesCaseSemanticsWithConnection() throws SQLException { + Connection conn = mock(Connection.class); + DatabaseMetaData meta = mock(DatabaseMetaData.class); + when(conn.getMetaData()).thenReturn(meta); + when(meta.getIdentifierQuoteString()).thenReturn("\""); + + // Database storing uppercase identifiers (e.g. HSQLDB, Oracle) + when(meta.storesUpperCaseIdentifiers()).thenReturn(true); + when(meta.storesLowerCaseIdentifiers()).thenReturn(false); + assertEquals("\"TARGET_APPEND\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("target_append", conn)); + assertEquals("\"MY_SCHEMA\".\"MY_TABLE\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("my_schema.my_table", conn)); + + // Already-quoted identifiers must not be folded or double-quoted + assertEquals("\"target_append\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("\"target_append\"", conn)); + assertEquals("`my-project.sales.orders`", + SqlIdentifierUtils.escapeAndQuoteIdentifier("`my-project.sales.orders`", conn)); + + // Database storing lowercase identifiers (e.g. PostgreSQL) + when(meta.storesUpperCaseIdentifiers()).thenReturn(false); + when(meta.storesLowerCaseIdentifiers()).thenReturn(true); + assertEquals("\"target_append\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("TARGET_APPEND", conn)); + + // Database preserving mixed case (e.g. MySQL) + when(meta.storesUpperCaseIdentifiers()).thenReturn(false); + when(meta.storesLowerCaseIdentifiers()).thenReturn(false); + assertEquals("\"Target_Append\"", + SqlIdentifierUtils.escapeAndQuoteIdentifier("Target_Append", conn)); + } +} 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..94272ac86 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(), product); 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], product); + 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], product); + 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..e43ffedaa 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,65 @@ 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 testBigQueryDocumentedBacktickTableNameClauseGeneration() { + JdbcTableSinkOperator sinkOp = new HsqldbTableSinkOperator( + "`my-project.sales.orders`", new String[]{"col1"}); + sinkOp.setMode("overwrite"); + assertEquals("CREATE TABLE `my-project.sales.orders` AS (", sinkOp.createSqlClause(null, null)); + + sinkOp.setMode("append"); + assertEquals("INSERT INTO `my-project.sales.orders`", sinkOp.createSqlClause(null, null)); + } + + @Test + void testBigQueryUnquotedTableNameClauseGeneration() { + JdbcTableSinkOperator sinkOp = new HsqldbTableSinkOperator( + "my-project.sales.orders", new String[]{"col1"}); + sinkOp.setMode("overwrite"); + assertEquals("CREATE TABLE \"my-project\".\"sales\".\"orders\" AS (", sinkOp.createSqlClause(null, null)); + + sinkOp.setMode("append"); + assertEquals("INSERT INTO \"my-project\".\"sales\".\"orders\"", 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)