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