yozefu_app/admin/
mod.rs

1use std::{collections::HashSet, time::Duration};
2
3use itertools::Itertools;
4use lib::{ConsumerGroupDetail, Error, TopicConfig, TopicDetail};
5use rdkafka::{
6    Offset, TopicPartitionList,
7    admin::{AdminClient as RDAdminClient, AdminOptions, ResourceSpecifier},
8    client::DefaultClientContext,
9    config::FromClientConfigAndContext,
10    consumer::{BaseConsumer, Consumer, StreamConsumer},
11};
12use tracing::warn;
13
14use crate::configuration::{Configuration, InternalConfig};
15
16pub struct AdminClient {
17    client: RDAdminClient<DefaultClientContext>,
18    config: InternalConfig,
19    options: AdminOptions,
20}
21
22impl AdminClient {
23    pub fn new(config: InternalConfig) -> Result<Self, Error> {
24        let client =
25            RDAdminClient::from_config_and_context(&config.client_config(), DefaultClientContext)?;
26        let options = AdminOptions::new();
27        Ok(Self {
28            client,
29            config,
30            options,
31        })
32    }
33    /// Loads the configuration details for the specified topic from the Kafka cluster.
34    pub async fn topic_config(&self, topic: &str) -> Result<Option<TopicConfig>, Error> {
35        let resource = ResourceSpecifier::Topic(topic);
36        let result = self
37            .client
38            .describe_configs(&[resource], &self.options)
39            .await?;
40
41        if result.is_empty() {
42            return Ok(None);
43        }
44
45        match result.first() {
46            Some(Ok(c)) => Ok(Some(c.into())),
47            Some(Err(e)) => Err(Error::Error(e.to_string())),
48            None => Ok(None),
49        }
50    }
51
52    /// Returns the topics details for a given list topics
53    /// This function is not ready yet
54    pub fn topic_details(&self, topics: HashSet<String>) -> Result<Vec<TopicDetail>, Error> {
55        let mut results = vec![];
56        for topic in topics {
57            let consumer: BaseConsumer = self.config.create_kafka_consumer()?;
58            let metadata = consumer.fetch_metadata(Some(&topic), Duration::from_secs(10))?;
59            let metadata = metadata.topics().first().unwrap();
60            let mut detail = TopicDetail {
61                name: topic.clone(),
62                replicas: metadata.partitions().first().unwrap().replicas().len(),
63                partitions: metadata.partitions().len(),
64                consumer_groups: vec![],
65                count: self.count_records_in_topic(&topic)?,
66                config: None,
67            };
68            let mut consumer_groups = vec![];
69            let metadata = consumer.fetch_group_list(None, Duration::from_secs(10))?;
70            for g in metadata.groups() {
71                consumer_groups.push(ConsumerGroupDetail {
72                    name: g.name().to_string(),
73                    members: vec![], //Self::parse_members(g, g.members())?,
74                    state: g.state().parse()?,
75                });
76            }
77            detail.consumer_groups = consumer_groups;
78            results.push(detail);
79        }
80
81        Ok(results)
82    }
83
84    pub fn estimate_number_of_records_to_read(
85        &self,
86        topic_partition_list: &TopicPartitionList,
87    ) -> Result<i64, Error> {
88        let client: StreamConsumer = self.config.create_kafka_consumer()?;
89        let mut count = 0;
90        for t in topic_partition_list.elements() {
91            // this function call be very slow
92            let watermarks: (i64, i64) =
93                match client.fetch_watermarks(t.topic(), t.partition(), Duration::from_secs(10)) {
94                    Ok(i) => i,
95                    Err(e) => {
96                        warn!(
97                            "I was not able to fetch watermarks of topic '{}', partition {}: {}",
98                            t.topic(),
99                            t.partition(),
100                            e
101                        );
102                        (0, 0)
103                    }
104                };
105            count += match t.offset() {
106                Offset::Beginning => watermarks.1 - watermarks.0,
107                Offset::End => 0,
108                Offset::Stored => 1,
109                Offset::Invalid => 1,
110                Offset::Offset(o) => watermarks.1 - o,
111                Offset::OffsetTail(o) => o,
112            }
113        }
114        Ok(count)
115    }
116
117    fn count_records_in_topic(&self, topic: &str) -> Result<i64, Error> {
118        let mut count = 0;
119        let consumer: BaseConsumer = self.config.create_kafka_consumer()?;
120        let metadata = consumer.fetch_metadata(Some(topic), Duration::from_secs(10))?;
121        let metadata_topic = metadata.topics().first();
122        if metadata_topic.is_none() {
123            return Ok(0);
124        }
125
126        let metadata_topic = metadata_topic.unwrap();
127        for partition in metadata_topic.partitions() {
128            let watermarks =
129                consumer.fetch_watermarks(topic, partition.id(), Duration::from_secs(10))?;
130            count += watermarks.1 - watermarks.0;
131        }
132
133        Ok(count)
134    }
135
136    pub fn list_topics(&self) -> Result<Vec<String>, Error> {
137        let consumer: StreamConsumer = self.config.create_kafka_consumer()?;
138        let metadata = consumer.fetch_metadata(None, Duration::from_secs(10))?;
139        let topics = metadata
140            .topics()
141            .iter()
142            .map(|t| t.name().to_string())
143            .collect_vec();
144        Ok(topics)
145    }
146}
yozefu_app/admin/mod.rs

yozefu_app/admin/
mod.rs
