You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iggy.apache.org by ma...@apache.org on 2025/11/19 12:26:27 UTC
(iggy) 01/01: fix(connectors): change connector state to binary representation
This is an automated email from the ASF dual-hosted git repository.
maciej pushed a commit to branch connectors-fix
in repository https://gitbox.apache.org/repos/asf/iggy.git
commit c081756f315729c9096615018d2d24b8abddd110
Author: Maciej Modzelewski <ma...@modzelewski.dev>
AuthorDate: Wed Nov 19 12:18:56 2025 +0100
fix(connectors): change connector state to binary representation
---
Cargo.lock | 1 -
core/connectors/runtime/src/source.rs | 32 +--
core/connectors/runtime/src/state.rs | 52 ++---
core/connectors/sdk/src/lib.rs | 236 +------------------
core/connectors/sdk/src/source.rs | 12 +-
.../sources/elasticsearch_source/src/lib.rs | 115 ++--------
.../elasticsearch_source/src/state_manager.rs | 252 ++++++++++++++++-----
core/connectors/sources/postgres_source/src/lib.rs | 7 +-
core/connectors/sources/random_source/Cargo.toml | 1 -
core/connectors/sources/random_source/src/lib.rs | 29 +--
10 files changed, 280 insertions(+), 457 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 0bf96ced7..dc43e362d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4915,7 +4915,6 @@ dependencies = [
"once_cell",
"rand 0.9.2",
"serde",
- "serde_json",
"simd-json",
"tokio",
"tracing",
diff --git a/core/connectors/runtime/src/source.rs b/core/connectors/runtime/src/source.rs
index 12d192e34..462ae4487 100644
--- a/core/connectors/runtime/src/source.rs
+++ b/core/connectors/runtime/src/source.rs
@@ -24,7 +24,8 @@ use iggy::prelude::{
DirectConfig, HeaderKey, HeaderValue, IggyClient, IggyDuration, IggyError, IggyMessage,
};
use iggy_connector_sdk::{
- DecodedMessage, Error, ProducedMessages, StreamEncoder, TopicMetadata, transforms::Transform,
+ ConnectorState, DecodedMessage, Error, ProducedMessages, StreamEncoder, TopicMetadata,
+ transforms::Transform,
};
use once_cell::sync::Lazy;
use std::{
@@ -32,7 +33,7 @@ use std::{
str::FromStr,
sync::{Arc, atomic::Ordering},
};
-use tracing::{debug, error, info, warn};
+use tracing::{debug, error, info, trace, warn};
use crate::configs::connectors::SourceConfig;
use crate::{
@@ -172,15 +173,13 @@ fn init_source(
container: &Container<SourceApi>,
plugin_config: &serde_json::Value,
id: u32,
- state: Option<serde_json::Value>,
+ state: Option<ConnectorState>,
) {
+ trace!("Initializing source plugin with config: {plugin_config:?} (ID: {id})");
let plugin_config =
serde_json::to_string(plugin_config).expect("Invalid source plugin config.");
- let state_bytes = state.as_ref().map_or(Vec::new(), |s| {
- serde_json::to_vec(s).expect("Failed to serialize state")
- });
- let state_ptr = state_bytes.as_ptr();
- let state_len = state_bytes.len();
+ let state_ptr = state.as_ref().map_or(std::ptr::null(), |s| s.0.as_ptr());
+ let state_len = state.as_ref().map_or(0, |s| s.0.len());
(container.open)(
id,
plugin_config.as_ptr(),
@@ -357,13 +356,16 @@ extern "C" fn handle_produced_messages(
unsafe {
if let Some(sender) = SOURCE_SENDERS.get(&plugin_id) {
let messages = std::slice::from_raw_parts(messages_ptr, messages_len);
- let Ok(messages) = postcard::from_bytes::<ProducedMessages>(messages) else {
- error!(
- "Failed to deserialize produced messages for source connector with ID: {plugin_id}"
- );
- return;
- };
- let _ = sender.send(messages);
+ match postcard::from_bytes::<ProducedMessages>(messages) {
+ Ok(messages) => {
+ let _ = sender.send(messages);
+ }
+ Err(err) => {
+ error!(
+ "Failed to deserialize produced messages for source connector with ID: {plugin_id}. {err}"
+ );
+ }
+ }
}
}
}
diff --git a/core/connectors/runtime/src/state.rs b/core/connectors/runtime/src/state.rs
index d583f1ddc..0b2456e0d 100644
--- a/core/connectors/runtime/src/state.rs
+++ b/core/connectors/runtime/src/state.rs
@@ -1,24 +1,25 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
use std::io::SeekFrom;
-use iggy_connector_sdk::Error;
-use serde_json::Value;
+use iggy_connector_sdk::{ConnectorState, Error};
use strum::Display;
use tokio::{
fs::{File, OpenOptions},
@@ -28,8 +29,8 @@ use tokio::{
use tracing::{debug, error, info};
pub trait StateProvider {
- async fn load(&self) -> Result<Option<Value>, Error>;
- async fn save(&self, state: Value) -> Result<(), Error>;
+ async fn load(&self) -> Result<Option<ConnectorState>, Error>;
+ async fn save(&self, state: ConnectorState) -> Result<(), Error>;
}
#[non_exhaustive]
@@ -55,7 +56,7 @@ impl FileStateProvider {
}
impl StateProvider for FileStateProvider {
- async fn load(&self) -> Result<Option<Value>, Error> {
+ async fn load(&self) -> Result<Option<ConnectorState>, Error> {
let mut file = OpenOptions::new()
.read(true)
.write(true)
@@ -76,12 +77,11 @@ impl StateProvider for FileStateProvider {
Ok(None)
} else {
info!("Loaded state file: {}", self.path);
- let state = serde_json::from_slice(&buffer).map_err(|_| Error::CannotReadStateFile)?;
- Ok(Some(state))
+ Ok(Some(ConnectorState(buffer)))
}
}
- async fn save(&self, state: Value) -> Result<(), Error> {
+ async fn save(&self, state: ConnectorState) -> Result<(), Error> {
let mut file = self.file.lock().await;
let Some(file) = file.as_mut() else {
return Err(Error::CannotReadStateFile);
@@ -97,9 +97,7 @@ impl StateProvider for FileStateProvider {
Error::CannotWriteStateFile
})?;
- let state_bytes = serde_json::to_vec(&state).map_err(|_| Error::CannotWriteStateFile)?;
-
- file.write_all(&state_bytes).await.map_err(|error| {
+ file.write_all(&state.0).await.map_err(|error| {
error!("Cannot write state file: {}. {error}.", self.path);
Error::CannotWriteStateFile
})?;
diff --git a/core/connectors/sdk/src/lib.rs b/core/connectors/sdk/src/lib.rs
index c290f9b50..98066f33d 100644
--- a/core/connectors/sdk/src/lib.rs
+++ b/core/connectors/sdk/src/lib.rs
@@ -65,26 +65,6 @@ pub trait Source: Send + Sync {
/// Invoked when the source is closed, allowing it to perform any necessary cleanup.
async fn close(&mut self) -> Result<(), Error>;
-
- /// Optional: Get the current state of the source (default implementation returns None)
- async fn get_state(&self) -> Result<Option<SourceState>, Error> {
- Ok(None)
- }
-
- /// Optional: Set the state of the source (default implementation does nothing)
- async fn set_state(&mut self, _state: SourceState) -> Result<(), Error> {
- Ok(())
- }
-
- /// Optional: Save the current state to persistent storage (default implementation does nothing)
- async fn save_state(&self) -> Result<(), Error> {
- Ok(())
- }
-
- /// Optional: Load the state from persistent storage (default implementation does nothing)
- async fn load_state(&mut self) -> Result<(), Error> {
- Ok(())
- }
}
/// The Sink trait defines the interface for a sink connector, responsible for consuming the messages from the configured topics.
@@ -104,220 +84,6 @@ pub trait Sink: Send + Sync {
/// Invoked when the sink is closed, allowing it to perform any necessary cleanup.
async fn close(&mut self) -> Result<(), Error>;
-
- /// Optional: Get the current state of the sink (default implementation returns None)
- async fn get_state(&self) -> Result<Option<SinkState>, Error> {
- Ok(None)
- }
-
- /// Optional: Set the state of the sink (default implementation does nothing)
- async fn set_state(&mut self, _state: SinkState) -> Result<(), Error> {
- Ok(())
- }
-
- /// Optional: Save the current state to persistent storage (default implementation does nothing)
- async fn save_state(&self) -> Result<(), Error> {
- Ok(())
- }
-
- /// Optional: Load the state from persistent storage (default implementation does nothing)
- async fn load_state(&mut self) -> Result<(), Error> {
- Ok(())
- }
-}
-
-/// State management for source connectors
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct SourceState {
- /// Unique identifier for this state
- pub id: String,
- /// Timestamp when this state was last updated
- pub last_updated: chrono::DateTime<chrono::Utc>,
- /// Version of the state format
- pub version: u32,
- /// Generic state data as JSON
- pub data: serde_json::Value,
- /// Optional metadata
- pub metadata: Option<serde_json::Value>,
-}
-
-/// State management for sink connectors
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct SinkState {
- /// Unique identifier for this state
- pub id: String,
- /// Timestamp when this state was last updated
- pub last_updated: chrono::DateTime<chrono::Utc>,
- /// Version of the state format
- pub version: u32,
- /// Generic state data as JSON
- pub data: serde_json::Value,
- /// Optional metadata
- pub metadata: Option<serde_json::Value>,
-}
-
-/// State storage backend trait
-#[async_trait]
-pub trait StateStorage: Send + Sync {
- /// Save source state to storage
- async fn save_source_state(&self, state: &SourceState) -> Result<(), Error>;
-
- /// Load source state from storage
- async fn load_source_state(&self, id: &str) -> Result<Option<SourceState>, Error>;
-
- /// Save sink state to storage
- async fn save_sink_state(&self, state: &SinkState) -> Result<(), Error>;
-
- /// Load sink state from storage
- async fn load_sink_state(&self, id: &str) -> Result<Option<SinkState>, Error>;
-
- /// Delete state from storage
- async fn delete_state(&self, id: &str) -> Result<(), Error>;
-
- /// List all state IDs
- async fn list_states(&self) -> Result<Vec<String>, Error>;
-}
-
-/// File-based state storage implementation
-pub struct FileStateStorage {
- base_path: std::path::PathBuf,
-}
-
-impl FileStateStorage {
- pub fn new<P: AsRef<std::path::Path>>(base_path: P) -> Self {
- Self {
- base_path: base_path.as_ref().to_path_buf(),
- }
- }
-
- fn get_state_path(&self, id: &str) -> std::path::PathBuf {
- self.base_path.join(format!("{id}.json"))
- }
-}
-
-#[async_trait]
-impl StateStorage for FileStateStorage {
- async fn save_source_state(&self, state: &SourceState) -> Result<(), Error> {
- use tokio::fs;
-
- // Ensure directory exists
- if let Some(parent) = self.base_path.parent() {
- fs::create_dir_all(parent)
- .await
- .map_err(|e| Error::Storage(format!("Failed to create state directory: {e}")))?;
- }
-
- let path = self.get_state_path(&state.id);
- let json = serde_json::to_string_pretty(state)
- .map_err(|e| Error::Serialization(format!("Failed to serialize source state: {e}")))?;
-
- fs::write(path, json)
- .await
- .map_err(|e| Error::Storage(format!("Failed to write state file: {e}")))?;
-
- Ok(())
- }
-
- async fn load_source_state(&self, id: &str) -> Result<Option<SourceState>, Error> {
- use tokio::fs;
-
- let path = self.get_state_path(id);
- if !path.exists() {
- return Ok(None);
- }
-
- let content = fs::read_to_string(path)
- .await
- .map_err(|e| Error::Storage(format!("Failed to read state file: {e}")))?;
-
- let state: SourceState = serde_json::from_str(&content).map_err(|e| {
- Error::Serialization(format!("Failed to deserialize source state: {e}"))
- })?;
-
- Ok(Some(state))
- }
-
- async fn save_sink_state(&self, state: &SinkState) -> Result<(), Error> {
- use tokio::fs;
-
- // Ensure directory exists
- if let Some(parent) = self.base_path.parent() {
- fs::create_dir_all(parent)
- .await
- .map_err(|e| Error::Storage(format!("Failed to create state directory: {e}")))?;
- }
-
- let path = self.get_state_path(&state.id);
- let json = serde_json::to_string_pretty(state)
- .map_err(|e| Error::Serialization(format!("Failed to serialize sink state: {e}")))?;
-
- fs::write(path, json)
- .await
- .map_err(|e| Error::Storage(format!("Failed to write state file: {e}")))?;
-
- Ok(())
- }
-
- async fn load_sink_state(&self, id: &str) -> Result<Option<SinkState>, Error> {
- use tokio::fs;
-
- let path = self.get_state_path(id);
- if !path.exists() {
- return Ok(None);
- }
-
- let content = fs::read_to_string(path)
- .await
- .map_err(|e| Error::Storage(format!("Failed to read state file: {e}")))?;
-
- let state: SinkState = serde_json::from_str(&content)
- .map_err(|e| Error::Serialization(format!("Failed to deserialize sink state: {e}")))?;
-
- Ok(Some(state))
- }
-
- async fn delete_state(&self, id: &str) -> Result<(), Error> {
- use tokio::fs;
-
- let path = self.get_state_path(id);
- if path.exists() {
- fs::remove_file(path)
- .await
- .map_err(|e| Error::Storage(format!("Failed to delete state file: {e}")))?;
- }
-
- Ok(())
- }
-
- async fn list_states(&self) -> Result<Vec<String>, Error> {
- use tokio::fs;
-
- let mut states = Vec::new();
-
- if !self.base_path.exists() {
- return Ok(states);
- }
-
- let mut entries = fs::read_dir(&self.base_path)
- .await
- .map_err(|e| Error::Storage(format!("Failed to read state directory: {e}")))?;
-
- while let Some(entry) = entries
- .next_entry()
- .await
- .map_err(|e| Error::Storage(format!("Failed to read directory entry: {e}")))?
- {
- if let Some(extension) = entry.path().extension()
- && extension == "json"
- && let Some(stem) = entry.path().file_stem()
- && let Some(id) = stem.to_str()
- {
- states.push(id.to_string());
- }
- }
-
- Ok(states)
- }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -455,7 +221,7 @@ pub struct ReceivedMessage {
pub struct ProducedMessages {
pub schema: Schema,
pub messages: Vec<ProducedMessage>,
- pub state: Option<serde_json::Value>,
+ pub state: Option<ConnectorState>,
}
#[repr(C)]
diff --git a/core/connectors/sdk/src/source.rs b/core/connectors/sdk/src/source.rs
index e86207e43..f4f2d1b39 100644
--- a/core/connectors/sdk/src/source.rs
+++ b/core/connectors/sdk/src/source.rs
@@ -1,4 +1,5 @@
-/* Licensed to the Apache Software Foundation (ASF) under one
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
@@ -16,7 +17,7 @@
* under the License.
*/
-use crate::{Error, Source, get_runtime};
+use crate::{ConnectorState, Error, Source, get_runtime};
use serde::de::DeserializeOwned;
use std::sync::Arc;
use tokio::{sync::watch, task::JoinHandle};
@@ -66,7 +67,7 @@ impl<T: Source + std::fmt::Debug + 'static> SourceContainer<T> {
factory: F,
) -> i32
where
- F: FnOnce(u32, C, Option<serde_json::Value>) -> T,
+ F: FnOnce(u32, C, Option<ConnectorState>) -> T,
C: DeserializeOwned,
{
unsafe {
@@ -89,10 +90,7 @@ impl<T: Source + std::fmt::Debug + 'static> SourceContainer<T> {
None
} else {
let state = std::slice::from_raw_parts(state_ptr, state_len);
- let Ok(state) = serde_json::from_slice(state) else {
- error!("Failed to parse state for source connector with ID: {id}",);
- return -1;
- };
+ let state = ConnectorState(state.to_vec());
Some(state)
};
diff --git a/core/connectors/sources/elasticsearch_source/src/lib.rs b/core/connectors/sources/elasticsearch_source/src/lib.rs
index ed52c3813..56a709365 100644
--- a/core/connectors/sources/elasticsearch_source/src/lib.rs
+++ b/core/connectors/sources/elasticsearch_source/src/lib.rs
@@ -16,7 +16,6 @@
* specific language governing permissions and limitations
* under the License.
*/
-
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use elasticsearch::{
@@ -25,8 +24,7 @@ use elasticsearch::{
http::{Url, transport::TransportBuilder},
};
use iggy_connector_sdk::{
- Error, FileStateStorage, ProducedMessage, ProducedMessages, Schema, Source, SourceState,
- StateStorage, source_connector,
+ ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector,
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
@@ -37,7 +35,8 @@ use tokio::{sync::Mutex, time::sleep};
use tracing::{info, warn};
mod state_manager;
-pub use state_manager::{StateInfo, StateManager, StateManagerExt, StateStats};
+use crate::state_manager::{FileStateStorage, SourceState, StateStorage};
+pub use state_manager::{StateInfo, StateManager, StateStats};
source_connector!(ElasticsearchSource);
@@ -80,7 +79,7 @@ pub struct StateConfig {
/// State storage type: "file", "elasticsearch", "redis", etc.
pub storage_type: Option<String>,
/// State storage configuration (depends on storage_type)
- pub storage_config: Option<serde_json::Value>,
+ pub storage_config: Option<Value>,
/// State ID for this connector instance
pub state_id: Option<String>,
/// Auto-save state interval (e.g., "30s", "5m")
@@ -95,7 +94,7 @@ pub struct ElasticsearchSourceConfig {
pub index: String,
pub username: Option<String>,
pub password: Option<String>,
- pub query: Option<serde_json::Value>,
+ pub query: Option<Value>,
pub polling_interval: Option<String>,
pub batch_size: Option<usize>,
pub timestamp_field: Option<String>,
@@ -113,11 +112,7 @@ pub struct ElasticsearchSource {
}
impl ElasticsearchSource {
- pub fn new(
- id: u32,
- config: ElasticsearchSourceConfig,
- _state: Option<serde_json::Value>,
- ) -> Self {
+ pub fn new(id: u32, config: ElasticsearchSourceConfig, _state: Option<ConnectorState>) -> Self {
let polling_interval = config
.polling_interval
.as_deref()
@@ -519,13 +514,16 @@ impl Source for ElasticsearchSource {
return Err(e);
}
};
-
- let state_guard = self.state.lock().await;
- let state_value = serde_json::to_value(&*state_guard).ok();
+ let state_value = {
+ let state = self.state.lock().await;
+ serde_json::to_vec(&*state).map_err(|err| {
+ Error::Serialization(format!("Failed to serialize state: {}", err))
+ })?
+ };
Ok(ProducedMessages {
schema: Schema::Json,
messages,
- state: state_value,
+ state: Some(ConnectorState(state_value)),
})
}
@@ -559,91 +557,4 @@ impl Source for ElasticsearchSource {
);
Ok(())
}
-
- async fn get_state(&self) -> Result<Option<SourceState>, Error> {
- if self
- .config
- .state
- .as_ref()
- .map(|s| s.enabled)
- .unwrap_or(false)
- {
- Ok(Some(self.internal_state_to_source_state().await?))
- } else {
- Ok(None)
- }
- }
-
- async fn set_state(&mut self, state: SourceState) -> Result<(), Error> {
- if self
- .config
- .state
- .as_ref()
- .map(|s| s.enabled)
- .unwrap_or(false)
- {
- self.source_state_to_internal_state(state).await
- } else {
- Ok(())
- }
- }
-
- async fn save_state(&self) -> Result<(), Error> {
- if !self
- .config
- .state
- .as_ref()
- .map(|s| s.enabled)
- .unwrap_or(false)
- {
- return Ok(());
- }
-
- let storage = self
- .create_state_storage()
- .ok_or_else(|| Error::Storage("State storage not configured".to_string()))?;
-
- let source_state = self.internal_state_to_source_state().await?;
- storage.save_source_state(&source_state).await?;
-
- info!(
- "Saved state for Elasticsearch source connector with ID: {}",
- self.id
- );
- Ok(())
- }
-
- async fn load_state(&mut self) -> Result<(), Error> {
- if !self
- .config
- .state
- .as_ref()
- .map(|s| s.enabled)
- .unwrap_or(false)
- {
- return Ok(());
- }
-
- let storage = self
- .create_state_storage()
- .ok_or_else(|| Error::Storage("State storage not configured".to_string()))?;
-
- let state_id = self.get_state_id();
- if let Some(source_state) = storage.load_source_state(&state_id).await? {
- self.source_state_to_internal_state(source_state).await?;
-
- let state = self.state.lock().await;
- info!(
- "Loaded state for Elasticsearch source connector with ID: {} - last poll: {:?}, total docs: {}, polls: {}",
- self.id, state.last_poll_timestamp, state.total_documents_fetched, state.poll_count
- );
- } else {
- info!(
- "No existing state found for Elasticsearch source connector with ID: {}, starting fresh",
- self.id
- );
- }
-
- Ok(())
- }
}
diff --git a/core/connectors/sources/elasticsearch_source/src/state_manager.rs b/core/connectors/sources/elasticsearch_source/src/state_manager.rs
index d310e7493..0bcb63cf9 100644
--- a/core/connectors/sources/elasticsearch_source/src/state_manager.rs
+++ b/core/connectors/sources/elasticsearch_source/src/state_manager.rs
@@ -18,13 +18,89 @@
*/
use crate::{ElasticsearchSource, StateConfig};
-use iggy_connector_sdk::{Error, FileStateStorage, Source, SourceState, StateStorage};
-use serde_json::Value;
+use async_trait::async_trait;
+use iggy_connector_sdk::Error;
+use serde::{Deserialize, Serialize};
use std::str::FromStr;
use std::sync::Arc;
use tokio::time::{Duration, interval};
use tracing::{error, info, warn};
+impl ElasticsearchSource {
+ async fn get_state(&self) -> Result<Option<SourceState>, Error> {
+ if self
+ .config
+ .state
+ .as_ref()
+ .map(|s| s.enabled)
+ .unwrap_or(false)
+ {
+ Ok(Some(self.internal_state_to_source_state().await?))
+ } else {
+ Ok(None)
+ }
+ }
+
+ pub(super) async fn save_state(&self) -> Result<(), Error> {
+ if !self
+ .config
+ .state
+ .as_ref()
+ .map(|s| s.enabled)
+ .unwrap_or(false)
+ {
+ return Ok(());
+ }
+
+ let storage = self
+ .create_state_storage()
+ .ok_or_else(|| Error::Storage("State storage not configured".to_string()))?;
+
+ let source_state = self.internal_state_to_source_state().await?;
+ storage.save_source_state(&source_state).await?;
+
+ info!(
+ "Saved state for Elasticsearch source connector with ID: {}",
+ self.id
+ );
+ Ok(())
+ }
+
+ pub(super) async fn load_state(&mut self) -> Result<(), Error> {
+ if !self
+ .config
+ .state
+ .as_ref()
+ .map(|s| s.enabled)
+ .unwrap_or(false)
+ {
+ return Ok(());
+ }
+
+ let storage = self
+ .create_state_storage()
+ .ok_or_else(|| Error::Storage("State storage not configured".to_string()))?;
+
+ let state_id = self.get_state_id();
+ if let Some(source_state) = storage.load_source_state(&state_id).await? {
+ self.source_state_to_internal_state(source_state).await?;
+
+ let state = self.state.lock().await;
+ info!(
+ "Loaded state for Elasticsearch source connector with ID: {} - last poll: {:?}, total docs: {}, polls: {}",
+ self.id, state.last_poll_timestamp, state.total_documents_fetched, state.poll_count
+ );
+ } else {
+ info!(
+ "No existing state found for Elasticsearch source connector with ID: {}, starting fresh",
+ self.id
+ );
+ }
+
+ Ok(())
+ }
+}
+
/// State manager for Elasticsearch source connector
pub struct StateManager {
storage: Arc<dyn StateStorage>,
@@ -81,7 +157,7 @@ impl StateManager {
}
/// Start auto-save background task
- pub async fn start_auto_save(&self, connector: std::sync::Arc<ElasticsearchSource>) {
+ pub async fn start_auto_save(&self, connector: Arc<ElasticsearchSource>) {
let interval_duration = self
.auto_save_interval
.unwrap_or_else(|| Duration::from_secs(60));
@@ -91,7 +167,7 @@ impl StateManager {
let mut interval = interval(interval_duration);
loop {
interval.tick().await;
- if let Ok(Some(state)) = Source::get_state(&*connector).await {
+ if let Ok(Some(state)) = connector.get_state().await {
if let Err(e) = storage.save_source_state(&state).await {
error!(
"Failed to auto-save state for {}: {}",
@@ -178,66 +254,136 @@ pub struct StateInfo {
pub connector_type: String,
}
-/// Extension trait for ElasticsearchSource to add state management utilities
-#[allow(async_fn_in_trait)]
-pub trait StateManagerExt {
- /// Get state manager for this connector
- fn get_state_manager(&self) -> Option<StateManager>;
+/// State management for source connectors
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SourceState {
+ /// Unique identifier for this state
+ pub id: String,
+ /// Timestamp when this state was last updated
+ pub last_updated: chrono::DateTime<chrono::Utc>,
+ /// Version of the state format
+ pub version: u32,
+ /// Generic state data as JSON
+ pub data: serde_json::Value,
+ /// Optional metadata
+ pub metadata: Option<serde_json::Value>,
+}
+
+/// State storage backend trait
+#[async_trait]
+pub trait StateStorage: Send + Sync {
+ /// Save source state to storage
+ async fn save_source_state(&self, state: &SourceState) -> Result<(), Error>;
- /// Export current state to JSON
- async fn export_state(&self) -> Result<Value, Error>;
+ /// Load source state from storage
+ async fn load_source_state(&self, id: &str) -> Result<Option<SourceState>, Error>;
- /// Import state from JSON
- async fn import_state(&mut self, state_json: Value) -> Result<(), Error>;
+ /// Delete state from storage
+ async fn delete_state(&self, id: &str) -> Result<(), Error>;
- /// Reset state (clear all state data)
- async fn reset_state(&mut self) -> Result<(), Error>;
+ /// List all state IDs
+ async fn list_states(&self) -> Result<Vec<String>, Error>;
}
-impl StateManagerExt for ElasticsearchSource {
- fn get_state_manager(&self) -> Option<StateManager> {
- self.config.state.as_ref().and_then(|config| {
- if config.enabled {
- StateManager::new(config.clone()).ok()
- } else {
- None
- }
- })
+/// File-based state storage implementation
+pub struct FileStateStorage {
+ base_path: std::path::PathBuf,
+}
+
+impl FileStateStorage {
+ pub fn new<P: AsRef<std::path::Path>>(base_path: P) -> Self {
+ Self {
+ base_path: base_path.as_ref().to_path_buf(),
+ }
}
- async fn export_state(&self) -> Result<Value, Error> {
- let state = Source::get_state(self).await?;
- serde_json::to_value(state).map_err(|e| Error::Storage(e.to_string()))
+ fn get_state_path(&self, id: &str) -> std::path::PathBuf {
+ self.base_path.join(format!("{id}.json"))
}
+}
+
+#[async_trait]
+impl StateStorage for FileStateStorage {
+ async fn save_source_state(&self, state: &SourceState) -> Result<(), Error> {
+ use tokio::fs;
+
+ // Ensure directory exists
+ if let Some(parent) = self.base_path.parent() {
+ fs::create_dir_all(parent)
+ .await
+ .map_err(|e| Error::Storage(format!("Failed to create state directory: {e}")))?;
+ }
- async fn import_state(&mut self, state_json: Value) -> Result<(), Error> {
- let source_state: SourceState =
- serde_json::from_value(state_json).map_err(|e| Error::Storage(e.to_string()))?;
- Source::set_state(self, source_state).await
+ let path = self.get_state_path(&state.id);
+ let json = serde_json::to_string_pretty(state)
+ .map_err(|e| Error::Serialization(format!("Failed to serialize source state: {e}")))?;
+
+ fs::write(path, json)
+ .await
+ .map_err(|e| Error::Storage(format!("Failed to write state file: {e}")))?;
+
+ Ok(())
}
- async fn reset_state(&mut self) -> Result<(), Error> {
- let mut state = self.state.lock().await;
- *state = crate::State {
- last_poll_timestamp: None,
- total_documents_fetched: 0,
- poll_count: 0,
- last_document_id: None,
- last_scroll_id: None,
- last_offset: None,
- error_count: 0,
- last_error: None,
- processing_stats: crate::ProcessingStats {
- total_bytes_processed: 0,
- avg_batch_processing_time_ms: 0.0,
- last_successful_poll: None,
- empty_polls_count: 0,
- successful_polls_count: 0,
- },
- };
- drop(state);
+ async fn load_source_state(&self, id: &str) -> Result<Option<SourceState>, Error> {
+ use tokio::fs;
+
+ let path = self.get_state_path(id);
+ if !path.exists() {
+ return Ok(None);
+ }
+
+ let content = fs::read_to_string(path)
+ .await
+ .map_err(|e| Error::Storage(format!("Failed to read state file: {e}")))?;
+
+ let state: SourceState = serde_json::from_str(&content).map_err(|e| {
+ Error::Serialization(format!("Failed to deserialize source state: {e}"))
+ })?;
+
+ Ok(Some(state))
+ }
+
+ async fn delete_state(&self, id: &str) -> Result<(), Error> {
+ use tokio::fs;
+
+ let path = self.get_state_path(id);
+ if path.exists() {
+ fs::remove_file(path)
+ .await
+ .map_err(|e| Error::Storage(format!("Failed to delete state file: {e}")))?;
+ }
+
+ Ok(())
+ }
+
+ async fn list_states(&self) -> Result<Vec<String>, Error> {
+ use tokio::fs;
+
+ let mut states = Vec::new();
+
+ if !self.base_path.exists() {
+ return Ok(states);
+ }
+
+ let mut entries = fs::read_dir(&self.base_path)
+ .await
+ .map_err(|e| Error::Storage(format!("Failed to read state directory: {e}")))?;
+
+ while let Some(entry) = entries
+ .next_entry()
+ .await
+ .map_err(|e| Error::Storage(format!("Failed to read directory entry: {e}")))?
+ {
+ if let Some(extension) = entry.path().extension()
+ && extension == "json"
+ && let Some(stem) = entry.path().file_stem()
+ && let Some(id) = stem.to_str()
+ {
+ states.push(id.to_string());
+ }
+ }
- // Save the reset state
- Source::save_state(self).await
+ Ok(states)
}
}
diff --git a/core/connectors/sources/postgres_source/src/lib.rs b/core/connectors/sources/postgres_source/src/lib.rs
index d38ca7ea5..27b8e3efa 100644
--- a/core/connectors/sources/postgres_source/src/lib.rs
+++ b/core/connectors/sources/postgres_source/src/lib.rs
@@ -1,4 +1,5 @@
-/* Licensed to the Apache Software Foundation (ASF) under one
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
@@ -21,7 +22,7 @@ use std::{collections::HashMap, str::FromStr, time::Duration};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use iggy_connector_sdk::{
- Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector,
+ ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector,
};
use serde::{Deserialize, Serialize};
use sqlx::{Column, Pool, Postgres, Row, TypeInfo, postgres::PgPoolOptions};
@@ -76,7 +77,7 @@ pub struct DatabaseRecord {
}
impl PostgresSource {
- pub fn new(id: u32, config: PostgresSourceConfig, _state: Option<serde_json::Value>) -> Self {
+ pub fn new(id: u32, config: PostgresSourceConfig, _state: Option<ConnectorState>) -> Self {
PostgresSource {
id,
pool: None,
diff --git a/core/connectors/sources/random_source/Cargo.toml b/core/connectors/sources/random_source/Cargo.toml
index 378ec2aa0..ad645601d 100644
--- a/core/connectors/sources/random_source/Cargo.toml
+++ b/core/connectors/sources/random_source/Cargo.toml
@@ -42,7 +42,6 @@ iggy_connector_sdk = { workspace = true }
once_cell = { workspace = true }
rand = { workspace = true }
serde = { workspace = true }
-serde_json = { workspace = true }
simd-json = { workspace = true }
tokio = { workspace = true }
tracing = { workspace = true }
diff --git a/core/connectors/sources/random_source/src/lib.rs b/core/connectors/sources/random_source/src/lib.rs
index 9c0497164..ef13b6413 100644
--- a/core/connectors/sources/random_source/src/lib.rs
+++ b/core/connectors/sources/random_source/src/lib.rs
@@ -1,4 +1,5 @@
-/* Licensed to the Apache Software Foundation (ASF) under one
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
@@ -16,17 +17,16 @@
* under the License.
*/
-use std::{str::FromStr, time::Duration};
-
use async_trait::async_trait;
use iggy_connector_sdk::{
- Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector,
+ ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector,
};
use rand::{
Rng,
distr::{Alphanumeric, Uniform},
};
use serde::{Deserialize, Serialize};
+use std::{str::FromStr, time::Duration};
use tokio::{sync::Mutex, time::sleep};
use tracing::{error, info};
use uuid::Uuid;
@@ -57,20 +57,23 @@ struct State {
}
impl RandomSource {
- pub fn new(id: u32, config: RandomSourceConfig, state: Option<serde_json::Value>) -> Self {
+ pub fn new(id: u32, config: RandomSourceConfig, state: Option<ConnectorState>) -> Self {
let interval = config.interval.unwrap_or("1s".to_string());
let interval = humantime::Duration::from_str(&interval)
.unwrap_or(humantime::Duration::from_str("1s").expect("Failed to parse interval"));
let current_number = if let Some(state) = state {
- serde_json::from_value::<usize>(state)
- .inspect_err(|error| {
- error!("Failed to convert state to current number. {error}");
- })
- .unwrap_or_default()
+ u64::from_le_bytes(
+ state.0[0..8]
+ .try_into()
+ .inspect_err(|error| {
+ error!("Failed to convert state to current number. {error}");
+ })
+ .unwrap_or_default(),
+ )
} else {
0
- };
+ } as usize;
RandomSource {
id,
@@ -152,7 +155,7 @@ impl Source for RandomSource {
return Ok(ProducedMessages {
schema: Schema::Json,
messages: vec![],
- state: Some(serde_json::to_value(state.current_number).unwrap()),
+ state: Some(ConnectorState(state.current_number.to_le_bytes().to_vec())),
});
}
@@ -166,7 +169,7 @@ impl Source for RandomSource {
Ok(ProducedMessages {
schema: Schema::Json,
messages,
- state: Some(serde_json::to_value(state.current_number).unwrap()),
+ state: Some(ConnectorState(state.current_number.to_le_bytes().to_vec())),
})
}