feat: Implement schema support for postgres (#1070)

getdozer · Mar 13, 2023 · 8c70609 · 8c70609
1 parent cfbd9bc
commit 8c70609
Show file tree

Hide file tree

Showing 30 changed files with 212 additions and 70 deletions.
diff --git a/dozer-ingestion/examples/postgres/main.rs b/dozer-ingestion/examples/postgres/main.rs
@@ -13,6 +13,7 @@ fn main() {
     let (ingestor, mut iterator) = Ingestor::initialize_channel(IngestionConfig::default());
     let tables = vec![TableInfo {
         name: "users".to_string(),
+        schema: Some("public".to_string()),
         columns: None,
     }];
     let postgres_config = PostgresConfig {

diff --git a/dozer-ingestion/src/connectors/delta_lake/schema_helper.rs b/dozer-ingestion/src/connectors/delta_lake/schema_helper.rs
@@ -45,6 +45,6 @@ impl SchemaHelper {
             .into();
         let schema_mapper = SchemaMapper::new(self.config.clone());
         let schema = schema_mapper.map_schema(id as u32, arrow_schema, table)?;
-        Ok(SourceSchema::new(table.name.clone(), schema, Nothing))
+        Ok(SourceSchema::new(table.name.clone(), None, schema, Nothing))
     }
 }
diff --git a/dozer-ingestion/src/connectors/delta_lake/test/deltalake_test.rs b/dozer-ingestion/src/connectors/delta_lake/test/deltalake_test.rs
@@ -23,6 +23,7 @@ fn get_schema_from_deltalake() {
     let connector = DeltaLakeConnector::new(1, config);
     let table_info = TableInfo {
         name: table_name.to_string(),
+        schema: Some("public".to_string()),
         columns: None,
     };
     let field = connector.get_schemas(Some(&vec![table_info])).unwrap()[0]
@@ -53,6 +54,7 @@ fn read_deltalake() {
     let (ingestor, iterator) = Ingestor::initialize_channel(config);
     let table = TableInfo {
         name: "test_table".to_string(),
+        schema: Some("public".to_string()),
         columns: None,
     };
     thread::spawn(move || {

diff --git a/dozer-ingestion/src/connectors/ethereum/log/connector.rs b/dozer-ingestion/src/connectors/ethereum/log/connector.rs
@@ -128,6 +128,7 @@ impl Connector for EthLogConnector {
     ) -> Result<Vec<SourceSchema>, ConnectorError> {
         let mut schemas = vec![SourceSchema::new(
             ETH_LOGS_TABLE.to_string(),
+            None,
             helper::get_eth_schema(),
             ReplicationChangesTrackingType::Nothing,
         )];

diff --git a/dozer-ingestion/src/connectors/ethereum/log/helper.rs b/dozer-ingestion/src/connectors/ethereum/log/helper.rs
@@ -51,6 +51,7 @@ pub fn get_contract_event_schemas(
 
             schemas.push(SourceSchema::new(
                 get_table_name(&contract_tuple, &event.name),
+                None,
                 Schema {
                     identifier: Some(SchemaIdentifier {
                         id: schema_id as u32,

diff --git a/dozer-ingestion/src/connectors/ethereum/log/tests/helper.rs b/dozer-ingestion/src/connectors/ethereum/log/tests/helper.rs
@@ -77,6 +77,7 @@ pub fn get_eth_producer(
     let schemas = eth_connector.get_schemas(None)?;
     for SourceSchema {
         name,
+        schema_name: _,
         schema,
         replication_type: _,
     } in schemas

diff --git a/dozer-ingestion/src/connectors/ethereum/trace/connector.rs b/dozer-ingestion/src/connectors/ethereum/trace/connector.rs
@@ -40,6 +40,7 @@ impl Connector for EthTraceConnector {
     ) -> Result<Vec<SourceSchema>, ConnectorError> {
         Ok(vec![SourceSchema::new(
             ETH_TRACE_TABLE.to_string(),
+            None,
             helper::get_trace_schema(),
             ReplicationChangesTrackingType::Nothing,
         )])

diff --git a/dozer-ingestion/src/connectors/grpc/adapter/arrow.rs b/dozer-ingestion/src/connectors/grpc/adapter/arrow.rs
@@ -44,11 +44,12 @@ impl ArrowAdapter {
 
             arrow_schemas.insert(id as u32, grpc_schema.schema);
 
-            schemas.push(SourceSchema {
-                name: grpc_schema.name,
+            schemas.push(SourceSchema::new(
+                grpc_schema.name,
+                None,
                 schema,
-                replication_type: grpc_schema.replication_type.clone(),
-            });
+                grpc_schema.replication_type.clone(),
+            ));
         }
         Ok((schemas, arrow_schemas))
     }

diff --git a/dozer-ingestion/src/connectors/kafka/debezium/no_schema_registry.rs b/dozer-ingestion/src/connectors/kafka/debezium/no_schema_registry.rs
@@ -53,6 +53,7 @@ impl NoSchemaRegistry {
 
                             schemas.push(SourceSchema::new(
                                 table.name.clone(),
+                                None,
                                 mapped_schema,
                                 ReplicationChangesTrackingType::FullChanges,
                             ));

diff --git a/dozer-ingestion/src/connectors/kafka/debezium/schema_registry.rs b/dozer-ingestion/src/connectors/kafka/debezium/schema_registry.rs
@@ -141,6 +141,7 @@ impl SchemaRegistry {
 
                                 schema_data = Some(Ok(vec![SourceSchema::new(
                                     table.name.clone(),
+                                    None,
                                     schema,
                                     ReplicationChangesTrackingType::FullChanges,
                                 )]));

diff --git a/dozer-ingestion/src/connectors/mod.rs b/dozer-ingestion/src/connectors/mod.rs
@@ -70,6 +70,7 @@ pub trait Connector: Send + Sync + Debug {
                         })
                         .collect(),
                 ),
+                schema: source_schema.schema_name,
             })
             .collect::<Vec<TableInfo>>())
     }
@@ -79,6 +80,7 @@ pub trait Connector: Send + Sync + Debug {
 #[serde(crate = "self::serde")]
 pub struct TableInfo {
     pub name: String,
+    pub schema: Option<String>,
     pub columns: Option<Vec<ColumnInfo>>,
 }
 

diff --git a/dozer-ingestion/src/connectors/object_store/schema_mapper.rs b/dozer-ingestion/src/connectors/object_store/schema_mapper.rs
@@ -69,6 +69,7 @@ impl<T: DozerObjectStore> Mapper<T> for SchemaMapper<T> {
                 .map(|t| TableInfo {
                     name: t.name.clone(),
                     columns: None,
+                    schema: None,
                 })
                 .collect()
         });
@@ -107,7 +108,7 @@ impl<T: DozerObjectStore> Mapper<T> for SchemaMapper<T> {
 
                 let schema = self.map_schema(id as u32, resolved_schema, table)?;
 
-                Ok(SourceSchema::new(table_name, schema, Nothing))
+                Ok(SourceSchema::new(table_name, None, schema, Nothing))
             })
             .collect()
     }

diff --git a/dozer-ingestion/src/connectors/object_store/tests/local_storage_tests.rs b/dozer-ingestion/src/connectors/object_store/tests/local_storage_tests.rs
@@ -75,6 +75,7 @@ fn test_read_parquet_file() {
 
     let table = TableInfo {
         name: "all_types_parquet".to_string(),
+        schema: None,
         columns: None,
     };
     thread::spawn(move || {
@@ -125,6 +126,7 @@ fn test_csv_read() {
 
     let table = TableInfo {
         name: "all_types_csv".to_string(),
+        schema: None,
         columns: None,
     };
 
@@ -200,6 +202,7 @@ fn test_missing_directory() {
         &ingestor,
         vec![TableInfo {
             name: table.name,
+            schema: None,
             columns: None,
         }],
     );

diff --git a/dozer-ingestion/src/connectors/postgres/connection/validator.rs b/dozer-ingestion/src/connectors/postgres/connection/validator.rs
@@ -428,6 +428,7 @@ mod tests {
 
             let tables = vec![TableInfo {
                 name: "not_existing".to_string(),
+                schema: Some("public".to_string()),
                 columns: None,
             }];
             let result = validate_connection("pg_test_conn", config, Some(&tables), None);
@@ -476,6 +477,7 @@ mod tests {
 
             let tables = vec![TableInfo {
                 name: "existing".to_string(),
+                schema: Some("public".to_string()),
                 columns: Some(columns),
             }];
 
@@ -654,6 +656,7 @@ mod tests {
         for (table_name, expected_result) in tables_with_result {
             let res = validate_tables_names(&vec![TableInfo {
                 name: table_name.to_string(),
+                schema: Some("public".to_string()),
                 columns: None,
             }]);
 
@@ -673,6 +676,7 @@ mod tests {
         for (column_name, expected_result) in columns_names_with_result {
             let res = validate_columns_names(&vec![TableInfo {
                 name: "column_test_table".to_string(),
+                schema: Some("public".to_string()),
                 columns: Some(vec![ColumnInfo {
                     name: column_name.to_string(),
                     data_type: None,
@@ -707,6 +711,7 @@ mod tests {
                 &mut pg_client,
                 &vec![TableInfo {
                     name: table_name,
+                    schema: Some("public".to_string()),
                     columns: None,
                 }],
             );
@@ -717,6 +722,7 @@ mod tests {
                 &mut pg_client,
                 &vec![TableInfo {
                     name: view_name,
+                    schema: Some("public".to_string()),
                     columns: None,
                 }],
             );

diff --git a/dozer-ingestion/src/connectors/postgres/connector.rs b/dozer-ingestion/src/connectors/postgres/connector.rs
@@ -44,7 +44,7 @@ impl PostgresConnector {
         let mut replication_conn_config = config.config.clone();
         replication_conn_config.replication_mode(ReplicationMode::Logical);
 
-        let helper = SchemaHelper::new(config.config.clone(), None);
+        let helper = SchemaHelper::new(config.config.clone());
 
         // conn_str - replication_conn_config
         // conn_str_plain- conn_config
@@ -101,6 +101,7 @@ impl Connector for PostgresConnector {
 
         let lsn = PostgresConnector::get_lsn_with_offset_from_seq(self.name.clone(), from_seq);
 
+        info!("TABLES iN STRT {:?}", tables);
         let iterator = PostgresIterator::new(
             self.id,
             self.name.clone(),
@@ -137,6 +138,7 @@ impl Connector for PostgresConnector {
             .map(|table_info| TableInfo {
                 name: table_info.name,
                 columns: Some(table_info.columns),
+                schema: Some(table_info.schema),
             })
             .collect())
     }
@@ -169,7 +171,18 @@ impl PostgresConnector {
         let table_str: String = match self.tables.as_ref() {
             None => "ALL TABLES".to_string(),
             Some(arr) => {
-                let table_names: Vec<String> = arr.iter().map(|t| t.name.clone()).collect();
+                let table_names: Vec<String> = arr
+                    .iter()
+                    .map(|t| {
+                        format!(
+                            "{}.{}",
+                            t.schema
+                                .as_ref()
+                                .map_or("public".to_string(), |s| s.clone()),
+                            t.name.clone()
+                        )
+                    })
+                    .collect();
                 format!("TABLE {}", table_names.join(" , "))
             }
         };

diff --git a/dozer-ingestion/src/connectors/postgres/iterator.rs b/dozer-ingestion/src/connectors/postgres/iterator.rs
@@ -171,6 +171,7 @@ impl<'a> PostgresIteratorHandler<'a> {
                 .map(|table_info| TableInfo {
                     name: table_info.name.clone(),
                     columns: Some(table_info.columns.clone()),
+                    schema: Some(table_info.schema.clone()),
                 })
                 .collect::<Vec<_>>();
             snapshotter.sync_tables(&tables)?;
-Original file line number
+Diff line change
@@ Expand Up / @@ -141,6 +141,7 @@ impl SchemaRegistry { @@
                                     schema_data = Some(Ok(vec![SourceSchema::new(
                                         table.name.clone(),
+                                        None,
                                         schema,
                                         ReplicationChangesTrackingType::FullChanges,
                                     )]));
@@ Expand Down @@