package com.cloudera.streaming.examples.flink.operators; import org.apache.flink.configuration.Configuration; import org.apache.flink.streaming.api.functions.async.ResultFuture; import org.apache.flink.streaming.api.functions.async.RichAsyncFunction; import com.cloudera.streaming.examples.flink.types.QueryResult; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.sql.Connection; import java.sql.DriverManager; import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; import java.util.Collections; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; /** * Enriches the item info by fetching the item name by item id. */ public class ItemInfoEnrichment extends RichAsyncFunction<QueryResult, QueryResult> { private static final Logger LOG = LoggerFactory.getLogger(ItemInfoEnrichment.class); private static final String ITEM_QUERY = "SELECT name FROM items WHERE itemId = ?;"; private final int threadPoolSize; private final String dbConnectionString; private transient Connection dbConnection; private transient ExecutorService executor; private transient PreparedStatement itemQuery; public ItemInfoEnrichment(int threadPoolSize, String dbConnectionString) { this.threadPoolSize = threadPoolSize; this.dbConnectionString = dbConnectionString; } @Override public void asyncInvoke(QueryResult queryResult, ResultFuture<QueryResult> resultFuture) throws Exception { executor.submit(() -> { try { itemQuery.setString(1, queryResult.itemInfo.itemId); ResultSet rs = itemQuery.executeQuery(); if (rs.next()) { queryResult.itemInfo.setItemName(rs.getString("name")); } resultFuture.complete(Collections.singletonList(queryResult)); } catch (SQLException t) { resultFuture.completeExceptionally(t); } }); } @Override public void open(Configuration configuration) throws Exception { executor = Executors.newFixedThreadPool(threadPoolSize); Class.forName("com.mysql.jdbc.Driver"); dbConnection = DriverManager.getConnection(dbConnectionString); itemQuery = dbConnection.prepareStatement(ITEM_QUERY); } @Override public void close() { try { executor.shutdownNow(); } catch (Throwable t) { LOG.error("Error while shutting down executor service.", t); } try { itemQuery.close(); } catch (Throwable t) { LOG.error("Error while closing connection.", t); } try { dbConnection.close(); } catch (Throwable t) { LOG.error("Error while closing connection.", t); } } }