DbStateBackend extends AbstractStateBackend { @Override public void initializeForJob(final Environment env, String operatorIdentifier, TypeSerializer<?> keySerializer) throws Exception { super.initializeForJob(env, operatorIdentifier, keySerializer); this.operatorIdentifier = operatorIdentifier; this.rnd = new Random(); this.env = env; this.jobId = env.getJobID().toString().substring(0, 16); connections = dbConfig.createShardedConnection(); connections.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED); if (nonPartitionedStateBackend == null) { insertStatement = retry(new Callable<PreparedStatement>() { public PreparedStatement call() throws SQLException { dbAdapter.createCheckpointsTable(jobId, getConnections().getFirst()); return dbAdapter.prepareCheckpointInsert(jobId, getConnections().getFirst()); } }, numSqlRetries, sqlRetrySleep); } else { nonPartitionedStateBackend.initializeForJob(env, operatorIdentifier, keySerializer); } if (LOG.isDebugEnabled()) { LOG.debug("Database state backend successfully initialized"); } } DbStateBackend(DbBackendConfig backendConfig);  DbStateBackend(DbBackendConfig backendConfig, AbstractStateBackend backend); ShardedConnection getConnections(); boolean isInitialized(); Environment getEnvironment(); DbBackendConfig getConfiguration(); @Override StateHandle<S> checkpointStateSerializable(final S state, final long checkpointID,
			final long timestamp); @Override CheckpointStateOutputStream createCheckpointStateOutputStream(long checkpointID, long timestamp); @Override void initializeForJob(final Environment env,
		String operatorIdentifier,
		TypeSerializer<?> keySerializer); @Override void close(); @Override void disposeAllStateForCurrentJob();  }
@Test @SuppressWarnings("unchecked,rawtypes") public void testListState() { File tempDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()); try { FsStateBackend fileBackend = new FsStateBackend(localFileUri(tempDir)); DbStateBackend backend = new DbStateBackend(conf, fileBackend); Environment env = new DummyEnvironment("test", 2, 0); backend.initializeForJob(env, "dummy_test_kv_list", IntSerializer.INSTANCE); ListStateDescriptor<String> kvId = new ListStateDescriptor<>("id", StringSerializer.INSTANCE); ListState<String> state = backend.getPartitionedState(null, VoidSerializer.INSTANCE, kvId); @SuppressWarnings("unchecked") KvState<Integer, Void, ListState<String>, ListStateDescriptor<String>, DbStateBackend> kv = (KvState<Integer, Void, ListState<String>, ListStateDescriptor<String>, DbStateBackend>) state; Joiner joiner = Joiner.on(","); kv.setCurrentKey(1); assertEquals("", joiner.join(state.get())); state.add("1"); kv.setCurrentKey(2); assertEquals("", joiner.join(state.get())); state.add("2"); kv.setCurrentKey(1); assertEquals("1", joiner.join(state.get())); KvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, DbStateBackend> snapshot1 = kv.snapshot(682375462378L, 2); kv.setCurrentKey(1); state.add("u1"); kv.setCurrentKey(2); state.add("u2"); kv.setCurrentKey(3); state.add("u3"); KvStateSnapshot<Integer, Void, ListState<String>, ListStateDescriptor<String>, DbStateBackend> snapshot2 = kv.snapshot(682375462379L, 4); kv.setCurrentKey(1); assertEquals("1,u1", joiner.join(state.get())); kv.setCurrentKey(2); assertEquals("2,u2", joiner.join(state.get())); kv.setCurrentKey(3); assertEquals("u3", joiner.join(state.get())); kv.dispose(); KvState<Integer, Void, ListState<String>, ListStateDescriptor<String>, DbStateBackend> restored2 = snapshot2.restoreState( backend, IntSerializer.INSTANCE, this.getClass().getClassLoader(), 20); @SuppressWarnings("unchecked") ListState<String> restored2State = (ListState<String>) restored2; restored2.setCurrentKey(1); assertEquals("1,u1", joiner.join(restored2State.get())); restored2.setCurrentKey(2); assertEquals("2,u2", joiner.join(restored2State.get())); restored2.setCurrentKey(3); assertEquals("u3", joiner.join(restored2State.get())); } catch (Exception e) { e.printStackTrace(); fail(e.getMessage()); } }
@Test @SuppressWarnings("unchecked,rawtypes") public void testReducingState() { File tempDir = new File(ConfigConstants.DEFAULT_TASK_MANAGER_TMP_PATH, UUID.randomUUID().toString()); try { FsStateBackend fileBackend = new FsStateBackend(localFileUri(tempDir)); DbStateBackend backend = new DbStateBackend(conf, fileBackend); Environment env = new DummyEnvironment("test", 2, 0); backend.initializeForJob(env, "dummy_test_kv_reduce", IntSerializer.INSTANCE); ReducingStateDescriptor<String> kvId = new ReducingStateDescriptor<>("id", new ReduceFunction<String>() { private static final long serialVersionUID = 1L; @Override public String reduce(String value1, String value2) throws Exception { return value1 + "," + value2; } }, StringSerializer.INSTANCE); ReducingState<String> state = backend.getPartitionedState(null, VoidSerializer.INSTANCE, kvId); @SuppressWarnings("unchecked") KvState<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, DbStateBackend> kv = (KvState<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, DbStateBackend>) state; Joiner joiner = Joiner.on(","); kv.setCurrentKey(1); assertEquals(null, state.get()); state.add("1"); kv.setCurrentKey(2); assertEquals(null, state.get()); state.add("2"); kv.setCurrentKey(1); assertEquals("1", state.get()); KvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, DbStateBackend> snapshot1 = kv.snapshot(682375462378L, 2); kv.setCurrentKey(1); state.add("u1"); kv.setCurrentKey(2); state.add("u2"); kv.setCurrentKey(3); state.add("u3"); KvStateSnapshot<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, DbStateBackend> snapshot2 = kv.snapshot(682375462379L, 4); kv.setCurrentKey(1); assertEquals("1,u1", state.get()); kv.setCurrentKey(2); assertEquals("2,u2", state.get()); kv.setCurrentKey(3); assertEquals("u3", state.get()); kv.dispose(); KvState<Integer, Void, ReducingState<String>, ReducingStateDescriptor<String>, DbStateBackend> restored2 = snapshot2.restoreState( backend, IntSerializer.INSTANCE, this.getClass().getClassLoader(), 20); @SuppressWarnings("unchecked") ReducingState<String> restored2State = (ReducingState<String>) restored2; restored2.setCurrentKey(1); assertEquals("1,u1", restored2State.get()); restored2.setCurrentKey(2); assertEquals("2,u2", restored2State.get()); restored2.setCurrentKey(3); assertEquals("u3", restored2State.get()); } catch (Exception e) { e.printStackTrace(); fail(e.getMessage()); } }