use crate::dialect::{ColumnType, Dialect, quote_qualified, validate_identifier};
use crate::Database;
use rustlavel_core::Result;
#[derive(Debug, Clone)]
enum Default {
Literal(String),
Now,
Uuid,
}
#[derive(Debug, Clone)]
pub struct Column {
name: String,
kind: ColumnType,
nullable: bool,
unique: bool,
primary: bool,
default: Option<Default>,
references: Option<(String, String)>,
on_delete: Option<&'static str>,
index: bool,
}
impl Column {
fn new(name: &str, kind: ColumnType) -> Self {
Column {
name: name.to_string(),
kind,
nullable: false,
unique: false,
primary: false,
default: None,
references: None,
on_delete: None,
index: false,
}
}
pub fn nullable(&mut self) -> &mut Self {
self.nullable = true;
self
}
pub fn unique(&mut self) -> &mut Self {
self.unique = true;
self
}
pub fn primary(&mut self) -> &mut Self {
self.primary = true;
self
}
pub fn index(&mut self) -> &mut Self {
self.index = true;
self
}
pub fn default(&mut self, value: &str) -> &mut Self {
self.default = Some(Default::Literal(format!("'{}'", value.replace('\'', "''"))));
self
}
pub fn default_int(&mut self, value: i64) -> &mut Self {
self.default = Some(Default::Literal(value.to_string()));
self
}
pub fn default_bool(&mut self, value: bool) -> &mut Self {
self.default = Some(Default::Literal(if value { "TRUE_LITERAL" } else { "FALSE_LITERAL" }.into()));
self
}
pub fn default_now(&mut self) -> &mut Self {
self.default = Some(Default::Now);
self
}
pub fn default_raw(&mut self, expression: &str) -> &mut Self {
self.default = Some(Default::Literal(expression.to_string()));
self
}
pub fn references(&mut self, table: &str, column: &str) -> &mut Self {
self.references = Some((table.to_string(), column.to_string()));
self
}
pub fn cascade_on_delete(&mut self) -> &mut Self {
self.on_delete = Some("cascade");
self
}
pub fn null_on_delete(&mut self) -> &mut Self {
self.on_delete = Some("set null");
self
}
fn to_sql(&self, dialect: &dyn Dialect) -> Result<String> {
validate_identifier(&self.name, dialect.max_identifier_length())?;
let mut sql = format!("{} {}", dialect.quote(&self.name), dialect.column_type(&self.kind));
if self.primary {
sql.push_str(" primary key");
} else if !self.nullable {
sql.push_str(" not null");
}
if self.unique && !self.primary {
sql.push_str(" unique");
}
if let Some(default) = &self.default {
let rendered = match default {
Default::Now => dialect.now().to_string(),
Default::Uuid => match dialect.uuid_default() {
Some(expression) => expression.to_string(),
None => return Err(rustlavel_core::Error::msg(format!(
"`{}` cannot default a uuid column: {} has no expression for it. \
Generate the id in the application instead.",
self.name,
dialect.name()
))),
},
Default::Literal(literal) if literal == "TRUE_LITERAL" => {
if dialect.booleans_are_integers() { "1".into() } else { "true".into() }
}
Default::Literal(literal) if literal == "FALSE_LITERAL" => {
if dialect.booleans_are_integers() { "0".into() } else { "false".into() }
}
Default::Literal(literal) => literal.clone(),
};
sql.push_str(&format!(" default {rendered}"));
}
Ok(sql)
}
fn foreign_key_sql(&self, dialect: &dyn Dialect, table: &str) -> Result<Option<String>> {
let Some((target, target_column)) = &self.references else { return Ok(None) };
let limit = dialect.max_identifier_length();
validate_identifier(target_column, limit)?;
let name = format!("{table}_{}_foreign", self.name);
validate_identifier(&name, limit)?;
let mut sql = format!(
"constraint {} foreign key ({}) references {} ({})",
dialect.quote(&name),
dialect.quote(&self.name),
quote_qualified(dialect, target)?,
dialect.quote(target_column)
);
if let Some(action) = self.on_delete {
sql.push_str(&format!(" on delete {action}"));
}
Ok(Some(sql))
}
}
#[derive(Default)]
pub struct Table {
columns: Vec<Column>,
indexes: Vec<(Vec<String>, bool)>,
drops: Vec<String>,
}
impl Table {
fn add(&mut self, column: Column) -> &mut Column {
self.columns.push(column);
self.columns.last_mut().expect("just pushed")
}
pub fn id(&mut self) -> &mut Column {
let mut column = Column::new("id", ColumnType::Id);
column.primary = true;
self.add(column)
}
pub fn uuid_id(&mut self) -> &mut Column {
let mut column = Column::new("id", ColumnType::UuidId);
column.primary = true;
column.default = Some(Default::Uuid);
self.add(column)
}
pub fn string(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::String { length: 255 }))
}
pub fn string_with(&mut self, name: &str, length: u32) -> &mut Column {
self.add(Column::new(name, ColumnType::String { length }))
}
pub fn text(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Text))
}
pub fn integer(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Integer))
}
pub fn big_integer(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::BigInteger))
}
pub fn float(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Float))
}
pub fn decimal(&mut self, name: &str, precision: u32, scale: u32) -> &mut Column {
self.add(Column::new(name, ColumnType::Decimal { precision, scale }))
}
pub fn boolean(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Boolean))
}
pub fn json(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Json))
}
pub fn uuid(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Uuid))
}
pub fn date(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Date))
}
pub fn timestamp(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Timestamp))
}
pub fn binary(&mut self, name: &str) -> &mut Column {
self.add(Column::new(name, ColumnType::Binary))
}
pub fn foreign_id(&mut self, singular: &str) -> &mut Column {
let name = format!("{singular}_id");
let table = crate::migration::pluralize(singular);
let mut column = Column::new(&name, ColumnType::BigInteger);
column.references = Some((table, "id".to_string()));
column.index = true;
self.add(column)
}
pub fn timestamps(&mut self) {
self.add(Column::new("created_at", ColumnType::Timestamp)).default_now();
self.add(Column::new("updated_at", ColumnType::Timestamp)).default_now();
}
pub fn soft_deletes(&mut self) {
self.add(Column::new("deleted_at", ColumnType::Timestamp)).nullable();
}
pub fn index(&mut self, columns: &[&str]) {
self.indexes.push((columns.iter().map(|c| (*c).to_string()).collect(), false));
}
pub fn unique(&mut self, columns: &[&str]) {
self.indexes.push((columns.iter().map(|c| (*c).to_string()).collect(), true));
}
pub fn drop_column(&mut self, name: &str) {
self.drops.push(name.to_string());
}
}
pub struct Schema<'a> {
db: &'a Database,
}
impl<'a> Schema<'a> {
pub fn new(db: &'a Database) -> Self {
Schema { db }
}
pub async fn create(&self, table: &str, define: impl FnOnce(&mut Table)) -> Result<()> {
for statement in create_statements(self.db.dialect(), table, define)? {
self.db.run(&statement).await?;
}
Ok(())
}
pub async fn create_if_missing(&self, table: &str, define: impl FnOnce(&mut Table)) -> Result<()> {
let dialect = self.db.dialect();
let mut statements = create_statements(dialect, table, define)?;
let create = statements.remove(0);
let create = match dialect.supports_if_not_exists_table() {
true => create.replacen("create table ", "create table if not exists ", 1),
false => create,
};
match self.db.run(&create).await {
Ok(_) => {}
Err(error) if is_already_exists(&error) => return Ok(()),
Err(error) => return Err(error),
}
for statement in statements {
if let Err(error) = self.db.run(&statement).await
&& !is_already_exists(&error)
{
return Err(error);
}
}
Ok(())
}
pub async fn alter(&self, table: &str, define: impl FnOnce(&mut Table)) -> Result<()> {
for statement in alter_statements(self.db.dialect(), table, define)? {
self.db.run(&statement).await?;
}
Ok(())
}
pub async fn drop(&self, table: &str) -> Result<()> {
let dialect = self.db.dialect();
let quoted = quote_qualified(dialect, table)?;
let sql = match dialect.name() {
"postgres" => format!("drop table if exists {quoted} cascade"),
"sqlserver" => format!("drop table if exists {quoted}"),
_ => format!("drop table if exists {quoted}"),
};
self.db.run(&sql).await?;
Ok(())
}
pub async fn rename(&self, from: &str, to: &str) -> Result<()> {
self.db
.run(&format!(
"alter table {} rename to {}",
quote_qualified(self.db.dialect(), from)?,
quote_qualified(self.db.dialect(), to)?
))
.await?;
Ok(())
}
pub async fn has_table(&self, table: &str) -> Result<bool> {
let dialect = self.db.dialect();
let found = self
.db
.scalar::<i64>(
&format!(
"select count(*) from information_schema.tables \
where table_schema = {} and table_name = {}",
dialect.current_schema_expression(),
dialect.placeholder(1)
),
&[crate::Value::from(table)],
)
.await?;
Ok(found.unwrap_or(0) > 0)
}
pub async fn has_column(&self, table: &str, column: &str) -> Result<bool> {
let dialect = self.db.dialect();
let found = self
.db
.scalar::<i64>(
&format!(
"select count(*) from information_schema.columns \
where table_schema = {} and table_name = {} and column_name = {}",
dialect.current_schema_expression(),
dialect.placeholder(1),
dialect.placeholder(2)
),
&[crate::Value::from(table), crate::Value::from(column)],
)
.await?;
Ok(found.unwrap_or(0) > 0)
}
}
pub fn create_statements(
dialect: &dyn Dialect,
table: &str,
define: impl FnOnce(&mut Table),
) -> Result<Vec<String>> {
let mut definition = Table::default();
define(&mut definition);
let quoted = quote_qualified(dialect, table)?;
let mut lines: Vec<String> = definition
.columns
.iter()
.map(|column| column.to_sql(dialect))
.collect::<Result<_>>()?;
for column in &definition.columns {
if let Some(constraint) = column.foreign_key_sql(dialect, table)? {
lines.push(constraint);
}
}
let mut statements = vec![format!("create table {quoted} (\n {}\n)", lines.join(",\n "))];
statements.extend(index_statements(dialect, table, &definition)?);
Ok(statements)
}
fn is_already_exists(error: &rustlavel_core::Error) -> bool {
let text = error.to_string().to_ascii_lowercase();
text.contains("already exists")
|| text.contains("42p07")
|| text.contains("pg_type_typname_nsp_index")
|| text.contains("there is already an object")
}
pub fn alter_statements(
dialect: &dyn Dialect,
table: &str,
define: impl FnOnce(&mut Table),
) -> Result<Vec<String>> {
let mut definition = Table::default();
define(&mut definition);
let quoted = quote_qualified(dialect, table)?;
let mut statements = Vec::new();
for column in &definition.columns {
statements.push(format!(
"alter table {quoted} {} {}",
dialect.add_column_clause(),
column.to_sql(dialect)?
));
if let Some(constraint) = column.foreign_key_sql(dialect, table)? {
statements.push(format!("alter table {quoted} add {constraint}"));
}
}
for name in &definition.drops {
validate_identifier(name, dialect.max_identifier_length())?;
statements.push(format!("alter table {quoted} drop column {}", dialect.quote(name)));
}
statements.extend(index_statements(dialect, table, &definition)?);
Ok(statements)
}
fn index_statements(
dialect: &dyn Dialect,
table: &str,
definition: &Table,
) -> Result<Vec<String>> {
let quoted = quote_qualified(dialect, table)?;
let guard = if dialect.supports_if_not_exists_index() { "if not exists " } else { "" };
let limit = dialect.max_identifier_length();
let mut statements = Vec::new();
for column in definition.columns.iter().filter(|c| c.index) {
validate_identifier(&column.name, limit)?;
let name = format!("{table}_{}_index", column.name);
validate_identifier(&name, limit)?;
statements.push(format!(
"create index {guard}{} on {quoted} ({})",
dialect.quote(&name),
dialect.quote(&column.name)
));
}
for (columns, unique) in &definition.indexes {
for column in columns {
validate_identifier(column, limit)?;
}
let name =
format!("{table}_{}_{}", columns.join("_"), if *unique { "unique" } else { "index" });
validate_identifier(&name, limit)?;
let quoted_columns: Vec<String> = columns.iter().map(|c| dialect.quote(c)).collect();
statements.push(format!(
"create {}index {guard}{} on {quoted} ({})",
if *unique { "unique " } else { "" },
dialect.quote(&name),
quoted_columns.join(", ")
));
}
Ok(statements)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::dialect::{MySql, Postgres, SqlServer};
#[test]
fn builds_a_create_table_statement() {
let statements = create_statements(&Postgres, "users", |t| {
t.id();
t.string("name");
t.string("email").unique();
t.boolean("active").default_bool(true);
t.text("bio").nullable();
t.timestamps();
})
.unwrap();
assert_eq!(
statements[0],
"create table \"users\" (\n \
\"id\" bigserial primary key,\n \
\"name\" varchar(255) not null,\n \
\"email\" varchar(255) not null unique,\n \
\"active\" boolean not null default true,\n \
\"bio\" text,\n \
\"created_at\" timestamptz not null default now(),\n \
\"updated_at\" timestamptz not null default now()\n)"
);
}
#[test]
fn a_foreign_id_points_at_the_pluralized_table_and_gets_an_index() {
let statements = create_statements(&Postgres, "posts", |t| {
t.id();
t.foreign_id("user").cascade_on_delete();
})
.unwrap();
assert!(statements[0].contains("\"user_id\" bigint not null,"), "{}", statements[0]);
assert!(
statements[0].contains(
"constraint \"posts_user_id_foreign\" foreign key (\"user_id\") \
references \"users\" (\"id\") on delete cascade"
),
"{}",
statements[0]
);
assert_eq!(
statements[1],
"create index if not exists \"posts_user_id_index\" on \"posts\" (\"user_id\")"
);
}
#[test]
fn composite_indexes_get_their_own_statements() {
let statements = create_statements(&Postgres, "memberships", |t| {
t.id();
t.integer("team_id");
t.integer("user_id");
t.unique(&["team_id", "user_id"]);
})
.unwrap();
assert_eq!(
statements[1],
"create unique index if not exists \"memberships_team_id_user_id_unique\" \
on \"memberships\" (\"team_id\", \"user_id\")"
);
}
#[test]
fn alter_adds_and_drops_columns() {
let statements = alter_statements(&Postgres, "users", |t| {
t.string("nickname").nullable();
t.drop_column("legacy_flag");
})
.unwrap();
assert_eq!(statements[0], "alter table \"users\" add column \"nickname\" varchar(255)");
assert_eq!(statements[1], "alter table \"users\" drop column \"legacy_flag\"");
let sqlserver = alter_statements(&SqlServer, "users", |t| {
t.string("nickname").nullable();
t.drop_column("legacy_flag");
})
.unwrap();
assert_eq!(sqlserver[0], "alter table [users] add [nickname] nvarchar(255)");
assert_eq!(sqlserver[1], "alter table [users] drop column [legacy_flag]");
}
#[test]
fn a_malicious_column_name_is_rejected() {
let error = create_statements(&Postgres, "users", |t| {
t.string("name\"; drop table users; --");
})
.unwrap_err();
assert!(error.to_string().contains("not a valid SQL identifier"));
}
#[test]
fn a_string_default_is_quoted_and_escaped() {
let statements = create_statements(&Postgres, "t", |t| {
t.string("motto").default("it's fine");
})
.unwrap();
assert!(statements[0].contains("default 'it''s fine'"));
}
#[test]
fn one_schema_definition_produces_correct_ddl_for_every_database() {
let define = |t: &mut Table| {
t.id();
t.string("email").unique();
t.boolean("active").default_bool(true);
t.timestamps();
};
let postgres = create_statements(&Postgres, "users", define).unwrap();
assert!(postgres[0].contains("\"id\" bigserial primary key"), "{}", postgres[0]);
assert!(postgres[0].contains("\"active\" boolean not null default true"));
assert!(postgres[0].contains("default now()"));
let mysql = create_statements(&MySql, "users", define).unwrap();
assert!(
mysql[0].contains("`id` bigint not null auto_increment primary key"),
"{}",
mysql[0]
);
assert!(mysql[0].contains("`active` tinyint(1) not null default 1"), "{}", mysql[0]);
assert!(mysql[0].contains("default current_timestamp(6)"));
let sqlserver = create_statements(&SqlServer, "users", define).unwrap();
assert!(sqlserver[0].contains("[id] bigint identity(1,1) primary key"), "{}", sqlserver[0]);
assert!(sqlserver[0].contains("[active] bit not null default 1"), "{}", sqlserver[0]);
assert!(sqlserver[0].contains("default sysutcdatetime()"));
}
#[test]
fn only_postgres_guards_an_index_with_if_not_exists() {
let define = |t: &mut Table| {
t.id();
t.integer("team_id").index();
};
assert!(create_statements(&Postgres, "m", define).unwrap()[1].contains("if not exists"));
assert!(!create_statements(&MySql, "m", define).unwrap()[1].contains("if not exists"));
assert!(!create_statements(&SqlServer, "m", define).unwrap()[1].contains("if not exists"));
}
#[test]
fn a_uuid_default_mysql_cannot_express_is_refused_rather_than_faked() {
let define = |t: &mut Table| {
t.uuid_id();
};
assert!(create_statements(&Postgres, "t", define).unwrap()[0].contains("gen_random_uuid()"));
assert!(create_statements(&SqlServer, "t", define).unwrap()[0].contains("newid()"));
let error = create_statements(&MySql, "t", define).unwrap_err().to_string();
assert!(error.contains("has no expression for it"), "{error}");
}
#[test]
fn soft_deletes_add_a_nullable_timestamp() {
let statements = create_statements(&Postgres, "posts", |t| {
t.id();
t.soft_deletes();
})
.unwrap();
assert!(statements[0].contains("\"deleted_at\" timestamptz"));
assert!(!statements[0].contains("\"deleted_at\" timestamptz not null"));
}
}