4195006582

Build # 4195006582

Build Type

github

Committed by GitHub

Commit Message

Merge 14ceac9a6 into dc166625f

Pull Request Pull Request #918: ObjectStore - Reduce code duplication & more clone prevention

Run Details

177 of 177 new or added lines in 5 files covered. (100.0%)

24368 of 37519 relevant lines covered (64.95%)

43913.78 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/dozer-ingestion/src/connectors/object_store/table_reader.rs

use crate::connectors::object_store::helper::{get_details, get_table, map_listing_options};
use crate::connectors::object_store::schema_helper::map_value_to_dozer_field;
use crate::connectors::{ColumnInfo, TableInfo};
use crate::errors::ObjectStoreConnectorError::TableReaderError;
use crate::errors::ObjectStoreObjectError::ListingPathParsingError;
use crate::errors::ObjectStoreTableReaderError::{
    ColumnsSelectFailed, StreamExecutionError, TableReadFailed,
};
use crate::errors::{ConnectorError, ObjectStoreConnectorError};
use crate::ingestion::Ingestor;
use datafusion::datasource::listing::{
    ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
};
use datafusion::prelude::SessionContext;
use dozer_types::ingestion_types::{IngestionMessage, LocalStorage, S3Storage, Table};
use dozer_types::parking_lot::RwLock;
use dozer_types::types::{Operation, Record, SchemaIdentifier};
use futures::StreamExt;
use object_store::aws::{AmazonS3, AmazonS3Builder};
use object_store::local::LocalFileSystem;
use object_store::ObjectStore;
use std::sync::Arc;
use tokio::runtime::Runtime;

pub struct TableReader<T: Clone + Send + Sync> {
    config: T,
}

impl<T: Clone + Send + Sync> TableReader<T> {
    pub fn new(config: T) -> TableReader<T> {
        Self { config }
    }

    pub async fn read(
        id: u32,
        ctx: SessionContext,
        table_path: ListingTableUrl,
        listing_options: ListingOptions,
        ingestor: Arc<RwLock<Ingestor>>,
        table: &TableInfo,
    ) -> Result<(), ObjectStoreConnectorError> {
        let resolved_schema = listing_options
            .infer_schema(&ctx.state(), &table_path)
            .await
            .map_err(|_| ObjectStoreConnectorError::InternalError)?;

        let mut idx = 0;
        let fields = resolved_schema.all_fields();

        let config = ListingTableConfig::new(table_path)
            .with_listing_options(listing_options)
            .with_schema(resolved_schema.clone());

        let provider = Arc::new(
            ListingTable::try_new(config)
                .map_err(ObjectStoreConnectorError::InternalDataFusionError)?,
        );

        let columns: Vec<ColumnInfo> = match &table.columns {
            Some(columns_list) if !columns_list.is_empty() => columns_list.clone(),
            _ => fields
                .iter()
                .map(|f| ColumnInfo {
                    name: f.name().to_string(),
                    data_type: Some(f.data_type().to_string()),
                })
                .collect(),
        };

        let cols: Vec<&str> = columns.iter().map(|c| c.name.as_str()).collect();
        let data = ctx
            .read_table(provider.clone())
            .map_err(|e| TableReaderError(TableReadFailed(e)))?
            .select_columns(&cols)
            .map_err(|e| TableReaderError(ColumnsSelectFailed(e)))?
            .execute_stream()
            .await
            .map_err(|e| TableReaderError(StreamExecutionError(e)))?;

        tokio::pin!(data);

        while let Some(Ok(batch)) = data.next().await {
            for row in 0..batch.num_rows() {
                let fields = batch
                    .columns()
                    .iter()
                    .enumerate()
                    .map(|(col, column)| {
                        map_value_to_dozer_field(column, &row, resolved_schema.field(col).name())
                    })
                    .collect::<Result<Vec<_>, _>>()?;

                ingestor
                    .write()
                    .handle_message((
                        (0_u64, idx),
                        IngestionMessage::OperationEvent(Operation::Insert {
                            new: Record {
                                schema_id: Some(SchemaIdentifier { id, version: 0 }),
                                values: fields,
                                version: None,
                            },
                        }),
                    ))
                    .unwrap();

                idx += 1;
            }
        }

        Ok(())
    }
}

pub trait DozerObjectStore: Clone + Send + Sync {
    type ObjectStore: ObjectStore;

    fn table_params<'a>(
        &'a self,
        table_name: &str,
    ) -> Result<DozerObjectStoreParams<'a, Self::ObjectStore>, ConnectorError>;
}

pub struct DozerObjectStoreParams<'a, T: ObjectStore> {
    pub scheme: &'static str,
    pub host: &'a str,
    pub object_store: T,
    pub table_path: String,
    pub data_fusion_table: &'a Table,
}

pub trait Reader<T> {
    fn read_tables(
        &self,
        tables: &[TableInfo],
        ingestor: &Arc<RwLock<Ingestor>>,
    ) -> Result<(), ConnectorError>;
}

impl<T: DozerObjectStore> Reader<T> for TableReader<T> {
    fn read_tables(
        &self,
        tables: &[TableInfo],
        ingestor: &Arc<RwLock<Ingestor>>,
    ) -> Result<(), ConnectorError> {
        for (id, table) in tables.iter().enumerate() {
            let params = self.config.table_params(&table.name)?;

            let table_path = ListingTableUrl::parse(params.table_path).map_err(|_| {
                ObjectStoreConnectorError::DataFusionStorageObjectError(ListingPathParsingError)
            })?;

            let listing_options = map_listing_options(params.data_fusion_table);

            let rt = Runtime::new().map_err(|_| ObjectStoreConnectorError::RuntimeCreationError)?;

            let ctx = SessionContext::new();

            ctx.runtime_env().register_object_store(
                params.scheme,
                params.host,
                Arc::new(params.object_store),
            );

            rt.block_on(Self::read(
                id as u32,
                ctx,
                table_path,
                listing_options,
                ingestor.clone(),
                table,
            ))?;
        }

        Ok(())
    }
}

impl DozerObjectStore for S3Storage {
    type ObjectStore = AmazonS3;

    fn table_params(
        &self,
        table_name: &str,
    ) -> Result<DozerObjectStoreParams<Self::ObjectStore>, ConnectorError> {
        let table = get_table(&self.tables, table_name)?;
        let details = get_details(&self.details)?;

        let object_store = AmazonS3Builder::new()
            .with_bucket_name(&details.bucket_name)
            .with_region(&details.region)
            .with_access_key_id(&details.access_key_id)
            .with_secret_access_key(&details.secret_access_key)
            .build()
            .map_err(|_| ConnectorError::InitializationError)?;

        Ok(DozerObjectStoreParams {
            scheme: "s3",
            host: &details.bucket_name,
            object_store,
            table_path: format!("s3://{}/{}/", details.bucket_name, table.prefix),
            data_fusion_table: table,
        })
    }
}

impl DozerObjectStore for LocalStorage {
    type ObjectStore = LocalFileSystem;

    fn table_params(
        &self,
        table_name: &str,
    ) -> Result<DozerObjectStoreParams<Self::ObjectStore>, ConnectorError> {
        let table = get_table(&self.tables, table_name)?;
        let path = get_details(&self.details)?.path.as_str();

        let object_store = LocalFileSystem::new_with_prefix(path)
            .map_err(|_| ConnectorError::InitializationError)?;

        Ok(DozerObjectStoreParams {
            scheme: "local",
            host: path,
            object_store,
            table_path: format!("s3://{path}/{}/", table.prefix),
            data_fusion_table: table,
        })
    }
}

1	use crate::connectors::object_store::helper::{get_details, get_table, map_listing_options};
2	use crate::connectors::object_store::schema_helper::map_value_to_dozer_field;
3	use crate::connectors::{ColumnInfo, TableInfo};
4	use crate::errors::ObjectStoreConnectorError::TableReaderError;
5	use crate::errors::ObjectStoreObjectError::ListingPathParsingError;
6	use crate::errors::ObjectStoreTableReaderError::{
7	ColumnsSelectFailed, StreamExecutionError, TableReadFailed,
8	};
9	use crate::errors::{ConnectorError, ObjectStoreConnectorError};
10	use crate::ingestion::Ingestor;
11	use datafusion::datasource::listing::{
12	ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
13	};
14	use datafusion::prelude::SessionContext;
15	use dozer_types::ingestion_types::{IngestionMessage, LocalStorage, S3Storage, Table};
16	use dozer_types::parking_lot::RwLock;
17	use dozer_types::types::{Operation, Record, SchemaIdentifier};
18	use futures::StreamExt;
19	use object_store::aws::{AmazonS3, AmazonS3Builder};
20	use object_store::local::LocalFileSystem;
21	use object_store::ObjectStore;
22	use std::sync::Arc;
23	use tokio::runtime::Runtime;
24
25	pub struct TableReader<T: Clone + Send + Sync> {
26	config: T,
27	}
28
29	impl<T: Clone + Send + Sync> TableReader<T> {
30	pub fn new(config: T) -> TableReader<T> {	×
31	Self { config }	×
32	}	×
33		×
34	pub async fn read(	×
35	id: u32,	×
36	ctx: SessionContext,	×
37	table_path: ListingTableUrl,	×
38	listing_options: ListingOptions,	×
39	ingestor: Arc<RwLock<Ingestor>>,	×
40	table: &TableInfo,	×
41	) -> Result<(), ObjectStoreConnectorError> {	×
42	let resolved_schema = listing_options	×
43	.infer_schema(&ctx.state(), &table_path)	×
44	.await	×
45	.map_err(\|_\| ObjectStoreConnectorError::InternalError)?;	×
46		×
47	let mut idx = 0;	×
48	let fields = resolved_schema.all_fields();	×
49		×
50	let config = ListingTableConfig::new(table_path)	×
51	.with_listing_options(listing_options)	×
52	.with_schema(resolved_schema.clone());	×
53		×
54	let provider = Arc::new(	×
55	ListingTable::try_new(config)	×
56	.map_err(ObjectStoreConnectorError::InternalDataFusionError)?,	×
57	);
58		×
59	let columns: Vec<ColumnInfo> = match &table.columns {	×
60	Some(columns_list) if !columns_list.is_empty() => columns_list.clone(),	×
61	_ => fields	×
62	.iter()	×
63	.map(\|f\| ColumnInfo {	×
64	name: f.name().to_string(),	×
65	data_type: Some(f.data_type().to_string()),	×
66	})	×
67	.collect(),	×
68	};
69		×
70	let cols: Vec<&str> = columns.iter().map(\|c\| c.name.as_str()).collect();	×
71	let data = ctx	×
72	.read_table(provider.clone())	×
73	.map_err(\|e\| TableReaderError(TableReadFailed(e)))?	×
74	.select_columns(&cols)	×
75	.map_err(\|e\| TableReaderError(ColumnsSelectFailed(e)))?	×
76	.execute_stream()	×
77	.await	×
78	.map_err(\|e\| TableReaderError(StreamExecutionError(e)))?;	×
79		×
80	tokio::pin!(data);	×
81		×
82	while let Some(Ok(batch)) = data.next().await {	×
83	for row in 0..batch.num_rows() {	×
84	let fields = batch	×
85	.columns()	×
86	.iter()	×
87	.enumerate()	×
88	.map(\|(col, column)\| {	×
89	map_value_to_dozer_field(column, &row, resolved_schema.field(col).name())	×
90	})	×
91	.collect::<Result<Vec<_>, _>>()?;	×
92
93	ingestor	×
94	.write()	×
95	.handle_message((	×
96	(0_u64, idx),	×
97	IngestionMessage::OperationEvent(Operation::Insert {	×
98	new: Record {	×
99	schema_id: Some(SchemaIdentifier { id, version: 0 }),	×
100	values: fields,	×
101	version: None,	×
102	},	×
103	}),	×
104	))	×
105	.unwrap();	×
106		×
107	idx += 1;	×
108	}	×
109	}
110
111	Ok(())	×
112	}	×
113	}	×
114		×
115	pub trait DozerObjectStore: Clone + Send + Sync {	×
116	type ObjectStore: ObjectStore;	×
117
118	fn table_params<'a>(
119	&'a self,
120	table_name: &str,
121	) -> Result<DozerObjectStoreParams<'a, Self::ObjectStore>, ConnectorError>;
122	}
123
124	pub struct DozerObjectStoreParams<'a, T: ObjectStore> {
125	pub scheme: &'static str,
126	pub host: &'a str,
127	pub object_store: T,
128	pub table_path: String,	×
129	pub data_fusion_table: &'a Table,	×
130	}	×
131		×
132	pub trait Reader<T> {	×
133	fn read_tables(	×
134	&self,	×
135	tables: &[TableInfo],	×
136	ingestor: &Arc<RwLock<Ingestor>>,	×
137	) -> Result<(), ConnectorError>;	×
138	}	×
139		×
140	impl<T: DozerObjectStore> Reader<T> for TableReader<T> {	×
141	fn read_tables(	×
142	&self,	×
143	tables: &[TableInfo],	×
144	ingestor: &Arc<RwLock<Ingestor>>,	×
145	) -> Result<(), ConnectorError> {	×
146	for (id, table) in tables.iter().enumerate() {	×
147	let params = self.config.table_params(&table.name)?;	×
148
149	let table_path = ListingTableUrl::parse(params.table_path).map_err(\|_\| {	×
150	ObjectStoreConnectorError::DataFusionStorageObjectError(ListingPathParsingError)	×
151	})?;	×
152		×
153	let listing_options = map_listing_options(params.data_fusion_table);	×
154		×
155	let rt = Runtime::new().map_err(\|_\| ObjectStoreConnectorError::RuntimeCreationError)?;	×
156		×
157	let ctx = SessionContext::new();	×
158		×
159	ctx.runtime_env().register_object_store(	×
160	params.scheme,	×
161	params.host,	×
162	Arc::new(params.object_store),	×
163	);	×
164		×
165	rt.block_on(Self::read(	×
166	id as u32,	×
167	ctx,	×
168	table_path,	×
169	listing_options,	×
170	ingestor.clone(),	×
171	table,	×
172	))?;	×
173	}	×
174		×
175	Ok(())	×
176	}	×
177	}
178		×
179	impl DozerObjectStore for S3Storage {	×
180	type ObjectStore = AmazonS3;	×
181		×
182	fn table_params(	×
183	&self,	×
184	table_name: &str,	×
185	) -> Result<DozerObjectStoreParams<Self::ObjectStore>, ConnectorError> {	×
186	let table = get_table(&self.tables, table_name)?;	×
187	let details = get_details(&self.details)?;	×
188		×
189	let object_store = AmazonS3Builder::new()	×
190	.with_bucket_name(&details.bucket_name)	×
191	.with_region(&details.region)	×
192	.with_access_key_id(&details.access_key_id)	×
193	.with_secret_access_key(&details.secret_access_key)	×
194	.build()	×
195	.map_err(\|_\| ConnectorError::InitializationError)?;	×
196		×
197	Ok(DozerObjectStoreParams {	×
198	scheme: "s3",	×
199	host: &details.bucket_name,	×
200	object_store,	×
201	table_path: format!("s3://{}/{}/", details.bucket_name, table.prefix),	×
202	data_fusion_table: table,	×
203	})	×
204	}	×
205	}	×
206		×
207	impl DozerObjectStore for LocalStorage {	×
208	type ObjectStore = LocalFileSystem;	×
209		×
210	fn table_params(	×
211	&self,	×
212	table_name: &str,	×
213	) -> Result<DozerObjectStoreParams<Self::ObjectStore>, ConnectorError> {	×
214	let table = get_table(&self.tables, table_name)?;	×
215	let path = get_details(&self.details)?.path.as_str();	×
216
217	let object_store = LocalFileSystem::new_with_prefix(path)	×
218	.map_err(\|_\| ConnectorError::InitializationError)?;	×
219		×
220	Ok(DozerObjectStoreParams {	×
221	scheme: "local",	×
222	host: path,	×
223	object_store,	×
224	table_path: format!("s3://{path}/{}/", table.prefix),	×
225	data_fusion_table: table,	×
226	})	×
227	}	×
228	}	×

getdozer / dozer / 4195006582

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous