chore: remove unnecessary Arc for schema provider

Author: benjamin-awd

lib/codecs/src/encoding/format/arrow.rs

@@ -55,10 +55,6 @@ pub struct ArrowStreamSerializerConfig {
     #[serde(default)]
     #[configurable(metadata(docs::examples = true))]
     pub allow_nullable_fields: bool,
-
-    /// Schema provider for lazy schema loading.
-    #[serde(skip)]
-    schema_provider: Option<Arc<dyn SchemaProvider>>,
 }
 
 impl std::fmt::Debug for ArrowStreamSerializerConfig {
@@ -72,10 +68,6 @@ impl std::fmt::Debug for ArrowStreamSerializerConfig {
                     .map(|s| format!("{} fields", s.fields().len())),
             )
             .field("allow_nullable_fields", &self.allow_nullable_fields)
-            .field(
-                "schema_provider",
-                &self.schema_provider.as_ref().map(|_| "<provider>"),
-            )
             .finish()
     }
 }
@@ -86,37 +78,20 @@ impl ArrowStreamSerializerConfig {
         Self {
             schema: Some(schema),
             allow_nullable_fields: false,
-            schema_provider: None,
         }
     }
 
-    /// Create a new ArrowStreamSerializerConfig with a schema provider
-    pub fn with_provider(mut self, provider: Arc<dyn SchemaProvider>) -> Self {
-        self.schema = None;
-        self.schema_provider = Some(provider);
-        self
-    }
-
-    /// Get the schema provider if one was configured
-    pub fn provider(&self) -> Option<&Arc<dyn SchemaProvider>> {
-        self.schema_provider.as_ref()
-    }
-
-    /// Resolve the schema from the provider if present.
-    pub async fn resolve(&mut self) -> Result<(), ArrowEncodingError> {
-        // If schema already exists, nothing to do
-        if self.schema.is_some() {
-            return Ok(());
-        }
-
-        // Fetch from provider if available
-        if let Some(provider) = &self.schema_provider {
-            let schema = provider.get_schema().await?;
-            self.schema = Some(schema);
-            Ok(())
-        } else {
-            Err(ArrowEncodingError::NoSchemaProvided)
-        }
+    /// Resolve the schema from a provider.
+    ///
+    /// This fetches the schema from the provider and stores it in this config,
+    /// preserving all other configuration fields (like `allow_nullable_fields`).
+    pub async fn resolve_with_provider(
+        &mut self,
+        provider: &dyn SchemaProvider,
+    ) -> Result<(), ArrowEncodingError> {
+        let schema = provider.get_schema().await?;
+        self.schema = Some(schema);
+        Ok(())
     }
 
     /// The data type of events that are accepted by `ArrowStreamEncoder`.

src/sinks/clickhouse/config.rs

@@ -1,6 +1,6 @@
 //! Configuration for the `Clickhouse` sink.
 
-use std::{fmt, sync::Arc};
+use std::fmt;
 
 use http::{Request, StatusCode, Uri};
 use hyper::Body;
@@ -290,21 +290,20 @@ impl ClickhouseConfig {
                 .into());
             }
 
-            let base_arrow_config = match batch_encoding {
+            let mut arrow_config = match batch_encoding {
                 BatchSerializerConfig::ArrowStream(config) => config.clone(),
             };
 
-            let resolved_config = self
-                .resolve_arrow_schema(
-                    client,
-                    endpoint.to_string(),
-                    database,
-                    auth,
-                    base_arrow_config,
-                )
-                .await?;
+            self.resolve_arrow_schema(
+                client,
+                endpoint.to_string(),
+                database,
+                auth,
+                &mut arrow_config,
+            )
+            .await?;
 
-            let resolved_batch_config = BatchSerializerConfig::ArrowStream(resolved_config);
+            let resolved_batch_config = BatchSerializerConfig::ArrowStream(arrow_config);
             let arrow_serializer = resolved_batch_config.build()?;
             let batch_serializer = BatchSerializer::Arrow(arrow_serializer);
             let encoder = EncoderKind::Batch(BatchEncoder::new(batch_serializer));
@@ -326,8 +325,8 @@ impl ClickhouseConfig {
         endpoint: String,
         database: &Template,
         auth: Option<&Auth>,
-        base_config: ArrowStreamSerializerConfig,
-    ) -> crate::Result<ArrowStreamSerializerConfig> {
+        config: &mut ArrowStreamSerializerConfig,
+    ) -> crate::Result<()> {
         use super::arrow;
 
         if self.table.is_dynamic() || database.is_dynamic() {
@@ -345,16 +344,15 @@ impl ClickhouseConfig {
             database_str, table_str
         );
 
-        let provider = Arc::new(arrow::ClickHouseSchemaProvider::new(
+        let provider = arrow::ClickHouseSchemaProvider::new(
             client.clone(),
             endpoint,
             database_str.to_string(),
             table_str.to_string(),
             auth.cloned(),
-        ));
+        );
 
-        let mut arrow_config = base_config.with_provider(provider);
-        arrow_config.resolve().await.map_err(|e| {
+        config.resolve_with_provider(&provider).await.map_err(|e| {
             format!(
                 "Failed to fetch schema for {}.{}: {}.",
                 database_str, table_str, e
@@ -363,14 +361,14 @@ impl ClickhouseConfig {
 
         debug!(
             "Successfully fetched Arrow schema with {} fields",
-            arrow_config
+            config
                 .schema
                 .as_ref()
                 .map(|s| s.fields().len())
                 .unwrap_or(0)
         );
 
-        Ok(arrow_config)
+        Ok(())
     }
 }