<commit_before>package command

import (
	"flag"
	"fmt"
	"sort"
	"strings"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/config"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
)

type Topics struct {
	Ui cli.Ui
}

func (this *Topics) Run(args []string) (exitCode int) {
	var (
		zone    string
		cluster string
		topic   string
	)
	cmdFlags := flag.NewFlagSet("brokers", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&topic, "t", "", "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if zone == "" {
		this.Ui.Error("empty zone not allowed")
		this.Ui.Output(this.Help())
		return 2
	}

	ensureZoneValid(zone)

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, config.ZonePath(zone)))
	if cluster != "" {
		this.displayTopicsOfCluster(cluster, zkzone)
		return
	}

	// all clusters
	zkzone.WithinClusters(func(cluster string, path string) {
		this.displayTopicsOfCluster(cluster, zkzone)
	})

	return
}

func (this *Topics) displayTopicsOfCluster(cluster string, zkzone *zk.ZkZone) {
	must := func(err error) {
		if err != nil {
			panic(err)
		}
	}

	this.Ui.Output(cluster)
	zkcluster := zkzone.NewCluster(cluster)

	// get all alive brokers within this cluster
	brokers := zkcluster.Brokers()
	if len(brokers) == 0 {
		this.Ui.Output(fmt.Sprintf("%4s%s", " ", color.Red("empty brokers")))
		return
	}

	sortedBrokerIds := make([]string, 0, len(brokers))
	for brokerId, _ := range brokers {
		sortedBrokerIds = append(sortedBrokerIds, brokerId)
	}
	sort.Strings(sortedBrokerIds)
	for _, brokerId := range sortedBrokerIds {
		this.Ui.Output(fmt.Sprintf("%4s%s %s", " ", color.Green(brokerId),
			brokers[brokerId]))
	}

	// find 1st broker in the cluster
	// each broker in the cluster has same metadata
	var broker0 *zk.Broker
	for _, broker := range brokers {
		broker0 = broker
		break
	}

	kfkClient, err := sarama.NewClient([]string{broker0.Addr()}, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(fmt.Sprintf("%5s%s %s", " ", broker0.Addr(),
			err.Error()))
		return
	}
	defer kfkClient.Close()

	topics, err := kfkClient.Topics()
	must(err)
	if len(topics) == 0 {
		this.Ui.Output(fmt.Sprintf("%5s%s", " ", color.Magenta("no topics")))
		return
	}

	this.Ui.Output(fmt.Sprintf("%80d topics", len(topics)))
	for _, topic := range topics {
		this.Ui.Output(strings.Repeat(" ", 4) + topic)

		// get partitions and check if some dead
		alivePartitions, err := kfkClient.WritablePartitions(topic)
		must(err)
		partions, err := kfkClient.Partitions(topic)
		must(err)
		if len(alivePartitions) != len(partions) {
			this.Ui.Output(fmt.Sprintf("topic[%s] has %s partitions: %+v/%+v",
				alivePartitions, color.Red("dead"), partions))
		}

		for _, partitionID := range alivePartitions {
			leader, err := kfkClient.Leader(topic, partitionID)
			must(err)<commit_msg>topics supports topic name pattern matching<commit_after>package command

import (
	"flag"
	"fmt"
	"regexp"
	"sort"
	"strings"

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/config"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
)

type Topics struct {
	Ui cli.Ui
}

func (this *Topics) Run(args []string) (exitCode int) {
	var (
		zone         string
		cluster      string
		topicPattern string
	)
	cmdFlags := flag.NewFlagSet("brokers", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", "", "")
	cmdFlags.StringVar(&topicPattern, "t", "", "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if zone == "" {
		this.Ui.Error("empty zone not allowed")
		this.Ui.Output(this.Help())
		return 2
	}

	ensureZoneValid(zone)

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, config.ZonePath(zone)))
	if cluster != "" {
		this.displayTopicsOfCluster(cluster, zkzone, topicPattern)
		return
	}

	// all clusters
	zkzone.WithinClusters(func(cluster string, path string) {
		this.displayTopicsOfCluster(cluster, zkzone, topicPattern)
	})

	return
}

func (this *Topics) displayTopicsOfCluster(cluster string, zkzone *zk.ZkZone,
	topicPattern string) {
	must := func(err error) {
		if err != nil {
			panic(err)
		}
	}

	this.Ui.Output(cluster)
	zkcluster := zkzone.NewCluster(cluster)

	// get all alive brokers within this cluster
	brokers := zkcluster.Brokers()
	if len(brokers) == 0 {
		this.Ui.Output(fmt.Sprintf("%4s%s", " ", color.Red("empty brokers")))
		return
	}

	if topicPattern == "" {
		sortedBrokerIds := make([]string, 0, len(brokers))
		for brokerId, _ := range brokers {
			sortedBrokerIds = append(sortedBrokerIds, brokerId)
		}
		sort.Strings(sortedBrokerIds)
		for _, brokerId := range sortedBrokerIds {
			this.Ui.Output(fmt.Sprintf("%4s%s %s", " ", color.Green(brokerId),
				brokers[brokerId]))
		}
	}

	// find 1st broker in the cluster
	// each broker in the cluster has same metadata
	var broker0 *zk.Broker
	for _, broker := range brokers {
		broker0 = broker
		break
	}

	kfkClient, err := sarama.NewClient([]string{broker0.Addr()}, sarama.NewConfig())
	if err != nil {
		this.Ui.Output(fmt.Sprintf("%5s%s %s", " ", broker0.Addr(),
			err.Error()))
		return
	}
	defer kfkClient.Close()

	topics, err := kfkClient.Topics()
	must(err)
	if len(topics) == 0 {
		if topicPattern == "" {
			this.Ui.Output(fmt.Sprintf("%5s%s", " ", color.Magenta("no topics")))
		}

		return
	}

	if topicPattern == "" {
		this.Ui.Output(fmt.Sprintf("%80d topics", len(topics)))
	}

	for _, topic := range topics {
		if topicPattern != "" {
			if matched, _ := regexp.MatchString(topicPattern, topic); !matched {
				continue
			}
		}

		this.Ui.Output(strings.Repeat(" ", 4) + color.Blue(topic))

		// get partitions and check if some dead
		alivePartitions, err := kfkClient.WritablePartitions(topic)
		must(err)
		partions, err := kfkClient.Partitions(topic)
		must(err)
		if len(alivePartitions) != len(partions) {
			this.Ui.Output(fmt.Sprintf("topic[%s] has %s partitions: %+v/%+v",
				alivePartitions, color.Red("dead"), partions))
		}

		for _, partitionID := range alivePartitions {
			leader, err := kfkClient.Leader(topic, partitionID)
			must(err)
{'fraction_non_alphanumeric': 0.1600189933523267, 'fraction_numerical': 0.0053814498258942705, 'mean_word_length': 4.067361668003207, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '92138', 'n_tokens_mistral': 2519, 'n_tokens_neox': 2330, 'n_words': 471}