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