use crate::data::column::ColumnMeta;
use crate::data::dataframe::DataFrame;
use crate::types::ColumnType;
use color_eyre::{eyre::eyre, Result};
use std::path::{Path, PathBuf};
const IDS_PER_STMT: usize = 500;
const DISPLAY_CAP: usize = 2000;
const BUSY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
pub(crate) fn open_sqlite(path: &Path) -> Result<rusqlite::Connection> {
let conn = rusqlite::Connection::open(path)?;
conn.busy_timeout(BUSY_TIMEOUT)?;
Ok(conn)
}
pub(crate) fn open_duckdb(path: &Path) -> Result<duckdb::Connection> {
duckdb::Connection::open(path).map_err(|e| {
let text = e.to_string();
if text.contains("Conflicting lock") {
eyre!(
"'{}' is open in another program — DuckDB allows one process at a time. \
Close it and try again.",
path.display()
)
} else {
eyre!(text)
}
})
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum DbKind {
Sqlite,
DuckDb,
}
impl DbKind {
pub fn name(&self) -> &'static str {
match self {
Self::Sqlite => "SQLite",
Self::DuckDb => "DuckDB",
}
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum DeclType {
Int,
Real,
Bool,
Text,
}
impl DeclType {
pub fn from_sql(decl: &str) -> Self {
let d = decl.to_ascii_uppercase();
if d.starts_with("BOOL") {
Self::Bool
} else if d.contains("INT") {
Self::Int
} else if d.contains("CHAR") || d.contains("CLOB") || d.contains("TEXT") {
Self::Text
} else if d.contains("REAL")
|| d.contains("FLOA")
|| d.contains("DOUB")
|| d.contains("DEC")
|| d.contains("NUMERIC")
{
Self::Real
} else {
Self::Text
}
}
pub fn name(&self) -> &'static str {
match self {
Self::Int => "integer",
Self::Real => "number",
Self::Bool => "boolean",
Self::Text => "text",
}
}
}
#[derive(Clone, Debug)]
pub struct DbColumn {
pub name: String,
pub decl_raw: String,
pub decl: DeclType,
pub notnull: bool,
pub pk: bool,
pub default_sql: Option<String>,
pub generated: bool,
}
#[derive(Clone)]
pub struct TableSource {
pub kind: DbKind,
pub db_path: PathBuf,
pub table: String,
pub key_col: String,
pub columns: Vec<DbColumn>,
pub original: polars::prelude::DataFrame,
}
impl TableSource {
pub fn at(&self, db_path: &std::path::Path) -> Self {
Self {
db_path: db_path.to_path_buf(),
..self.clone()
}
}
pub fn column(&self, name: &str) -> Option<&DbColumn> {
self.columns.iter().find(|c| c.name == name)
}
}
#[derive(Clone, Default, Debug, serde::Serialize, serde::Deserialize)]
pub struct DbRows {
pub ids: Vec<Option<i64>>,
pub deleted: Vec<i64>,
}
impl DbRows {
pub fn new(ids: Vec<Option<i64>>) -> Self {
Self {
ids,
deleted: Vec::new(),
}
}
}
#[derive(Clone, PartialEq, Debug)]
pub enum Val {
Null,
Int(i64),
Real(f64),
Bool(bool),
Text(String),
}
impl Val {
fn literal(&self) -> String {
match self {
Self::Null => "NULL".to_string(),
Self::Int(i) => i.to_string(),
Self::Real(f) => f.to_string(),
Self::Bool(b) => if *b { "TRUE" } else { "FALSE" }.to_string(),
Self::Text(s) => format!("'{}'", s.replace('\'', "''")),
}
}
fn parse(raw: Option<&str>, col: &DbColumn) -> Result<Self, String> {
let Some(s) = raw else {
return Ok(Self::Null);
};
if s.contains('\0') {
return Err("contains a NUL byte".to_string());
}
if matches!(col.decl, DeclType::Text) {
return Ok(Self::Text(s.to_string()));
}
if s.is_empty() {
return Ok(Self::Null);
}
let t = s.trim();
match col.decl {
DeclType::Int => t
.parse::<i64>()
.map(Self::Int)
.map_err(|_| format!("'{}' is not an integer ({} column)", s, col.decl_raw)),
DeclType::Real => t
.parse::<f64>()
.map(Self::Real)
.map_err(|_| format!("'{}' is not a number ({} column)", s, col.decl_raw)),
DeclType::Bool => match t.to_ascii_lowercase().as_str() {
"true" | "t" | "1" | "yes" => Ok(Self::Bool(true)),
"false" | "f" | "0" | "no" => Ok(Self::Bool(false)),
_ => Err(format!(
"'{}' is not a boolean ({} column)",
s, col.decl_raw
)),
},
DeclType::Text => unreachable!("handled above"),
}
}
}
impl rusqlite::ToSql for Val {
fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
use rusqlite::types::{ToSqlOutput, Value as V, ValueRef};
Ok(match self {
Self::Null => ToSqlOutput::Borrowed(ValueRef::Null),
Self::Int(i) => ToSqlOutput::Owned(V::Integer(*i)),
Self::Real(f) => ToSqlOutput::Owned(V::Real(*f)),
Self::Bool(b) => ToSqlOutput::Owned(V::Integer(i64::from(*b))),
Self::Text(s) => ToSqlOutput::Borrowed(ValueRef::Text(s.as_bytes())),
})
}
}
impl duckdb::ToSql for Val {
fn to_sql(&self) -> duckdb::Result<duckdb::types::ToSqlOutput<'_>> {
use duckdb::types::{ToSqlOutput, Value as V, ValueRef};
Ok(match self {
Self::Null => ToSqlOutput::Borrowed(ValueRef::Null),
Self::Int(i) => ToSqlOutput::Owned(V::BigInt(*i)),
Self::Real(f) => ToSqlOutput::Owned(V::Double(*f)),
Self::Bool(b) => ToSqlOutput::Owned(V::Boolean(*b)),
Self::Text(s) => ToSqlOutput::Borrowed(ValueRef::Text(s.as_bytes())),
})
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum StmtKind {
Update,
Insert,
Delete,
Schema,
}
#[derive(Clone, Debug)]
pub struct Stmt {
pub sql: String,
pub params: Vec<Val>,
pub display: String,
pub kind: StmtKind,
}
#[derive(Clone, Debug)]
pub struct RowCheck {
pub id: i64,
pub values: Vec<Option<String>>,
}
#[derive(Clone, Debug, Default)]
pub struct WritePlan {
pub stmts: Vec<Stmt>,
pub checks: Vec<RowCheck>,
pub create: bool,
pub rebuild: bool,
pub warnings: Vec<String>,
pub schema: usize,
pub updates: usize,
pub inserts: usize,
pub deletes: usize,
}
impl WritePlan {
pub fn is_empty(&self) -> bool {
self.stmts.is_empty()
}
fn push(&mut self, sql: String, display: String, params: Vec<Val>, kind: StmtKind) {
let display = if self.stmts.len() < DISPLAY_CAP {
display
} else {
String::new()
};
self.stmts.push(Stmt {
sql,
display,
params,
kind,
});
}
pub fn hidden_stmts(&self) -> usize {
self.stmts.iter().filter(|s| s.display.is_empty()).count()
}
pub fn summary(&self) -> String {
let mut parts = Vec::new();
for (n, label) in [
(self.schema, "SCHEMA"),
(self.updates, "UPDATE"),
(self.inserts, "INSERT"),
(self.deletes, "DELETE"),
] {
if n > 0 {
parts.push(format!("{} {}", n, label));
}
}
if parts.is_empty() {
"no changes".to_string()
} else {
parts.join(" · ")
}
}
}
fn push_val(sql: &mut String, display: &mut String, params: &mut Vec<Val>, v: Val) {
sql.push('?');
display.push_str(&v.literal());
params.push(v);
}
fn quote_ident(name: &str) -> String {
name.replace('"', "\"\"")
}
#[derive(Debug, Default)]
pub struct SchemaPlan {
pub drops: Vec<String>,
pub renames: Vec<(String, String)>,
pub adds: Vec<usize>,
pub retypes: Vec<(usize, ColumnType)>,
pub reorder: bool,
}
impl SchemaPlan {
pub fn is_empty(&self) -> bool {
self.drops.is_empty()
&& self.renames.is_empty()
&& self.adds.is_empty()
&& self.retypes.is_empty()
&& !self.reorder
}
}
pub fn declared_type(kind: DbKind, t: ColumnType) -> &'static str {
match t {
ColumnType::Integer | ColumnType::FileSize => match kind {
DbKind::Sqlite => "INTEGER",
DbKind::DuckDb => "BIGINT",
},
ColumnType::Float | ColumnType::Percentage | ColumnType::Currency => match kind {
DbKind::Sqlite => "REAL",
DbKind::DuckDb => "DOUBLE",
},
ColumnType::Boolean => "BOOLEAN",
_ => match kind {
DbKind::Sqlite => "TEXT",
DbKind::DuckDb => "VARCHAR",
},
}
}
pub fn apply_declared_types(df: &mut DataFrame, columns: &[DbColumn]) {
use polars::prelude::DataType;
for (i, col) in columns.iter().enumerate() {
let (target, as_type) = match col.decl {
DeclType::Int => (DataType::Int64, ColumnType::Integer),
DeclType::Real => (DataType::Float64, ColumnType::Float),
DeclType::Bool | DeclType::Text => continue,
};
let Some(column) = df.df.columns().get(i) else {
continue;
};
let Ok(cast) = column.as_materialized_series().strict_cast(&target) else {
continue;
};
if df.df.with_column(cast.into()).is_ok() {
df.columns[i].col_type = as_type;
}
}
df.calc_widths(40, 1000);
}
fn same_value(a: Option<&str>, b: Option<&str>, col: &DbColumn) -> bool {
if a == b {
return true;
}
if !matches!(col.decl, DeclType::Int | DeclType::Real) {
return false;
}
match (Val::parse(a, col), Val::parse(b, col)) {
(Ok(Val::Real(x)), Ok(Val::Real(y))) => x == y || (x.is_nan() && y.is_nan()),
(Ok(x), Ok(y)) => x == y,
_ => false,
}
}
fn is_binary(col: &DbColumn) -> bool {
col.decl_raw.to_ascii_uppercase().contains("BLOB")
}
fn is_parsed_out_of_text(t: ColumnType, live: &DbColumn) -> bool {
matches!(t, ColumnType::Date | ColumnType::Datetime) && live.decl == DeclType::Text
}
fn is_display_only_type(t: ColumnType) -> bool {
matches!(
t,
ColumnType::Percentage | ColumnType::Currency | ColumnType::FileSize
)
}
impl TableSource {
pub fn writeback_status<'a>(&self, df: &'a DataFrame) -> Result<&'a DbRows, String> {
let refuse = |why: String| {
Err(format!(
"Cannot write back into '{}': {}. Save to a different file instead.",
self.table, why
))
};
let Some(rows) = df.db_rows.as_ref() else {
return refuse(
"row identity was lost (window column, transpose, pivot, join or group)".into(),
);
};
if rows.ids.len() != df.df.height() {
return refuse("row identity no longer lines up with the table".into());
}
if df.columns.is_empty() {
return refuse("the sheet has no columns".into());
}
if df.columns.len() != df.df.width() {
return refuse("the column metadata no longer matches the data".into());
}
for (meta, actual) in df.columns.iter().zip(df.df.get_column_names()) {
if meta.name.as_str() != actual.as_str() {
return refuse(format!(
"column '{}' does not line up with the data behind it",
meta.name
));
}
}
let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
for meta in &df.columns {
if let Some(origin) = meta.db_origin.as_deref() {
if self.column(origin).is_none() {
return refuse(format!("column '{}' is no longer in the table", origin));
}
if !seen.insert(origin) {
return refuse(format!("two columns both claim to be '{}'", origin));
}
}
if let Some(t) = meta.db_retype {
if is_display_only_type(t) {
return refuse(format!(
"column '{}' is shown as a {} — that is a display format, not a \
storage type (a percentage is stored divided by 100), so writing \
it back would silently rescale the column",
meta.name,
t.name()
));
}
if meta
.db_origin
.as_deref()
.and_then(|o| self.column(o))
.is_some_and(|c| c.generated)
{
return refuse(format!(
"column '{}' is generated by the database and its type is not ours \
to change",
meta.name
));
}
}
}
Ok(rows)
}
pub fn schema_plan(&self, df: &DataFrame) -> SchemaPlan {
let mut plan = SchemaPlan::default();
let claimed: std::collections::HashSet<&str> = df
.columns
.iter()
.filter_map(|c| c.db_origin.as_deref())
.collect();
for col in &self.columns {
if !claimed.contains(col.name.as_str()) {
plan.drops.push(col.name.clone());
}
}
for (i, meta) in df.columns.iter().enumerate() {
match meta.db_origin.as_deref() {
None => plan.adds.push(i),
Some(origin) => {
if origin != meta.name {
plan.renames.push((origin.to_string(), meta.name.clone()));
}
}
}
if let Some(t) = meta.db_retype {
if meta.db_origin.is_some()
&& DeclType::from_sql(declared_type(self.kind, t))
!= DeclType::from_sql(self.decl_of(meta))
{
plan.retypes.push((i, t));
}
}
}
let live: Vec<&str> = df
.columns
.iter()
.filter_map(|c| c.db_origin.as_deref())
.collect();
let table: Vec<&str> = self
.columns
.iter()
.map(|c| c.name.as_str())
.filter(|n| live.contains(n))
.collect();
plan.reorder = live != table;
plan
}
fn decl_of(&self, meta: &ColumnMeta) -> &str {
meta.db_origin
.as_deref()
.and_then(|o| self.column(o))
.map(|c| c.decl_raw.as_str())
.unwrap_or_default()
}
pub fn live_columns(&self, df: &DataFrame) -> Vec<DbColumn> {
df.columns
.iter()
.map(|meta| {
let retyped = meta.db_retype.filter(|t| !is_display_only_type(*t));
match meta.db_origin.as_deref().and_then(|o| self.column(o)) {
Some(existing) => {
let mut c = existing.clone();
c.name = meta.name.clone();
if let Some(t) = retyped {
c.decl_raw = declared_type(self.kind, t).to_string();
c.decl = DeclType::from_sql(&c.decl_raw);
}
c
}
None => {
let decl_raw =
declared_type(self.kind, retyped.unwrap_or(meta.col_type)).to_string();
DbColumn {
name: meta.name.clone(),
decl: DeclType::from_sql(&decl_raw),
decl_raw,
notnull: false,
pk: false,
default_sql: None,
generated: false,
}
}
}
})
.collect()
}
}
fn cell(pdf: &polars::prelude::DataFrame, row: usize, col: usize) -> Option<String> {
match pdf.columns().get(col)?.get(row) {
Ok(polars::prelude::AnyValue::Null) | Err(_) => None,
Ok(v) => Some(DataFrame::anyvalue_to_string_fmt(&v)),
}
}
pub fn build_plan(src: &TableSource, df: &DataFrame) -> Result<WritePlan> {
let rows = src.writeback_status(df).map_err(|e| eyre!(e))?;
let deleted: std::collections::HashSet<i64> = rows.deleted.iter().copied().collect();
let mut phys_of_id: std::collections::HashMap<i64, usize> = std::collections::HashMap::new();
for (i, id) in rows.ids.iter().enumerate() {
if let Some(id) = id {
phys_of_id.insert(*id, i);
}
}
let mut display_of_phys: std::collections::HashMap<usize, usize> =
std::collections::HashMap::new();
for (d, p) in df.row_order.iter().enumerate() {
display_of_phys.insert(*p, d + 1);
}
let live_cols = src.live_columns(df);
let schema = src.schema_plan(df);
let snapshot_of: std::collections::HashMap<String, usize> = src
.original
.get_column_names()
.iter()
.enumerate()
.map(|(i, n)| (n.to_string(), i))
.collect();
enum ColDiff {
Added,
Plain(usize),
Retyped(usize),
}
let mut diff_cols: Vec<(usize, ColDiff)> = Vec::new();
let mut unwritable: Vec<String> = Vec::new();
for (i, meta) in df.columns.iter().enumerate() {
match meta.db_origin.as_deref() {
None => diff_cols.push((i, ColDiff::Added)),
Some(origin) => {
let Some(&s) = snapshot_of.get(origin) else {
continue;
};
if let Some(t) = meta.db_retype {
if is_parsed_out_of_text(t, &live_cols[i]) {
let changed = (0..df.df.height())
.any(|r| cell(&df.df, r, i) != cell(&src.original, r, s));
if changed {
unwritable.push(format!(
"column '{}' is shown as a {} but stored as text, so it \
will not be written — press t to put it back to edit it",
meta.name,
t.name()
));
}
continue;
}
diff_cols.push((i, ColDiff::Retyped(s)));
continue;
}
if schema.retypes.iter().any(|(c, _)| *c == i) {
diff_cols.push((i, ColDiff::Retyped(s)));
} else {
match (src.original.columns().get(s), df.df.columns().get(i)) {
(Some(a), Some(b)) if a.equals_missing(b) => {}
_ => diff_cols.push((i, ColDiff::Plain(s))),
}
}
}
}
}
let mut groups: indexmap::IndexMap<Vec<(usize, Option<String>)>, Vec<i64>> =
indexmap::IndexMap::new();
let mut updated_ids: Vec<i64> = Vec::new();
for (phys, id) in rows.ids.iter().enumerate() {
let Some(id) = *id else { continue };
if deleted.contains(&id) {
continue;
}
let mut change = Vec::new();
for (i, how) in &diff_cols {
let now = cell(&df.df, phys, *i);
let differs = match how {
ColDiff::Added => now.is_some(),
ColDiff::Plain(s) => now != cell(&src.original, phys, *s),
ColDiff::Retyped(s) => {
let before = cell(&src.original, phys, *s);
Val::parse(before.as_deref(), &live_cols[*i]).ok()
!= Val::parse(now.as_deref(), &live_cols[*i]).ok()
}
};
if differs {
change.push((*i, now));
}
}
if change.is_empty() {
continue;
}
updated_ids.push(id);
groups.entry(change).or_default().push(id);
}
let mut plan = WritePlan {
warnings: unwritable,
..Default::default()
};
let table = quote_ident(&src.table);
let key = quote_ident(&src.key_col);
push_schema_stmts(src, df, &schema, &live_cols, &mut plan)?;
for (change, ids) in &groups {
let mut sets: Vec<(String, Val)> = Vec::with_capacity(change.len());
for (c, raw) in change {
let col = &live_cols[*c];
if col.generated {
return Err(eyre!(
"Column '{}' is generated by the database and cannot be edited",
col.name
));
}
if is_binary(col) {
return Err(eyre!(
"Column '{}' holds binary data: it can be read but not edited. \
Undo the change there and the rest of the row still saves",
col.name
));
}
let val = Val::parse(raw.as_deref(), col).map_err(|why| {
let where_ = ids
.first()
.and_then(|id| phys_of_id.get(id))
.and_then(|p| display_of_phys.get(p))
.map(|d| format!("row {}", d))
.unwrap_or_else(|| "a row".to_string());
eyre!("Column '{}', {}: {}", col.name, where_, why)
})?;
sets.push((col.name.clone(), val));
}
for chunk in ids.chunks(IDS_PER_STMT) {
let mut sql = format!("UPDATE \"{}\" SET ", table);
let mut display = sql.clone();
let mut params: Vec<Val> = Vec::new();
for (i, (name, val)) in sets.iter().enumerate() {
if i > 0 {
sql.push_str(", ");
display.push_str(", ");
}
let assign = format!("\"{}\" = ", quote_ident(name));
sql.push_str(&assign);
display.push_str(&assign);
push_val(&mut sql, &mut display, &mut params, val.clone());
}
write_where(&mut sql, &mut display, &mut params, &key, chunk);
plan.push(sql, display, params, StmtKind::Update);
}
}
plan.updates = updated_ids.len();
let insertable: Vec<(usize, &DbColumn)> = live_cols
.iter()
.enumerate()
.filter(|(_, c)| !c.generated)
.collect();
for (phys, id) in rows.ids.iter().enumerate() {
if id.is_some() {
continue;
}
let mut sql = format!("INSERT INTO \"{}\" (", table);
for (i, (_, col)) in insertable.iter().enumerate() {
if i > 0 {
sql.push_str(", ");
}
sql.push_str(&format!("\"{}\"", quote_ident(&col.name)));
}
sql.push_str(") VALUES (");
let mut display = sql.clone();
let mut params: Vec<Val> = Vec::new();
for (i, (c, col)) in insertable.iter().enumerate() {
if i > 0 {
sql.push_str(", ");
display.push_str(", ");
}
let raw = cell(&df.df, phys, *c);
let val = Val::parse(raw.as_deref(), col).map_err(|why| {
let where_ = display_of_phys
.get(&phys)
.map(|d| format!("row {}", d))
.unwrap_or_else(|| "a new row".to_string());
eyre!("Column '{}', {}: {}", col.name, where_, why)
})?;
if is_binary(col) && !matches!(val, Val::Null) {
return Err(eyre!(
"Column '{}' holds binary data: a new row can only leave it empty",
col.name
));
}
push_val(&mut sql, &mut display, &mut params, val);
}
sql.push(')');
display.push(')');
plan.push(sql, display, params, StmtKind::Insert);
plan.inserts += 1;
}
for chunk in rows.deleted.chunks(IDS_PER_STMT) {
let mut sql = format!("DELETE FROM \"{}\"", table);
let mut display = sql.clone();
let mut params: Vec<Val> = Vec::new();
write_where(&mut sql, &mut display, &mut params, &key, chunk);
plan.push(sql, display, params, StmtKind::Delete);
}
plan.deletes = rows.deleted.len();
push_post_stmts(src, df, &schema, &live_cols, &mut plan)?;
for id in updated_ids.iter().chain(rows.deleted.iter()) {
let Some(&phys) = phys_of_id.get(id) else {
continue;
};
plan.checks.push(RowCheck {
id: *id,
values: (0..src.columns.len())
.map(|c| cell(&src.original, phys, c))
.collect(),
});
}
Ok(plan)
}
const PARAMS_PER_STMT: usize = 500;
pub fn new_columns(kind: DbKind, df: &DataFrame) -> Result<Vec<DbColumn>> {
if df.columns.is_empty() {
return Err(eyre!("This sheet has no columns to make a table from"));
}
df.columns
.iter()
.enumerate()
.map(|(i, meta)| {
if meta.name.trim().is_empty() {
return Err(eyre!(
"Column {} has no name — rename it with 'ze' before creating a table",
i + 1
));
}
let decl_raw = declared_type(kind, meta.col_type).to_string();
Ok(DbColumn {
name: meta.name.clone(),
decl: DeclType::from_sql(&decl_raw),
decl_raw,
notnull: false,
pk: false,
default_sql: None,
generated: false,
})
})
.collect()
}
pub fn validate_table_name(name: &str) -> Result<(), &'static str> {
let trimmed = name.trim();
if trimmed.is_empty() {
Err("a table needs a name")
} else if trimmed.contains('\0') || trimmed.contains('\n') {
Err("a table name cannot contain a newline")
} else if trimmed.to_ascii_lowercase().starts_with("sqlite_") {
Err("names starting with 'sqlite_' are reserved by the engine")
} else {
Ok(())
}
}
pub fn table_exists(kind: DbKind, path: &Path, name: &str) -> bool {
if !path.exists() {
return false;
}
match kind {
DbKind::Sqlite => open_sqlite(path)
.and_then(|conn| {
let mut stmt = conn.prepare(
"SELECT 1 FROM sqlite_master WHERE type IN ('table', 'view') AND name = ?1",
)?;
Ok(stmt.exists([name])?)
})
.unwrap_or(false),
DbKind::DuckDb => open_duckdb(path)
.and_then(|conn| {
let mut stmt = conn.prepare(
"SELECT 1 FROM duckdb_tables() WHERE table_name = ? \
UNION ALL SELECT 1 FROM duckdb_views() WHERE view_name = ?",
)?;
let mut rows = stmt.query([name, name])?;
Ok(rows.next()?.is_some())
})
.unwrap_or(false),
}
}
pub fn is_view(kind: DbKind, path: &Path, name: &str) -> bool {
if !path.exists() {
return false;
}
match kind {
DbKind::Sqlite => open_sqlite(path)
.and_then(|conn| {
let mut stmt =
conn.prepare("SELECT 1 FROM sqlite_master WHERE type = 'view' AND name = ?1")?;
Ok(stmt.exists([name])?)
})
.unwrap_or(false),
DbKind::DuckDb => open_duckdb(path)
.and_then(|conn| {
let mut stmt = conn.prepare("SELECT 1 FROM duckdb_views() WHERE view_name = ?")?;
let mut rows = stmt.query([name])?;
Ok(rows.next()?.is_some())
})
.unwrap_or(false),
}
}
pub fn existing_tables(kind: DbKind, path: &Path) -> Vec<String> {
if !path.exists() {
return Vec::new();
}
match kind {
DbKind::Sqlite => super::sqlite::sqlite_table_names(path).unwrap_or_default(),
DbKind::DuckDb => super::duckdb::duckdb_table_names(path).unwrap_or_default(),
}
}
fn preflight_replace(kind: DbKind, path: &Path, table: &str) -> Result<Vec<String>> {
let mut warnings = Vec::new();
let refuse = |other: &str| -> Result<Vec<String>> {
Err(eyre!(
"Table '{}' has a foreign key into '{}', which replacing it would leave \
pointing at nothing. Choose another table name.",
other,
table
))
};
match kind {
DbKind::Sqlite => {
let conn = open_sqlite(path)?;
let mut objs = conn.prepare(
"SELECT type, name FROM sqlite_master \
WHERE tbl_name = ?1 AND sql IS NOT NULL AND type IN ('index', 'trigger') \
ORDER BY type, name",
)?;
let mut rows = objs.query([table])?;
while let Some(row) = rows.next()? {
warnings.push(format!(
"{} '{}' will be lost",
row.get::<_, String>(0)?,
row.get::<_, String>(1)?
));
}
let mut views =
conn.prepare("SELECT name, sql FROM sqlite_master WHERE type = 'view'")?;
let mut rows = views.query([])?;
while let Some(row) = rows.next()? {
if mentions_word(&row.get::<_, String>(1)?, table) {
warnings.push(format!(
"view '{}' is built on this table and will break",
row.get::<_, String>(0)?
));
}
}
for other in super::sqlite::sqlite_table_names(path)? {
if other.eq_ignore_ascii_case(table) {
continue;
}
let mut fks = conn.prepare(&format!(
"PRAGMA foreign_key_list(\"{}\")",
quote_ident(&other)
))?;
let mut rows = fks.query([])?;
while let Some(row) = rows.next()? {
if row.get::<_, String>(2)?.eq_ignore_ascii_case(table) {
return refuse(&other);
}
}
}
}
DbKind::DuckDb => {
let conn = open_duckdb(path)?;
let mut idx =
conn.prepare("SELECT index_name FROM duckdb_indexes() WHERE table_name = ?")?;
let mut rows = idx.query([table])?;
while let Some(row) = rows.next()? {
warnings.push(format!("index '{}' will be lost", row.get::<_, String>(0)?));
}
let mut views = conn.prepare("SELECT view_name, sql FROM duckdb_views()")?;
let mut rows = views.query([])?;
while let Some(row) = rows.next()? {
if mentions_word(&row.get::<_, String>(1)?, table) {
warnings.push(format!(
"view '{}' is built on this table and will break",
row.get::<_, String>(0)?
));
}
}
let mut fks = conn.prepare(
"SELECT table_name, constraint_text FROM duckdb_constraints() \
WHERE constraint_type = 'FOREIGN KEY'",
)?;
let mut rows = fks.query([])?;
while let Some(row) = rows.next()? {
let other: String = row.get(0)?;
if !other.eq_ignore_ascii_case(table)
&& mentions_word(&row.get::<_, String>(1)?, table)
{
return refuse(&other);
}
}
}
}
Ok(warnings)
}
pub fn create_plan(
kind: DbKind,
path: &Path,
table: &str,
df: &DataFrame,
) -> Result<(WritePlan, TableSource)> {
let cols = new_columns(kind, df)?;
let quoted = quote_ident(table);
let mut plan = WritePlan {
create: true,
..Default::default()
};
let push_schema = |sql: String, plan: &mut WritePlan| {
plan.stmts.push(Stmt {
display: sql.clone(),
sql,
params: Vec::new(),
kind: StmtKind::Schema,
});
plan.schema += 1;
};
if table_exists(kind, path, table) {
plan.warnings = preflight_replace(kind, path, table)?;
push_schema(format!("DROP TABLE \"{}\"", quoted), &mut plan);
plan.rebuild = true;
}
let defs: Vec<String> = cols.iter().map(|c| column_def(c, false)).collect();
push_schema(
format!("CREATE TABLE \"{}\" ({})", quoted, defs.join(", ")),
&mut plan,
);
let names: Vec<String> = cols
.iter()
.map(|c| format!("\"{}\"", quote_ident(&c.name)))
.collect();
let prefix = format!("INSERT INTO \"{}\" ({}) VALUES ", quoted, names.join(", "));
let per_stmt = (PARAMS_PER_STMT / cols.len()).clamp(1, IDS_PER_STMT);
for chunk in df.row_order.chunks(per_stmt) {
let mut sql = prefix.clone();
let mut display = prefix.clone();
let mut params: Vec<Val> = Vec::with_capacity(chunk.len() * cols.len());
for (r, &phys) in chunk.iter().enumerate() {
if r > 0 {
sql.push_str(", ");
display.push_str(", ");
}
sql.push('(');
display.push('(');
for (c, col) in cols.iter().enumerate() {
if c > 0 {
sql.push_str(", ");
display.push_str(", ");
}
let raw = cell(&df.df, phys, c);
let val = Val::parse(raw.as_deref(), col).map_err(|why| {
let shown = df.row_order.iter().position(|p| *p == phys).unwrap_or(phys) + 1;
eyre!("Column '{}', row {}: {}", col.name, shown, why)
})?;
push_val(&mut sql, &mut display, &mut params, val);
}
sql.push(')');
display.push(')');
}
plan.push(sql, display, params, StmtKind::Insert);
}
plan.inserts = df.row_order.len();
let source = TableSource {
kind,
db_path: path.to_path_buf(),
table: table.to_string(),
key_col: "rowid".to_string(),
columns: cols,
original: polars::prelude::DataFrame::empty(),
};
Ok((plan, source))
}
pub fn remove_new_file(path: &Path) {
let _ = std::fs::remove_file(path);
let name = path.to_string_lossy().into_owned();
for extra in [
format!("{}-wal", name),
format!("{}-shm", name),
format!("{}.wal", name),
] {
let _ = std::fs::remove_file(extra);
}
}
pub fn create_table(kind: DbKind, path: &Path, table: &str, df: &DataFrame) -> Result<()> {
let existed = path.exists();
let (plan, src) = create_plan(kind, path, table, df)?;
let result = apply(&src, &plan);
if result.is_err() && !existed {
remove_new_file(path);
}
result
}
pub fn kind_of_file(path: &Path) -> Option<DbKind> {
use std::io::Read;
let mut head = [0u8; 16];
let mut file = std::fs::File::open(path).ok()?;
if file.read_exact(&mut head).is_err() {
return None;
}
if head.starts_with(b"SQLite format 3\0") {
Some(DbKind::Sqlite)
} else if &head[8..12] == b"DUCK" {
Some(DbKind::DuckDb)
} else {
None
}
}
pub fn kind_for_path(path: &Path) -> DbKind {
if let Some(kind) = kind_of_file(path) {
return kind;
}
match path
.extension()
.and_then(|e| e.to_str())
.unwrap_or_default()
.to_ascii_lowercase()
.as_str()
{
"duckdb" | "ddb" => DbKind::DuckDb,
_ => DbKind::Sqlite,
}
}
pub fn is_db_ext(path: &Path) -> bool {
is_db_name(
path.extension()
.and_then(|e| e.to_str())
.unwrap_or_default(),
)
}
pub fn is_db_name(ext: &str) -> bool {
matches!(
ext.to_ascii_lowercase().as_str(),
"db" | "sqlite" | "sqlite3" | "duckdb" | "ddb"
)
}
pub fn same_file(a: &Path, b: &Path) -> bool {
match (a.canonicalize(), b.canonicalize()) {
(Ok(a), Ok(b)) => a == b,
_ => a == b,
}
}
pub fn apply(src: &TableSource, plan: &WritePlan) -> Result<()> {
match src.kind {
DbKind::Sqlite => apply_sqlite(src, plan),
DbKind::DuckDb => apply_duckdb(src, plan),
}
}
fn check_query(src: &TableSource, n_ids: usize) -> String {
let cast = |name: &str| match src.kind {
DbKind::Sqlite => format!("\"{}\"", quote_ident(name)),
DbKind::DuckDb => format!("CAST(\"{}\" AS VARCHAR)", quote_ident(name)),
};
let cols: Vec<String> = src.columns.iter().map(|c| cast(&c.name)).collect();
let placeholders = vec!["?"; n_ids].join(", ");
format!(
"SELECT \"{key}\", {cols} FROM \"{table}\" WHERE \"{key}\" IN ({placeholders})",
key = quote_ident(&src.key_col),
cols = cols.join(", "),
table = quote_ident(&src.table),
)
}
fn compare_drift(
src: &TableSource,
checks: &[RowCheck],
actual: &std::collections::HashMap<i64, Vec<Option<String>>>,
) -> Result<()> {
let stale = |detail: String| {
eyre!(
"{} — '{}' changed since it was opened; something else has written to it. \
Nothing was written; reopen the table and redo the edit.",
detail,
src.table
)
};
for check in checks {
let Some(now) = actual.get(&check.id) else {
return Err(stale(format!("Row {} is gone", check.id)));
};
for (i, expected) in check.values.iter().enumerate() {
let Some(col) = src.columns.get(i) else {
continue;
};
let found = now.get(i).and_then(|v| v.as_deref());
if !same_value(found, expected.as_deref(), col) {
let name = col.name.as_str();
return Err(stale(format!(
"Row {}, column '{}': the database has {:?}, expected {:?}",
check.id,
name,
now.get(i).cloned().flatten(),
expected
)));
}
}
}
Ok(())
}
fn check_shape(src: &TableSource, actual: Vec<String>) -> Result<()> {
let expected: Vec<&str> = src.columns.iter().map(|c| c.name.as_str()).collect();
if actual != expected {
return Err(eyre!(
"The columns of '{}' changed since it was opened — the table now has [{}] \
where it had [{}]. Nothing was written; reopen the table and redo the edit.",
src.table,
actual.join(", "),
expected.join(", ")
));
}
Ok(())
}
fn apply_sqlite(src: &TableSource, plan: &WritePlan) -> Result<()> {
use super::sqlite::value_to_opt_string;
let mut conn = open_sqlite(&src.db_path)?;
if plan.rebuild {
conn.pragma_update(None, "foreign_keys", false)?;
}
let result = apply_sqlite_inner(src, plan, &mut conn, value_to_opt_string);
if plan.rebuild {
let _ = conn.pragma_update(None, "foreign_keys", true);
}
result
}
fn apply_sqlite_inner(
src: &TableSource,
plan: &WritePlan,
conn: &mut rusqlite::Connection,
value_to_opt_string: fn(rusqlite::types::Value) -> Option<String>,
) -> Result<()> {
let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
if !plan.create {
let mut stmt = tx.prepare(&format!(
"PRAGMA table_xinfo(\"{}\")",
quote_ident(&src.table)
))?;
let mut rows = stmt.query([])?;
let mut actual = Vec::new();
while let Some(row) = rows.next()? {
actual.push(row.get::<_, String>(1)?);
}
check_shape(src, actual)?;
}
for chunk in plan.checks.chunks(IDS_PER_STMT) {
let mut stmt = tx.prepare(&check_query(src, chunk.len()))?;
let mut rows = stmt.query(rusqlite::params_from_iter(chunk.iter().map(|c| c.id)))?;
let mut actual = std::collections::HashMap::new();
while let Some(row) = rows.next()? {
let id: i64 = row.get(0)?;
let mut vals = Vec::with_capacity(src.columns.len());
for i in 0..src.columns.len() {
vals.push(value_to_opt_string(row.get(i + 1)?));
}
actual.insert(id, vals);
}
compare_drift(src, chunk, &actual)?;
}
for stmt in &plan.stmts {
if stmt.params.is_empty() {
tx.execute_batch(&stmt.sql)?;
} else {
tx.execute(&stmt.sql, rusqlite::params_from_iter(stmt.params.iter()))?;
}
}
if plan.rebuild {
let mut check = tx.prepare("PRAGMA foreign_key_check")?;
let mut rows = check.query([])?;
if let Some(row) = rows.next()? {
return Err(eyre!(
"Rebuilding '{}' would leave a dangling reference from '{}'. Nothing was \
written.",
src.table,
row.get::<_, String>(0).unwrap_or_default()
));
}
}
tx.commit()?;
Ok(())
}
fn apply_duckdb(src: &TableSource, plan: &WritePlan) -> Result<()> {
let mut conn = open_duckdb(&src.db_path)?;
let tx = conn.transaction()?;
if !plan.create {
let mut stmt = tx.prepare(&format!(
"PRAGMA table_info('{}')",
src.table.replace('\'', "''")
))?;
let mut rows = stmt.query([])?;
let mut actual = Vec::new();
while let Some(row) = rows.next()? {
actual.push(row.get::<_, String>(1)?);
}
check_shape(src, actual)?;
}
for chunk in plan.checks.chunks(IDS_PER_STMT) {
let mut stmt = tx.prepare(&check_query(src, chunk.len()))?;
let mut rows = stmt.query(duckdb::params_from_iter(chunk.iter().map(|c| c.id)))?;
let mut actual = std::collections::HashMap::new();
while let Some(row) = rows.next()? {
let id: i64 = row.get(0)?;
let mut vals = Vec::with_capacity(src.columns.len());
for i in 0..src.columns.len() {
vals.push(row.get::<_, Option<String>>(i + 1)?);
}
actual.insert(id, vals);
}
compare_drift(src, chunk, &actual)?;
}
for stmt in &plan.stmts {
tx.execute(&stmt.sql, duckdb::params_from_iter(stmt.params.iter()))?;
}
tx.commit()?;
Ok(())
}
pub fn copy_db(src: &TableSource, dest: &Path) -> Result<()> {
if dest.exists() {
return Err(eyre!(
"{} already exists. Saving a database elsewhere writes a fresh copy of the \
whole file — remove it first or choose another name.",
dest.display()
));
}
match src.kind {
DbKind::Sqlite => {
let conn = open_sqlite(&src.db_path)?;
conn.execute("VACUUM INTO ?1", [dest.to_string_lossy().as_ref()])?;
Ok(())
}
DbKind::DuckDb => {
{
let conn = open_duckdb(&src.db_path)?;
conn.execute_batch("CHECKPOINT;")?;
}
std::fs::copy(&src.db_path, dest)?;
let wal = dest.with_extension(format!(
"{}.wal",
dest.extension().unwrap_or_default().to_string_lossy()
));
let _ = std::fs::remove_file(wal);
Ok(())
}
}
}
fn preflight_schema(src: &TableSource, schema: &SchemaPlan) -> Result<()> {
if schema.drops.is_empty() {
return Ok(());
}
for name in &schema.drops {
if src.column(name).is_some_and(|c| c.pk) {
return Err(eyre!(
"'{}' is the primary key of '{}' and cannot be dropped. Save to a \
different file instead.",
name,
src.table
));
}
}
if src.kind != DbKind::Sqlite {
return Ok(());
}
let conn = open_sqlite(&src.db_path)?;
let table = quote_ident(&src.table);
let mut list = conn.prepare(&format!("PRAGMA index_list(\"{}\")", table))?;
let mut indexes = list.query([])?;
while let Some(row) = indexes.next()? {
let index: String = row.get(1)?;
let mut info = conn.prepare(&format!("PRAGMA index_info(\"{}\")", quote_ident(&index)))?;
let mut cols = info.query([])?;
while let Some(c) = cols.next()? {
let col: Option<String> = c.get(2)?;
if let Some(col) = col {
if schema.drops.contains(&col) {
return Err(eyre!(
"'{}' is used by index '{}' and cannot be dropped. Drop the index \
first, or save to a different file.",
col,
index
));
}
}
}
}
let mut objs = conn.prepare(
"SELECT type, name, sql FROM sqlite_master \
WHERE type IN ('view', 'trigger') AND sql IS NOT NULL",
)?;
let mut rows = objs.query([])?;
while let Some(row) = rows.next()? {
let kind: String = row.get(0)?;
let name: String = row.get(1)?;
let sql: String = row.get(2)?;
for col in &schema.drops {
if mentions_word(&sql, col) {
return Err(eyre!(
"'{}' is referenced by {} '{}' and cannot be dropped. Save to a \
different file instead.",
col,
kind,
name
));
}
}
}
Ok(())
}
pub(crate) fn mentions_word(sql: &str, word: &str) -> bool {
let hay = sql.to_ascii_lowercase();
let needle = word.to_ascii_lowercase();
let boundary = |c: Option<char>| !c.is_some_and(|c| c.is_alphanumeric() || c == '_');
let mut from = 0;
while let Some(at) = hay[from..].find(&needle) {
let start = from + at;
let end = start + needle.len();
if boundary(hay[..start].chars().next_back()) && boundary(hay[end..].chars().next()) {
return true;
}
from = end;
}
false
}
const SWAP_SUFFIX: &str = "__tuitab_swap";
const REBUILD_SUFFIX: &str = "__tuitab_rebuild";
fn needs_rebuild(kind: DbKind, schema: &SchemaPlan) -> bool {
schema.reorder || (kind == DbKind::Sqlite && !schema.retypes.is_empty())
}
fn column_def(col: &DbColumn, inline_pk: bool) -> String {
let mut def = format!("\"{}\" {}", quote_ident(&col.name), col.decl_raw);
if col.notnull {
def.push_str(" NOT NULL");
}
if let Some(d) = &col.default_sql {
def.push_str(&format!(" DEFAULT {}", d));
}
if inline_pk && col.pk {
def.push_str(" PRIMARY KEY");
}
def
}
fn preflight_rebuild_sqlite(
src: &TableSource,
schema: &SchemaPlan,
live_cols: &[DbColumn],
) -> Result<()> {
let conn = open_sqlite(&src.db_path)?;
let table = quote_ident(&src.table);
let refuse = |why: String| -> Result<()> {
Err(eyre!(
"'{}' has to be rebuilt to do this, and tuitab will not rebuild it: {}. \
Save to a different file instead.",
src.table,
why
))
};
if live_cols.iter().filter(|c| c.pk).count() > 1 {
return refuse("its primary key spans several columns".into());
}
if live_cols.iter().any(|c| c.generated) {
return refuse("it has a generated column".into());
}
let mut fks = conn.prepare(&format!("PRAGMA foreign_key_list(\"{}\")", table))?;
if fks.query([])?.next()?.is_some() {
return refuse("it has a foreign key of its own".into());
}
let mut list = conn.prepare(&format!("PRAGMA index_list(\"{}\")", table))?;
let mut rows = list.query([])?;
while let Some(row) = rows.next()? {
if row.get::<_, String>(3)?.as_str() != "c" {
return refuse(format!(
"index '{}' comes from a UNIQUE or PRIMARY KEY constraint in the table \
definition",
row.get::<_, String>(1)?
));
}
}
let ddl: String = conn.query_row(
"SELECT sql FROM sqlite_master WHERE type = 'table' AND name = ?1",
[&src.table],
|r| r.get(0),
)?;
for word in [
"CHECK",
"COLLATE",
"AUTOINCREMENT",
"WITHOUT ROWID",
"STRICT",
"GENERATED",
] {
if mentions_word(&ddl, word) {
return refuse(format!("its definition uses {}", word));
}
}
for other in super::sqlite::sqlite_table_names(&src.db_path)? {
if other == src.table {
continue;
}
let mut stmt = conn.prepare(&format!(
"PRAGMA foreign_key_list(\"{}\")",
quote_ident(&other)
))?;
let mut rows = stmt.query([])?;
while let Some(row) = rows.next()? {
if row.get::<_, String>(2)?.eq_ignore_ascii_case(&src.table) {
return refuse(format!("table '{}' has a foreign key into it", other));
}
}
}
if !schema.renames.is_empty() {
let mut objs = conn.prepare(
"SELECT type, name, sql FROM sqlite_master \
WHERE tbl_name = ?1 AND sql IS NOT NULL AND type IN ('index', 'trigger')",
)?;
let mut rows = objs.query([&src.table])?;
while let Some(row) = rows.next()? {
let sql: String = row.get(2)?;
for (from, _) in &schema.renames {
if mentions_word(&sql, from) {
return refuse(format!(
"{} '{}' names the column '{}' that is being renamed",
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
from
));
}
}
}
}
let mut views = conn
.prepare("SELECT name, sql FROM sqlite_master WHERE type = 'view' AND sql IS NOT NULL")?;
let mut rows = views.query([])?;
while let Some(row) = rows.next()? {
let sql: String = row.get(1)?;
if mentions_word(&sql, &src.table) {
return refuse(format!(
"view '{}' is built on it",
row.get::<_, String>(0)?
));
}
}
if conn
.query_row(
"SELECT COUNT(*) FROM sqlite_master WHERE name = ?1",
[&format!("{}{}", src.table, REBUILD_SUFFIX)],
|r| r.get::<_, i64>(0),
)
.unwrap_or(0)
> 0
{
return refuse("a leftover rebuild table from an earlier run is in the way".into());
}
Ok(())
}
fn preflight_rebuild_duckdb(src: &TableSource) -> Result<()> {
let conn = open_duckdb(&src.db_path)?;
let refuse = |why: String| -> Result<()> {
Err(eyre!(
"'{}' has to be rebuilt to reorder its columns, and tuitab will not rebuild \
it: {}. Save to a different file instead.",
src.table,
why
))
};
let mut stmt = conn.prepare(
"SELECT constraint_type FROM duckdb_constraints() \
WHERE table_name = ? AND constraint_type <> 'NOT NULL'",
)?;
let mut rows = stmt.query([&src.table])?;
while let Some(row) = rows.next()? {
let kind: String = row.get(0)?;
if kind != "PRIMARY KEY" {
return refuse(format!("it has a {} constraint", kind));
}
}
let mut idx = conn.prepare("SELECT index_name FROM duckdb_indexes() WHERE table_name = ?")?;
let mut rows = idx.query([&src.table])?;
if let Some(row) = rows.next()? {
return refuse(format!(
"index '{}' would be lost",
row.get::<_, String>(0)?
));
}
let ddl: String = conn.query_row(
"SELECT sql FROM duckdb_tables() WHERE table_name = ?",
[&src.table],
|r| r.get(0),
)?;
for word in ["GENERATED", "CHECK", "COLLATE"] {
if mentions_word(&ddl, word) {
return refuse(format!("its definition uses {}", word));
}
}
let mut views = conn.prepare("SELECT view_name, sql FROM duckdb_views()")?;
let mut rows = views.query([])?;
while let Some(row) = rows.next()? {
if mentions_word(&row.get::<_, String>(1)?, &src.table) {
return refuse(format!(
"view '{}' is built on it",
row.get::<_, String>(0)?
));
}
}
let mut fks = conn.prepare(
"SELECT table_name, constraint_text FROM duckdb_constraints() \
WHERE constraint_type = 'FOREIGN KEY'",
)?;
let mut rows = fks.query([])?;
while let Some(row) = rows.next()? {
let other: String = row.get(0)?;
if !other.eq_ignore_ascii_case(&src.table)
&& mentions_word(&row.get::<_, String>(1)?, &src.table)
{
return refuse(format!("table '{}' has a foreign key into it", other));
}
}
let scratch = format!("{}{}", src.table, REBUILD_SUFFIX);
let mut left = conn.prepare("SELECT 1 FROM duckdb_tables() WHERE table_name = ?")?;
if left.query([&scratch])?.next()?.is_some() {
return refuse("a leftover rebuild table from an earlier run is in the way".into());
}
Ok(())
}
fn push_rebuild_stmts(
src: &TableSource,
df: &DataFrame,
live_cols: &[DbColumn],
plan: &mut WritePlan,
) -> Result<()> {
let table = quote_ident(&src.table);
let scratch = quote_ident(&format!("{}{}", src.table, REBUILD_SUFFIX));
let mut push = |sql: String| {
plan.stmts.push(Stmt {
display: sql.clone(),
sql,
params: Vec::new(),
kind: StmtKind::Schema,
});
plan.schema += 1;
};
let inline_pk = live_cols.iter().filter(|c| c.pk).count() == 1;
let defs: Vec<String> = live_cols.iter().map(|c| column_def(c, inline_pk)).collect();
push(format!(
"CREATE TABLE \"{}\" ({})",
scratch,
defs.join(", ")
));
let post_alter = src.kind == DbKind::DuckDb;
let mut targets: Vec<String> = Vec::new();
let mut sources: Vec<String> = Vec::new();
if src.kind == DbKind::Sqlite {
targets.push("rowid".to_string());
sources.push("rowid".to_string());
}
for (meta, col) in df.columns.iter().zip(live_cols) {
let source = match (post_alter, meta.db_origin.as_deref()) {
(true, _) => col.name.as_str(),
(false, Some(origin)) => origin,
(false, None) => continue,
};
targets.push(format!("\"{}\"", quote_ident(&col.name)));
sources.push(format!("\"{}\"", quote_ident(source)));
}
let order = if post_alter {
" ORDER BY rowid".to_string()
} else {
String::new()
};
push(format!(
"INSERT INTO \"{}\" ({}) SELECT {} FROM \"{}\"{}",
scratch,
targets.join(", "),
sources.join(", "),
table,
order
));
push(format!("DROP TABLE \"{}\"", table));
if src.kind == DbKind::Sqlite {
push("PRAGMA legacy_alter_table = ON".to_string());
}
push(format!(
"ALTER TABLE \"{}\" RENAME TO \"{}\"",
scratch, table
));
if src.kind == DbKind::Sqlite {
push("PRAGMA legacy_alter_table = OFF".to_string());
let conn = open_sqlite(&src.db_path)?;
let mut stmt = conn.prepare(
"SELECT sql FROM sqlite_master \
WHERE tbl_name = ?1 AND sql IS NOT NULL AND type IN ('index', 'trigger') \
ORDER BY type, name",
)?;
let mut rows = stmt.query([&src.table])?;
while let Some(row) = rows.next()? {
push(row.get::<_, String>(0)?);
}
}
plan.rebuild = true;
Ok(())
}
fn push_schema_stmts(
src: &TableSource,
df: &DataFrame,
schema: &SchemaPlan,
live_cols: &[DbColumn],
plan: &mut WritePlan,
) -> Result<()> {
if schema.is_empty() {
return Ok(());
}
preflight_schema(src, schema)?;
if src.kind == DbKind::Sqlite && needs_rebuild(src.kind, schema) {
preflight_rebuild_sqlite(src, schema, live_cols)?;
return push_rebuild_stmts(src, df, live_cols, plan);
}
let table = quote_ident(&src.table);
let mut push = |sql: String| {
plan.stmts.push(Stmt {
display: sql.clone(),
sql,
params: Vec::new(),
kind: StmtKind::Schema,
});
plan.schema += 1;
};
for name in &schema.drops {
push(format!(
"ALTER TABLE \"{}\" DROP COLUMN \"{}\"",
table,
quote_ident(name)
));
}
let mut pending: Vec<(String, String)> = schema.renames.clone();
while !pending.is_empty() {
let occupied: std::collections::HashSet<&str> =
pending.iter().map(|(from, _)| from.as_str()).collect();
let ready: Vec<usize> = pending
.iter()
.enumerate()
.filter(|(_, (from, to))| from == to || !occupied.contains(to.as_str()))
.map(|(i, _)| i)
.collect();
if ready.is_empty() {
let taken: std::collections::HashSet<&str> = src
.columns
.iter()
.chain(live_cols.iter())
.map(|c| c.name.as_str())
.collect();
let from = pending[0].0.clone();
let mut scratch = format!("{}{}", from, SWAP_SUFFIX);
for n in 2.. {
if !taken.contains(scratch.as_str()) {
break;
}
scratch = format!("{}{}{}", from, SWAP_SUFFIX, n);
}
push(format!(
"ALTER TABLE \"{}\" RENAME COLUMN \"{}\" TO \"{}\"",
table,
quote_ident(&from),
quote_ident(&scratch)
));
pending[0].0 = scratch;
continue;
}
for i in ready.iter().rev() {
let (from, to) = pending.remove(*i);
push(format!(
"ALTER TABLE \"{}\" RENAME COLUMN \"{}\" TO \"{}\"",
table,
quote_ident(&from),
quote_ident(&to)
));
}
}
for &i in &schema.adds {
let col = &live_cols[i];
push(format!(
"ALTER TABLE \"{}\" ADD COLUMN \"{}\" {}",
table,
quote_ident(&col.name),
col.decl_raw
));
}
for &(i, _) in &schema.retypes {
let col = &live_cols[i];
push(format!(
"ALTER TABLE \"{}\" ALTER COLUMN \"{}\" TYPE {}",
table,
quote_ident(&col.name),
col.decl_raw
));
}
Ok(())
}
fn push_post_stmts(
src: &TableSource,
df: &DataFrame,
schema: &SchemaPlan,
live_cols: &[DbColumn],
plan: &mut WritePlan,
) -> Result<()> {
if src.kind != DbKind::DuckDb || !needs_rebuild(src.kind, schema) {
return Ok(());
}
preflight_rebuild_duckdb(src)?;
push_rebuild_stmts(src, df, live_cols, plan)
}
fn write_where(
sql: &mut String,
display: &mut String,
params: &mut Vec<Val>,
key: &str,
ids: &[i64],
) {
if let [only] = ids {
sql.push_str(&format!(" WHERE \"{}\" = ", key));
display.push_str(&format!(" WHERE \"{}\" = ", key));
push_val(sql, display, params, Val::Int(*only));
} else {
sql.push_str(&format!(" WHERE \"{}\" IN (", key));
display.push_str(&format!(" WHERE \"{}\" IN (", key));
for (i, id) in ids.iter().enumerate() {
if i > 0 {
sql.push_str(", ");
display.push_str(", ");
}
push_val(sql, display, params, Val::Int(*id));
}
sql.push(')');
display.push(')');
}
}