Examples with JdbcWriterCommands org.apache.gobblin.writer.commands.JdbcWriterCommands used on opensource projects

Example 6 with JdbcWriterCommands

use of org.apache.gobblin.writer.commands.JdbcWriterCommands in project incubator-gobblin by apache.

the class AvroToJdbcEntryConverterTest method testFieldNameConversion.

@Test
public void testFieldNameConversion() throws IOException, SchemaConversionException, SQLException {
    Map<String, JdbcType> dateColums = new HashMap<>();
    dateColums.put("last_updated", JdbcType.TIMESTAMP);
    final String db = "db";
    final String table = "users";
    JdbcWriterCommands mockWriterCommands = mock(JdbcWriterCommands.class);
    when(mockWriterCommands.retrieveDateColumns(db, table)).thenReturn(dateColums);
    JdbcWriterCommandsFactory factory = mock(JdbcWriterCommandsFactory.class);
    when(factory.newInstance(any(State.class), any(Connection.class))).thenReturn(mockWriterCommands);
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.appendToListProp(JdbcPublisher.JDBC_PUBLISHER_FINAL_TABLE_NAME, table);
    String fieldPairJson = "{\"userId\":\"user_id\" , \"memberId\":\"member_id\" , \"businessUnit\":\"business_unit\", \"geoRegion\":\"geo_region\", \"superRegion\":\"super_region\", \"subRegion\":\"sub_region\"}";
    workUnitState.appendToListProp(ConfigurationKeys.CONVERTER_AVRO_JDBC_ENTRY_FIELDS_PAIRS, fieldPairJson);
    workUnitState.appendToListProp(ConfigurationKeys.WRITER_DESTINATION_TYPE_KEY, DestinationType.MYSQL.name());
    AvroToJdbcEntryConverter converter = new AvroToJdbcEntryConverter(workUnitState);
    Schema inputSchema = new Schema.Parser().parse(getClass().getResourceAsStream("/converter/user.avsc"));
    List<JdbcEntryMetaDatum> jdbcEntryMetaData = new ArrayList<>();
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("user_id", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("member_id", JdbcType.BIGINT));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("business_unit", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("level", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("geo_region", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("super_region", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("sub_region", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("currency", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("segment", JdbcType.VARCHAR));
    jdbcEntryMetaData.add(new JdbcEntryMetaDatum("vertical", JdbcType.VARCHAR));
    JdbcEntrySchema expected = new JdbcEntrySchema(jdbcEntryMetaData);
    Map<String, JdbcType> dateColumnMapping = Maps.newHashMap();
    workUnitState.appendToListProp(AvroToJdbcEntryConverter.CONVERTER_AVRO_JDBC_DATE_FIELDS, new Gson().toJson(dateColumnMapping));
    JdbcEntrySchema actual = converter.convertSchema(inputSchema, workUnitState);
    Assert.assertEquals(expected, actual);
}

Example 7 with JdbcWriterCommands

use of org.apache.gobblin.writer.commands.JdbcWriterCommands in project incubator-gobblin by apache.

the class JdbcWriterTest method writeFailRollbackTest.

@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);
}

Example 8 with JdbcWriterCommands

use of org.apache.gobblin.writer.commands.JdbcWriterCommands in project incubator-gobblin by apache.

the class JdbcWriterTest method writeAndCommitTest.

@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();
}

Example 9 with JdbcWriterCommands

use of org.apache.gobblin.writer.commands.JdbcWriterCommands in project incubator-gobblin by apache.

the class JdbcPublisher method publishData.

/**
 * 1. Truncate destination table if requested
 * 2. Move data from staging to destination
 * 3. Update Workunit state
 *
 * TODO: Research on running this in parallel. While testing publishing it in parallel, it turns out delete all from the table locks the table
 * so that copying table threads wait until transaction lock times out and throwing exception(MySQL). Is there a way to avoid this?
 *
 * {@inheritDoc}
 * @see org.apache.gobblin.publisher.DataPublisher#publishData(java.util.Collection)
 */
@Override
public void publishData(Collection<? extends WorkUnitState> states) throws IOException {
    LOG.info("Start publishing data");
    int branches = this.state.getPropAsInt(ConfigurationKeys.FORK_BRANCHES_KEY, 1);
    Set<String> emptiedDestTables = Sets.newHashSet();
    final Connection conn = createConnection();
    final JdbcWriterCommands commands = this.jdbcWriterCommandsFactory.newInstance(this.state, conn);
    try {
        conn.setAutoCommit(false);
        for (int i = 0; i < branches; i++) {
            final String destinationTable = this.state.getProp(ForkOperatorUtils.getPropertyNameForBranch(JDBC_PUBLISHER_FINAL_TABLE_NAME, branches, i));
            final String databaseName = this.state.getProp(ForkOperatorUtils.getPropertyNameForBranch(JDBC_PUBLISHER_DATABASE_NAME, branches, i));
            Preconditions.checkNotNull(destinationTable);
            if (this.state.getPropAsBoolean(ForkOperatorUtils.getPropertyNameForBranch(JDBC_PUBLISHER_REPLACE_FINAL_TABLE, branches, i), false) && !emptiedDestTables.contains(destinationTable)) {
                LOG.info("Deleting table " + destinationTable);
                commands.deleteAll(databaseName, destinationTable);
                emptiedDestTables.add(destinationTable);
            }
            Map<String, List<WorkUnitState>> stagingTables = getStagingTables(states, branches, i);
            for (Map.Entry<String, List<WorkUnitState>> entry : stagingTables.entrySet()) {
                String stagingTable = entry.getKey();
                LOG.info("Copying data from staging table " + stagingTable + " into destination table " + destinationTable);
                commands.copyTable(databaseName, stagingTable, destinationTable);
                for (WorkUnitState workUnitState : entry.getValue()) {
                    workUnitState.setWorkingState(WorkUnitState.WorkingState.COMMITTED);
                }
            }
        }
        LOG.info("Commit publish data");
        conn.commit();
    } catch (Exception e) {
        try {
            LOG.error("Failed publishing. Rolling back.");
            conn.rollback();
        } catch (SQLException se) {
            LOG.error("Failed rolling back.", se);
        }
        throw new RuntimeException("Failed publishing", e);
    } finally {
        try {
            conn.close();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }
}