Escape backslashes in queries passed to UNLOAD

We need to escape backslashes that appear in the query passed to the `UNLOAD` command. This patch adds a regression test plus additional backslash escaping. Fixes #215. Author: Josh Rosen <[email protected]> Closes #228 from JoshRosen/unload-backslash-escape.
databricks · Jul 8, 2016 · 127e5fc · 127e5fc
1 parent 32c08d6
commit 127e5fc
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/src/it/scala/com/databricks/spark/redshift/RedshiftIntegrationSuite.scala b/src/it/scala/com/databricks/spark/redshift/RedshiftIntegrationSuite.scala
@@ -186,6 +186,19 @@ class RedshiftIntegrationSuite extends IntegrationSuiteBase {
  )
  }
 
+ test("backslashes in queries/subqueries are escaped (regression test for #215)") {
+ val loadedDf = sqlContext.read
+ .format("com.databricks.spark.redshift")
+ .option("url", jdbcUrl)
+ .option("query", s"select replace(teststring, '\\\\', '') as col from $test_table")
+ .option("tempdir", tempDir)
+ .load()
+ checkAnswer(
+ loadedDf.filter("col = 'asdf'"),
+ Seq(Row("asdf"))
+ )
+ }
+
  test("Can load output when 'dbtable' is a subquery wrapped in parentheses") {
  // scalastyle:off
  val query =

diff --git a/src/main/scala/com/databricks/spark/redshift/RedshiftRelation.scala b/src/main/scala/com/databricks/spark/redshift/RedshiftRelation.scala
@@ -174,8 +174,8 @@ private[redshift] case class RedshiftRelation(
  val credsString: String = AWSCredentialsUtils.getRedshiftCredentialsString(params, creds)
  val query = {
  // Since the query passed to UNLOAD will be enclosed in single quotes, we need to escape
- // any single quotes that appear in the query itself
- val escapedTableNameOrSubqury = tableNameOrSubquery.replace("'", "\\'")
+ // any backslashes and single quotes that appear in the query itself
+ val escapedTableNameOrSubqury = tableNameOrSubquery.replace("\\", "\\\\").replace("'", "\\'")
  s"SELECT $columnList FROM $escapedTableNameOrSubqury $whereClause"
  }
  // We need to remove S3 credentials from the unload path URI because they will conflict with

diff --git a/src/test/scala/com/databricks/spark/redshift/RedshiftSourceSuite.scala b/src/test/scala/com/databricks/spark/redshift/RedshiftSourceSuite.scala
@@ -182,10 +182,10 @@ class RedshiftSourceSuite
  |UNLOAD \('SELECT "testbyte", "testbool" FROM
  | \(select testbyte, testbool
  | from test_table
- | where teststring = \\'Unicode\\'\\'s樂趣\\'\) '\)
+ | where teststring = \\'\\\\\\\\Unicode\\'\\'s樂趣\\'\) '\)
  """.stripMargin.lines.map(_.trim).mkString(" ").trim.r
  val query =
- """select testbyte, testbool from test_table where teststring = 'Unicode''s樂趣'"""
+ """select testbyte, testbool from test_table where teststring = '\\Unicode''s樂趣'"""
  // scalastyle:on
  val querySchema =
  StructType(Seq(StructField("testbyte", ByteType), StructField("testbool", BooleanType)))