rust_pgdatadiff/diff/table/query/
table_query.rs

1use crate::diff::table::query::table_types::{
2    IncludedExcludedTables, TableMode, TableName, TableOffset, TablePosition, TablePrimaryKeys,
3};
4use crate::diff::types::SchemaName;
5use std::fmt::Display;
6
7pub enum TableQuery {
8    AllTablesForSchema(SchemaName, IncludedExcludedTables),
9    CountRowsForTable(SchemaName, TableName),
10    FindPrimaryKeyForTable(TableName),
11    HashQuery(
12        SchemaName,
13        TableName,
14        TablePrimaryKeys,
15        TablePosition,
16        TableOffset,
17    ),
18}
19
20impl Display for TableQuery {
21    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22        match self {
23            Self::AllTablesForSchema(schema_name, included_excluded_tables) => {
24                let inclusion_exclusion_statement = match included_excluded_tables.table_mode() {
25                    None => "".to_string(),
26                    Some(table_mode) => match table_mode {
27                        TableMode::Include => included_excluded_tables.inclusion_statement(),
28                        TableMode::Exclude => included_excluded_tables.exclusion_statement(),
29                    },
30                };
31
32                write!(
33                    f,
34                    r#"
35                SELECT table_name
36                FROM information_schema.tables
37                WHERE table_schema = '{}'
38                {}
39                "#,
40                    schema_name.name(),
41                    inclusion_exclusion_statement
42                )
43            }
44            // https://stackoverflow.com/questions/7943233/fast-way-to-discover-the-row-count-of-a-table-in-postgresql
45            TableQuery::CountRowsForTable(schema_name, table_name) => {
46                write!(
47                    f,
48                    "SELECT count(*) FROM {}.{}",
49                    schema_name.name(),
50                    table_name.name()
51                )
52            }
53            TableQuery::FindPrimaryKeyForTable(table_name) => write!(
54                f,
55                // language=postgresql
56                r#"
57                SELECT a.attname
58                FROM   pg_index i
59                JOIN   pg_attribute a ON a.attrelid = i.indrelid
60                                     AND a.attnum = ANY(i.indkey)
61                WHERE  i.indrelid = '{}'::regclass
62                AND    i.indisprimary"#,
63                table_name.name()
64            ),
65            TableQuery::HashQuery(
66                schema_name,
67                table_name,
68                table_primary_keys,
69                table_position,
70                table_offset,
71            ) => {
72                write!(
73                    f,
74                    r#"
75                    SELECT md5(array_agg(md5((t.*)::varchar))::varchar)
76                    FROM (
77                        SELECT *
78                        FROM {}.{}
79                        ORDER BY {} limit {} offset {}
80                    ) AS t
81                    "#,
82                    schema_name.name(),
83                    table_name.name(),
84                    table_primary_keys.keys(),
85                    table_offset.offset(),
86                    table_position.position(),
87                )
88            }
89        }
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use pretty_assertions::assert_eq;
97
98    #[test]
99    fn test_display_all_tables_for_schema_with_included_tables() {
100        let schema_name = SchemaName::new("public");
101        let included_tables = vec!["table1".to_string(), "table2".to_string()];
102        let excluded_tables: Vec<String> = vec![];
103        let included_excluded_tables =
104            IncludedExcludedTables::new(included_tables, excluded_tables);
105        let query = TableQuery::AllTablesForSchema(schema_name, included_excluded_tables);
106        let expected = r#"
107                SELECT table_name
108                FROM information_schema.tables
109                WHERE table_schema = 'public'
110                AND table_name IN ('table1','table2')
111                "#;
112        assert_eq!(expected, query.to_string());
113    }
114
115    #[test]
116    fn test_display_all_tables_for_schema_with_excluded_tables() {
117        let schema_name = SchemaName::new("public");
118        let included_tables: Vec<String> = vec![];
119        let excluded_tables = vec!["table1", "table2"];
120        let included_excluded_tables =
121            IncludedExcludedTables::new(included_tables, excluded_tables);
122        let query = TableQuery::AllTablesForSchema(schema_name, included_excluded_tables);
123        let expected = r#"
124                SELECT table_name
125                FROM information_schema.tables
126                WHERE table_schema = 'public'
127                AND table_name NOT IN ('table1','table2')
128                "#;
129        assert_eq!(expected, query.to_string());
130    }
131
132    #[test]
133    fn test_display_count_rows_for_table() {
134        let schema_name = SchemaName::new("public".to_string());
135        let table_name = TableName::new("table1".to_string());
136        let query = TableQuery::CountRowsForTable(schema_name, table_name);
137        let expected = "SELECT count(*) FROM public.table1";
138        assert_eq!(expected, query.to_string());
139    }
140
141    #[test]
142    fn test_display_find_primary_key_for_table() {
143        let table_name = TableName::new("table1".to_string());
144        let query = TableQuery::FindPrimaryKeyForTable(table_name);
145        let expected = r#"
146                SELECT a.attname
147                FROM   pg_index i
148                JOIN   pg_attribute a ON a.attrelid = i.indrelid
149                                     AND a.attnum = ANY(i.indkey)
150                WHERE  i.indrelid = 'table1'::regclass
151                AND    i.indisprimary"#;
152        assert_eq!(expected, query.to_string());
153    }
154
155    #[test]
156    fn test_display_hash_query() {
157        let schema_name = SchemaName::new("public".to_string());
158        let table_name = TableName::new("table1".to_string());
159        let table_primary_keys = TablePrimaryKeys::new("id".to_string());
160        let table_position = TablePosition::new(0);
161        let table_offset = TableOffset::new(100);
162        let query = TableQuery::HashQuery(
163            schema_name,
164            table_name,
165            table_primary_keys,
166            table_position,
167            table_offset,
168        );
169        let expected = r#"
170                    SELECT md5(array_agg(md5((t.*)::varchar))::varchar)
171                    FROM (
172                        SELECT *
173                        FROM public.table1
174                        ORDER BY id limit 100 offset 0
175                    ) AS t
176                    "#;
177        assert_eq!(expected, query.to_string());
178    }
179}