laminar_connectors/kafka/schema_registry/
mod.rs

1//! Confluent Schema Registry client.
2//!
3//! [`SchemaRegistryClient`] provides a lightweight async REST client for
4//! the Confluent Schema Registry API, with in-memory caching, arrow
5//! schema conversion, and compatibility checking.
6
7use std::time::{Duration, Instant};
8
9use quick_cache::sync::Cache;
10
11use arrow_schema::SchemaRef;
12use reqwest::Client;
13use serde::{Deserialize, Serialize};
14
15use crate::error::{ConnectorError, SerdeError};
16use crate::kafka::config::{CompatibilityLevel, SrAuth};
17
18mod resolution;
19
20const SCHEMA_REGISTRY_CONNECT_TIMEOUT: Duration = Duration::from_secs(2);
21const SCHEMA_REGISTRY_READ_TIMEOUT: Duration = Duration::from_secs(2);
22const SCHEMA_REGISTRY_REQUEST_TIMEOUT: Duration = Duration::from_secs(3);
23
24#[derive(Clone, Copy)]
25struct SchemaRegistryHttpTimeouts {
26    connect: Duration,
27    read: Duration,
28    request: Duration,
29}
30
31const SCHEMA_REGISTRY_HTTP_TIMEOUTS: SchemaRegistryHttpTimeouts = SchemaRegistryHttpTimeouts {
32    connect: SCHEMA_REGISTRY_CONNECT_TIMEOUT,
33    read: SCHEMA_REGISTRY_READ_TIMEOUT,
34    request: SCHEMA_REGISTRY_REQUEST_TIMEOUT,
35};
36
37fn http_client_builder(timeouts: SchemaRegistryHttpTimeouts) -> reqwest::ClientBuilder {
38    Client::builder()
39        .redirect(reqwest::redirect::Policy::none())
40        .connect_timeout(timeouts.connect)
41        .read_timeout(timeouts.read)
42        .timeout(timeouts.request)
43}
44
45fn build_http_client(timeouts: SchemaRegistryHttpTimeouts) -> Result<Client, ConnectorError> {
46    http_client_builder(timeouts).build().map_err(|e| {
47        ConnectorError::ConfigurationError(format!(
48            "failed to build Schema Registry HTTP client: {e}"
49        ))
50    })
51}
52
53fn schema_registry_http_error(
54    operation: &str,
55    status: reqwest::StatusCode,
56    detail: &str,
57) -> ConnectorError {
58    let _ = detail;
59    let message = format!("schema registry {operation} failed: {status}");
60    if status == reqwest::StatusCode::REQUEST_TIMEOUT
61        || status == reqwest::StatusCode::TOO_MANY_REQUESTS
62        || status.is_server_error()
63    {
64        ConnectorError::ConnectionFailed(message)
65    } else {
66        ConnectorError::ConfigurationError(message)
67    }
68}
69
70fn schema_registry_request_error(operation: &str, error: reqwest::Error) -> ConnectorError {
71    let builder_error = error.is_builder();
72    let message = format!(
73        "schema registry {operation} failed: {}",
74        error.without_url()
75    );
76    if builder_error {
77        ConnectorError::ConfigurationError(message)
78    } else {
79        ConnectorError::ConnectionFailed(message)
80    }
81}
82
83/// Schema type as reported by the Schema Registry.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum SchemaType {
86    /// Apache Avro schema.
87    Avro,
88    /// Protocol Buffers schema.
89    Protobuf,
90    /// JSON Schema.
91    Json,
92}
93
94impl std::str::FromStr for SchemaType {
95    type Err = ConnectorError;
96
97    fn from_str(s: &str) -> Result<Self, Self::Err> {
98        match s.to_uppercase().as_str() {
99            "AVRO" => Ok(SchemaType::Avro),
100            "PROTOBUF" => Ok(SchemaType::Protobuf),
101            "JSON" => Ok(SchemaType::Json),
102            other => Err(ConnectorError::ConfigurationError(format!(
103                "unknown schema type: '{other}'"
104            ))),
105        }
106    }
107}
108
109impl std::fmt::Display for SchemaType {
110    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
111        match self {
112            SchemaType::Avro => write!(f, "AVRO"),
113            SchemaType::Protobuf => write!(f, "PROTOBUF"),
114            SchemaType::Json => write!(f, "JSON"),
115        }
116    }
117}
118
119fn require_avro_mutation(schema_type: SchemaType) -> Result<(), ConnectorError> {
120    if schema_type == SchemaType::Avro {
121        Ok(())
122    } else {
123        Err(ConnectorError::ConfigurationError(format!(
124            "Kafka Schema Registry registration supports AVRO only, got {schema_type}"
125        )))
126    }
127}
128
129/// Configuration for the Schema Registry cache.
130#[derive(Debug, Clone)]
131pub struct SchemaRegistryCacheConfig {
132    /// Maximum number of cached schemas. Default: 1000.
133    pub max_entries: usize,
134    /// TTL for cache entries. `None` means no expiry. Default: 1 hour.
135    pub ttl: Option<Duration>,
136}
137
138impl Default for SchemaRegistryCacheConfig {
139    fn default() -> Self {
140        Self {
141            max_entries: 1000,
142            ttl: Some(Duration::from_secs(3600)),
143        }
144    }
145}
146
147/// A cached schema entry from the Schema Registry.
148#[derive(Debug, Clone)]
149pub struct CachedSchema {
150    /// Schema with transitive named references resolved for the selected Arrow codec.
151    pub resolved_schema_str: String,
152    /// Complete, concrete transitive reference definitions.
153    pub references: Vec<crate::schema::resolution::NativeSchemaArtifact>,
154    /// Schema Registry schema ID.
155    pub id: i32,
156    /// Schema version within its subject.
157    pub version: i32,
158    /// The schema type.
159    pub schema_type: SchemaType,
160    /// Raw schema string (e.g., Avro JSON).
161    pub schema_str: String,
162    /// Derived Arrow schema for `RecordBatch` construction.
163    pub arrow_schema: SchemaRef,
164    /// When this entry was inserted or last accessed.
165    inserted_at: Instant,
166}
167
168/// Result of a compatibility check.
169#[derive(Debug, Clone)]
170pub struct CompatibilityResult {
171    /// Whether the schema is compatible.
172    pub is_compatible: bool,
173    /// Incompatibility reasons (if any).
174    pub messages: Vec<String>,
175}
176
177/// Async client for the Confluent Schema Registry REST API.
178///
179/// Provides schema lookup by ID and subject, caching with LRU eviction
180/// and TTL, compatibility checking, and Avro-to-Arrow schema conversion.
181pub struct SchemaRegistryClient {
182    client: Client,
183    base_url: String,
184    auth: Option<SrAuth>,
185    /// Cache by schema ID (`quick_cache`, S3-FIFO-style eviction).
186    cache: Cache<i32, CachedSchema>,
187    /// Cache by subject name (latest version).
188    subject_cache: Cache<String, CachedSchema>,
189    /// Cache configuration.
190    cache_config: SchemaRegistryCacheConfig,
191    // INVARIANT: serialize cache misses per authenticated registry client. Reference fetches
192    // use get_json directly, so this guard never recursively acquires itself.
193    lookup_gate: tokio::sync::Mutex<()>,
194}
195
196// -- Schema Registry REST API response types --
197
198#[derive(Deserialize)]
199struct SchemaByIdResponse {
200    #[serde(default)]
201    references: Vec<resolution::SchemaReference>,
202    schema: String,
203    #[serde(default = "default_schema_type")]
204    #[serde(rename = "schemaType")]
205    schema_type: String,
206}
207
208#[derive(Deserialize)]
209struct SchemaVersionResponse {
210    #[serde(default)]
211    references: Vec<resolution::SchemaReference>,
212    id: i32,
213    version: i32,
214    schema: String,
215    #[serde(default = "default_schema_type")]
216    #[serde(rename = "schemaType")]
217    schema_type: String,
218}
219
220#[derive(Deserialize)]
221struct CompatibilityResponse {
222    is_compatible: bool,
223    #[serde(default)]
224    messages: Vec<String>,
225}
226
227#[derive(Deserialize)]
228struct ConfigResponse {
229    #[serde(rename = "compatibilityLevel")]
230    compatibility_level: String,
231}
232
233#[derive(Serialize)]
234struct CompatibilityRequest {
235    schema: String,
236    #[serde(rename = "schemaType")]
237    schema_type: String,
238}
239
240#[derive(Serialize)]
241struct ConfigUpdateRequest {
242    compatibility: String,
243}
244
245#[derive(Serialize)]
246struct RegisterSchemaRequest {
247    schema: String,
248    #[serde(rename = "schemaType")]
249    schema_type: String,
250}
251
252#[derive(Deserialize)]
253struct RegisterSchemaResponse {
254    id: i32,
255}
256
257fn default_schema_type() -> String {
258    "AVRO".to_string()
259}
260
261impl SchemaRegistryClient {
262    /// Creates a new Schema Registry client with default cache config.
263    ///
264    /// # Errors
265    ///
266    /// Returns `ConnectorError::ConfigurationError` if the HTTP client cannot be built.
267    pub fn new(base_url: impl Into<String>, auth: Option<SrAuth>) -> Result<Self, ConnectorError> {
268        Self::with_cache_config(base_url, auth, SchemaRegistryCacheConfig::default())
269    }
270
271    /// Creates a TLS client (CA cert only). Delegates to [`Self::with_tls_mtls`].
272    ///
273    /// # Errors
274    ///
275    /// Returns `ConnectorError::ConfigurationError` if the CA cert cannot be read.
276    pub fn with_tls(
277        base_url: impl Into<String>,
278        auth: Option<SrAuth>,
279        ca_cert_path: &str,
280    ) -> Result<Self, ConnectorError> {
281        Self::with_tls_mtls(base_url, auth, Some(ca_cert_path), None, None)
282    }
283
284    /// Creates a client with full TLS/mTLS support. Without `ca_cert_path` the server is
285    /// verified against the default trust roots; a client certificate applies either way.
286    ///
287    /// # Errors
288    ///
289    /// Returns `ConnectorError::ConfigurationError` if any cert/key file
290    /// cannot be read or parsed.
291    pub fn with_tls_mtls(
292        base_url: impl Into<String>,
293        auth: Option<SrAuth>,
294        ca_cert_path: Option<&str>,
295        client_cert_path: Option<&str>,
296        client_key_path: Option<&str>,
297    ) -> Result<Self, ConnectorError> {
298        let mut builder = http_client_builder(SCHEMA_REGISTRY_HTTP_TIMEOUTS);
299        if let Some(ca_cert_path) = ca_cert_path {
300            let pem = std::fs::read(ca_cert_path).map_err(|e| {
301                ConnectorError::ConfigurationError(format!(
302                    "failed to read SR CA cert at '{ca_cert_path}': {e}"
303                ))
304            })?;
305            let cert = reqwest::tls::Certificate::from_pem(&pem).map_err(|e| {
306                ConnectorError::ConfigurationError(format!(
307                    "invalid PEM CA cert at '{ca_cert_path}': {e}"
308                ))
309            })?;
310            builder = builder.add_root_certificate(cert);
311        }
312
313        if client_cert_path.is_some() != client_key_path.is_some() {
314            return Err(ConnectorError::ConfigurationError(
315                "mTLS requires both client cert and key — only one was provided".into(),
316            ));
317        }
318        if let (Some(cert_path), Some(key_path)) = (client_cert_path, client_key_path) {
319            let mut identity_pem = std::fs::read(cert_path).map_err(|e| {
320                ConnectorError::ConfigurationError(format!(
321                    "failed to read SR client cert at '{cert_path}': {e}"
322                ))
323            })?;
324            let key_pem = std::fs::read(key_path).map_err(|e| {
325                ConnectorError::ConfigurationError(format!(
326                    "failed to read SR client key at '{key_path}': {e}"
327                ))
328            })?;
329            // reqwest Identity expects cert + key concatenated in PEM format.
330            identity_pem.extend_from_slice(&key_pem);
331            let identity = reqwest::tls::Identity::from_pem(&identity_pem).map_err(|e| {
332                ConnectorError::ConfigurationError(format!("invalid client cert/key PEM: {e}"))
333            })?;
334            builder = builder.identity(identity);
335        }
336
337        let client = builder.build().map_err(|e| {
338            ConnectorError::ConfigurationError(format!("failed to build TLS client: {e}"))
339        })?;
340
341        Ok(Self::from_http_client(
342            base_url,
343            auth,
344            SchemaRegistryCacheConfig::default(),
345            client,
346        ))
347    }
348
349    /// Creates a new Schema Registry client with custom cache config.
350    ///
351    /// # Errors
352    ///
353    /// Returns `ConnectorError::ConfigurationError` if the HTTP client cannot be built.
354    pub fn with_cache_config(
355        base_url: impl Into<String>,
356        auth: Option<SrAuth>,
357        cache_config: SchemaRegistryCacheConfig,
358    ) -> Result<Self, ConnectorError> {
359        Self::with_cache_config_and_timeouts(
360            base_url,
361            auth,
362            cache_config,
363            SCHEMA_REGISTRY_HTTP_TIMEOUTS,
364        )
365    }
366
367    fn with_cache_config_and_timeouts(
368        base_url: impl Into<String>,
369        auth: Option<SrAuth>,
370        cache_config: SchemaRegistryCacheConfig,
371        timeouts: SchemaRegistryHttpTimeouts,
372    ) -> Result<Self, ConnectorError> {
373        let client = build_http_client(timeouts)?;
374        Ok(Self::from_http_client(base_url, auth, cache_config, client))
375    }
376
377    fn from_http_client(
378        base_url: impl Into<String>,
379        auth: Option<SrAuth>,
380        cache_config: SchemaRegistryCacheConfig,
381        client: Client,
382    ) -> Self {
383        let cache = Cache::new(cache_config.max_entries.clamp(1, 64));
384        // Subject cache is small — one entry per subject
385        let subject_cache = Cache::new(16);
386        Self {
387            client,
388            base_url: base_url.into().trim_end_matches('/').to_string(),
389            auth,
390            cache,
391            subject_cache,
392            cache_config,
393            lookup_gate: tokio::sync::Mutex::new(()),
394        }
395    }
396
397    /// Returns the base URL of the Schema Registry.
398    #[must_use]
399    pub fn base_url(&self) -> &str {
400        &self.base_url
401    }
402
403    /// Returns `true` if authentication is configured.
404    #[must_use]
405    pub fn has_auth(&self) -> bool {
406        self.auth.is_some()
407    }
408
409    /// Returns the cache configuration.
410    #[must_use]
411    pub fn cache_config(&self) -> &SchemaRegistryCacheConfig {
412        &self.cache_config
413    }
414
415    /// Inserts a schema into the cache.
416    ///
417    /// `quick_cache` handles eviction internally (S3-FIFO-style).
418    fn cache_insert(&self, id: i32, mut schema: CachedSchema) {
419        schema.inserted_at = Instant::now();
420        self.cache.insert(id, schema);
421    }
422
423    /// Gets from cache, returning `None` if expired.
424    ///
425    /// TTL is checked lazily on access — expired entries are removed
426    /// and treated as cache misses.
427    fn cache_get(&self, id: i32) -> Option<CachedSchema> {
428        let schema = self.cache.get(&id)?;
429        if let Some(ttl) = self.cache_config.ttl {
430            if schema.inserted_at.elapsed() > ttl {
431                self.cache.remove(&id);
432                return None;
433            }
434        }
435        // quick_cache's get() already promotes the entry in the eviction policy
436        Some(schema)
437    }
438
439    /// Fetches a schema by its global ID.
440    ///
441    /// Results are cached for subsequent lookups.
442    ///
443    /// # Errors
444    ///
445    /// Returns `ConnectorError` if the HTTP request fails or the schema
446    /// cannot be parsed.
447    pub async fn get_schema_by_id(&self, id: i32) -> Result<CachedSchema, ConnectorError> {
448        if let Some(cached) = self.cache_get(id) {
449            return Ok(cached);
450        }
451
452        if id <= 0 {
453            return Err(ConnectorError::SchemaMismatch(
454                "registry schema ID must be positive".into(),
455            ));
456        }
457        tokio::time::timeout(std::time::Duration::from_secs(30), async {
458            let _lookup = self.lookup_gate.lock().await;
459            if let Some(cached) = self.cache_get(id) {
460                return Ok(cached);
461            }
462            let url = format!("{}/schemas/ids/{}", self.base_url, id);
463            let operation = format!("fetch schema ID {id}");
464            let resp: SchemaByIdResponse = self.get_json(&url, &operation).await?;
465            let cached = self
466                .complete_schema(id, 0, resp.schema, &resp.schema_type, resp.references)
467                .await?;
468            self.cache_insert(id, cached.clone());
469            Ok(cached)
470        })
471        .await
472        .map_err(|_| ConnectorError::Timeout(30_000))?
473    }
474
475    pub(super) async fn fetch_registered_schema(
476        &self,
477        id: i32,
478    ) -> Result<CachedSchema, ConnectorError> {
479        // Registration echoes only an ID. Verify its definition instead of trusting
480        // the optimistic cache populated from our request body.
481        self.cache.remove(&id);
482        self.get_schema_by_id(id).await
483    }
484
485    /// Fetches the latest schema version for a subject.
486    ///
487    /// # Errors
488    ///
489    /// Returns `ConnectorError` if the HTTP request fails.
490    pub async fn get_latest_schema(&self, subject: &str) -> Result<CachedSchema, ConnectorError> {
491        tokio::time::timeout(std::time::Duration::from_secs(30), async {
492            let _lookup = self.lookup_gate.lock().await;
493            let url = format!(
494                "{}/subjects/{}/versions/latest",
495                self.base_url,
496                resolution::encoded_subject(subject)
497            );
498            let operation = format!("fetch latest schema for subject '{subject}'");
499            let resp: SchemaVersionResponse = self.get_json(&url, &operation).await?;
500
501            if resp.version <= 0 {
502                return Err(ConnectorError::SchemaMismatch(
503                    "registry returned a different or invalid concrete version".into(),
504                ));
505            }
506            let cached = self
507                .complete_schema(
508                    resp.id,
509                    resp.version,
510                    resp.schema,
511                    &resp.schema_type,
512                    resp.references,
513                )
514                .await?;
515            self.cache_insert(cached.id, cached.clone());
516            self.subject_cache
517                .insert(subject.to_string(), cached.clone());
518            Ok(cached)
519        })
520        .await
521        .map_err(|_| ConnectorError::Timeout(30_000))?
522    }
523
524    /// Fetches a specific schema version for a subject.
525    ///
526    /// # Errors
527    ///
528    /// Returns `ConnectorError` if the HTTP request fails.
529    pub async fn get_schema_version(
530        &self,
531        subject: &str,
532        version: i32,
533    ) -> Result<CachedSchema, ConnectorError> {
534        tokio::time::timeout(std::time::Duration::from_secs(30), async {
535            let _lookup = self.lookup_gate.lock().await;
536            let url = format!(
537                "{}/subjects/{}/versions/{}",
538                self.base_url,
539                resolution::encoded_subject(subject),
540                version
541            );
542            let operation = format!("fetch version {version} for subject '{subject}'");
543            let resp: SchemaVersionResponse = self.get_json(&url, &operation).await?;
544
545            if resp.version <= 0 || resp.version != version {
546                return Err(ConnectorError::SchemaMismatch(
547                    "registry returned a different or invalid concrete version".into(),
548                ));
549            }
550            let cached = self
551                .complete_schema(
552                    resp.id,
553                    resp.version,
554                    resp.schema,
555                    &resp.schema_type,
556                    resp.references,
557                )
558                .await?;
559            self.cache_insert(cached.id, cached.clone());
560            Ok(cached)
561        })
562        .await
563        .map_err(|_| ConnectorError::Timeout(30_000))?
564    }
565
566    /// Checks compatibility of an Avro schema against the latest version.
567    ///
568    /// # Errors
569    ///
570    /// Returns `ConnectorError` if the HTTP request fails.
571    pub async fn check_compatibility(
572        &self,
573        subject: &str,
574        schema_str: &str,
575    ) -> Result<CompatibilityResult, ConnectorError> {
576        let url = format!(
577            "{}/compatibility/subjects/{}/versions/latest",
578            self.base_url,
579            resolution::encoded_subject(subject)
580        );
581
582        let body = CompatibilityRequest {
583            schema: schema_str.to_string(),
584            schema_type: SchemaType::Avro.to_string(),
585        };
586
587        let mut req = self.client.post(&url).json(&body);
588        if let Some(ref auth) = self.auth {
589            req = req.basic_auth(&auth.username, Some(&auth.password));
590        }
591
592        let operation = format!("compatibility check for subject '{subject}'");
593        let resp = req
594            .send()
595            .await
596            .map_err(|error| schema_registry_request_error(&operation, error))?;
597
598        if !resp.status().is_success() {
599            let status = resp.status();
600            let text = String::new();
601            if status == reqwest::StatusCode::NOT_FOUND {
602                return Ok(CompatibilityResult {
603                    is_compatible: true,
604                    messages: Vec::new(),
605                });
606            }
607            return Err(schema_registry_http_error(&operation, status, &text));
608        }
609
610        let result: CompatibilityResponse =
611            Self::bounded_json(resp, &operation).await.map_err(|e| {
612                ConnectorError::Internal(format!(
613                    "schema registry {operation} returned an invalid response: {e}"
614                ))
615            })?;
616
617        Ok(CompatibilityResult {
618            is_compatible: result.is_compatible,
619            messages: result.messages,
620        })
621    }
622
623    /// Gets the compatibility level for a subject.
624    ///
625    /// # Errors
626    ///
627    /// Returns `ConnectorError` if the HTTP request fails.
628    pub async fn get_compatibility_level(
629        &self,
630        subject: &str,
631    ) -> Result<CompatibilityLevel, ConnectorError> {
632        let url = format!(
633            "{}/config/{}",
634            self.base_url,
635            resolution::encoded_subject(subject)
636        );
637        let operation = format!("fetch compatibility config for subject '{subject}'");
638        let resp: ConfigResponse = self.get_json(&url, &operation).await?;
639        resp.compatibility_level.parse()
640    }
641
642    /// Sets the compatibility level for a subject.
643    ///
644    /// # Errors
645    ///
646    /// Returns `ConnectorError` if the HTTP request fails.
647    pub async fn set_compatibility_level(
648        &self,
649        subject: &str,
650        level: CompatibilityLevel,
651    ) -> Result<(), ConnectorError> {
652        let url = format!(
653            "{}/config/{}",
654            self.base_url,
655            resolution::encoded_subject(subject)
656        );
657        let body = ConfigUpdateRequest {
658            compatibility: level.as_str().to_string(),
659        };
660
661        let mut req = self.client.put(&url).json(&body);
662        if let Some(ref auth) = self.auth {
663            req = req.basic_auth(&auth.username, Some(&auth.password));
664        }
665
666        let operation = format!("update compatibility config for subject '{subject}'");
667        let resp = req
668            .send()
669            .await
670            .map_err(|error| schema_registry_request_error(&operation, error))?;
671
672        if !resp.status().is_success() {
673            let status = resp.status();
674            let text = String::new();
675            return Err(schema_registry_http_error(&operation, status, &text));
676        }
677
678        Ok(())
679    }
680
681    /// Resolves a Confluent schema ID, returning from cache if available.
682    ///
683    /// This is the hot-path method called during Avro deserialization to
684    /// look up schemas by the 4-byte ID in the Confluent wire format.
685    ///
686    /// # Errors
687    ///
688    /// Returns `ConnectorError` if the schema cannot be fetched.
689    pub async fn resolve_confluent_id(&self, id: i32) -> Result<CachedSchema, ConnectorError> {
690        self.get_schema_by_id(id).await
691    }
692
693    /// Registers an Avro schema with the Schema Registry under the given subject.
694    ///
695    /// Returns the schema ID assigned by the registry. Caches the result
696    /// so subsequent calls with the same subject return immediately.
697    ///
698    /// # Errors
699    ///
700    /// Returns `ConnectorError` if `schema_type` is not Avro, the HTTP request
701    /// fails, or the response is malformed.
702    pub async fn register_schema(
703        &self,
704        subject: &str,
705        schema_str: &str,
706        schema_type: SchemaType,
707    ) -> Result<i32, ConnectorError> {
708        require_avro_mutation(schema_type)?;
709
710        // Check subject cache — only return cached ID if schema hasn't changed.
711        if let Some(cached) = self.subject_cache.get(subject) {
712            if cached.schema_str == schema_str {
713                return Ok(cached.id);
714            }
715        }
716
717        let url = format!(
718            "{}/subjects/{}/versions",
719            self.base_url,
720            resolution::encoded_subject(subject)
721        );
722        let body = RegisterSchemaRequest {
723            schema: schema_str.to_string(),
724            schema_type: schema_type.to_string(),
725        };
726
727        let mut req = self.client.post(&url).json(&body);
728        if let Some(ref auth) = self.auth {
729            req = req.basic_auth(&auth.username, Some(&auth.password));
730        }
731
732        let operation = format!("register schema for subject '{subject}'");
733        let resp = req
734            .send()
735            .await
736            .map_err(|error| schema_registry_request_error(&operation, error))?;
737
738        if !resp.status().is_success() {
739            let status = resp.status();
740            let text = String::new();
741            return Err(schema_registry_http_error(&operation, status, &text));
742        }
743
744        let result: RegisterSchemaResponse =
745            Self::bounded_json(resp, &operation).await.map_err(|e| {
746                ConnectorError::Internal(format!(
747                    "schema registry {operation} returned an invalid response: {e}"
748                ))
749            })?;
750
751        if result.id <= 0 {
752            return Err(ConnectorError::SchemaMismatch(
753                "registry returned a non-positive registration ID".into(),
754            ));
755        }
756        let arrow_schema = avro_to_arrow_schema(schema_str)?;
757        let cached = CachedSchema {
758            resolved_schema_str: schema_str.to_string(),
759            references: Vec::new(),
760            id: result.id,
761            version: 0,
762            schema_type,
763            schema_str: schema_str.to_string(),
764            arrow_schema,
765            inserted_at: Instant::now(),
766        };
767        self.cache_insert(result.id, cached.clone());
768        self.subject_cache.insert(subject.to_string(), cached);
769
770        Ok(result.id)
771    }
772
773    /// Validates compatibility and registers a schema in one step.
774    ///
775    /// If the subject already has schemas registered, checks compatibility
776    /// first. Returns `SerdeError::SchemaIncompatible` if the new schema
777    /// is not compatible with the existing versions.
778    ///
779    /// # Errors
780    ///
781    /// Returns `ConnectorError::Serde(SchemaIncompatible)` if incompatible,
782    /// or `ConnectorError` for HTTP/network errors.
783    pub async fn validate_and_register_schema(
784        &self,
785        subject: &str,
786        schema_str: &str,
787        schema_type: SchemaType,
788    ) -> Result<i32, ConnectorError> {
789        // Reject unsupported mutation types before the compatibility request.
790        // `register_schema` repeats this guard because it is also public.
791        require_avro_mutation(schema_type)?;
792
793        let result = self.check_compatibility(subject, schema_str).await?;
794        if !result.is_compatible {
795            let message = if result.messages.is_empty() {
796                "new schema is not compatible with existing version".to_string()
797            } else {
798                result.messages.join("; ")
799            };
800            return Err(ConnectorError::Serde(SerdeError::SchemaIncompatible {
801                subject: subject.to_string(),
802                message,
803            }));
804        }
805
806        self.register_schema(subject, schema_str, schema_type).await
807    }
808
809    /// Returns `true` if the schema ID is in the local cache.
810    #[must_use]
811    pub fn is_cached(&self, id: i32) -> bool {
812        self.cache.contains_key(&id)
813    }
814
815    /// Returns the number of cached schemas.
816    #[must_use]
817    pub fn cache_size(&self) -> usize {
818        self.cache.len()
819    }
820
821    /// Helper to perform a GET request and deserialize JSON.
822    ///
823    /// Retries transient failures (408, 429, 5xx, and transport errors) up to
824    /// 3 attempts with exponential backoff (100ms, 500ms).
825    async fn bounded_json<T: serde::de::DeserializeOwned>(
826        mut response: reqwest::Response,
827        operation: &str,
828    ) -> Result<T, ConnectorError> {
829        const LIMIT: usize = 2 * 1024 * 1024;
830        if response
831            .content_length()
832            .is_some_and(|length| length > LIMIT as u64)
833        {
834            return Err(ConnectorError::SchemaMismatch(
835                "registry response exceeds 2 MiB".into(),
836            ));
837        }
838        let mut bytes = Vec::new();
839        while let Some(chunk) = response
840            .chunk()
841            .await
842            .map_err(|error| schema_registry_request_error(operation, error))?
843        {
844            if bytes.len().saturating_add(chunk.len()) > LIMIT {
845                return Err(ConnectorError::SchemaMismatch(
846                    "registry response exceeds 2 MiB".into(),
847                ));
848            }
849            bytes.extend_from_slice(&chunk);
850        }
851        serde_json::from_slice(&bytes).map_err(|_| {
852            ConnectorError::SchemaMismatch(format!(
853                "registry {operation} returned malformed metadata"
854            ))
855        })
856    }
857
858    async fn get_json<T: serde::de::DeserializeOwned>(
859        &self,
860        url: &str,
861        operation: &str,
862    ) -> Result<T, ConnectorError> {
863        let backoffs = [
864            std::time::Duration::from_millis(100),
865            std::time::Duration::from_millis(500),
866        ];
867        let mut last_err = None;
868
869        for (attempt, backoff) in std::iter::once(&std::time::Duration::ZERO)
870            .chain(backoffs.iter())
871            .enumerate()
872        {
873            if attempt > 0 {
874                tokio::time::sleep(*backoff).await;
875            }
876
877            let mut req = self.client.get(url);
878            if let Some(ref auth) = self.auth {
879                req = req.basic_auth(&auth.username, Some(&auth.password));
880            }
881
882            let resp = match req.send().await {
883                Ok(r) => r,
884                Err(e) => {
885                    let error = schema_registry_request_error(operation, e);
886                    if !error.is_transient() {
887                        return Err(error);
888                    }
889                    tracing::warn!(
890                        attempt = attempt + 1,
891                        error = %error,
892                        "schema registry request failed, retrying"
893                    );
894                    last_err = Some(error);
895                    continue;
896                }
897            };
898
899            let status = resp.status();
900            if status.is_success() {
901                return Self::bounded_json(resp, operation).await;
902            }
903
904            let transient = status == reqwest::StatusCode::REQUEST_TIMEOUT
905                || status == reqwest::StatusCode::TOO_MANY_REQUESTS
906                || status.is_server_error();
907            if !transient {
908                let text = String::new();
909                return Err(schema_registry_http_error(operation, status, &text));
910            }
911
912            let text = String::new();
913            tracing::warn!(
914                attempt = attempt + 1,
915                status = %status,
916                "schema registry server error, retrying"
917            );
918            last_err = Some(schema_registry_http_error(operation, status, &text));
919        }
920
921        Err(last_err.unwrap_or_else(|| {
922            ConnectorError::ConnectionFailed(format!(
923                "schema registry {operation} exhausted all retries"
924            ))
925        }))
926    }
927}
928
929impl std::fmt::Debug for SchemaRegistryClient {
930    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
931        f.debug_struct("SchemaRegistryClient")
932            .field("base_url", &self.base_url)
933            .field("has_auth", &self.auth.is_some())
934            .field("cached_schemas", &self.cache.len())
935            .field("cached_subjects", &self.subject_cache.len())
936            .finish_non_exhaustive()
937    }
938}
939
940/// Dispatch a Schema Registry payload to the right Arrow converter.
941/// Only Avro is wired today; JSON Schema and Protobuf return an
942/// actionable error until a maintained conversion library lands.
943fn schema_to_arrow(schema_type: SchemaType, schema_str: &str) -> Result<SchemaRef, ConnectorError> {
944    let name = match schema_type {
945        SchemaType::Avro => return avro_to_arrow_schema(schema_str),
946        SchemaType::Json => "JSON Schema Registry",
947        SchemaType::Protobuf => "Protobuf Schema Registry",
948    };
949    Err(ConnectorError::SchemaMismatch(format!(
950        "{name} subjects are not yet supported for auto-discovery \
951         — declare columns explicitly or use an Avro subject"
952    )))
953}
954
955/// Converts an Avro JSON schema string to an Arrow [`SchemaRef`] via `arrow-avro`'s Decoder.
956///
957/// # Errors
958///
959/// Returns `ConnectorError::SchemaMismatch` if the JSON is invalid or conversion fails.
960pub fn avro_to_arrow_schema(avro_schema_str: &str) -> Result<SchemaRef, ConnectorError> {
961    use arrow_avro::reader::ReaderBuilder;
962    use arrow_avro::schema::{AvroSchema, Fingerprint, FingerprintAlgorithm, SchemaStore};
963
964    let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
965    let avro_schema = AvroSchema::new(avro_schema_str.to_string());
966    let fp = Fingerprint::Id(0);
967    store
968        .set(fp, avro_schema)
969        .map_err(|e| ConnectorError::SchemaMismatch(format!("invalid Avro schema: {e}")))?;
970
971    let decoder = ReaderBuilder::new()
972        .with_writer_schema_store(store)
973        .with_active_fingerprint(fp)
974        .build_decoder()
975        .map_err(|e| ConnectorError::SchemaMismatch(format!("Avro→Arrow conversion: {e}")))?;
976
977    Ok(decoder.schema())
978}
979
980/// Converts an Arrow [`SchemaRef`] to an Avro JSON schema string.
981///
982/// Generates a record schema named `"record"` with fields mapped from
983/// Arrow data types to Avro primitives.
984///
985/// # Errors
986///
987/// Returns `SerdeError` if an Arrow type has no Avro equivalent.
988pub fn arrow_to_avro_schema(schema: &SchemaRef, record_name: &str) -> Result<String, SerdeError> {
989    let mut metadata = schema.metadata().clone();
990    metadata
991        .entry("avro.name".into())
992        .or_insert_with(|| record_name.replace('-', "_"));
993    let native_schema = schema.as_ref().clone().with_metadata(metadata);
994    arrow_avro::schema::AvroSchema::try_from(&native_schema)
995        .map(|schema| schema.json_string)
996        .map_err(|error| SerdeError::UnsupportedFormat(format!("Arrow to Avro: {error}")))
997}
998
999#[cfg(test)]
1000mod tests;
laminar_connectors/kafka/schema_registry/mod.rs

laminar_connectors/kafka/schema_registry/
mod.rs
