connectorx/sources/
mod.rs

1//! This module defines four traits [`Source`], [`SourcePartition`], [`PartitionParser`], and [`Produce`]  to define a source.
2//! This module also contains source implementations for various databases.
3
4#[cfg(feature = "src_bigquery")]
5pub mod bigquery;
6#[cfg(feature = "src_clickhouse")]
7pub mod clickhouse;
8#[cfg(feature = "src_csv")]
9pub mod csv;
10#[cfg(feature = "src_dummy")]
11pub mod dummy;
12#[cfg(feature = "src_mssql_common")]
13pub mod mssql;
14#[cfg(feature = "src_mysql")]
15pub mod mysql;
16#[cfg(feature = "src_oracle")]
17pub mod oracle;
18#[cfg(feature = "src_postgres")]
19pub mod postgres;
20#[cfg(feature = "src_sqlite")]
21pub mod sqlite;
22#[cfg(feature = "src_trino")]
23pub mod trino;
24
25use crate::data_order::DataOrder;
26use crate::errors::ConnectorXError;
27use crate::sql::CXQuery;
28use crate::typesystem::{TypeAssoc, TypeSystem};
29use std::fmt::Debug;
30
31/// A `Source` connects to a database, fetches the metadata of a query and splits the work into
32/// [`SourcePartition`]s, one per partitioned query.
33///
34/// [`Dispatcher`](crate::dispatcher::Dispatcher) drives a `Source` in this order:
35/// [`set_data_order`](Source::set_data_order), [`set_queries`](Source::set_queries),
36/// [`set_origin_query`](Source::set_origin_query), [`fetch_metadata`](Source::fetch_metadata),
37/// then [`schema`](Source::schema) and [`names`](Source::names), optionally
38/// [`result_rows`](Source::result_rows), and finally [`partition`](Source::partition).
39pub trait Source {
40    /// Supported data orders, ordering by preference.
41    const DATA_ORDERS: &'static [DataOrder];
42    /// The type system this `Source` associated with.
43    type TypeSystem: TypeSystem;
44    /// The partition type produced by [`partition`](Source::partition). It is sent to a worker
45    /// thread, so it must be `Send`.
46    type Partition: SourcePartition<TypeSystem = Self::TypeSystem, Error = Self::Error> + Send;
47    /// The error type returned by this source and its partitions.
48    type Error: From<ConnectorXError> + Send + Debug;
49
50    /// Set the data order negotiated between this source and the destination, one of
51    /// [`DATA_ORDERS`](Source::DATA_ORDERS). Return an error if the order is not supported.
52    fn set_data_order(&mut self, data_order: DataOrder) -> Result<(), Self::Error>;
53
54    /// Set the partitioned queries. Each query becomes one [`SourcePartition`].
55    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]);
56
57    /// Set the original, unpartitioned query, if there is one. Sources can use it to
58    /// count the rows of the whole result in [`result_rows`](Source::result_rows).
59    fn set_origin_query(&mut self, query: Option<String>);
60
61    /// Set queries to run on each connection before the partitioned queries.
62    /// The default implementation panics; sources that support it override it.
63    fn set_pre_execution_queries(&mut self, _pre_execution_queries: Option<&[String]>) {
64        unimplemented!("pre_execution_queries is not implemented in this source type");
65    }
66
67    /// Fetch the result schema (column names and types) of the queries, e.g. by preparing the
68    /// statement or running a `LIMIT 0` query. Called before [`names`](Source::names) and
69    /// [`schema`](Source::schema).
70    fn fetch_metadata(&mut self) -> Result<(), Self::Error>;
71    /// Get total number of rows if available.
72    ///
73    /// Return `Ok(None)` when the total cannot be derived; the dispatcher then counts each
74    /// partition with [`SourcePartition::result_rows`] instead.
75    fn result_rows(&mut self) -> Result<Option<usize>, Self::Error>;
76
77    /// The column names of the result, available after [`fetch_metadata`](Source::fetch_metadata).
78    fn names(&self) -> Vec<String>;
79
80    /// The column types of the result, available after [`fetch_metadata`](Source::fetch_metadata).
81    fn schema(&self) -> Vec<Self::TypeSystem>;
82
83    /// Consume the source and create one partition per query given to
84    /// [`set_queries`](Source::set_queries), in the same order.
85    fn partition(self) -> Result<Vec<Self::Partition>, Self::Error>;
86}
87
88/// In general, a `DataSource` abstracts the data source as a stream, which can produce
89/// a sequence of values of variate types by repetitively calling the function `produce`.
90pub trait SourcePartition {
91    /// The type system of the parent [`Source`].
92    type TypeSystem: TypeSystem;
93    /// The parser that reads the rows of this partition, borrowing from it.
94    type Parser<'a>: PartitionParser<'a, TypeSystem = Self::TypeSystem, Error = Self::Error>
95    where
96        Self: 'a;
97    /// The error type of the parent [`Source`].
98    type Error: From<ConnectorXError> + Send + Debug;
99
100    /// Count total number of rows in each partition.
101    ///
102    /// The count is stored in the partition and reported by [`nrows`](SourcePartition::nrows).
103    /// It is only called when the destination needs a count and
104    /// [`Source::result_rows`] returned `None`.
105    fn result_rows(&mut self) -> Result<(), Self::Error>;
106
107    /// Run the partition's query and return a parser over its result.
108    fn parser(&mut self) -> Result<Self::Parser<'_>, Self::Error>;
109
110    /// Number of rows this `DataSource` got.
111    /// Sometimes it is not possible for the source to know how many rows it gets before reading the whole data.
112    fn nrows(&self) -> usize;
113
114    /// Number of cols this `DataSource` got.
115    fn ncols(&self) -> usize;
116}
117
118/// A `PartitionParser` reads the values of one partition's result, one cell at a time.
119///
120/// The dispatcher calls [`fetch_next`](PartitionParser::fetch_next) to buffer a batch of rows,
121/// then [`parse`](PartitionParser::parse) once per cell in the negotiated
122/// [`DataOrder`]. Each cell's type comes from the source schema.
123pub trait PartitionParser<'a>: Send {
124    /// The type system of the parent [`Source`].
125    type TypeSystem: TypeSystem;
126    /// The error type of the parent [`Source`].
127    type Error: From<ConnectorXError> + Send + Debug;
128
129    /// Read a value `T` by calling `Produce<T>::produce`. Usually this function does not need to be
130    /// implemented.
131    fn parse<'r, T>(&'r mut self) -> Result<T, <Self as PartitionParser<'a>>::Error>
132    where
133        T: TypeAssoc<Self::TypeSystem>,
134        Self: Produce<'r, T, Error = <Self as PartitionParser<'a>>::Error>,
135    {
136        self.produce()
137    }
138
139    /// Fetch next batch of rows from database, return (number of rows fetched to local, whether all rows are fechted from database).
140    /// There might be rows that are not consumed yet when calling the next fetch_next.
141    /// The function might be called even after the last batch is fetched.
142    fn fetch_next(&mut self) -> Result<(usize, bool), Self::Error>;
143}
144
145/// A type implemented `Produce<T>` means that it can produce a value `T` by consuming part of it's raw data buffer.
146pub trait Produce<'r, T> {
147    /// The error returned when the next value cannot be produced as `T`.
148    type Error: From<ConnectorXError> + Send;
149
150    /// Produce the next value from the buffer as `T` and advance past it.
151    fn produce(&'r mut self) -> Result<T, Self::Error>;
152}
connectorx/sources/mod.rs

connectorx/sources/
mod.rs
