dms-cdc-operator 0.1.26

The dms-cdc-operator is a Rust-based utility for comparing the state of a list of tables in an Amazon RDS database with data stored in Parquet files on Amazon S3, particularly useful for change data capture (CDC) scenarios
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

#[cfg(test)]
mod tests {
    use indexmap::IndexMap;
    use mockall::predicate::*;
    use polars::prelude::*;

    use crate::postgres::postgres_operator::{
        InsertDataframePayload, MockPostgresOperator, PostgresOperator, UpsertDataframePayload,
    };

    #[tokio::test]
    async fn test_get_table_columns() {
        let mut postgres_operator = MockPostgresOperator::new();
        postgres_operator
            .expect_get_table_columns()
            .times(1)
            .with(eq("schema"), eq("table"))
            .returning(|_, _| {
                let mut columns = IndexMap::new();
                columns.insert("column1".to_string(), "text".to_string());
                columns.insert("column2".to_string(), "text".to_string());
                Ok(columns)
            });

        let result = postgres_operator
            .get_table_columns("schema", "table")
            .await
            .unwrap();
        assert_eq!(result.len(), 2);
        assert_eq!(result.get("column1").unwrap(), "text");
        assert_eq!(result.get("column2").unwrap(), "text");
    }

    #[tokio::test]
    async fn test_get_primary_key() {
        let mut postgres_operator = MockPostgresOperator::new();
        postgres_operator
            .expect_get_primary_key()
            .times(1)
            .with(eq("table"), eq("schema"))
            .returning(|_, _| Ok(vec!["primary_key".to_string()]));

        let result = postgres_operator
            .get_primary_key("table", "schema")
            .await
            .unwrap();
        assert_eq!(result, vec!["primary_key"]);
    }

    #[tokio::test]
    async fn test_create_table() {
        let mut postgres_operator = MockPostgresOperator::new();
        postgres_operator
            .expect_create_table()
            .times(1)
            .returning(|_, _, _, _| Ok(()));

        let mut column_data_types = IndexMap::new();
        column_data_types.insert("column1".to_string(), "text".to_string());
        column_data_types.insert("column2".to_string(), "text".to_string());

        postgres_operator
            .create_table(
                &column_data_types,
                vec!["primary_key".to_string()].as_slice(),
                "schema",
                "table",
            )
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_insert_dataframe_in_target_db() {
        let mut postgres_operator = MockPostgresOperator::new();
        postgres_operator
            .expect_insert_dataframe_in_target_db()
            .times(1)
            .returning(|_, _| Ok(()));

        let df = DataFrame::new(vec![Series::new("column1".into(), &[1, 2, 3]).into()]).unwrap();
        let payload = InsertDataframePayload {
            database_name: "database".to_string(),
            schema_name: "schema".to_string(),
            table_name: "table".to_string(),
        };

        postgres_operator
            .insert_dataframe_in_target_db(&df, &payload)
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_upsert_dataframe_in_target_db() {
        let mut postgres_operator = MockPostgresOperator::new();
        postgres_operator
            .expect_upsert_dataframe_in_target_db()
            .times(1)
            .returning(|_, _| Ok(()));

        let df = DataFrame::new(vec![Series::new("column1".into(), &[1, 2, 3]).into()]).unwrap();
        let payload = UpsertDataframePayload {
            database_name: "database".to_string(),
            schema_name: "schema".to_string(),
            table_name: "table".to_string(),
            primary_key: "primary_key".to_string(),
        };
        postgres_operator
            .upsert_dataframe_in_target_db(&df, &payload)
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_close_connection_pool() {
        let mut postgres_operator = MockPostgresOperator::new();
        postgres_operator
            .expect_close_connection_pool()
            .times(1)
            .return_const(());

        postgres_operator.close_connection_pool().await;
    }
}