rust_pgdatadiff/diff/table/query/
table_query.rs1use 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 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 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}