retry on pragma.data_version error in sqlite instead of failing (#9825)

GitOrigin-RevId: ec31b9d781ea4bd365368b93b68771416be59842

Author: zxqfd555

src/connectors/data_storage.rs

@@ -8,18 +8,15 @@ use std::any::type_name;
 use std::borrow::Borrow;
 use std::borrow::Cow;
 use std::collections::HashMap;
-use std::collections::HashSet;
-use std::collections::VecDeque;
 use std::fmt;
 use std::fmt::{Debug, Display};
 use std::io;
 use std::io::BufRead;
 use std::io::BufWriter;
 use std::io::Write;
 use std::mem::take;
-use std::str::{from_utf8, Utf8Error};
+use std::str::Utf8Error;
 use std::sync::Arc;
-use std::thread::sleep;
 use std::time::Duration;
 
 use arcstr::ArcStr;
@@ -60,18 +57,15 @@ use crate::connectors::data_format::{
     FormatterError, COMMIT_LITERAL,
 };
 use crate::connectors::data_lake::buffering::IncorrectSnapshotError;
-use crate::connectors::metadata::{KafkaMetadata, SQLiteMetadata, SourceMetadata};
-use crate::connectors::offset::EMPTY_OFFSET;
+use crate::connectors::metadata::{KafkaMetadata, SourceMetadata};
 use crate::connectors::posix_like::PosixLikeReader;
 use crate::connectors::scanner::s3::S3CommandName;
 use crate::connectors::{Offset, OffsetKey, OffsetValue, SPECIAL_FIELD_DIFF, SPECIAL_FIELD_TIME};
 use crate::engine::error::limit_length;
 use crate::engine::error::DynResult;
 use crate::engine::error::STANDARD_OBJECT_LENGTH_LIMIT;
 use crate::engine::time::DateTime;
-use crate::engine::DateTimeNaive;
-use crate::engine::Type;
-use crate::engine::{Duration as EngineDuration, Key, Value};
+use crate::engine::{DateTimeNaive, Duration as EngineDuration, Key, Type, Value};
 use crate::persistence::backends::Error as PersistenceBackendError;
 use crate::persistence::frontier::OffsetAntichain;
 use crate::persistence::tracker::WorkerPersistentStorage;
@@ -107,8 +101,6 @@ use rdkafka::producer::{BaseRecord, DefaultProducerContext, Producer, ThreadedPr
 use rdkafka::topic_partition_list::Offset as KafkaOffset;
 use rdkafka::Message;
 use rdkafka::TopicPartitionList;
-use rusqlite::types::ValueRef as SqliteValue;
-use rusqlite::Connection as SqliteConnection;
 use rusqlite::Error as SqliteError;
 use serde::{Deserialize, Serialize};
 
@@ -119,6 +111,7 @@ pub use super::data_lake::iceberg::IcebergReader;
 pub use super::data_lake::LakeWriter;
 pub use super::nats::NatsReader;
 pub use super::nats::NatsWriter;
+pub use super::sqlite::SqliteReader;
 
 #[derive(Clone, Debug, Eq, PartialEq, Copy)]
 pub enum DataEventType {
@@ -1932,207 +1925,6 @@ impl Writer for NullWriter {
     }
 }
 
-const SQLITE_DATA_VERSION_PRAGMA: &str = "data_version";
-
-pub struct SqliteReader {
-    connection: SqliteConnection,
-    table_name: String,
-    schema: Vec<(String, Type)>,
-
-    last_saved_data_version: Option<i64>,
-    stored_state: HashMap<i64, ValuesMap>,
-    queued_updates: VecDeque<ReadResult>,
-}
-
-impl SqliteReader {
-    pub fn new(
-        connection: SqliteConnection,
-        table_name: String,
-        schema: Vec<(String, Type)>,
-    ) -> Self {
-        Self {
-            connection,
-            table_name,
-            schema,
-
-            last_saved_data_version: None,
-            queued_updates: VecDeque::new(),
-            stored_state: HashMap::new(),
-        }
-    }
-
-    /// Data version is required to check if there was an update in the database.
-    /// There are also hooks, but they only work for changes happened in the same
-    /// connection.
-    /// More details why hooks don't help here: <https://sqlite.org/forum/forumpost/3174b39eeb79b6a4>
-    pub fn data_version(&self) -> i64 {
-        let version: ::rusqlite::Result<i64> = self.connection.pragma_query_value(
-            Some(::rusqlite::DatabaseName::Main),
-            SQLITE_DATA_VERSION_PRAGMA,
-            |row| row.get(0),
-        );
-        version.expect("pragma.data_version request should not fail")
-    }
-
-    /// Convert raw `SQLite` field into one of internal value types
-    /// There are only five supported types: null, integer, real, text, blob
-    /// See also: <https://www.sqlite.org/datatype3.html>
-    fn convert_to_value(
-        orig_value: SqliteValue<'_>,
-        field_name: &str,
-        dtype: &Type,
-    ) -> Result<Value, Box<ConversionError>> {
-        let value = match (dtype, orig_value) {
-            (Type::Optional(_) | Type::Any, SqliteValue::Null) => Some(Value::None),
-            (Type::Optional(arg), value) => Self::convert_to_value(value, field_name, arg).ok(),
-            (Type::Int | Type::Any, SqliteValue::Integer(val)) => Some(Value::Int(val)),
-            (Type::Float | Type::Any, SqliteValue::Real(val)) => Some(Value::Float(val.into())),
-            (Type::String | Type::Any, SqliteValue::Text(val)) => from_utf8(val)
-                .ok()
-                .map(|parsed_string| Value::String(parsed_string.into())),
-            (Type::Json, SqliteValue::Text(val)) => from_utf8(val)
-                .ok()
-                .and_then(|parsed_string| {
-                    serde_json::from_str::<serde_json::Value>(parsed_string).ok()
-                })
-                .map(Value::from),
-            (Type::Bytes | Type::Any, SqliteValue::Blob(val)) => Some(Value::Bytes(val.into())),
-            _ => None,
-        };
-        if let Some(value) = value {
-            Ok(value)
-        } else {
-            let value_repr = limit_length(format!("{orig_value:?}"), STANDARD_OBJECT_LENGTH_LIMIT);
-            Err(Box::new(ConversionError::new(
-                value_repr,
-                field_name.to_owned(),
-                dtype.clone(),
-                None,
-            )))
-        }
-    }
-
-    fn load_table(&mut self) -> Result<(), ReadError> {
-        let column_names: Vec<&str> = self
-            .schema
-            .iter()
-            .map(|(name, _dtype)| name.as_str())
-            .collect();
-        let query = format!(
-            "SELECT {},_rowid_ FROM {}",
-            column_names.join(","),
-            self.table_name
-        );
-
-        let mut statement = self.connection.prepare(&query)?;
-        let mut rows = statement.query([])?;
-
-        let mut present_rowids = HashSet::new();
-        while let Some(row) = rows.next()? {
-            let rowid: i64 = row.get(self.schema.len())?;
-            let mut values = HashMap::with_capacity(self.schema.len());
-            for (column_idx, (column_name, column_dtype)) in self.schema.iter().enumerate() {
-                let value =
-                    Self::convert_to_value(row.get_ref(column_idx)?, column_name, column_dtype);
-                values.insert(column_name.clone(), value);
-            }
-            let values: ValuesMap = values.into();
-            self.stored_state
-                .entry(rowid)
-                .and_modify(|current_values| {
-                    if current_values != &values {
-                        let key = vec![Value::Int(rowid)];
-                        self.queued_updates.push_back(ReadResult::Data(
-                            ReaderContext::from_diff(
-                                DataEventType::Delete,
-                                Some(key.clone()),
-                                take(current_values),
-                            ),
-                            EMPTY_OFFSET,
-                        ));
-                        self.queued_updates.push_back(ReadResult::Data(
-                            ReaderContext::from_diff(
-                                DataEventType::Insert,
-                                Some(key),
-                                values.clone(),
-                            ),
-                            EMPTY_OFFSET,
-                        ));
-                        current_values.clone_from(&values);
-                    }
-                })
-                .or_insert_with(|| {
-                    let key = vec![Value::Int(rowid)];
-                    self.queued_updates.push_back(ReadResult::Data(
-                        ReaderContext::from_diff(DataEventType::Insert, Some(key), values.clone()),
-                        EMPTY_OFFSET,
-                    ));
-                    values
-                });
-            present_rowids.insert(rowid);
-        }
-
-        self.stored_state.retain(|rowid, values| {
-            if present_rowids.contains(rowid) {
-                true
-            } else {
-                let key = vec![Value::Int(*rowid)];
-                self.queued_updates.push_back(ReadResult::Data(
-                    ReaderContext::from_diff(DataEventType::Delete, Some(key), take(values)),
-                    EMPTY_OFFSET,
-                ));
-                false
-            }
-        });
-
-        if !self.queued_updates.is_empty() {
-            self.queued_updates.push_back(ReadResult::FinishedSource {
-                commit_possibility: CommitPossibility::Possible,
-            });
-        }
-
-        Ok(())
-    }
-
-    fn wait_period() -> Duration {
-        Duration::from_millis(500)
-    }
-}
-
-impl Reader for SqliteReader {
-    fn seek(&mut self, _frontier: &OffsetAntichain) -> Result<(), ReadError> {
-        todo!("seek is not supported for Sqlite source: persistent history of changes unavailable")
-    }
-
-    fn read(&mut self) -> Result<ReadResult, ReadError> {
-        loop {
-            if let Some(queued_update) = self.queued_updates.pop_front() {
-                return Ok(queued_update);
-            }
-
-            let current_data_version = self.data_version();
-            if self.last_saved_data_version != Some(current_data_version) {
-                self.load_table()?;
-                self.last_saved_data_version = Some(current_data_version);
-                return Ok(ReadResult::NewSource(
-                    SQLiteMetadata::new(current_data_version).into(),
-                ));
-            }
-            // Sleep to avoid non-stop pragma requests of a table
-            // that did not change
-            sleep(Self::wait_period());
-        }
-    }
-
-    fn short_description(&self) -> Cow<'static, str> {
-        format!("SQLite({})", self.table_name).into()
-    }
-
-    fn storage_type(&self) -> StorageType {
-        StorageType::Sqlite
-    }
-}
-
 #[derive(Debug)]
 enum BufferedMongoEvent {
     Insert {

src/connectors/mod.rs

@@ -29,6 +29,7 @@ pub mod nats;
 pub mod offset;
 pub mod posix_like;
 pub mod scanner;
+pub mod sqlite;
 pub mod synchronization;
 
 use crate::connectors::monitoring::ConnectorMonitor;