/src/main/java/org/apache/kafka/connect/mongodb/MongodbReader.java

https://github.com/DataReply/kafka-connect-mongodb · Java · 99 lines · 79 code · 13 blank · 7 comment · 7 complexity · 9c22c2763c55cb6ba26a89ae8ec01528 MD5 · raw file

  1. package org.apache.kafka.connect.mongodb;
  2. import org.bson.Document;
  3. import java.util.ArrayList;
  4. import java.util.Collections;
  5. import java.util.List;
  6. import java.util.Map;
  7. import java.util.concurrent.ConcurrentLinkedQueue;
  8. /**
  9. * Class that creates a new DatabaseReader thread for every db
  10. *
  11. * @author Andrea Patelli
  12. */
  13. public class MongodbReader {
  14. private final ConcurrentLinkedQueue<Document> messages;
  15. private final List<String> dbs;
  16. private final String uri;
  17. private final String host;
  18. private final Integer port;
  19. private final Integer batchSize;
  20. private final Map<Map<String, String>, Map<String, Object>> start;
  21. private final List<Thread> threads;
  22. public MongodbReader(String uri, List<String> dbs, Integer batchSize, Map<Map<String, String>, Map<String, Object>> start) {
  23. this.uri = uri;
  24. this.host = null;
  25. this.port = null;
  26. this.batchSize = batchSize;
  27. this.dbs = new ArrayList<>(0);
  28. this.dbs.addAll(dbs);
  29. this.threads = new ArrayList<>(0);
  30. this.start = start;
  31. this.messages = new ConcurrentLinkedQueue<>();
  32. }
  33. public MongodbReader(String host, Integer port, List<String> dbs, Integer batchSize, Map<Map<String, String>, Map<String, Object>> start) {
  34. this.uri = null;
  35. this.host = host;
  36. this.port = port;
  37. this.batchSize = batchSize;
  38. this.dbs = new ArrayList<>(0);
  39. this.dbs.addAll(dbs);
  40. this.threads = new ArrayList<>(0);
  41. this.start = start;
  42. this.messages = new ConcurrentLinkedQueue<>();
  43. }
  44. public void run() {
  45. // for every database to watch
  46. for (String db : dbs) {
  47. final String start = getStartOffset(db);
  48. DatabaseReader reader;
  49. if(uri != null){
  50. reader = new DatabaseReader(uri, db, start, batchSize, messages);
  51. }
  52. else{
  53. reader = new DatabaseReader(host, port, db, start, batchSize, messages);
  54. }
  55. final Thread thread = new Thread(reader);
  56. thread.start();
  57. threads.add(thread);
  58. }
  59. }
  60. private String getStartOffset(String db){
  61. String start;
  62. final List<Map<String, String>> partitions = new ArrayList<>();
  63. final Map<String, String> partition = Collections.singletonMap("mongodb", db);
  64. partitions.add(partition);
  65. // get the last message that was read
  66. Map<String, Object> dbOffset = this.start.get(Collections.singletonMap("mongodb", db));
  67. if (dbOffset == null || dbOffset.isEmpty()){
  68. start = "0";
  69. }
  70. else{
  71. start = (String) this.start.get(Collections.singletonMap("mongodb", db)).get(db);
  72. }
  73. return start;
  74. }
  75. public void stop(){
  76. for(Thread thread : threads){
  77. thread.interrupt();
  78. }
  79. }
  80. public boolean isEmpty(){
  81. return messages.isEmpty();
  82. }
  83. public Document pool() {
  84. return messages.poll();
  85. }
  86. }