public void testPublish() throws IOException, SQLException { publisher.publish(workUnitStates); InOrder inOrder = inOrder(conn, commands, workUnitState); inOrder.verify(conn, times(1)).setAutoCommit(false); inOrder.verify(commands, times(1)).copyTable(database, stagingTable, destinationTable); inOrder.verify(workUnitState, times(1)).setWorkingState(WorkUnitState.WorkingState.COMMITTED); inOrder.verify(conn, times(1)).commit(); inOrder.verify(conn, times(1)).close(); verify(commands, never()).deleteAll(database, destinationTable); }
public void userCreatedStagingTable() throws SQLException { this.state.setProp(ConfigurationKeys.WRITER_STAGING_TABLE, STAGING_TABLE); when(this.commands.isEmpty(DB, STAGING_TABLE)).thenReturn(Boolean.TRUE); this.initializer.initialize(); Assert.assertEquals(STAGING_TABLE, this.workUnit.getProp(ConfigurationKeys.WRITER_STAGING_TABLE)); verify(this.commands, never()).createTableStructure(anyString(), anyString(), anyString()); verify(this.commands, never()).truncate(anyString(), anyString()); verify(this.commands, never()).drop(anyString(), anyString()); }
/** * Flushes JdbcWriterCommands and commit. * {@inheritDoc} * @see org.apache.gobblin.writer.DataWriter#commit() */ @Override public void commit() throws IOException { try { LOG.info("Flushing pending insert."); this.commands.flush(); LOG.info("Commiting transaction."); this.conn.commit(); } catch (Exception e) { this.failed = true; throw new RuntimeException(e); } }
/** * Drop table if it's created by this instance. * Truncate staging tables passed by user. * {@inheritDoc} * @see org.apache.gobblin.Initializer#close() */ @Override public void close() { LOG.info("Closing " + this.getClass().getSimpleName()); try (Connection conn = createConnection()) { JdbcWriterCommands commands = createJdbcWriterCommands(conn); if (!this.createdStagingTables.isEmpty()) { for (String stagingTable : this.createdStagingTables) { LOG.info("Dropping staging table " + this.createdStagingTables); commands.drop(database, stagingTable); } } if (this.userCreatedStagingTable != null) { LOG.info("Truncating staging table " + this.userCreatedStagingTable); commands.truncate(database, this.userCreatedStagingTable); } } catch (SQLException e) { throw new RuntimeException("Failed to close", e); } }
LOG.info("Staging table " + tmp + " does not exist. Creating."); try { commands.createTableStructure(database, destinationTable, tmp); LOG.info("Test if staging table can be dropped. Test by dropping and Creating staging table."); commands.drop(database, tmp); commands.createTableStructure(database, destinationTable, tmp); stagingTable = tmp; break;
public void skipStagingTableTruncateDestTable() throws SQLException { this.state.setProp(ConfigurationKeys.JOB_COMMIT_POLICY_KEY, "partial"); this.state.setProp(ConfigurationKeys.PUBLISH_DATA_AT_JOB_LEVEL, Boolean.toString(false)); this.state.setProp(JdbcPublisher.JDBC_PUBLISHER_REPLACE_FINAL_TABLE, Boolean.toString(true)); this.initializer.initialize(); Assert.assertEquals(DEST_TABLE, this.workUnit.getProp(ConfigurationKeys.WRITER_STAGING_TABLE)); verify(this.commands, never()).createTableStructure(anyString(), anyString(), anyString()); InOrder inOrder = inOrder(this.commands); inOrder.verify(this.commands, times(1)).truncate(DB, DEST_TABLE); this.initializer.close(); inOrder.verify(this.commands, never()).truncate(anyString(), anyString()); verify(this.commands, never()).drop(anyString(), anyString()); }
LOG.info("User chose to replace final table " + publishTable + " on branch " + this.branchId + " workunit " + i); commands.truncate(database, publishTable); ConfigurationKeys.WRITER_TRUNCATE_STAGING_TABLE, this.branches, this.branchId), false)) { LOG.info("Truncating staging table " + stagingTable + " as requested."); commands.truncate(database, stagingTable); if (!commands.isEmpty(database, stagingTable)) { LOG.error("Staging table " + stagingTable + " is not empty. Failing."); throw new IllegalArgumentException("Staging table " + stagingTable + " should be empty.");
@Test public void writeAndCommitTest() throws SQLException, IOException { final String database = "db"; final String table = "users"; final int writeCount = 25; JdbcWriterCommands writerCommands = mock(JdbcWriterCommands.class); Connection conn = mock(Connection.class); try (JdbcWriter writer = new JdbcWriter(writerCommands, new State(), database, table, conn)) { for(int i = 0; i < writeCount; i++) { writer.write(null); } writer.commit(); Assert.assertEquals(writer.recordsWritten(), writeCount); } verify(writerCommands, times(writeCount)).insert(anyString(), anyString(), any(JdbcEntryData.class)); verify(conn, times(1)).commit(); verify(conn, never()).rollback(); verify(writerCommands, times(1)).flush(); verify(conn, times(1)).close(); }
/** * AvroToJdbcEntryConverter list of date columns existing in the table. As we don't want each converter * making a connection against database to get the same information. Here, ConverterInitializer will * retrieve it and store it into WorkUnit so that AvroToJdbcEntryConverter will use it later. * * {@inheritDoc} * @see org.apache.gobblin.initializer.Initializer#initialize() */ @Override public void initialize() { String table = Preconditions.checkNotNull(this.state.getProp(ForkOperatorUtils .getPropertyNameForBranch(JdbcPublisher.JDBC_PUBLISHER_FINAL_TABLE_NAME, this.branches, this.branchId))); String db = Preconditions.checkNotNull(this.state.getProp(ForkOperatorUtils .getPropertyNameForBranch(JdbcPublisher.JDBC_PUBLISHER_DATABASE_NAME, this.branches, this.branchId))); try (Connection conn = createConnection()) { JdbcWriterCommands commands = this.jdbcWriterCommandsFactory.newInstance(this.state, conn); Map<String, JdbcType> dateColumnMapping = commands.retrieveDateColumns(db, table); LOG.info("Date column mapping: " + dateColumnMapping); final String dateFieldsKey = ForkOperatorUtils.getPropertyNameForBranch( AvroToJdbcEntryConverter.CONVERTER_AVRO_JDBC_DATE_FIELDS, this.branches, this.branchId); for (WorkUnit wu : this.workUnits) { wu.setProp(dateFieldsKey, new Gson().toJson(dateColumnMapping)); } } catch (SQLException e) { throw new RuntimeException(e); } }
/** * Invokes JdbcWriterCommands.insert * {@inheritDoc} * @see org.apache.gobblin.writer.DataWriter#write(java.lang.Object) */ @Override public void write(JdbcEntryData record) throws IOException { if (LOG.isDebugEnabled()) { LOG.debug("Writing " + record); } try { this.commands.insert(this.databaseName, this.tableName, record); this.recordWrittenCount++; } catch (Exception e) { this.failed = true; throw new RuntimeException(e); } }
public JdbcWriter(JdbcWriterBuilder builder) { this.state = builder.destination.getProperties(); this.state.setProp(ConfigurationKeys.FORK_BRANCH_ID_KEY, Integer.toString(builder.branch)); String databaseTableKey = ForkOperatorUtils.getPropertyNameForBranch(JdbcPublisher.JDBC_PUBLISHER_DATABASE_NAME, builder.branches, builder.branch); this.databaseName = Preconditions.checkNotNull(this.state.getProp(databaseTableKey), "Staging table is missing with key " + databaseTableKey); String stagingTableKey = ForkOperatorUtils.getPropertyNameForBranch(ConfigurationKeys.WRITER_STAGING_TABLE, builder.branches, builder.branch); this.tableName = Preconditions.checkNotNull(this.state.getProp(stagingTableKey), "Staging table is missing with key " + stagingTableKey); try { this.conn = createConnection(); this.commands = new JdbcWriterCommandsFactory().newInstance(this.state, this.conn); this.commands.setConnectionParameters(this.state.getProperties(), this.conn); } catch (SQLException e) { throw new RuntimeException(e); } }
public void skipStagingTable() throws SQLException { this.state.setProp(ConfigurationKeys.JOB_COMMIT_POLICY_KEY, "partial"); this.state.setProp(ConfigurationKeys.PUBLISH_DATA_AT_JOB_LEVEL, Boolean.toString(false)); this.initializer.initialize(); this.initializer.close(); Assert.assertEquals(DEST_TABLE, this.workUnit.getProp(ConfigurationKeys.WRITER_STAGING_TABLE)); verify(this.commands, never()).createTableStructure(anyString(), anyString(), anyString()); verify(this.commands, never()).truncate(anyString(), anyString()); verify(this.commands, never()).drop(anyString(), anyString()); }
/** * Drop table if it's created by this instance. * Truncate staging tables passed by user. * {@inheritDoc} * @see org.apache.gobblin.Initializer#close() */ @Override public void close() { LOG.info("Closing " + this.getClass().getSimpleName()); try (Connection conn = createConnection()) { JdbcWriterCommands commands = createJdbcWriterCommands(conn); if (!this.createdStagingTables.isEmpty()) { for (String stagingTable : this.createdStagingTables) { LOG.info("Dropping staging table " + this.createdStagingTables); commands.drop(database, stagingTable); } } if (this.userCreatedStagingTable != null) { LOG.info("Truncating staging table " + this.userCreatedStagingTable); commands.truncate(database, this.userCreatedStagingTable); } } catch (SQLException e) { throw new RuntimeException("Failed to close", e); } }
LOG.info("Staging table " + tmp + " does not exist. Creating."); try { commands.createTableStructure(database, destinationTable, tmp); LOG.info("Test if staging table can be dropped. Test by dropping and Creating staging table."); commands.drop(database, tmp); commands.createTableStructure(database, destinationTable, tmp); stagingTable = tmp; break;
LOG.info("User chose to replace final table " + publishTable + " on branch " + this.branchId + " workunit " + i); commands.truncate(database, publishTable); ConfigurationKeys.WRITER_TRUNCATE_STAGING_TABLE, this.branches, this.branchId), false)) { LOG.info("Truncating staging table " + stagingTable + " as requested."); commands.truncate(database, stagingTable); if (!commands.isEmpty(database, stagingTable)) { LOG.error("Staging table " + stagingTable + " is not empty. Failing."); throw new IllegalArgumentException("Staging table " + stagingTable + " should be empty.");
when(mockWriterCommands.retrieveDateColumns(db, table)).thenReturn(dateColums);
@Test public void writeFailRollbackTest() throws SQLException, IOException { final String database = "db"; final String table = "users"; JdbcWriterCommands writerCommands = mock(JdbcWriterCommands.class); Connection conn = mock(Connection.class); doThrow(RuntimeException.class).when(writerCommands).insert(anyString(), anyString(), any(JdbcEntryData.class)); JdbcWriter writer = new JdbcWriter(writerCommands, new State(), database, table, conn); try { writer.write(null); Assert.fail("Test case didn't throw Exception."); } catch (RuntimeException e) { Assert.assertTrue(e instanceof RuntimeException); } writer.close(); verify(writerCommands, times(1)).insert(anyString(), anyString(), any(JdbcEntryData.class)); verify(conn, times(1)).rollback(); verify(conn, never()).commit(); verify(conn, times(1)).close(); Assert.assertEquals(writer.recordsWritten(), 0L); } }
public JdbcWriter(JdbcWriterBuilder builder) { this.state = builder.destination.getProperties(); this.state.setProp(ConfigurationKeys.FORK_BRANCH_ID_KEY, Integer.toString(builder.branch)); String databaseTableKey = ForkOperatorUtils.getPropertyNameForBranch(JdbcPublisher.JDBC_PUBLISHER_DATABASE_NAME, builder.branches, builder.branch); this.databaseName = Preconditions.checkNotNull(this.state.getProp(databaseTableKey), "Staging table is missing with key " + databaseTableKey); String stagingTableKey = ForkOperatorUtils.getPropertyNameForBranch(ConfigurationKeys.WRITER_STAGING_TABLE, builder.branches, builder.branch); this.tableName = Preconditions.checkNotNull(this.state.getProp(stagingTableKey), "Staging table is missing with key " + stagingTableKey); try { this.conn = createConnection(); this.commands = new JdbcWriterCommandsFactory().newInstance(this.state, this.conn); this.commands.setConnectionParameters(this.state.getProperties(), this.conn); } catch (SQLException e) { throw new RuntimeException(e); } }
public void userCreatedStagingTableTruncate() throws SQLException { this.state.setProp(ConfigurationKeys.WRITER_STAGING_TABLE, STAGING_TABLE); this.state.setProp(ConfigurationKeys.WRITER_TRUNCATE_STAGING_TABLE, Boolean.toString(true)); when(this.commands.isEmpty(DB, STAGING_TABLE)).thenReturn(Boolean.TRUE); this.initializer.initialize(); Assert.assertEquals(STAGING_TABLE, this.workUnit.getProp(ConfigurationKeys.WRITER_STAGING_TABLE)); InOrder inOrder = inOrder(this.commands); inOrder.verify(this.commands, times(1)).truncate(DB, STAGING_TABLE); this.initializer.close(); inOrder.verify(this.commands, times(1)).truncate(DB, STAGING_TABLE); verify(this.commands, never()).createTableStructure(anyString(), anyString(), anyString()); verify(this.commands, never()).drop(anyString(), anyString()); }
public void testPublishReplaceOutput() throws IOException, SQLException { state.setProp(JdbcPublisher.JDBC_PUBLISHER_REPLACE_FINAL_TABLE, Boolean.toString(true)); publisher.publish(workUnitStates); InOrder inOrder = inOrder(conn, commands, workUnitState); inOrder.verify(conn, times(1)).setAutoCommit(false); inOrder.verify(commands, times(1)).deleteAll(database, destinationTable); inOrder.verify(commands, times(1)).copyTable(database, stagingTable, destinationTable); inOrder.verify(workUnitState, times(1)).setWorkingState(WorkUnitState.WorkingState.COMMITTED); inOrder.verify(conn, times(1)).commit(); inOrder.verify(conn, times(1)).close(); }