diff --git a/amazon-redshift-plugin/src/main/java/io/cdap/plugin/amazon/redshift/RedshiftConnector.java b/amazon-redshift-plugin/src/main/java/io/cdap/plugin/amazon/redshift/RedshiftConnector.java
index fb8cac4a7..57a253ebb 100644
--- a/amazon-redshift-plugin/src/main/java/io/cdap/plugin/amazon/redshift/RedshiftConnector.java
+++ b/amazon-redshift-plugin/src/main/java/io/cdap/plugin/amazon/redshift/RedshiftConnector.java
@@ -82,7 +82,7 @@ protected SchemaReader getSchemaReader(String sessionID) {
@Override
protected String getTableName(String database, String schema, String table) {
- return String.format("\"%s\".\"%s\"", schema, table);
+ return quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
}
@Override
diff --git a/amazon-redshift-plugin/src/test/java/io/cdap/plugin/amazon/redshift/RedshiftConnectorUnitTest.java b/amazon-redshift-plugin/src/test/java/io/cdap/plugin/amazon/redshift/RedshiftConnectorUnitTest.java
index 94ea6004b..16c94a3f0 100644
--- a/amazon-redshift-plugin/src/test/java/io/cdap/plugin/amazon/redshift/RedshiftConnectorUnitTest.java
+++ b/amazon-redshift-plugin/src/test/java/io/cdap/plugin/amazon/redshift/RedshiftConnectorUnitTest.java
@@ -41,6 +41,15 @@ public void getTableNameTest() {
CONNECTOR.getTableName("db", "schema", "table"));
}
+ /**
+ * Unit test for getTableName() with identifiers containing double quotes
+ */
+ @Test
+ public void getTableNameEscapesQuotesTest() {
+ Assert.assertEquals("\"my\"\"schema\".\"t\"\"; DELETE FROM x; --\"",
+ CONNECTOR.getTableName("db", "my\"schema", "t\"; DELETE FROM x; --"));
+ }
+
@Test
public void getRandomQuery() {
Assert.assertEquals("SELECT * FROM TestData\n" +
diff --git a/cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLConnector.java b/cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLConnector.java
index 0751bd160..dccb8b047 100644
--- a/cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLConnector.java
+++ b/cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLConnector.java
@@ -82,7 +82,7 @@ protected SchemaReader getSchemaReader(String sessionID) {
@Override
protected String getTableName(String database, String schema, String table) {
- return String.format("`%s`.`%s`", database, table);
+ return quoteIdentifier(database, '`') + "." + quoteIdentifier(table, '`');
}
@Override
diff --git a/cloudsql-postgresql-plugin/src/main/java/io/cdap/plugin/cloudsql/postgres/CloudSQLPostgreSQLConnector.java b/cloudsql-postgresql-plugin/src/main/java/io/cdap/plugin/cloudsql/postgres/CloudSQLPostgreSQLConnector.java
index 348a1b94c..8f1921478 100644
--- a/cloudsql-postgresql-plugin/src/main/java/io/cdap/plugin/cloudsql/postgres/CloudSQLPostgreSQLConnector.java
+++ b/cloudsql-postgresql-plugin/src/main/java/io/cdap/plugin/cloudsql/postgres/CloudSQLPostgreSQLConnector.java
@@ -86,7 +86,7 @@ protected SchemaReader getSchemaReader(String sessionID) {
@Override
protected String getTableName(String database, String schema, String table) {
- return String.format("\"%s\".\"%s\"", schema, table);
+ return quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
}
@Override
diff --git a/database-commons/pom.xml b/database-commons/pom.xml
index 98bcb3803..2d92db70d 100644
--- a/database-commons/pom.xml
+++ b/database-commons/pom.xml
@@ -77,6 +77,10 @@
dev.failsafe
failsafe
+
+ com.github.jsqlparser
+ jsqlparser
+
diff --git a/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java b/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java
index e0407203b..f1b0fdf6f 100644
--- a/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java
+++ b/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java
@@ -150,8 +150,17 @@ protected String getConnectionString(String database) {
}
protected String getTableName(String database, String schema, String table) {
- return schema == null ? String.format("\"%s\".\"%s\"", database, table)
- : String.format("\"%s\".\"%s\".\"%s\"", database, schema, table);
+ return schema == null ? quoteIdentifier(database, '"') + "." + quoteIdentifier(table, '"')
+ : quoteIdentifier(database, '"') + "." + quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
+ }
+
+ /**
+ * Wraps the given identifier in the given quote character, doubling any embedded quote characters so that a
+ * user supplied database, schema, table or column name can never break out of the quoted identifier.
+ */
+ public static String quoteIdentifier(String identifier, char quote) {
+ String quoteString = String.valueOf(quote);
+ return quoteString + identifier.replace(quoteString, quoteString + quoteString) + quoteString;
}
protected String getTableQuery(String database, String schema, String table) {
diff --git a/database-commons/src/main/java/io/cdap/plugin/db/source/AbstractDBSource.java b/database-commons/src/main/java/io/cdap/plugin/db/source/AbstractDBSource.java
index 908d27eb9..53f26c617 100644
--- a/database-commons/src/main/java/io/cdap/plugin/db/source/AbstractDBSource.java
+++ b/database-commons/src/main/java/io/cdap/plugin/db/source/AbstractDBSource.java
@@ -52,6 +52,7 @@
import io.cdap.plugin.util.DriverCleanup;
import io.cdap.plugin.util.RetryPolicyUtil;
import io.cdap.plugin.util.RetryUtils;
+import io.cdap.plugin.util.SelectQueryValidator;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.mapreduce.MRJobConfig;
import org.apache.hadoop.mapreduce.lib.db.DBConfiguration;
@@ -106,6 +107,7 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
StageConfigurer stageConfigurer = pipelineConfigurer.getStageConfigurer();
FailureCollector collector = stageConfigurer.getFailureCollector();
sourceConfig.validate(collector);
+ boolean queriesValid = validateQueries(collector);
if (sourceConfig.getSchema() != null) {
stageConfigurer.setOutputSchema(sourceConfig.getSchema());
return;
@@ -119,7 +121,8 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
Class driverClass = DBUtils.getDriverClass(
pipelineConfigurer, sourceConfig, ConnectionConfig.JDBC_PLUGIN_TYPE);
- if (sourceConfig.canConnect()) {
+ // never run a query that failed the read-only validation, not even to fetch its schema
+ if (sourceConfig.canConnect() && queriesValid) {
try {
stageConfigurer.setOutputSchema(getSchema(driverClass));
} catch (IllegalAccessException | InstantiationException e) {
@@ -228,6 +231,43 @@ protected String getErrorDetailsProviderClassName() {
return null;
}
+ /**
+ * Returns whether square brackets quote identifiers in the database's SQL dialect, as in SQL Server.
+ * Override this method to allow such identifiers in the import and bounding queries.
+ */
+ protected boolean useSquareBracketQuotes() {
+ return false;
+ }
+
+ /**
+ * Validates that the import and bounding queries are single read-only SELECT statements, since they are executed
+ * with the credential of the connection. Queries that are macros are validated at runtime, once they are evaluated.
+ *
+ * @return true if the queries are valid or cannot be validated yet
+ */
+ private boolean validateQueries(FailureCollector collector) {
+ boolean importQueryValid = validateQuery(DBSourceConfig.IMPORT_QUERY, "Import Query",
+ sourceConfig.getImportQuery(), collector);
+ boolean boundingQueryValid = validateQuery(DBSourceConfig.BOUNDING_QUERY, "Bounding Query",
+ sourceConfig.getBoundingQuery(), collector);
+ return importQueryValid && boundingQueryValid;
+ }
+
+ private boolean validateQuery(String property, String label, @Nullable String query, FailureCollector collector) {
+ if (sourceConfig.containsMacro(property) || Strings.isNullOrEmpty(query)) {
+ return true;
+ }
+ try {
+ SelectQueryValidator.validate(query, useSquareBracketQuotes());
+ return true;
+ } catch (IllegalArgumentException e) {
+ collector.addFailure(String.format("Invalid %s: %s", label, e.getMessage()),
+ "Provide a single SELECT statement that only reads data.")
+ .withConfigProperty(property);
+ return false;
+ }
+ }
+
/**
* Returns the DBErrorDetailsProvider instance.
* Override this method to provide a custom DBErrorDetailsProvider instance.
@@ -272,6 +312,7 @@ private Connection getConnection() {
public void prepareRun(BatchSourceContext context) throws Exception {
FailureCollector collector = context.getFailureCollector();
sourceConfig.validate(collector);
+ validateQueries(collector);
collector.getOrThrowException();
LOG.debug("pluginType = {}; pluginName = {}; connectionString = {}; importQuery = {}; " +
diff --git a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java
index c78ceca6f..5997cf627 100644
--- a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java
+++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java
@@ -116,15 +116,15 @@ protected SchemaReader getSchemaReader(String sessionID) {
@Override
protected String getTableName(String database, String schema, String table) {
if (database == null && schema == null) {
- return String.format("`%s`", table);
+ return quoteIdentifier(table, '`');
}
if (database == null) {
- return String.format("`%s`.`%s`", schema, table);
+ return quoteIdentifier(schema, '`') + "." + quoteIdentifier(table, '`');
}
if (schema == null) {
- return String.format("`%s`.`%s`", database, table);
+ return quoteIdentifier(database, '`') + "." + quoteIdentifier(table, '`');
}
- return String.format("`%s`.`%s`.`%s`", database, schema, table);
+ return quoteIdentifier(database, '`') + "." + quoteIdentifier(schema, '`') + "." + quoteIdentifier(table, '`');
}
@Override
@@ -136,6 +136,7 @@ protected String getRandomQuery(String tableName, int limit) {
@Override
protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) {
+ String strataColumn = quoteIdentifier(strata, '`');
return String.format("WITH t_%s AS (\n" +
" SELECT *,\n" +
" ROW_NUMBER() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" +
@@ -146,8 +147,8 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
"WHERE MOD(sqn_%s, GREATEST(1, CAST(c_%s / %d AS BIGINT))) = 1\n" +
"ORDER BY %s\n" +
"LIMIT %d",
- sessionID, strata, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
- limit, strata, limit);
+ sessionID, strataColumn, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
+ limit, strataColumn, limit);
}
@Override
diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java
index a04dfb398..27b830c76 100644
--- a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java
+++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java
@@ -38,6 +38,21 @@ public void testGetTableName() {
CONNECTOR.getTableName(null, null, "my_table"));
}
+ @Test
+ public void testGetTableNameEscapesBackticks() {
+ Assert.assertEquals("`main`.`default`.`t``; DELETE FROM x; --`",
+ CONNECTOR.getTableName("main", "default", "t`; DELETE FROM x; --"));
+ Assert.assertEquals("`my``schema`.`my``table`",
+ CONNECTOR.getTableName(null, "my`schema", "my`table"));
+ }
+
+ @Test
+ public void testGetStratifiedQueryEscapesStrata() {
+ String query = CONNECTOR.getStratifiedQuery("`main`.`default`.`my_table`", 10, "id`; DROP TABLE t; --", "s1");
+ Assert.assertTrue(query.contains("ORDER BY `id``; DROP TABLE t; --`, RAND()"));
+ Assert.assertTrue(query.contains("ORDER BY `id``; DROP TABLE t; --`\n"));
+ }
+
@Test
public void testGetRandomQuery() {
Assert.assertEquals("SELECT * FROM `main`.`default`.`my_table`\n" +
@@ -49,13 +64,13 @@ public void testGetRandomQuery() {
public void testGetStratifiedQuery() {
Assert.assertEquals("WITH t_s1 AS (\n" +
" SELECT *,\n" +
- " ROW_NUMBER() OVER (ORDER BY id, RAND()) AS sqn_s1,\n" +
+ " ROW_NUMBER() OVER (ORDER BY `id`, RAND()) AS sqn_s1,\n" +
" COUNT(*) OVER () AS c_s1\n" +
" FROM `main`.`default`.`my_table`\n" +
" )\n" +
"SELECT * FROM t_s1\n" +
"WHERE MOD(sqn_s1, GREATEST(1, CAST(c_s1 / 10 AS BIGINT))) = 1\n" +
- "ORDER BY id\n" +
+ "ORDER BY `id`\n" +
"LIMIT 10",
CONNECTOR.getStratifiedQuery("`main`.`default`.`my_table`", 10, "id", "s1"));
}
diff --git a/generic-database-plugin/src/main/java/io/cdap/plugin/jdbc/DatabaseSource.java b/generic-database-plugin/src/main/java/io/cdap/plugin/jdbc/DatabaseSource.java
index b35b025b7..1c07b2d8a 100644
--- a/generic-database-plugin/src/main/java/io/cdap/plugin/jdbc/DatabaseSource.java
+++ b/generic-database-plugin/src/main/java/io/cdap/plugin/jdbc/DatabaseSource.java
@@ -46,6 +46,14 @@ protected String createConnectionString() {
return databaseSourceConfig.connectionString;
}
+ @Override
+ protected boolean useSquareBracketQuotes() {
+ // SQL Server quotes identifiers with square brackets
+ String connectionString = databaseSourceConfig.connectionString;
+ return connectionString != null && (connectionString.toLowerCase().startsWith("jdbc:sqlserver:")
+ || connectionString.toLowerCase().startsWith("jdbc:jtds:sqlserver:"));
+ }
+
/**
* Generic database source configuration.
*/
diff --git a/generic-database-plugin/src/test/java/io/cdap/plugin/db/sink/DBSourceTestRun.java b/generic-database-plugin/src/test/java/io/cdap/plugin/db/sink/DBSourceTestRun.java
index bedb720a2..f11342461 100644
--- a/generic-database-plugin/src/test/java/io/cdap/plugin/db/sink/DBSourceTestRun.java
+++ b/generic-database-plugin/src/test/java/io/cdap/plugin/db/sink/DBSourceTestRun.java
@@ -20,16 +20,19 @@
import io.cdap.cdap.api.common.Bytes;
import io.cdap.cdap.api.data.format.StructuredRecord;
import io.cdap.cdap.api.dataset.table.Table;
+import io.cdap.cdap.datapipeline.SmartWorkflow;
import io.cdap.cdap.etl.api.batch.BatchSource;
import io.cdap.cdap.etl.mock.batch.MockSink;
import io.cdap.cdap.etl.proto.v2.ETLBatchConfig;
import io.cdap.cdap.etl.proto.v2.ETLPlugin;
import io.cdap.cdap.etl.proto.v2.ETLStage;
+import io.cdap.cdap.proto.ProgramRunStatus;
import io.cdap.cdap.proto.artifact.AppRequest;
import io.cdap.cdap.proto.id.ApplicationId;
import io.cdap.cdap.proto.id.NamespaceId;
import io.cdap.cdap.test.ApplicationManager;
import io.cdap.cdap.test.DataSetManager;
+import io.cdap.cdap.test.WorkflowManager;
import io.cdap.plugin.GenericDatabasePluginTestBase;
import io.cdap.plugin.common.Constants;
import io.cdap.plugin.db.ConnectionConfig;
@@ -41,7 +44,11 @@
import java.math.BigDecimal;
import java.math.MathContext;
import java.nio.ByteBuffer;
+import java.sql.Connection;
import java.sql.Date;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
import java.sql.Time;
import java.text.SimpleDateFormat;
import java.time.LocalDate;
@@ -52,6 +59,7 @@
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.TimeUnit;
/**
* Test for ETL using databases.
@@ -329,6 +337,61 @@ public void testUserNamePasswordCombinations() throws Exception {
deployApplication(appId, appRequest);
}
+ @Test
+ public void testNonSelectImportQueryFailsDeployment() throws Exception {
+ int rowCount = countRows("my_table");
+ String importQuery = "SELECT * FROM \"my_table\" WHERE $CONDITIONS; DELETE FROM \"my_table\"";
+ ETLPlugin sourceConfig = new ETLPlugin(
+ DatabaseConstants.PLUGIN_NAME,
+ BatchSource.PLUGIN_TYPE,
+ ImmutableMap.builder()
+ .put(ConnectionConfig.CONNECTION_STRING, getConnectionURL())
+ .put(ConnectionConfig.JDBC_PLUGIN_NAME, JDBC_DRIVER_NAME)
+ .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, importQuery)
+ .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, "SELECT MIN(ID),MAX(ID) from \"my_table\"")
+ .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, "ID")
+ .put(Constants.Reference.REFERENCE_NAME, "DBNonSelectQueryTest")
+ .build(),
+ null);
+ ETLBatchConfig etlConfig = getETLBatchConfig(sourceConfig, MockSink.getPlugin("nonSelectQueryOutput"));
+ assertDeploymentFailure(NamespaceId.DEFAULT.app("dbSourceNonSelectQueryTest"), etlConfig, DATAPIPELINE_ARTIFACT,
+ "Deploying DB Source with a non SELECT import query should have failed.");
+ Assert.assertEquals(rowCount, countRows("my_table"));
+ }
+
+ @Test
+ public void testNonSelectImportQueryMacroFailsRun() throws Exception {
+ int rowCount = countRows("my_table");
+ ETLPlugin sourceConfig = new ETLPlugin(
+ DatabaseConstants.PLUGIN_NAME,
+ BatchSource.PLUGIN_TYPE,
+ ImmutableMap.builder()
+ .put(ConnectionConfig.CONNECTION_STRING, getConnectionURL())
+ .put(ConnectionConfig.JDBC_PLUGIN_NAME, JDBC_DRIVER_NAME)
+ .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, "${importQuery}")
+ .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, "SELECT MIN(ID),MAX(ID) from \"my_table\"")
+ .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, "ID")
+ .put(Constants.Reference.REFERENCE_NAME, "DBNonSelectQueryMacroTest")
+ .build(),
+ null);
+ ApplicationManager appManager = deployETL(sourceConfig, MockSink.getPlugin("nonSelectQueryMacroOutput"),
+ DATAPIPELINE_ARTIFACT, "dbSourceNonSelectQueryMacroTest");
+ WorkflowManager workflowManager = appManager.getWorkflowManager(SmartWorkflow.NAME);
+ workflowManager.start(ImmutableMap.of(
+ "importQuery", "SELECT * FROM \"my_table\" WHERE $CONDITIONS; DELETE FROM \"my_table\""));
+ workflowManager.waitForRun(ProgramRunStatus.FAILED, 5, TimeUnit.MINUTES);
+ Assert.assertEquals(rowCount, countRows("my_table"));
+ }
+
+ private int countRows(String table) throws SQLException {
+ try (Connection connection = getConnection();
+ Statement statement = connection.createStatement();
+ ResultSet resultSet = statement.executeQuery(String.format("SELECT COUNT(*) FROM \"%s\"", table))) {
+ resultSet.next();
+ return resultSet.getInt(1);
+ }
+ }
+
@Test
public void testNonExistentDBTable() throws Exception {
// source
diff --git a/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerConnector.java b/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerConnector.java
index 7e86d97ce..4f3ba8963 100644
--- a/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerConnector.java
+++ b/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerConnector.java
@@ -117,7 +117,8 @@ public StructuredRecord transform(LongWritable longWritable, SqlServerSourceDBRe
@Override
protected String getTableName(String database, String schema, String table) {
- return String.format("\"%s\".\"%s\".\"%s\"", database, schema, table);
+ return quoteIdentifier(database, '"') + "." + quoteIdentifier(schema, '"') + "."
+ + quoteIdentifier(table, '"');
}
@Override
@@ -137,6 +138,7 @@ protected String getRandomQuery(String tableName, int limit) {
@Override
protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) {
+ String strataColumn = quoteIdentifier(strata, '"');
return String.format("WITH t_%s AS (\n" +
" SELECT *,\n" +
" ROW_NUMBER() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" +
@@ -146,8 +148,8 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
"SELECT TOP %d * FROM t_%s\n" +
"WHERE sqn_%s %% CAST(0.5 * ((c_%s / %d + 1) + ABS(c_%s / %d - 1)) AS bigint) = 1\n" +
"ORDER BY %s",
- sessionID, strata, sessionID, sessionID, tableName, limit, sessionID, sessionID, sessionID,
- limit, sessionID, limit, strata);
+ sessionID, strataColumn, sessionID, sessionID, tableName, limit, sessionID, sessionID,
+ sessionID, limit, sessionID, limit, strataColumn);
}
@Override
diff --git a/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerSource.java b/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerSource.java
index d8d4cf39e..b6bcb9ff9 100644
--- a/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerSource.java
+++ b/mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerSource.java
@@ -81,6 +81,11 @@ protected String getErrorDetailsProviderClassName() {
return SqlServerErrorDetailsProvider.class.getName();
}
+ @Override
+ protected boolean useSquareBracketQuotes() {
+ return true;
+ }
+
@Override
protected DBErrorDetailsProvider getErrorDetailsProvider() {
if (dbErrorDetailsProvider == null) {
diff --git a/mssql-plugin/src/test/java/io/cdap/plugin/mssql/SqlServerConnectorUnitTest.java b/mssql-plugin/src/test/java/io/cdap/plugin/mssql/SqlServerConnectorUnitTest.java
index 0163cd55e..231d072b9 100644
--- a/mssql-plugin/src/test/java/io/cdap/plugin/mssql/SqlServerConnectorUnitTest.java
+++ b/mssql-plugin/src/test/java/io/cdap/plugin/mssql/SqlServerConnectorUnitTest.java
@@ -34,6 +34,27 @@ public class SqlServerConnectorUnitTest {
// private static final SqlServerConnector CONNECTOR = new SqlServerConnector(new SqlServerConnectorConfig());
+ /**
+ * Unit test for getTableName() with identifiers containing double quotes
+ */
+ @Test
+ public void getTableNameEscapesQuotesTest() {
+ Assert.assertEquals("\"db\".\"dbo\".\"t\"\"; DELETE FROM x; --\"",
+ CONNECTOR.getTableName("db", "dbo", "t\"; DELETE FROM x; --"));
+ Assert.assertEquals("\"my\"\"db\".\"my\"\"schema\".\"my\"\"table\"",
+ CONNECTOR.getTableName("my\"db", "my\"schema", "my\"table"));
+ }
+
+ /**
+ * Unit test for getStratifiedQuery() with a strata column containing double quotes
+ */
+ @Test
+ public void getStratifiedQueryEscapesStrataTest() {
+ String query = CONNECTOR.getStratifiedQuery("\"db\".\"dbo\".\"table\"", 10, "id\"; DROP TABLE t; --", "s1");
+ Assert.assertTrue(query.contains("ORDER BY \"id\"\"; DROP TABLE t; --\", RAND()"));
+ Assert.assertTrue(query.endsWith("ORDER BY \"id\"\"; DROP TABLE t; --\""));
+ }
+
/**
* Unit tests for getTableQuery()
*/
diff --git a/mysql-plugin/src/main/java/io/cdap/plugin/mysql/MysqlConnector.java b/mysql-plugin/src/main/java/io/cdap/plugin/mysql/MysqlConnector.java
index 214d8650e..0f854bc88 100644
--- a/mysql-plugin/src/main/java/io/cdap/plugin/mysql/MysqlConnector.java
+++ b/mysql-plugin/src/main/java/io/cdap/plugin/mysql/MysqlConnector.java
@@ -98,7 +98,7 @@ protected void setConnectorSpec(ConnectorSpecRequest request, DBConnectorPath pa
@Override
protected String getTableName(String database, String schema, String table) {
- return String.format("`%s`.`%s`", database, table);
+ return quoteIdentifier(database, '`') + "." + quoteIdentifier(table, '`');
}
@Override
@@ -112,6 +112,7 @@ protected String getRandomQuery(String tableName, int limit) {
@Override
protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) {
+ String strataColumn = quoteIdentifier(strata, '`');
return String.format("WITH t_%s AS (\n" +
" SELECT *,\n" +
" ROW_NUMBER() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" +
@@ -122,8 +123,8 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
"WHERE MOD(sqn_%s, GREATEST(1, CAST(c_%s / %d AS UNSIGNED))) = 1\n" +
"ORDER BY %s\n" +
"LIMIT %d",
- sessionID, strata, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
- limit, strata, limit);
+ sessionID, strataColumn, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
+ limit, strataColumn, limit);
}
@Override
diff --git a/mysql-plugin/src/test/java/io/cdap/plugin/mysql/MysqlConnectorUnitTest.java b/mysql-plugin/src/test/java/io/cdap/plugin/mysql/MysqlConnectorUnitTest.java
index df731d194..01820842e 100644
--- a/mysql-plugin/src/test/java/io/cdap/plugin/mysql/MysqlConnectorUnitTest.java
+++ b/mysql-plugin/src/test/java/io/cdap/plugin/mysql/MysqlConnectorUnitTest.java
@@ -42,6 +42,25 @@ public void getTableNameTest() {
CONNECTOR.getTableName("db", "schema", "table"));
}
+ /**
+ * Unit test for getTableName() with identifiers containing backticks
+ */
+ @Test
+ public void getTableNameEscapesBackticksTest() {
+ Assert.assertEquals("`my``db`.`t``; DELETE FROM x; --`",
+ CONNECTOR.getTableName("my`db", "schema", "t`; DELETE FROM x; --"));
+ }
+
+ /**
+ * Unit test for getStratifiedQuery() with a strata column containing backticks
+ */
+ @Test
+ public void getStratifiedQueryEscapesStrataTest() {
+ String query = CONNECTOR.getStratifiedQuery("`db`.`table`", 10, "id`; DROP TABLE t; --", "s1");
+ Assert.assertTrue(query.contains("ORDER BY `id``; DROP TABLE t; --`, RAND()"));
+ Assert.assertTrue(query.contains("ORDER BY `id``; DROP TABLE t; --`\n"));
+ }
+
/**
* Unit tests for getTableQuery()
*/
diff --git a/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleConnector.java b/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleConnector.java
index 469ca45f2..a59c0688e 100644
--- a/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleConnector.java
+++ b/oracle-plugin/src/main/java/io/cdap/plugin/oracle/OracleConnector.java
@@ -135,7 +135,7 @@ protected String getConnectionString(@Nullable String database) {
@Override
protected String getTableName(String database, String schema, String table) {
- return String.format("\"%s\".\"%s\"", schema, table);
+ return quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
}
@Override
@@ -157,6 +157,7 @@ protected String getRandomQuery(String tableName, int limit) {
@Override
protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) {
+ String strataColumn = quoteIdentifier(strata, '"');
return String.format("WITH t_%s AS (\n" +
" SELECT %s.*,\n" +
" ROW_NUMBER() OVER (ORDER BY %s, dbms_random.random) AS s_%s,\n" +
@@ -166,8 +167,8 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
" )\n" +
"SELECT * FROM t_%s\n" +
"WHERE MOD(s_%s, GREATEST(1, CAST(c_%s / %d AS INTEGER))) = 1 AND ROWNUM <= %d",
- sessionID, tableName, strata, sessionID, tableName, sessionID, tableName,
- strata, sessionID, sessionID, sessionID, limit, limit);
+ sessionID, tableName, strataColumn, sessionID, tableName, sessionID, tableName,
+ strataColumn, sessionID, sessionID, sessionID, limit, limit);
}
@Override
diff --git a/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleConnectorUnitTest.java b/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleConnectorUnitTest.java
index f89a6c1d7..2441862d0 100644
--- a/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleConnectorUnitTest.java
+++ b/oracle-plugin/src/test/java/io/cdap/plugin/oracle/OracleConnectorUnitTest.java
@@ -37,6 +37,25 @@ public void getTableNameTest() {
CONNECTOR.getTableName("db", "schema", "table"));
}
+ /**
+ * Unit test for getTableName() with identifiers containing double quotes
+ */
+ @Test
+ public void getTableNameEscapesQuotesTest() {
+ Assert.assertEquals("\"MY\"\"SCHEMA\".\"T\"\" WHERE 1=1 --\"",
+ CONNECTOR.getTableName("db", "MY\"SCHEMA", "T\" WHERE 1=1 --"));
+ }
+
+ /**
+ * Unit test for getStratifiedQuery() with a strata column containing double quotes
+ */
+ @Test
+ public void getStratifiedQueryEscapesStrataTest() {
+ String query = CONNECTOR.getStratifiedQuery("\"SCHEMA\".\"TABLE\"", 10, "ID\", 1) --", "s1");
+ Assert.assertTrue(query.contains("ORDER BY \"ID\"\", 1) --\", dbms_random.random"));
+ Assert.assertTrue(query.contains("ORDER BY \"ID\"\", 1) --\" )"));
+ }
+
/**
* Unit tests for getTableQuery()
*/
diff --git a/pom.xml b/pom.xml
index c739d8855..a04552767 100644
--- a/pom.xml
+++ b/pom.xml
@@ -310,6 +310,12 @@
failsafe
3.3.2
+
+
+ com.github.jsqlparser
+ jsqlparser
+ 4.9
+
org.hsqldb
hsqldb
diff --git a/postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresConnector.java b/postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresConnector.java
index 5b23ce6cd..4ad703b61 100644
--- a/postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresConnector.java
+++ b/postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresConnector.java
@@ -117,7 +117,7 @@ public StructuredRecord transform(LongWritable longWritable, PostgresDBRecord re
@Override
protected String getTableName(String database, String schema, String table) {
- return String.format("\"%s\".\"%s\"", schema, table);
+ return quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
}
@Override
@@ -129,6 +129,7 @@ protected String getRandomQuery(String tableName, int limit) {
@Override
protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) {
+ String strataColumn = quoteIdentifier(strata, '"');
// this query times out on larger datasets; would benefit from optimization
return String.format("WITH t_%s AS (\n" +
" SELECT *,\n" +
@@ -140,8 +141,8 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
"WHERE mod(sqn_%s, GREATEST(1, CAST(c_%s / %d AS bigint))) = 1\n" +
"ORDER BY %s\n" +
"LIMIT %d",
- sessionID, strata, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
- limit, strata, limit);
+ sessionID, strataColumn, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
+ limit, strataColumn, limit);
}
@Override
diff --git a/postgresql-plugin/src/test/java/io/cdap/plugin/postgres/PostgresConnectorUnitTest.java b/postgresql-plugin/src/test/java/io/cdap/plugin/postgres/PostgresConnectorUnitTest.java
index f286cf574..57e6399d4 100644
--- a/postgresql-plugin/src/test/java/io/cdap/plugin/postgres/PostgresConnectorUnitTest.java
+++ b/postgresql-plugin/src/test/java/io/cdap/plugin/postgres/PostgresConnectorUnitTest.java
@@ -41,6 +41,28 @@ public void getTableNameTest() {
CONNECTOR.getTableName("db", "schema", "table"));
}
+ /**
+ * Unit test for getTableName() with identifiers containing double quotes
+ */
+ @Test
+ public void getTableNameEscapesQuotesTest() {
+ Assert.assertEquals("\"pg_catalog\".\"pg_class\"\" LIMIT 0; UPDATE t SET v='x'; SELECT * FROM \"\"pg_class\"",
+ CONNECTOR.getTableName("db", "pg_catalog",
+ "pg_class\" LIMIT 0; UPDATE t SET v='x'; SELECT * FROM \"pg_class"));
+ Assert.assertEquals("\"my\"\"schema\".\"my\"\"table\"",
+ CONNECTOR.getTableName("db", "my\"schema", "my\"table"));
+ }
+
+ /**
+ * Unit test for getStratifiedQuery() with a strata column containing double quotes
+ */
+ @Test
+ public void getStratifiedQueryEscapesStrataTest() {
+ String query = CONNECTOR.getStratifiedQuery("\"schema\".\"table\"", 10, "id\"; DROP TABLE t; --", "s1");
+ Assert.assertTrue(query.contains("ORDER BY \"id\"\"; DROP TABLE t; --\", random()"));
+ Assert.assertTrue(query.contains("ORDER BY \"id\"\"; DROP TABLE t; --\"\n"));
+ }
+
/**
* Unit tests for getTableQuery()
*/