risingwave_connector/parser/avro/
pulsar_resolver.rs1use 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#[derive(Debug)]
29pub struct PulsarSchemaCache {
30 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 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 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}