zio.kafka.consumer.Subscription.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of zio-kafka_3 Show documentation
zio-kafka
The newest version!
package zio.kafka.consumer

import zio.{ NonEmptyChunk, Task, ZIO }

import scala.util.matching.Regex
import java.util.regex.{ Pattern => JPattern }
import org.apache.kafka.common.TopicPartition

sealed trait Subscription
object Subscription {
  final case class Topics(topics: Set[String])                  extends Subscription
  final case class Pattern(pattern: Regex)                      extends Subscription
  final case class Manual(topicPartitions: Set[TopicPartition]) extends Subscription

  /**
   * Create a subscription for one or more topics
   *
   * @param name
   *   Topic to subscribe to
   * @param names
   *   Additional topics to subscribe to
   */
  def topics(name: String, names: String*): Subscription =
    Topics(Set(name) ++ names.toSet)

  /**
   * Create a subscription for all topics matching the given pattern
   *
   * @param pattern
   *   Pattern
   */
  def pattern(pattern: Regex): Subscription =
    Pattern(pattern)

  /**
   * Create a subscription for all topics matching the given pattern
   *
   * @param pattern
   *   Pattern to be converted to a Regex
   * @return
   *   The created subscription or failure when the pattern is invalid
   */
  def pattern(pattern: String): Task[Subscription] =
    ZIO.attempt(new Regex(pattern)).map(Pattern.apply)

  /**
   * Create a subscription for all topics matching the given pattern
   *
   * @param pattern
   *   Pattern
   */
  def pattern(pattern: JPattern): Subscription =
    Pattern(new Regex(pattern.pattern()))

  /**
   * Create a manual subscription to a fixed set of topic-partitions
   *
   * A consumer with this type of subscription does not perform consumer rebalancing
   *
   * @param topicPartitions
   *   Tuples of topic and partition
   */
  def manual(topicPartitions: (String, Int)*): Manual =
    Manual(topicPartitions.map { case (topic, partition) => new TopicPartition(topic, partition) }.toSet)

  /**
   * Create a subscription to a single topic-partition
   *
   * A consumer with this type of subscription does not perform consumer rebalancing
   *
   * @param topic
   * @param partition
   * @return
   */
  def manual(topic: String, partition: Int): Manual =
    manual((topic, partition))

  private[consumer] def union(s1: Subscription, s2: Subscription): Option[Subscription] = (s1, s2) match {
    case (Pattern(p1), Pattern(p2)) => Some(Pattern(s"$p1|$p2".r))
    case (Topics(t1), Topics(t2))   => Some(Topics(t1 ++ t2))
    case (Manual(tp1), Manual(tp2)) => Some(Manual(tp1 ++ tp2))
    case _ =>
      None // Although we could combine Pattern and Topics, the Kafka consumer will throw exceptions when changing the subscription type
  }

  private[consumer] def unionAll(subscriptions: NonEmptyChunk[Subscription]): Option[Subscription] =
    subscriptions.foldLeft(Option(subscriptions.head)) { case (s1, s2) => s1.flatMap(union(_, s2)) }

  private[consumer] def subscriptionMatches(subscription: Subscription, tp: TopicPartition): Boolean =
    subscription match {
      case Topics(topics)          => topics.contains(tp.topic())
      case Pattern(pattern)        => pattern.findFirstIn(tp.topic()).isDefined
      case Manual(topicPartitions) => topicPartitions.contains(tp)
    }
}