Skip to content

Validate and escape table identifiers in JDBC sink operators #839

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AliRana30 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Validate and escape table identifiers in JDBC sink operators #839

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/operators/TableSink.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -51,6 +52,12 @@ public TableSink(Properties props, String mode, String tableName, String... colu

public TableSink(Properties props, String mode, String tableName, String[] columnNames, DataSetType<T> type) {
super(type);
if (tableName != null && !tableName.trim().isEmpty()) {
SqlIdentifierUtils.validateTableName(tableName);
}
if (columnNames != null && columnNames.length > 0) {
SqlIdentifierUtils.validateColumnNames(columnNames);
}
this.tableName = tableName;
this.columnNames = columnNames == null ? null : java.util.Arrays.copyOf(columnNames, columnNames.length);
this.props = new Properties();
Expand Down Expand Up @@ -78,6 +85,7 @@ public String getTableName() {
}

protected void setColumnNames(String[] columnNames) {
SqlIdentifierUtils.validateColumnNames(columnNames);
this.columnNames = columnNames == null ? null : java.util.Arrays.copyOf(columnNames, columnNames.length);
}

Expand Down
2 changes: 2 additions & 0 deletions wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/DatabaseProduct.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,5 +30,7 @@ public enum DatabaseProduct {
H2,
DERBY,
MSSQL,
HSQLDB,
BIGQUERY,
UNKNOWN
}
374 changes: 374 additions & 0 deletions ...g-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/SqlIdentifierUtils.java
View file
Open in desktop

Large diffs are not rendered by default.

7 changes: 6 additions & 1 deletion wayang-commons/wayang-basic/src/main/java/org/apache/wayang/basic/util/SqlTypeUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;

/**
Expand Down Expand Up @@ -80,7 +81,7 @@ public class SqlTypeUtils {
public static DatabaseProduct detectProduct(String url) {
if (url == null)
return DatabaseProduct.UNKNOWN;
String lowerUrl = url.toLowerCase();
String lowerUrl = url.toLowerCase(Locale.ROOT);
if (lowerUrl.contains("postgresql") || lowerUrl.contains("postgres"))
return DatabaseProduct.POSTGRESQL;
if (lowerUrl.contains("mysql"))
Expand All @@ -90,12 +91,16 @@ public static DatabaseProduct detectProduct(String url) {
if (lowerUrl.contains("sqlite")) {
return DatabaseProduct.SQLITE;
}
if (lowerUrl.contains("hsqldb"))
return DatabaseProduct.HSQLDB;
if (lowerUrl.contains("h2"))
return DatabaseProduct.H2;
if (lowerUrl.contains("derby"))
return DatabaseProduct.DERBY;
if (lowerUrl.contains("mssql") || lowerUrl.contains("sqlserver"))
return DatabaseProduct.MSSQL;
if (lowerUrl.contains("bigquery"))
return DatabaseProduct.BIGQUERY;
return DatabaseProduct.UNKNOWN;
}

Expand Down
354 changes: 354 additions & 0 deletions ...mmons/wayang-basic/src/test/java/org/apache/wayang/basic/util/SqlIdentifierUtilsTest.java
View file
Open in desktop

Large diffs are not rendered by default.

21 changes: 21 additions & 0 deletions wayang-commons/wayang-basic/src/test/java/org/apache/wayang/basic/util/SqlTypeUtilsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.Locale;

import static org.junit.jupiter.api.Assertions.assertEquals;

Expand All @@ -43,9 +44,29 @@ public void testDetectProduct() {
SqlTypeUtils.detectProduct("jdbc:derby:memory:test;create=true"));
assertEquals(DatabaseProduct.MSSQL,
SqlTypeUtils.detectProduct("jdbc:sqlserver://localhost:1433;databaseName=db"));
assertEquals(DatabaseProduct.HSQLDB,
SqlTypeUtils.detectProduct("jdbc:hsqldb:mem:test"));
assertEquals(DatabaseProduct.BIGQUERY,
SqlTypeUtils.detectProduct("jdbc:bigquery://https://www.googleapis.com/bigquery/v2:443"));
assertEquals(DatabaseProduct.UNKNOWN, SqlTypeUtils.detectProduct("jdbc:unknown:db"));
}

@Test
public void testDetectProductTurkishLocale() {
Locale defaultLocale = Locale.getDefault();
try {
Locale.setDefault(Locale.forLanguageTag("tr-TR"));
assertEquals(DatabaseProduct.POSTGRESQL,
SqlTypeUtils.detectProduct("JDBC:POSTGRESQL://localhost:5432/db"));
assertEquals(DatabaseProduct.SQLITE,
SqlTypeUtils.detectProduct("JDBC:SQLITE:test.db"));
assertEquals(DatabaseProduct.BIGQUERY,
SqlTypeUtils.detectProduct("JDBC:BIGQUERY://https://www.googleapis.com/bigquery/v2:443"));
} finally {
Locale.setDefault(defaultLocale);
}
}

@Test
public void testGetSqlTypeDefault() {
assertEquals("INT", SqlTypeUtils.getSqlType(Integer.class, DatabaseProduct.UNKNOWN));
Expand Down
6 changes: 6 additions & 0 deletions wayang-platforms/wayang-java/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,12 @@
<version>5.11.0</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.hsqldb</groupId>
<artifactId>hsqldb</artifactId>
<version>2.7.1</version>
<scope>test</scope>
</dependency>
</dependencies>

</project>
20 changes: 9 additions & 11 deletions ...g-platforms/wayang-java/src/main/java/org/apache/wayang/java/operators/JavaTableSink.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -158,35 +159,32 @@ public Tuple<Collection<ExecutionLineageNode>, Collection<ChannelInstance>> 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 = ", ";
}
sb.append(")");
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 (");
Expand Down
Loading
Loading