use serde::Serialize;
use sqlx::Row;
use crate::sql::config::Engine;
use crate::sql::pool::ConnectionPool;
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct Column {
pub name: String,
#[serde(rename = "type")]
pub type_: String,
}
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct Table {
pub name: String,
pub columns: Vec<Column>,
}
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct Schema {
pub engine: String,
pub tables: Vec<Table>,
}
pub async fn introspect(engine: Engine, pool: &ConnectionPool) -> Result<Schema, String> {
let tables = match (engine, pool) {
(Engine::Postgres, ConnectionPool::Postgres(p)) => {
let rows = sqlx::query(
"SELECT table_name, column_name, data_type \
FROM information_schema.columns \
WHERE table_schema = 'public' \
ORDER BY table_name, ordinal_position",
)
.fetch_all(p)
.await
.map_err(|e| format!("postgres introspect: {e}"))?;
group_columns(rows.iter().map(|r| {
(
r.get::<String, _>("table_name"),
r.get::<String, _>("column_name"),
r.get::<String, _>("data_type"),
)
}))
}
(Engine::Mysql, ConnectionPool::Mysql(p)) => {
let rows = sqlx::query(
"SELECT table_name, column_name, data_type \
FROM information_schema.columns \
WHERE table_schema = DATABASE() \
ORDER BY table_name, ordinal_position",
)
.fetch_all(p)
.await
.map_err(|e| format!("mysql introspect: {e}"))?;
group_columns(rows.iter().map(|r| {
(
r.get::<String, _>("table_name"),
r.get::<String, _>("column_name"),
r.get::<String, _>("data_type"),
)
}))
}
(Engine::Sqlite, ConnectionPool::Sqlite(p)) => {
let names = sqlx::query(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name",
)
.fetch_all(p)
.await
.map_err(|e| format!("sqlite tables: {e}"))?;
let mut tables = Vec::new();
for row in names {
let table: String = row.get("name");
let escaped = table.replace('"', "\"\"");
let cols = sqlx::query(&format!("PRAGMA table_info(\"{escaped}\")"))
.fetch_all(p)
.await
.map_err(|e| format!("sqlite columns: {e}"))?;
let columns = cols
.iter()
.map(|c| Column {
name: c.get("name"),
type_: c.get("type"),
})
.collect();
tables.push(Table {
name: table,
columns,
});
}
tables
}
_ => return Err("engine/pool mismatch".to_string()),
};
Ok(Schema {
engine: engine_str(engine).to_string(),
tables,
})
}
fn engine_str(engine: Engine) -> &'static str {
match engine {
Engine::Postgres => "postgres",
Engine::Mysql => "mysql",
Engine::Sqlite => "sqlite",
}
}
fn group_columns(rows: impl Iterator<Item = (String, String, String)>) -> Vec<Table> {
let mut tables: Vec<Table> = Vec::new();
for (table, column, type_) in rows {
if tables.last().map(|t| &t.name) != Some(&table) {
tables.push(Table {
name: table.clone(),
columns: Vec::new(),
});
}
tables.last_mut().unwrap().columns.push(Column {
name: column,
type_,
});
}
tables
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn groups_columns_by_table() {
let rows = vec![
("users".into(), "id".into(), "integer".into()),
("users".into(), "email".into(), "text".into()),
("orders".into(), "id".into(), "integer".into()),
];
let tables = group_columns(rows.into_iter());
assert_eq!(tables.len(), 2);
assert_eq!(tables[0].name, "users");
assert_eq!(tables[0].columns.len(), 2);
}
#[tokio::test]
async fn introspects_sqlite() {
let pool = sqlx::sqlite::SqlitePool::connect("sqlite::memory:")
.await
.unwrap();
sqlx::query("CREATE TABLE users (id INTEGER, email TEXT)")
.execute(&pool)
.await
.unwrap();
sqlx::query("CREATE TABLE orders (id INTEGER)")
.execute(&pool)
.await
.unwrap();
let cp = crate::sql::pool::ConnectionPool::Sqlite(pool);
let schema = introspect(Engine::Sqlite, &cp).await.unwrap();
assert_eq!(schema.engine, "sqlite");
let users = schema.tables.iter().find(|t| t.name == "users").unwrap();
assert_eq!(users.columns.len(), 2);
assert!(users.columns.iter().any(|c| c.name == "email"));
}
#[tokio::test]
async fn introspects_sqlite_quoted_table_name() {
let pool = sqlx::sqlite::SqlitePool::connect("sqlite::memory:")
.await
.unwrap();
sqlx::query(r#"CREATE TABLE "weird name" (id INTEGER)"#)
.execute(&pool)
.await
.unwrap();
let cp = crate::sql::pool::ConnectionPool::Sqlite(pool);
let schema = introspect(Engine::Sqlite, &cp).await.unwrap();
let t = schema
.tables
.iter()
.find(|t| t.name == "weird name")
.unwrap();
assert_eq!(t.columns.len(), 1);
}
}