-
Notifications
You must be signed in to change notification settings - Fork 346
Simplify use of transactions when performing overwrites and when creating new tables #157
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.
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
Changes from all commits
77323b3
deb2345
a2e241b
3f4c1b9
3a6e9fd
8edc588
d7af870
41c080c
9045bda
b2804bf
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -26,7 +26,6 @@ import org.apache.spark.TaskContext | |
| import org.slf4j.LoggerFactory | ||
|
|
||
| import scala.collection.mutable | ||
| import scala.util.Random | ||
| import scala.util.control.NonFatal | ||
|
|
||
| import com.databricks.spark.redshift.Parameters.MergedParameters | ||
|
|
@@ -48,18 +47,15 @@ import org.apache.spark.sql.types._ | |
| * non-empty. After the write operation completes, we use this to construct a list of non-empty | ||
| * Avro partition files. | ||
| * | ||
| * - Use JDBC to issue any CREATE TABLE commands, if required. | ||
| * | ||
| * - If there is data to be written (i.e. not all partitions were empty), then use the list of | ||
| * non-empty Avro files to construct a JSON manifest file to tell Redshift to load those files. | ||
| * This manifest is written to S3 alongside the Avro files themselves. We need to use an | ||
| * explicit manifest, as opposed to simply passing the name of the directory containing the | ||
| * Avro files, in order to work around a bug related to parsing of empty Avro files (see #96). | ||
| * | ||
| * - Use JDBC to issue a COPY command in order to instruct Redshift to load the Avro data into | ||
| * the appropriate table. If the Overwrite SaveMode is being used, then by default the data | ||
| * will be loaded into a temporary staging table, which later will atomically replace the | ||
| * original table via a transaction. | ||
| * - Start a new JDBC transaction and disable auto-commit. Depending on the SaveMode, issue | ||
| * DELETE TABLE or CREATE TABLE commands, then use the COPY command to instruct Redshift to load | ||
| * the Avro data into the appropriate table. | ||
| */ | ||
| private[redshift] class RedshiftWriter( | ||
| jdbcWrapper: JDBCWrapper, | ||
|
|
@@ -101,43 +97,6 @@ private[redshift] class RedshiftWriter( | |
| s"AVRO 'auto' manifest ${params.extraCopyOptions}" | ||
| } | ||
|
|
||
| /** | ||
| * Sets up a staging table then runs the given action, passing the temporary table name | ||
| * as a parameter. | ||
| */ | ||
| private def withStagingTable( | ||
| conn: Connection, | ||
| table: TableName, | ||
| action: (String) => Unit) { | ||
| val randomSuffix = Math.abs(Random.nextInt()).toString | ||
| val tempTable = | ||
| table.copy(unescapedTableName = s"${table.unescapedTableName}_staging_$randomSuffix") | ||
| val backupTable = | ||
| table.copy(unescapedTableName = s"${table.unescapedTableName}_backup_$randomSuffix") | ||
| log.info("Loading new Redshift data to: " + tempTable) | ||
| log.info("Existing data will be backed up in: " + backupTable) | ||
|
|
||
| try { | ||
| action(tempTable.toString) | ||
|
|
||
| if (jdbcWrapper.tableExists(conn, table.toString)) { | ||
| jdbcWrapper.executeInterruptibly(conn.prepareStatement( | ||
| s""" | ||
| | BEGIN; | ||
| | ALTER TABLE $table RENAME TO ${backupTable.escapedTableName}; | ||
| | ALTER TABLE $tempTable RENAME TO ${table.escapedTableName}; | ||
| | DROP TABLE $backupTable; | ||
| | END; | ||
| """.stripMargin.trim)) | ||
| } else { | ||
| jdbcWrapper.executeInterruptibly(conn.prepareStatement( | ||
| s"ALTER TABLE $tempTable RENAME TO ${table.escapedTableName}")) | ||
| } | ||
| } finally { | ||
| jdbcWrapper.executeInterruptibly(conn.prepareStatement(s"DROP TABLE IF EXISTS $tempTable")) | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Generate COMMENT SQL statements for the table and columns. | ||
| */ | ||
|
|
@@ -151,23 +110,15 @@ private[redshift] class RedshiftWriter( | |
| } | ||
|
|
||
| /** | ||
| * Perform the Redshift load, including deletion of existing data in the case of an overwrite, | ||
| * and creating the table if it doesn't already exist. | ||
| * Perform the Redshift load by issuing a COPY statement. | ||
| */ | ||
| private def doRedshiftLoad( | ||
| conn: Connection, | ||
| data: DataFrame, | ||
| saveMode: SaveMode, | ||
| params: MergedParameters, | ||
| creds: AWSCredentials, | ||
| manifestUrl: Option[String]): Unit = { | ||
|
|
||
| // Overwrites must drop the table, in case there has been a schema update | ||
| if (saveMode == SaveMode.Overwrite) { | ||
| jdbcWrapper.executeInterruptibly( | ||
| conn.prepareStatement(s"DROP TABLE IF EXISTS ${params.table.get}")) | ||
| } | ||
|
|
||
| // If the table doesn't exist, we need to create it first, using JDBC to infer column types | ||
| val createStatement = createTableSql(data, params) | ||
| log.info(createStatement) | ||
|
|
@@ -190,8 +141,11 @@ private[redshift] class RedshiftWriter( | |
| jdbcWrapper.executeInterruptibly(conn.prepareStatement(copyStatement)) | ||
| } catch { | ||
| case e: SQLException => | ||
| log.error("SQLException thrown while running COPY query; will attempt to retrieve " + | ||
| "more information by querying the STL_LOAD_ERRORS table", e) | ||
| // Try to query Redshift's STL_LOAD_ERRORS table to figure out why the load failed. | ||
| // See http://docs.aws.amazon.com/redshift/latest/dg/r_STL_LOAD_ERRORS.html for details. | ||
| conn.rollback() | ||
| val errorLookupQuery = | ||
| """ | ||
| | SELECT * | ||
|
|
@@ -374,6 +328,12 @@ private[redshift] class RedshiftWriter( | |
| "For save operations you must specify a Redshift table name with the 'dbtable' parameter") | ||
| } | ||
|
|
||
| if (!params.useStagingTable) { | ||
| log.warn("Setting useStagingTable=false is deprecated; instead, we recommend that you " + | ||
| "drop the target table yourself. For more details on this deprecation, see" + | ||
| "https://github.com/databricks/spark-redshift/pull/157") | ||
| } | ||
|
|
||
| val creds: AWSCredentials = | ||
| AWSCredentialsUtils.load(params, sqlContext.sparkContext.hadoopConfiguration) | ||
|
|
||
|
|
@@ -382,19 +342,35 @@ private[redshift] class RedshiftWriter( | |
|
|
||
| Utils.checkThatBucketHasObjectLifecycleConfiguration(params.rootTempDir, s3ClientFactory(creds)) | ||
|
|
||
| // Save the table's rows to S3: | ||
| val manifestUrl = unloadData(sqlContext, data, params.createPerQueryTempDir()) | ||
| val conn = jdbcWrapper.getConnector(params.jdbcDriver, params.jdbcUrl, params.credentials) | ||
|
|
||
| conn.setAutoCommit(false) | ||
| try { | ||
| val tempDir = params.createPerQueryTempDir() | ||
| val manifestUrl = unloadData(sqlContext, data, tempDir) | ||
| if (saveMode == SaveMode.Overwrite && params.useStagingTable) { | ||
| withStagingTable(conn, params.table.get, stagingTable => { | ||
| val updatedParams = MergedParameters(params.parameters.updated("dbtable", stagingTable)) | ||
| doRedshiftLoad(conn, data, saveMode, updatedParams, creds, manifestUrl) | ||
| }) | ||
| } else { | ||
| doRedshiftLoad(conn, data, saveMode, params, creds, manifestUrl) | ||
| val table: TableName = params.table.get | ||
| if (saveMode == SaveMode.Overwrite) { | ||
| // Overwrites must drop the table in case there has been a schema update | ||
| jdbcWrapper.executeInterruptibly(conn.prepareStatement(s"DROP TABLE IF EXISTS $table;")) | ||
| if (!params.useStagingTable) { | ||
| // If we're not using a staging table, commit now so that Redshift doesn't have to | ||
| // maintain a snapshot of the old table during the COPY; this sacrifices atomicity for | ||
| // performance. | ||
| conn.commit() | ||
| } | ||
| } | ||
| log.info(s"Loading new Redshift data to: $table") | ||
| doRedshiftLoad(conn, data, params, creds, manifestUrl) | ||
| conn.commit() | ||
| } catch { | ||
| case NonFatal(e) => | ||
| try { | ||
| log.error("Exception thrown during Redshift load; will roll back transaction", e) | ||
| conn.rollback() | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Might a log.info or log.warn be good here, to inform anyone testing their application code that the load failed and what they're now waiting for is a rollback to finish?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done. |
||
| } catch { | ||
| case NonFatal(e2) => | ||
| log.error("Exception while rolling back transaction", e2) | ||
| } | ||
| throw e | ||
| } finally { | ||
| conn.close() | ||
| } | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -297,20 +297,12 @@ class RedshiftSourceSuite | |
| "usestagingtable" -> "true") | ||
|
|
||
| val expectedCommands = Seq( | ||
| "DROP TABLE IF EXISTS \"PUBLIC\".\"test_table_staging_.*\"".r, | ||
| "CREATE TABLE IF NOT EXISTS \"PUBLIC\".\"test_table_staging_.*\"".r, | ||
| "DELETE FROM \"PUBLIC\".\"test_table_staging_.*\" WHERE id < 100".r, | ||
| "DELETE FROM \"PUBLIC\".\"test_table_staging_.*\" WHERE id > 100".r, | ||
| "DELETE FROM \"PUBLIC\".\"test_table_staging_.*\" WHERE id = -1".r, | ||
| "COPY \"PUBLIC\".\"test_table_staging_.*\"".r, | ||
| """ | ||
| | BEGIN; | ||
| | ALTER TABLE "PUBLIC"\."test_table" RENAME TO "test_table_backup_.*"; | ||
| | ALTER TABLE "PUBLIC"\."test_table_staging_.*" RENAME TO "test_table"; | ||
| | DROP TABLE "PUBLIC"\."test_table_backup_.*"; | ||
| | END; | ||
| """.stripMargin.trim.r, | ||
| "DROP TABLE IF EXISTS \"PUBLIC\"\\.\"test_table_staging_.*\"".r) | ||
| "DROP TABLE IF EXISTS \"PUBLIC\".\"test_table.*\"".r, | ||
| "CREATE TABLE IF NOT EXISTS \"PUBLIC\".\"test_table.*\"".r, | ||
| "DELETE FROM \"PUBLIC\".\"test_table.*\" WHERE id < 100".r, | ||
| "DELETE FROM \"PUBLIC\".\"test_table.*\" WHERE id > 100".r, | ||
| "DELETE FROM \"PUBLIC\".\"test_table.*\" WHERE id = -1".r, | ||
| "COPY \"PUBLIC\".\"test_table.*\"".r) | ||
|
|
||
| source.createRelation(testSqlContext, SaveMode.Overwrite, params, expectedDataDF) | ||
| mockRedshift.verifyThatExpectedQueriesWereIssued(expectedCommands) | ||
|
|
@@ -324,19 +316,11 @@ class RedshiftSourceSuite | |
| "distkey" -> "testint") | ||
|
|
||
| val expectedCommands = Seq( | ||
| "DROP TABLE IF EXISTS \"PUBLIC\"\\.\"test_table_staging_.*\"".r, | ||
| ("CREATE TABLE IF NOT EXISTS \"PUBLIC\"\\.\"test_table_staging.*" + | ||
| "DROP TABLE IF EXISTS \"PUBLIC\"\\.\"test_table.*\"".r, | ||
| ("CREATE TABLE IF NOT EXISTS \"PUBLIC\"\\.\"test_table.*" + | ||
| " DISTSTYLE KEY DISTKEY \\(testint\\).*").r, | ||
| "COPY \"PUBLIC\"\\.\"test_table_staging_.*\"".r, | ||
| "GRANT SELECT ON \"PUBLIC\"\\.\"test_table_staging.+\" TO jeremy".r, | ||
| """ | ||
| | BEGIN; | ||
| | ALTER TABLE "PUBLIC"\."test_table" RENAME TO "test_table_backup_.*"; | ||
| | ALTER TABLE "PUBLIC"\."test_table_staging_.*" RENAME TO "test_table"; | ||
| | DROP TABLE "PUBLIC"\."test_table_backup_.*"; | ||
| | END; | ||
| """.stripMargin.trim.r, | ||
| "DROP TABLE IF EXISTS \"PUBLIC\"\\.\"test_table_staging_.*\"".r) | ||
| "COPY \"PUBLIC\"\\.\"test_table.*\"".r, | ||
| "GRANT SELECT ON \"PUBLIC\"\\.\"test_table\" TO jeremy".r) | ||
|
|
||
| val mockRedshift = new MockRedshift( | ||
| defaultParams("url"), | ||
|
|
@@ -374,6 +358,8 @@ class RedshiftSourceSuite | |
| testSqlContext, df, SaveMode.Append, Parameters.mergeParameters(defaultParams)) | ||
| } | ||
| mockRedshift.verifyThatConnectionsWereClosed() | ||
| mockRedshift.verifyThatCommitWasNotCalled() | ||
| mockRedshift.verifyThatRollbackWasCalled() | ||
| mockRedshift.verifyThatExpectedQueriesWereIssued(Seq.empty) | ||
| } | ||
|
|
||
|
|
@@ -383,23 +369,22 @@ class RedshiftSourceSuite | |
| val mockRedshift = new MockRedshift( | ||
| defaultParams("url"), | ||
| Map(TableName.parseFromEscaped("test_table").toString -> TestUtils.testSchema), | ||
| jdbcQueriesThatShouldFail = Seq("COPY \"PUBLIC\".\"test_table_staging_.*\"".r)) | ||
| jdbcQueriesThatShouldFail = Seq("COPY \"PUBLIC\".\"test_table.*\"".r)) | ||
|
|
||
| val expectedCommands = Seq( | ||
| "DROP TABLE IF EXISTS \"PUBLIC\".\"test_table_staging_.*\"".r, | ||
| "CREATE TABLE IF NOT EXISTS \"PUBLIC\".\"test_table_staging_.*\"".r, | ||
| "COPY \"PUBLIC\".\"test_table_staging_.*\"".r, | ||
| ".*FROM stl_load_errors.*".r, | ||
| "DROP TABLE IF EXISTS \"PUBLIC\".\"test_table_staging_.*\"".r | ||
| "DROP TABLE IF EXISTS \"PUBLIC\".\"test_table.*\"".r, | ||
| "CREATE TABLE IF NOT EXISTS \"PUBLIC\".\"test_table.*\"".r, | ||
| "COPY \"PUBLIC\".\"test_table.*\"".r, | ||
| ".*FROM stl_load_errors.*".r | ||
| ) | ||
|
|
||
| val source = new DefaultSource(mockRedshift.jdbcWrapper, _ => mockS3Client) | ||
| intercept[Exception] { | ||
| source.createRelation(testSqlContext, SaveMode.Overwrite, params, expectedDataDF) | ||
| mockRedshift.verifyThatConnectionsWereClosed() | ||
| mockRedshift.verifyThatExpectedQueriesWereIssued(expectedCommands) | ||
| } | ||
| mockRedshift.verifyThatConnectionsWereClosed() | ||
| mockRedshift.verifyThatCommitWasNotCalled() | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should this also verify that rollback was called?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done. |
||
| mockRedshift.verifyThatRollbackWasCalled() | ||
| mockRedshift.verifyThatExpectedQueriesWereIssued(expectedCommands) | ||
| } | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Does this need to handle InterruptedException? i.e. can this block for a really long time?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It has been my experience that virtually any Redshift command,
COMMITincluded, can block for several minutes under certain circumstances. I think the most likely cause is that the cluster has WLM parameters configured to put the connected client on a limited pool of some type, such that all commands will be queued when all slots are taken by other queries.I have a feeling that means though, if you plan to send a
ROLLBACKcommand in response to the interruption, that will also block for many minutes...There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not sure that it's safe for us to wrap this in order to catch InterruptedException since I don't think that it's safe to call
.rollback()while the other thread is in the middle of executingcommit(). Therefore, I'm going to leave this as-is for now and will revisit later if this turns out to be a problem in practice.