lezeh_db/psql/
db_metadata.rs

1use std::cell::RefCell;
2use std::collections::HashMap;
3use std::rc::Rc;
4
5use itertools::Itertools;
6use postgres::types::ToSql;
7use postgres::Row;
8
9use crate::psql::connection::PsqlConnection;
10use crate::psql::dto::*;
11use lezeh_common::types::ResultAnyError;
12
13pub type PsqlParamValue = Box<dyn ToSql + Sync>;
14
15const TABLE_WITH_FK_QUERY: &'static str = "
16    SELECT
17      tc.constraint_name,
18      tc.table_schema,
19      tc.table_name,
20      kcu.column_name,
21      c.data_type AS column_data_type,
22      ccu.table_schema AS foreign_table_schema,
23      ccu.table_name AS foreign_table_name,
24      ccu.column_name AS foreign_column_name,
25      foreign_c_meta.data_type AS foreign_column_data_type
26    FROM
27      information_schema.table_constraints AS tc
28        JOIN information_schema.key_column_usage AS kcu ON
29          tc.constraint_name = kcu.constraint_name AND
30          tc.table_schema = kcu.table_schema
31        JOIN information_schema.constraint_column_usage AS ccu ON
32          ccu.constraint_name = tc.constraint_name
33        JOIN information_schema.columns as c ON
34          c.table_name = tc.table_name AND
35          c.column_name = kcu.column_name
36        JOIN information_schema.columns as foreign_c_meta ON
37          foreign_c_meta.table_schema = ccu.table_schema AND
38          foreign_c_meta.table_name = ccu.table_name AND
39          foreign_c_meta.column_name = ccu.column_name
40    WHERE tc.constraint_type = 'FOREIGN KEY';
41";
42
43#[derive(PartialEq, Debug)]
44pub struct ForeignKeyInformationRow {
45  constraint_name: String,
46
47  // From table X
48  table_schema: String,
49  table_name: String,
50  column_name: String,
51  column_data_type: String,
52
53  // referencing to table Y
54  foreign_table_schema: String,
55  foreign_table_name: String,
56  foreign_column_name: String,
57  foreign_column_data_type: String,
58}
59
60pub struct Query {
61  connection: Rc<RefCell<PsqlConnection>>,
62}
63
64impl Query {
65  fn fetch_fk_info(&mut self, _schema: &str) -> ResultAnyError<Vec<ForeignKeyInformationRow>> {
66    // First try to build the UML for all of the tables
67    // we'll query from psql information_schema tables.
68    let rows: Vec<Row> = self
69      .connection
70      .borrow_mut()
71      .get()
72      .query(TABLE_WITH_FK_QUERY, &[])?;
73
74    let fk_info_rows: Vec<ForeignKeyInformationRow> = rows
75      .into_iter()
76      .map(|row: Row| -> ForeignKeyInformationRow {
77        return ForeignKeyInformationRow {
78          table_schema: row.get("table_schema"),
79          constraint_name: row.get("constraint_name"),
80          table_name: row.get("table_name"),
81          column_name: row.get("column_name"),
82          column_data_type: row.get("column_data_type"),
83          foreign_table_schema: row.get("foreign_table_schema"),
84          foreign_table_name: row.get("foreign_table_name"),
85          foreign_column_name: row.get("foreign_column_name"),
86          foreign_column_data_type: row.get("foreign_column_data_type"),
87        };
88      })
89      .collect();
90
91    return Ok(fk_info_rows);
92  }
93
94  fn get_table_by_id(&mut self) -> ResultAnyError<HashMap<PsqlTableIdentity, PsqlTable>> {
95    let rows: Vec<Row> = self.connection.borrow_mut().get().query(
96      "
97      SELECT
98        tc.constraint_name,
99        tc.table_schema,
100        tc.table_name,
101        kcu.column_name as primary_column_name,
102        c.data_type AS primary_column_data_type
103      FROM
104        information_schema.table_constraints AS tc
105          JOIN information_schema.key_column_usage AS kcu
106            ON tc.constraint_name = kcu.constraint_name
107            AND tc.table_schema = kcu.table_schema
108          JOIN information_schema.columns as c
109            ON c.table_schema = tc.table_schema
110            AND c.table_name = tc.table_name
111            AND c.column_name = kcu.column_name
112      WHERE tc.constraint_type = 'PRIMARY KEY' and
113       tc.table_schema not in ('pg_catalog', 'information_schema')
114      ",
115      &[],
116    )?;
117
118    let psql_table_by_id: HashMap<PsqlTableIdentity, PsqlTable> = rows
119      .into_iter()
120      .map(|row| {
121        let psql_table = PsqlTable::new(
122          row.get::<_, String>("table_schema"),
123          row.get::<_, String>("table_name"),
124          PsqlTableColumn::new(
125            row.get::<_, String>("primary_column_name"),
126            row.get::<_, String>("primary_column_data_type"),
127          ),
128          Default::default(),
129          Default::default(),
130          Default::default(),
131        );
132
133        return (psql_table.id.clone(), psql_table);
134      })
135      .collect();
136
137    return Ok(psql_table_by_id);
138  }
139}
140
141pub struct DbMetadata {
142  /// We know that we own this query so it's ok
143  /// to directl borrow_mut() without checking ownership
144  query: RefCell<Query>,
145}
146
147impl DbMetadata {
148  pub fn new(psql_connection: Rc<RefCell<PsqlConnection>>) -> DbMetadata {
149    return DbMetadata {
150      query: RefCell::new(Query {
151        connection: psql_connection,
152      }),
153    };
154  }
155}
156
157impl DbMetadata {
158  pub fn load_table_structure(
159    &self,
160    schema: &str,
161  ) -> ResultAnyError<HashMap<PsqlTableIdentity, PsqlTable>> {
162    let fk_info_rows = self.query.borrow_mut().fetch_fk_info(schema)?;
163
164    let mut table_by_id = self.query.borrow_mut().get_table_by_id()?;
165
166    psql_table_map_from_foreign_key_info_rows(&mut table_by_id, &fk_info_rows);
167
168    return Ok(table_by_id);
169  }
170}
171
172fn psql_table_map_from_foreign_key_info_rows(
173  table_by_id: &mut HashMap<PsqlTableIdentity, PsqlTable>,
174  rows: &Vec<ForeignKeyInformationRow>,
175) {
176  let fk_info_rows_by_foreign_table_id: HashMap<PsqlTableIdentity, Vec<&ForeignKeyInformationRow>> =
177    rows.iter().into_group_map_by(|row| {
178      return PsqlTableIdentity::new(&row.foreign_table_schema, &row.foreign_table_name);
179    });
180
181  let fk_info_rows_by_table_id: HashMap<PsqlTableIdentity, Vec<&ForeignKeyInformationRow>> =
182    rows.iter().into_group_map_by(|row| {
183      return PsqlTableIdentity::new(&row.table_schema, &row.table_name);
184    });
185
186  for (table_id, table) in table_by_id.into_iter() {
187    let referencing_fk_rows = fk_info_rows_by_table_id.get(&table_id);
188
189    if referencing_fk_rows.is_some() {
190      let referencing_fk_rows = referencing_fk_rows.unwrap();
191
192      table.referencing_fk_by_constraint_name = referencing_fk_rows
193        .iter()
194        .map(|fk_row| {
195          return (
196            fk_row.constraint_name.clone(),
197            PsqlForeignKey::new(
198              fk_row.constraint_name.clone(),
199              PsqlTableColumn::new(fk_row.column_name.clone(), fk_row.column_data_type.clone()),
200              fk_row.foreign_table_schema.clone(),
201              fk_row.foreign_table_name.clone(),
202            ),
203          );
204        })
205        .collect();
206    }
207
208    let referenced_fk_rows = fk_info_rows_by_foreign_table_id.get(&table_id);
209
210    if referenced_fk_rows.is_some() {
211      let referenced_fk_rows = referenced_fk_rows.unwrap();
212
213      table.referenced_fk_by_constraint_name = referenced_fk_rows
214        .iter()
215        .map(|fk_row| {
216          return (
217            fk_row.constraint_name.clone(),
218            PsqlForeignKey::new(
219              fk_row.constraint_name.clone(),
220              PsqlTableColumn::new(fk_row.column_name.clone(), fk_row.column_data_type.clone()),
221              fk_row.table_schema.clone(),
222              fk_row.table_name.clone(),
223            ),
224          );
225        })
226        .collect();
227    }
228  }
229}
230
231#[cfg(test)]
232mod test {
233  use super::*;
234  use std::borrow::Cow;
235  use std::collections::HashSet;
236
237  impl PsqlTable {
238    fn basic<'a, S>(schema: S, name: S, primary_column: PsqlTableColumn) -> PsqlTable
239    where
240      S: Into<Cow<'a, str>>,
241    {
242      return PsqlTable {
243        id: PsqlTableIdentity::new(schema, name),
244        primary_column,
245        columns: Default::default(),
246        referenced_fk_by_constraint_name: Default::default(),
247        referencing_fk_by_constraint_name: Default::default(),
248      };
249    }
250  }
251
252  mod psql_tables_from_foreign_key_info_rows {
253    use super::*;
254    use lezeh_common::macros::hashmap_literal;
255
256    #[test]
257    fn it_should_load_rows() {
258      // Db diagram view https://dbdiagram.io/d/6205540d85022f4ee57331e2
259      let fk_info_rows = vec![
260        ForeignKeyInformationRow {
261          table_schema: "public".into(),
262          constraint_name: "orders_store_id_foreign".into(),
263          table_name: "orders".into(),
264          column_name: "store_id".into(),
265          column_data_type: "integer".into(),
266          foreign_table_schema: "public".into(),
267          foreign_table_name: "stores".into(),
268          foreign_column_name: "id".into(),
269          foreign_column_data_type: "integer".into(),
270        },
271        ForeignKeyInformationRow {
272          table_schema: "public".into(),
273          constraint_name: "order_statuses_store_id_foreign".into(),
274          table_name: "order_statuses".into(),
275          column_name: "store_id".into(),
276          column_data_type: "integer".into(),
277          foreign_table_schema: "public".into(),
278          foreign_table_name: "stores".into(),
279          foreign_column_name: "id".into(),
280          foreign_column_data_type: "integer".into(),
281        },
282        ForeignKeyInformationRow {
283          table_schema: "public".into(),
284          constraint_name: "product_images_product_id_foreign".into(),
285          table_name: "product_images".into(),
286          column_name: "product_id".into(),
287          column_data_type: "integer".into(),
288          foreign_table_schema: "public".into(),
289          foreign_table_name: "products".into(),
290          foreign_column_name: "id".into(),
291          foreign_column_data_type: "integer".into(),
292        },
293        ForeignKeyInformationRow {
294          table_schema: "public".into(),
295          constraint_name: "product_stock_ledgers_product_id_foreign".into(),
296          table_name: "product_stock_ledgers".into(),
297          column_name: "product_id".into(),
298          column_data_type: "integer".into(),
299          foreign_table_schema: "public".into(),
300          foreign_table_name: "products".into(),
301          foreign_column_name: "id".into(),
302          foreign_column_data_type: "integer".into(),
303        },
304        ForeignKeyInformationRow {
305          table_schema: "public".into(),
306          constraint_name: "store_customers_store_id_foreign".into(),
307          table_name: "store_customers".into(),
308          column_name: "store_id".into(),
309          column_data_type: "integer".into(),
310          foreign_table_schema: "public".into(),
311          foreign_table_name: "stores".into(),
312          foreign_column_name: "id".into(),
313          foreign_column_data_type: "integer".into(),
314        },
315        ForeignKeyInformationRow {
316          table_schema: "public".into(),
317          constraint_name: "store_staffs_stores_store_staff_role_id_foreign".into(),
318          table_name: "store_staffs_stores".into(),
319          column_name: "store_staff_role_id".into(),
320          column_data_type: "uuid".into(),
321          foreign_table_schema: "public".into(),
322          foreign_table_name: "store_staff_roles".into(),
323          foreign_column_name: "id".into(),
324          foreign_column_data_type: "uuid".into(),
325        },
326        ForeignKeyInformationRow {
327          table_schema: "public".into(),
328          constraint_name: "store_staffs_stores_store_staff_id_foreign".into(),
329          table_name: "store_staffs_stores".into(),
330          column_name: "store_staff_id".into(),
331          column_data_type: "integer".into(),
332          foreign_table_schema: "public".into(),
333          foreign_table_name: "store_staffs".into(),
334          foreign_column_name: "id".into(),
335          foreign_column_data_type: "integer".into(),
336        },
337        ForeignKeyInformationRow {
338          table_schema: "public".into(),
339          constraint_name: "store_staffs_stores_store_id_foreign".into(),
340          table_name: "store_staffs_stores".into(),
341          column_name: "store_id".into(),
342          column_data_type: "integer".into(),
343          foreign_table_schema: "public".into(),
344          foreign_table_name: "stores".into(),
345          foreign_column_name: "id".into(),
346          foreign_column_data_type: "integer".into(),
347        },
348        ForeignKeyInformationRow {
349          table_schema: "public".into(),
350          constraint_name: "products_store_id_foreign".into(),
351          table_name: "products".into(),
352          column_name: "store_id".into(),
353          column_data_type: "integer".into(),
354          foreign_table_schema: "public".into(),
355          foreign_table_name: "stores".into(),
356          foreign_column_name: "id".into(),
357          foreign_column_data_type: "integer".into(),
358        },
359        ForeignKeyInformationRow {
360          table_schema: "public".into(),
361          constraint_name: "order_items_order_id_foreign".into(),
362          table_name: "order_items".into(),
363          column_name: "order_id".into(),
364          column_data_type: "integer".into(),
365          foreign_table_schema: "public".into(),
366          foreign_table_name: "orders".into(),
367          foreign_column_name: "id".into(),
368          foreign_column_data_type: "integer".into(),
369        },
370        ForeignKeyInformationRow {
371          table_schema: "public".into(),
372          constraint_name: "order_items_product_id_foreign".into(),
373          table_name: "order_items".into(),
374          column_name: "product_id".into(),
375          column_data_type: "integer".into(),
376          foreign_table_schema: "public".into(),
377          foreign_table_name: "products".into(),
378          foreign_column_name: "id".into(),
379          foreign_column_data_type: "integer".into(),
380        },
381      ];
382
383      let mut psql_table_by_id: HashMap<PsqlTableIdentity, PsqlTable> = hashmap_literal! {
384        PsqlTableIdentity::new("public", "stores") => PsqlTable::basic("public", "stores", PsqlTableColumn{
385          name: "id".into(),
386          data_type: "integer".into(),
387        }),
388        PsqlTableIdentity::new("public", "orders") => PsqlTable::basic("public", "orders", PsqlTableColumn{
389          name: "id".into(),
390          data_type: "integer".into(),
391        }),
392        PsqlTableIdentity::new("public", "order_items") => PsqlTable::basic("public", "order_items", PsqlTableColumn{
393          name: "id".into(),
394          data_type: "integer".into(),
395        }),
396        PsqlTableIdentity::new("public", "order_statuses") => PsqlTable::basic("public", "order_statuses", PsqlTableColumn{
397          name: "id".into(),
398          data_type: "integer".into(),
399        }),
400        PsqlTableIdentity::new("public", "products") => PsqlTable::basic("public", "products", PsqlTableColumn{
401          name: "id".into(),
402          data_type: "integer".into(),
403        }),
404        PsqlTableIdentity::new("public", "product_images") => PsqlTable::basic("public", "product_images", PsqlTableColumn{
405          name: "id".into(),
406          data_type: "integer".into(),
407        }),
408        PsqlTableIdentity::new("public", "product_stock_ledgers") => PsqlTable::basic("public", "product_stock_ledgers", PsqlTableColumn{
409          name: "id".into(),
410          data_type: "integer".into(),
411        }),
412        PsqlTableIdentity::new("public", "store_customers") => PsqlTable::basic("public", "store_customers", PsqlTableColumn{
413          name: "id".into(),
414          data_type: "integer".into(),
415        }),
416        PsqlTableIdentity::new("public", "store_staffs_stores") => PsqlTable::basic("public", "store_staffs_stores", PsqlTableColumn{
417          name: "id".into(),
418          data_type: "uuid".into(),
419        }),
420        PsqlTableIdentity::new("public", "store_staff_roles") => PsqlTable::basic("public", "store_staff_roles", PsqlTableColumn{
421          name: "id".into(),
422          data_type: "uuid".into(),
423        }),
424        PsqlTableIdentity::new("public", "store_staffs") => PsqlTable::basic("public", "store_staffs", PsqlTableColumn{
425          name: "id".into(),
426          data_type: "integer".into(),
427        }),
428      };
429
430      // TODO: Need to prefil psql tables
431      psql_table_map_from_foreign_key_info_rows(&mut psql_table_by_id, &fk_info_rows);
432
433      // Make sure relations are set correctly
434      // -------------------------------------------
435      // table: order_items
436      let order_items_table: &PsqlTable = psql_table_by_id
437        .get(&PsqlTableIdentity::new("public", "order_items"))
438        .unwrap();
439
440      assert_eq!(
441        order_items_table.id,
442        PsqlTableIdentity::new("public", "order_items")
443      );
444      assert_eq!(order_items_table.referencing_fk_by_constraint_name.len(), 2);
445      assert_eq!(order_items_table.referenced_fk_by_constraint_name.len(), 0);
446
447      let fk_to_orders_table_from_order_items = order_items_table
448        .referencing_fk_by_constraint_name
449        .get("order_items_order_id_foreign");
450
451      assert!(fk_to_orders_table_from_order_items.is_some());
452
453      // table: store_staffs_stores
454      let store_staffs_stores_table: &PsqlTable = psql_table_by_id
455        .get(&PsqlTableIdentity::new("public", "store_staffs_stores"))
456        .ok_or_else(|| "could not get store_staffs_stores")
457        .unwrap();
458
459      assert_eq!(
460        store_staffs_stores_table.id,
461        PsqlTableIdentity::new("public", "store_staffs_stores")
462      );
463      assert_eq!(
464        store_staffs_stores_table
465          .referencing_fk_by_constraint_name
466          .len(),
467        3
468      );
469
470      // table: store_staffs_stores
471      let products_table: &PsqlTable = psql_table_by_id
472        .get(&PsqlTableIdentity::new("public", "products"))
473        .unwrap();
474
475      assert_eq!(
476        products_table.id,
477        PsqlTableIdentity::new("public", "products")
478      );
479
480      assert_eq!(products_table.referencing_fk_by_constraint_name.len(), 1);
481      assert_eq!(products_table.referenced_fk_by_constraint_name.len(), 3);
482
483      // Make sure created tables have equal size
484      // with unique table names in fk info rows
485      // -------------------------------------------
486      let _available_tables: HashSet<&String> =
487        fk_info_rows.iter().map(|row| &row.table_name).collect();
488
489      assert_eq!(psql_table_by_id.len(), 11)
490    }
491  }
492}
lezeh_db/psql/db_metadata.rs

lezeh_db/psql/
db_metadata.rs
