Advertisement
Guest User

Untitled

a guest
Dec 9th, 2016
75
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 1.79 KB | None | 0 0
  1. class ResultSource[T](dataSource: DataSource, sql: String, prepare: PreparedStatement => PreparedStatement,
  2. extractor: ResultSet => T) extends GraphStage[SourceShape[T]] {
  3.  
  4. private val out: Outlet[T] = Outlet("EnvisiaResultSource.out")
  5.  
  6. override val shape: SourceShape[T] = SourceShape.of(out)
  7.  
  8. override def createLogic(inheritedAttributes: Attributes): GraphStageLogic = new GraphStageLogic(shape) with OutHandler {
  9. private var connection: Connection = null
  10. private var preparedStatement: PreparedStatement = null
  11. private var resultSet: ResultSet = null
  12.  
  13. override def preStart(): Unit = {
  14. try {
  15. connection = dataSource.getConnection
  16. preparedStatement = prepare(connection.prepareStatement(sql))
  17. resultSet = preparedStatement.executeQuery()
  18. if (!resultSet.next()) {
  19. completeStage()
  20. }
  21. } catch {
  22. case e: Exception => fail(out, e)
  23. }
  24. }
  25.  
  26. override def onPull(): Unit = {
  27. if (resultSet != null) {
  28. try {
  29. push(out, extractor(resultSet))
  30. } catch {
  31. case NonFatal(e) => fail(out, e)
  32. }
  33. // if there is no next result we just fail
  34. if (!resultSet.next()) {
  35. completeStage()
  36. }
  37. } else {
  38. fail(out, new Exception("ResultSet can't be null at this point"))
  39. }
  40. }
  41.  
  42. private def release() = {
  43. if (resultSet != null) {
  44. resultSet.close()
  45. }
  46. resultSet = null
  47. if (preparedStatement != null) {
  48. preparedStatement.close()
  49. }
  50. preparedStatement = null
  51. if (connection != null) {
  52. connection.close()
  53. }
  54. connection = null
  55. }
  56.  
  57. override def onDownstreamFinish() = {
  58. release()
  59. }
  60.  
  61. override def postStop() = {
  62. release()
  63. }
  64.  
  65. setHandler(out, this)
  66. }
  67.  
  68. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement