JdbcIO { public static <T> Write<T> write() { return new Write(); } private  JdbcIO(); static Read<T> read(); @Experimental(Kind.SCHEMAS) static ReadRows readRows(); static ReadAll<ParameterT, OutputT> readAll(); static Write<T> write(); static WriteVoid<T> writeVoid();  }
@Test public void testWrite() throws Exception { String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE"); DatabaseTestHelper.createTable(DATA_SOURCE, tableName); try { ArrayList<KV<Integer, String>> data = getDataToWrite(EXPECTED_ROW_COUNT); pipeline.apply(Create.of(data)).apply(getJdbcWrite(tableName)); pipeline.run(); assertRowCount(tableName, EXPECTED_ROW_COUNT); } finally { DatabaseTestHelper.deleteTable(DATA_SOURCE, tableName); } }
@Test public void testWriteWithBackoff() throws Exception { String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_BACKOFF"); DatabaseTestHelper.createTable(DATA_SOURCE, tableName); Connection connection = DATA_SOURCE.getConnection(); Statement lockStatement = connection.createStatement(); lockStatement.execute("ALTER TABLE " + tableName + " LOCKSIZE TABLE"); lockStatement.execute("LOCK TABLE " + tableName + " IN EXCLUSIVE MODE"); connection.setAutoCommit(false); PreparedStatement insertStatement = connection.prepareStatement("insert into " + tableName + " values(?, ?)"); insertStatement.setInt(1, 1); insertStatement.setString(2, "TEST"); insertStatement.execute(); pipeline .apply(Create.of(Collections.singletonList(KV.of(1, "TEST")))) .apply( JdbcIO.<KV<Integer, String>>write() .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION) .withStatement(String.format("insert into %s values(?, ?)", tableName)) .withRetryStrategy( (JdbcIO.RetryStrategy) e -> { return "40XL1" .equals(e.getSQLState()); }) .withPreparedStatementSetter( (element, statement) -> { statement.setInt(1, element.getKey()); statement.setString(2, element.getValue()); })); Thread commitThread = new Thread( () -> { try { Thread.sleep(10000); connection.commit(); } catch (Exception e) { } }); commitThread.start(); pipeline.run(); commitThread.join(); expectedLogs.verifyWarn("Deadlock detected, retrying"); assertRowCount(tableName, 2); }
@Test public void testWriteWithoutPreparedStatement() throws Exception { final int rowsToAdd = 10; Schema.Builder schemaBuilder = Schema.builder(); schemaBuilder.addField(Schema.Field.of("column_boolean", Schema.FieldType.BOOLEAN)); schemaBuilder.addField(Schema.Field.of("column_string", Schema.FieldType.STRING)); schemaBuilder.addField(Schema.Field.of("column_int", Schema.FieldType.INT32)); schemaBuilder.addField(Schema.Field.of("column_long", Schema.FieldType.INT64)); schemaBuilder.addField(Schema.Field.of("column_float", Schema.FieldType.FLOAT)); schemaBuilder.addField(Schema.Field.of("column_double", Schema.FieldType.DOUBLE)); schemaBuilder.addField(Schema.Field.of("column_bigdecimal", Schema.FieldType.DECIMAL)); schemaBuilder.addField(Schema.Field.of("column_date", LogicalTypes.JDBC_DATE_TYPE)); schemaBuilder.addField(Schema.Field.of("column_time", LogicalTypes.JDBC_TIME_TYPE)); schemaBuilder.addField( Schema.Field.of("column_timestamptz", LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE)); schemaBuilder.addField(Schema.Field.of("column_timestamp", Schema.FieldType.DATETIME)); schemaBuilder.addField(Schema.Field.of("column_short", Schema.FieldType.INT16)); Schema schema = schemaBuilder.build(); String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_PS"); StringBuilder stmt = new StringBuilder("CREATE TABLE "); stmt.append(tableName); stmt.append(" ("); stmt.append("column_boolean BOOLEAN,"); stmt.append("column_string VARCHAR(254),"); stmt.append("column_int INTEGER,"); stmt.append("column_long BIGINT,"); stmt.append("column_float REAL,"); stmt.append("column_double DOUBLE PRECISION,"); stmt.append("column_bigdecimal DECIMAL(13,0),"); stmt.append("column_date DATE,"); stmt.append("column_time TIME,"); stmt.append("column_timestamptz TIMESTAMP,"); stmt.append("column_timestamp TIMESTAMP,"); stmt.append("column_short SMALLINT"); stmt.append(" )"); DatabaseTestHelper.createTableWithStatement(DATA_SOURCE, stmt.toString()); try { ArrayList<Row> data = getRowsToWrite(rowsToAdd, schema); pipeline .apply(Create.of(data)) .setRowSchema(schema) .apply( JdbcIO.<Row>write() .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION) .withBatchSize(10L) .withTable(tableName)); pipeline.run(); assertRowCount(tableName, rowsToAdd); } finally { DatabaseTestHelper.deleteTable(DATA_SOURCE, tableName); } }
@Test public void testWriteWithoutPsWithNonNullableTableField() throws Exception { final int rowsToAdd = 10; Schema.Builder schemaBuilder = Schema.builder(); schemaBuilder.addField(Schema.Field.of("column_boolean", Schema.FieldType.BOOLEAN)); schemaBuilder.addField(Schema.Field.of("column_string", Schema.FieldType.STRING)); Schema schema = schemaBuilder.build(); String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE"); StringBuilder stmt = new StringBuilder("CREATE TABLE "); stmt.append(tableName); stmt.append(" ("); stmt.append("column_boolean BOOLEAN,"); stmt.append("column_int INTEGER NOT NULL"); stmt.append(" )"); DatabaseTestHelper.createTableWithStatement(DATA_SOURCE, stmt.toString()); try { ArrayList<Row> data = getRowsToWrite(rowsToAdd, schema); pipeline .apply(Create.of(data)) .setRowSchema(schema) .apply( JdbcIO.<Row>write() .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION) .withBatchSize(10L) .withTable(tableName)); pipeline.run(); } finally { DatabaseTestHelper.deleteTable(DATA_SOURCE, tableName); thrown.expect(RuntimeException.class); } }
@Test public void testWriteWithoutPreparedStatementAndNonRowType() throws Exception { final int rowsToAdd = 10; String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_PS_NON_ROW"); DatabaseTestHelper.createTableForRowWithSchema(DATA_SOURCE, tableName); try { List<RowWithSchema> data = getRowsWithSchemaToWrite(rowsToAdd); pipeline .apply(Create.of(data)) .apply( JdbcIO.<RowWithSchema>write() .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION) .withBatchSize(10L) .withTable(tableName)); pipeline.run(); assertRowCount(tableName, rowsToAdd); } finally { DatabaseTestHelper.deleteTable(DATA_SOURCE, tableName); } }
@Test public void testWriteWithEmptyPCollection() { pipeline .apply(Create.empty(KvCoder.of(VarIntCoder.of(), StringUtf8Coder.of()))) .apply( JdbcIO.<KV<Integer, String>>write() .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION) .withStatement("insert into BEAM values(?, ?)") .withPreparedStatementSetter( (element, statement) -> { statement.setInt(1, element.getKey()); statement.setString(2, element.getValue()); })); pipeline.run(); }
@Test public void testCustomFluentBackOffConfiguration() throws Exception { String tableName = DatabaseTestHelper.getTestTableName("UT_FLUENT_BACKOFF"); DatabaseTestHelper.createTable(DATA_SOURCE, tableName); Connection connection = DATA_SOURCE.getConnection(); Statement lockStatement = connection.createStatement(); lockStatement.execute("ALTER TABLE " + tableName + " LOCKSIZE TABLE"); lockStatement.execute("LOCK TABLE " + tableName + " IN EXCLUSIVE MODE"); connection.setAutoCommit(false); PreparedStatement insertStatement = connection.prepareStatement("insert into " + tableName + " values(?, ?)"); insertStatement.setInt(1, 1); insertStatement.setString(2, "TEST"); insertStatement.execute(); pipeline .apply(Create.of(Collections.singletonList(KV.of(1, "TEST")))) .apply( JdbcIO.<KV<Integer, String>>write() .withDataSourceConfiguration(DATA_SOURCE_CONFIGURATION) .withStatement(String.format("insert into %s values(?, ?)", tableName)) .withRetryStrategy( (JdbcIO.RetryStrategy) e -> { return "40XL1" .equals(e.getSQLState()); }) .withRetryConfiguration( JdbcIO.RetryConfiguration.create(2, null, Duration.standardSeconds(1))) .withPreparedStatementSetter( (element, statement) -> { statement.setInt(1, element.getKey()); statement.setString(2, element.getValue()); })); PipelineExecutionException exception = assertThrows( PipelineExecutionException.class, () -> { pipeline.run().waitUntilFinish(); }); connection.commit(); assertThat( exception.getMessage(), containsString( "java.sql.BatchUpdateException: A lock could not be obtained within the time requested")); expectedLogs.verifyLogRecords( new TypeSafeMatcher<Iterable<LogRecord>>() { @Override public void describeTo(Description description) {} @Override protected boolean matchesSafely(Iterable<LogRecord> logRecords) { int count = 0; for (LogRecord logRecord : logRecords) { if (logRecord.getMessage().contains("Deadlock detected, retrying")) { count += 1; } } return count == 3; } }); assertRowCount(tableName, 1); }
@Test public void testWrite() throws Exception { final long rowsToAdd = 1000L; String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE"); DatabaseTestHelper.createTable(dataSource, tableName); try { ArrayList<KV<Integer, String>> data = getDataToWrite(rowsToAdd); pipeline.apply(Create.of(data)).apply(getJdbcWrite(tableName)); pipeline.run(); assertRowCount(tableName, EXPECTED_ROW_COUNT); } finally { DatabaseTestHelper.deleteTable(dataSource, tableName); } }
@Test public void testWriteWithBackoff() throws Exception { String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_BACKOFF"); DatabaseTestHelper.createTable(dataSource, tableName); Connection connection = dataSource.getConnection(); Statement lockStatement = connection.createStatement(); lockStatement.execute("ALTER TABLE " + tableName + " LOCKSIZE TABLE"); lockStatement.execute("LOCK TABLE " + tableName + " IN EXCLUSIVE MODE"); connection.setAutoCommit(false); PreparedStatement insertStatement = connection.prepareStatement("insert into " + tableName + " values(?, ?)"); insertStatement.setInt(1, 1); insertStatement.setString(2, "TEST"); insertStatement.execute(); pipeline .apply(Create.of(Collections.singletonList(KV.of(1, "TEST")))) .apply( JdbcIO.<KV<Integer, String>>write() .withDataSourceConfiguration( JdbcIO.DataSourceConfiguration.create( "org.apache.derby.jdbc.ClientDriver", "jdbc:derby: .withStatement(String.format("insert into %s values(?, ?)", tableName)) .withRetryStrategy( (JdbcIO.RetryStrategy) e -> { return "XJ208" .equals(e.getSQLState()); }) .withPreparedStatementSetter( (element, statement) -> { statement.setInt(1, element.getKey()); statement.setString(2, element.getValue()); })); Thread commitThread = new Thread( () -> { try { Thread.sleep(10000); connection.commit(); } catch (Exception e) { } }); commitThread.start(); pipeline.run(); commitThread.join(); expectedLogs.verifyWarn("Deadlock detected, retrying"); assertRowCount(tableName, 2); }
@Test public void testWriteWithoutPreparedStatement() throws Exception { final int rowsToAdd = 10; Schema.Builder schemaBuilder = Schema.builder(); schemaBuilder.addField(Schema.Field.of("column_boolean", Schema.FieldType.BOOLEAN)); schemaBuilder.addField(Schema.Field.of("column_string", Schema.FieldType.STRING)); schemaBuilder.addField(Schema.Field.of("column_int", Schema.FieldType.INT32)); schemaBuilder.addField(Schema.Field.of("column_long", Schema.FieldType.INT64)); schemaBuilder.addField(Schema.Field.of("column_float", Schema.FieldType.FLOAT)); schemaBuilder.addField(Schema.Field.of("column_double", Schema.FieldType.DOUBLE)); schemaBuilder.addField(Schema.Field.of("column_bigdecimal", Schema.FieldType.DECIMAL)); schemaBuilder.addField(Schema.Field.of("column_date", LogicalTypes.JDBC_DATE_TYPE)); schemaBuilder.addField(Schema.Field.of("column_time", LogicalTypes.JDBC_TIME_TYPE)); schemaBuilder.addField( Schema.Field.of("column_timestamptz", LogicalTypes.JDBC_TIMESTAMP_WITH_TIMEZONE_TYPE)); schemaBuilder.addField(Schema.Field.of("column_timestamp", Schema.FieldType.DATETIME)); schemaBuilder.addField(Schema.Field.of("column_short", Schema.FieldType.INT16)); Schema schema = schemaBuilder.build(); String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_PS"); StringBuilder stmt = new StringBuilder("CREATE TABLE "); stmt.append(tableName); stmt.append(" ("); stmt.append("column_boolean BOOLEAN,"); stmt.append("column_string VARCHAR(254),"); stmt.append("column_int INTEGER,"); stmt.append("column_long BIGINT,"); stmt.append("column_float REAL,"); stmt.append("column_double DOUBLE PRECISION,"); stmt.append("column_bigdecimal DECIMAL(13,0),"); stmt.append("column_date DATE,"); stmt.append("column_time TIME,"); stmt.append("column_timestamptz TIMESTAMP,"); stmt.append("column_timestamp TIMESTAMP,"); stmt.append("column_short SMALLINT"); stmt.append(" )"); DatabaseTestHelper.createTableWithStatement(dataSource, stmt.toString()); try { ArrayList<Row> data = getRowsToWrite(rowsToAdd, schema); pipeline .apply(Create.of(data)) .setRowSchema(schema) .apply( JdbcIO.<Row>write() .withDataSourceConfiguration( JdbcIO.DataSourceConfiguration.create( "org.apache.derby.jdbc.ClientDriver", "jdbc:derby: .withBatchSize(10L) .withTable(tableName)); pipeline.run(); assertRowCount(tableName, rowsToAdd); } finally { DatabaseTestHelper.deleteTable(dataSource, tableName); } }
@Test public void testWriteWithoutPsWithNonNullableTableField() throws Exception { final int rowsToAdd = 10; Schema.Builder schemaBuilder = Schema.builder(); schemaBuilder.addField(Schema.Field.of("column_boolean", Schema.FieldType.BOOLEAN)); schemaBuilder.addField(Schema.Field.of("column_string", Schema.FieldType.STRING)); Schema schema = schemaBuilder.build(); String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE"); StringBuilder stmt = new StringBuilder("CREATE TABLE "); stmt.append(tableName); stmt.append(" ("); stmt.append("column_boolean BOOLEAN,"); stmt.append("column_int INTEGER NOT NULL"); stmt.append(" )"); DatabaseTestHelper.createTableWithStatement(dataSource, stmt.toString()); try { ArrayList<Row> data = getRowsToWrite(rowsToAdd, schema); pipeline .apply(Create.of(data)) .setRowSchema(schema) .apply( JdbcIO.<Row>write() .withDataSourceConfiguration( JdbcIO.DataSourceConfiguration.create( "org.apache.derby.jdbc.ClientDriver", "jdbc:derby: .withBatchSize(10L) .withTable(tableName)); pipeline.run(); } finally { DatabaseTestHelper.deleteTable(dataSource, tableName); thrown.expect(RuntimeException.class); } }
@Test public void testWriteWithoutPreparedStatementAndNonRowType() throws Exception { final int rowsToAdd = 10; String tableName = DatabaseTestHelper.getTestTableName("UT_WRITE_PS_NON_ROW"); DatabaseTestHelper.createTableForRowWithSchema(dataSource, tableName); try { List<RowWithSchema> data = getRowsWithSchemaToWrite(rowsToAdd); pipeline .apply(Create.of(data)) .apply( JdbcIO.<RowWithSchema>write() .withDataSourceConfiguration( JdbcIO.DataSourceConfiguration.create( "org.apache.derby.jdbc.ClientDriver", "jdbc:derby: .withBatchSize(10L) .withTable(tableName)); pipeline.run(); assertRowCount(tableName, rowsToAdd); } finally { DatabaseTestHelper.deleteTable(dataSource, tableName); } }
@Test public void testWriteWithEmptyPCollection() { pipeline .apply(Create.empty(KvCoder.of(VarIntCoder.of(), StringUtf8Coder.of()))) .apply( JdbcIO.<KV<Integer, String>>write() .withDataSourceConfiguration( JdbcIO.DataSourceConfiguration.create( "org.apache.derby.jdbc.ClientDriver", "jdbc:derby: .withStatement("insert into BEAM values(?, ?)") .withPreparedStatementSetter( (element, statement) -> { statement.setInt(1, element.getKey()); statement.setString(2, element.getValue()); })); pipeline.run(); }