Skip to main content

risingwave_connector/parser/avro/
pulsar_resolver.rs

1// Copyright 2026 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use anyhow::Context;
18use apache_avro::Schema;
19use moka::future::Cache;
20use risingwave_common::bail;
21
22use crate::error::ConnectorResult;
23use crate::schema::pulsar_schema::{Client, PulsarSchemaInfo};
24use crate::source::pulsar::topic::{Topic, parse_topic};
25
26/// Fetch schemas of a single Pulsar topic and cache writer schemas by version.
27/// Each parser configuration owns a resolver and shares it with its clones through `Arc`.
28#[derive(Debug)]
29pub struct PulsarSchemaCache {
30    /// The topic without partition suffix, as schemas are shared by all partitions.
31    topic: Topic,
32    writer_schemas: Cache<i64, Arc<Schema>>,
33    pulsar_client: Client,
34}
35
36impl PulsarSchemaCache {
37    pub fn new(pulsar_client: Client, topic: &str) -> ConnectorResult<Self> {
38        let mut topic = parse_topic(topic)?;
39        topic.topic = topic.topic_str_without_partition()?;
40        topic.partition_index = None;
41        Ok(Self {
42            topic,
43            writer_schemas: Cache::new(u64::MAX),
44            pulsar_client,
45        })
46    }
47
48    fn parse_schema(raw_schema: PulsarSchemaInfo) -> ConnectorResult<Arc<Schema>> {
49        if !raw_schema.r#type.eq_ignore_ascii_case("AVRO") {
50            bail!("expected Pulsar AVRO schema, got {}", raw_schema.r#type);
51        }
52        let schema =
53            Schema::parse_str(&raw_schema.data).context("failed to parse Pulsar Avro schema")?;
54        Ok(Arc::new(schema))
55    }
56
57    /// Fetches the current latest schema for a parser's reader schema and caches its version.
58    pub async fn get_latest(&self) -> ConnectorResult<Arc<Schema>> {
59        let raw_schema = self.pulsar_client.get_schema(&self.topic, None).await?;
60        let version = raw_schema.version;
61        let schema = Self::parse_schema(raw_schema)?;
62        self.writer_schemas
63            .insert(version, Arc::clone(&schema))
64            .await;
65        Ok(schema)
66    }
67
68    /// Gets a specific writer schema. Concurrent misses share one lookup; failed lookups can
69    /// be retried because errors are not cached.
70    pub async fn get_by_version(&self, version: i64) -> ConnectorResult<Arc<Schema>> {
71        self.writer_schemas
72            .try_get_with(version, async {
73                let raw_schema = self
74                    .pulsar_client
75                    .get_schema(&self.topic, Some(version))
76                    .await?;
77                if raw_schema.version != version {
78                    bail!(
79                        "expected Pulsar schema version {version}, got {}",
80                        raw_schema.version
81                    );
82                }
83                Self::parse_schema(raw_schema)
84            })
85            .await
86            .map_err(Into::into)
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use std::collections::BTreeMap;
93    #[cfg(not(madsim))]
94    use std::io::{Read, Write};
95    #[cfg(not(madsim))]
96    use std::net::TcpListener;
97    #[cfg(not(madsim))]
98    use std::thread;
99
100    use super::*;
101    use crate::schema::pulsar_schema::{
102        PULSAR_SCHEMA_AUTH_TOKEN_KEY, PULSAR_SCHEMA_URL_KEY, PulsarSchemaConfig,
103    };
104
105    fn config(url: &str, token: Option<&str>) -> PulsarSchemaConfig {
106        let mut options = BTreeMap::from([(PULSAR_SCHEMA_URL_KEY.to_owned(), url.to_owned())]);
107        if let Some(token) = token {
108            options.insert(PULSAR_SCHEMA_AUTH_TOKEN_KEY.to_owned(), token.to_owned());
109        }
110        PulsarSchemaConfig::from_options(&options).unwrap().unwrap()
111    }
112
113    #[tokio::test]
114    async fn resolver_clones_share_schemas_but_new_resolvers_are_independent() {
115        let config = config("https://cache-lifetime-test:8443", Some("token-a"));
116        let first =
117            Arc::new(PulsarSchemaCache::new(Client::new(&config).unwrap(), "events").unwrap());
118        let cloned = Arc::clone(&first);
119        let second = PulsarSchemaCache::new(Client::new(&config).unwrap(), "events").unwrap();
120        let schema =
121            Arc::new(Schema::parse_str(r#"{"type":"record","name":"Event","fields":[]}"#).unwrap());
122        first.writer_schemas.insert(0, Arc::clone(&schema)).await;
123        assert!(Arc::ptr_eq(
124            &schema,
125            &cloned.get_by_version(0).await.unwrap()
126        ));
127        assert!(second.writer_schemas.get(&0_i64).await.is_none());
128
129        let weak = Arc::downgrade(&first);
130        drop(first);
131        assert!(weak.upgrade().is_some());
132        drop(cloned);
133        assert!(weak.upgrade().is_none());
134    }
135
136    #[test]
137    fn topic_is_normalized_once_to_the_non_partitioned_topic() {
138        let config = config("http://topic-cache-test:8080", None);
139        for topic in [
140            "events",
141            "public/default/events",
142            "persistent://public/default/events",
143            "persistent://public/default/events-partition-2",
144        ] {
145            let cache = PulsarSchemaCache::new(Client::new(&config).unwrap(), topic).unwrap();
146            assert_eq!(
147                cache.topic.to_string(),
148                "persistent://public/default/events"
149            );
150            assert_eq!(cache.topic.partition_index, None);
151        }
152    }
153
154    #[cfg(not(madsim))]
155    fn spawn_schema_server(
156        responses: Vec<(u16, String)>,
157    ) -> (String, thread::JoinHandle<Vec<String>>) {
158        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
159        let url = format!("http://{}", listener.local_addr().unwrap());
160        let handle = thread::spawn(move || {
161            let mut requests = Vec::new();
162            for (status, body) in responses {
163                let (mut stream, _) = listener.accept().unwrap();
164                let mut request = Vec::new();
165                while !request.ends_with(b"\r\n\r\n") {
166                    let mut byte = [0];
167                    stream.read_exact(&mut byte).unwrap();
168                    request.push(byte[0]);
169                }
170                requests.push(String::from_utf8(request).unwrap());
171                write!(
172                    stream,
173                    "HTTP/1.1 {status} Test\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
174                    body.len()
175                )
176                .unwrap();
177            }
178            requests
179        });
180        (url, handle)
181    }
182
183    #[cfg(not(madsim))]
184    fn schema_response(version: i64, data: &str) -> (u16, String) {
185        (
186            200,
187            serde_json::json!({"version": version, "type": "AVRO", "data": data}).to_string(),
188        )
189    }
190
191    #[cfg(not(madsim))]
192    #[tokio::test]
193    async fn concurrent_misses_share_a_lookup_and_failures_can_be_retried() {
194        let data = r#"{"type":"record","name":"Event","fields":[]}"#;
195        let (url, server) =
196            spawn_schema_server(vec![(404, String::new()), schema_response(0, data)]);
197        let config = config(&url, None);
198        let first = Arc::new(
199            PulsarSchemaCache::new(
200                Client::new(&config).unwrap(),
201                "persistent://public/default/events-partition-0",
202            )
203            .unwrap(),
204        );
205        let second = Arc::clone(&first);
206        assert!(first.get_by_version(0).await.is_err());
207
208        let (a, b) = tokio::join!(first.get_by_version(0), second.get_by_version(0));
209        assert!(Arc::ptr_eq(&a.unwrap(), &b.unwrap()));
210        let requests = server.join().unwrap();
211        assert_eq!(requests.len(), 2);
212        for request in requests {
213            assert!(request.starts_with("GET /admin/v2/schemas/public/default/events/schema/0 "));
214        }
215    }
216
217    #[cfg(not(madsim))]
218    #[tokio::test]
219    async fn old_writer_schema_resolves_to_the_latest_reader() {
220        use apache_avro::types::Value;
221        use apache_avro::{from_avro_datum, to_avro_datum};
222
223        let v0 = r#"{"type":"record","name":"Event","fields":[{"name":"id","type":"int"}]}"#;
224        let v1 = r#"{"type":"record","name":"Event","fields":[{"name":"id","type":"int"},{"name":"name","type":"string","default":"unknown"}]}"#;
225        let (url, server) =
226            spawn_schema_server(vec![schema_response(1, v1), schema_response(0, v0)]);
227        let cache =
228            PulsarSchemaCache::new(Client::new(&config(&url, None)).unwrap(), "events").unwrap();
229        let reader = cache.get_latest().await.unwrap();
230        let writer = cache.get_by_version(0).await.unwrap();
231        assert!(Arc::ptr_eq(
232            &writer,
233            &cache.get_by_version(0).await.unwrap()
234        ));
235        assert!(Arc::ptr_eq(
236            &reader,
237            &cache.get_by_version(1).await.unwrap()
238        ));
239        let bytes =
240            to_avro_datum(&writer, Value::Record(vec![("id".into(), Value::Int(42))])).unwrap();
241        assert_eq!(
242            from_avro_datum(&writer, &mut bytes.as_slice(), Some(&reader)).unwrap(),
243            Value::Record(vec![
244                ("id".into(), Value::Int(42)),
245                ("name".into(), Value::String("unknown".into())),
246            ])
247        );
248        let requests = server.join().unwrap();
249        assert_eq!(requests.len(), 2);
250        assert!(requests[0].starts_with("GET /admin/v2/schemas/public/default/events/schema "));
251        assert!(requests[1].starts_with("GET /admin/v2/schemas/public/default/events/schema/0 "));
252    }
253}