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

use anyhow::{Context, Result};
use dbui_core::database::column::Column;
use dbui_core::database::table::Table;
use dbui_core::{Project, ResponseMessage};
use itertools::Itertools;
use uuid::Uuid;

fn run_sql(mut conn: postgres::Client, id: Uuid, sql: &str) -> Result<Vec<postgres::row::Row>> {
  let stmt = conn.prepare(sql).with_context(|| format!("Cannot parse query [{}]", id))?;
  conn.query(&stmt, &[]).with_context(|| format!("Cannot run query [{}]", id))
}

/// Called to refresh a cached schema
pub fn refresh_schema(project: &Project, conn: postgres::Client, log: &slog::Logger, full: bool) -> Result<ResponseMessage> {
  if full {
    refresh_full(project, conn, log)
  } else {
    refresh_quick(project, conn, log)
  }
}

/// Refreshes a schema, skipping detailed column information
fn refresh_quick(project: &Project, conn: postgres::Client, log: &slog::Logger) -> Result<ResponseMessage> {
  slog::info!(log, "Refreshing schema for [{}]", project.name());
  let id = Uuid::new_v4();
  let sql = crate::named::LIST_TABLES.to_string();
  let rows = run_sql(conn, id, &sql)?;

  let (mut tables, mut views, types) = (Vec::new(), Vec::new(), Vec::new());
  for row in rows {
    let oid: u32 = crate::get(&row, "oid")?;
    let schema: String = crate::get(&row, "schema")?;
    let name: String = crate::get(&row, "name")?;
    let typ: String = crate::get(&row, "type")?;
    let owner: String = crate::get(&row, "owner")?;
    let size: i64 = crate::get(&row, "size")?;
    let description: Option<String> = crate::get(&row, "description")?;
    match typ.as_ref() {
      "table" => {
        tables.push(dbui_core::Table::new(oid, schema, name, owner, Vec::new(), size, description));
        Ok(())
      }
      "view" => {
        views.push(dbui_core::View::new(oid, schema, name, owner, Vec::new(), size, description));
        Ok(())
      }
      x => Err(anyhow::anyhow!("Unhandled object type [{}]", x))
    }?
  }

  let schema = dbui_core::Schema::new(project.conn().database().into(), tables, views, types);
  Ok(ResponseMessage::SchemaResponse { schema })
}

/// Refreshes a schema, including detailed column information
fn refresh_full(project: &Project, conn: postgres::Client, log: &slog::Logger) -> Result<ResponseMessage> {
  slog::info!(log, "Refreshing full schema for [{}]", project.name());
  let id = Uuid::new_v4();
  let sql = crate::named::LIST_COLUMNS.to_string();
  let rows = run_sql(conn, id, &sql)?;

  let (okr, er): (Vec<_>, Vec<_>) = rows
    .iter()
    .map(|row| Ok(crate::schema::column::ColumnDefinition::from_row(&row)?))
    .partition(Result::is_ok);
  let ok: Vec<_> = okr
    .iter()
    .filter_map(|x: &Result<crate::schema::column::ColumnDefinition>| x.as_ref().ok())
    .collect();
  let e: Vec<_> = er.iter().filter_map(|x| x.as_ref().err()).collect();
  for err in e {
    slog::warn!(log, "Error parsing column: {}", err);
  }
  let grouped = ok.iter().group_by(|x| (x.table_schema(), x.table_name()));

  let mut tables: Vec<Table> = Vec::new();

  for (name, columns) in &grouped {
    let mut cols: Vec<Column> = Vec::new();
    for c in columns {
      cols.push(c.column())
    }
    tables.push(Table::new(0, name.0.into(), name.1.into(), "???".into(), cols, 0, None));
  }

  let schema = dbui_core::Schema::new(project.conn().database().into(), tables, Vec::new(), Vec::new());
  Ok(ResponseMessage::SchemaResponse { schema })
}