kafka-console-consumer.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  1. package main
  2. import (
  3. "flag"
  4. "fmt"
  5. "log"
  6. "os"
  7. "os/signal"
  8. "strconv"
  9. "strings"
  10. "sync"
  11. "github.com/Shopify/sarama"
  12. "github.com/Shopify/sarama/tools/tls"
  13. )
  14. var (
  15. brokerList = flag.String("brokers", os.Getenv("KAFKA_PEERS"), "The comma separated list of brokers in the Kafka cluster")
  16. topic = flag.String("topic", "", "REQUIRED: the topic to consume")
  17. partitions = flag.String("partitions", "all", "The partitions to consume, can be 'all' or comma-separated numbers")
  18. offset = flag.String("offset", "newest", "The offset to start with. Can be `oldest`, `newest`")
  19. verbose = flag.Bool("verbose", false, "Whether to turn on sarama logging")
  20. tlsEnabled = flag.Bool("tls-enabled", false, "Whether to enable TLS")
  21. tlsSkipVerify = flag.Bool("tls-skip-verify", false, "Whether skip TLS server cert verification")
  22. tlsClientCert = flag.String("tls-client-cert", "", "Client cert for client authentication (use with -tls-enabled and -tls-client-key)")
  23. tlsClientKey = flag.String("tls-client-key", "", "Client key for client authentication (use with tls-enabled and -tls-client-cert)")
  24. bufferSize = flag.Int("buffer-size", 256, "The buffer size of the message channel.")
  25. logger = log.New(os.Stderr, "", log.LstdFlags)
  26. )
  27. func main() {
  28. flag.Parse()
  29. if *brokerList == "" {
  30. printUsageErrorAndExit("You have to provide -brokers as a comma-separated list, or set the KAFKA_PEERS environment variable.")
  31. }
  32. if *topic == "" {
  33. printUsageErrorAndExit("-topic is required")
  34. }
  35. if *verbose {
  36. sarama.Logger = logger
  37. }
  38. var initialOffset int64
  39. switch *offset {
  40. case "oldest":
  41. initialOffset = sarama.OffsetOldest
  42. case "newest":
  43. initialOffset = sarama.OffsetNewest
  44. default:
  45. printUsageErrorAndExit("-offset should be `oldest` or `newest`")
  46. }
  47. config := sarama.NewConfig()
  48. if *tlsEnabled {
  49. tlsConfig, err := tls.NewConfig(*tlsClientCert, *tlsClientKey)
  50. if err != nil {
  51. printErrorAndExit(69, "Failed to create TLS config: %s", err)
  52. }
  53. config.Net.TLS.Enable = true
  54. config.Net.TLS.Config = tlsConfig
  55. config.Net.TLS.Config.InsecureSkipVerify = *tlsSkipVerify
  56. }
  57. c, err := sarama.NewConsumer(strings.Split(*brokerList, ","), config)
  58. if err != nil {
  59. printErrorAndExit(69, "Failed to start consumer: %s", err)
  60. }
  61. partitionList, err := getPartitions(c)
  62. if err != nil {
  63. printErrorAndExit(69, "Failed to get the list of partitions: %s", err)
  64. }
  65. var (
  66. messages = make(chan *sarama.ConsumerMessage, *bufferSize)
  67. closing = make(chan struct{})
  68. wg sync.WaitGroup
  69. )
  70. go func() {
  71. signals := make(chan os.Signal, 1)
  72. signal.Notify(signals, os.Kill, os.Interrupt)
  73. <-signals
  74. logger.Println("Initiating shutdown of consumer...")
  75. close(closing)
  76. }()
  77. for _, partition := range partitionList {
  78. pc, err := c.ConsumePartition(*topic, partition, initialOffset)
  79. if err != nil {
  80. printErrorAndExit(69, "Failed to start consumer for partition %d: %s", partition, err)
  81. }
  82. go func(pc sarama.PartitionConsumer) {
  83. <-closing
  84. pc.AsyncClose()
  85. }(pc)
  86. wg.Add(1)
  87. go func(pc sarama.PartitionConsumer) {
  88. defer wg.Done()
  89. for message := range pc.Messages() {
  90. messages <- message
  91. }
  92. }(pc)
  93. }
  94. go func() {
  95. for msg := range messages {
  96. fmt.Printf("Partition:\t%d\n", msg.Partition)
  97. fmt.Printf("Offset:\t%d\n", msg.Offset)
  98. fmt.Printf("Key:\t%s\n", string(msg.Key))
  99. fmt.Printf("Value:\t%s\n", string(msg.Value))
  100. fmt.Println()
  101. }
  102. }()
  103. wg.Wait()
  104. logger.Println("Done consuming topic", *topic)
  105. close(messages)
  106. if err := c.Close(); err != nil {
  107. logger.Println("Failed to close consumer: ", err)
  108. }
  109. }
  110. func getPartitions(c sarama.Consumer) ([]int32, error) {
  111. if *partitions == "all" {
  112. return c.Partitions(*topic)
  113. }
  114. tmp := strings.Split(*partitions, ",")
  115. var pList []int32
  116. for i := range tmp {
  117. val, err := strconv.ParseInt(tmp[i], 10, 32)
  118. if err != nil {
  119. return nil, err
  120. }
  121. pList = append(pList, int32(val))
  122. }
  123. return pList, nil
  124. }
  125. func printErrorAndExit(code int, format string, values ...interface{}) {
  126. fmt.Fprintf(os.Stderr, "ERROR: %s\n", fmt.Sprintf(format, values...))
  127. fmt.Fprintln(os.Stderr)
  128. os.Exit(code)
  129. }
  130. func printUsageErrorAndExit(format string, values ...interface{}) {
  131. fmt.Fprintf(os.Stderr, "ERROR: %s\n", fmt.Sprintf(format, values...))
  132. fmt.Fprintln(os.Stderr)
  133. fmt.Fprintln(os.Stderr, "Available command line options:")
  134. flag.PrintDefaults()
  135. os.Exit(64)
  136. }