Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(torii-sqlite): support enum upgrade of variants #2930

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open

feat(torii-sqlite): support enum upgrade of variants #2930

Changes from 8 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
d44e11b
refactor(torii): support enum upgrade
Larkooo Jan 20, 2025
f7c1ce8
fmt
Larkooo Jan 20, 2025
8087be0
Merge branch 'main' into upgrade-enum
Larkooo Jan 21, 2025
9423c97
refactor: check constraint
Larkooo Jan 21, 2025
6e2c333
fix diff propagation in upgrade
Larkooo Jan 22, 2025
353804d
fmt
Larkooo Jan 22, 2025
415daca
update with temporary table
Larkooo Jan 22, 2025
95cebf7
f
Larkooo Jan 22, 2025
4bc0b36
fix: new enum diff
Larkooo Jan 24, 2025
f993071
fix: recursive diff
Larkooo Jan 27, 2025
1df0eaa
fmt
Larkooo Jan 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 80 additions & 15 deletions crates/torii/sqlite/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -847,20 +847,44 @@
));
};

let modify_column =
|alter_table_queries: &mut Vec<String>, name: &str, sql_type: &str, sql_value: &str| {
// SQLite doesn't support ALTER COLUMN directly, so we need to:
// 1. Create a temporary table to store the current values
// 2. Drop the old column & index
// 3. Create new column with new type/constraint
// 4. Copy values back & create new index
alter_table_queries.push(format!(
"CREATE TEMPORARY TABLE tmp_values_{name} AS SELECT internal_id, [{name}] FROM \
[{table_id}]"
));
alter_table_queries.push(format!("DROP INDEX IF EXISTS [idx_{table_id}_{name}]"));
alter_table_queries.push(format!("ALTER TABLE [{table_id}] DROP COLUMN [{name}]"));
alter_table_queries
.push(format!("ALTER TABLE [{table_id}] ADD COLUMN [{name}] {sql_type}"));
alter_table_queries.push(format!("UPDATE [{table_id}] SET [{name}] = {sql_value}"));
alter_table_queries.push(format!("DROP TABLE tmp_values_{name}"));
alter_table_queries.push(format!(
"CREATE INDEX IF NOT EXISTS [idx_{table_id}_{name}] ON [{table_id}] ([{name}]);"
));
};

Check warning on line 870 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L851-L870 

Added lines #L851 - L870 were not covered by tests
Comment on lines +850 to +870
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Ohayo sensei! Fix potential data loss in column modification.

The modify_column closure creates a temporary table but doesn't copy the values back to the new column. The UPDATE statement on line 865 sets a static sql_value instead of restoring the original values.

Apply this diff to fix the data loss:

             alter_table_queries.push(format!(
                 "CREATE TEMPORARY TABLE tmp_values_{name} AS SELECT internal_id, [{name}] FROM \
                  [{table_id}]"
             ));
             alter_table_queries.push(format!("DROP INDEX IF EXISTS [idx_{table_id}_{name}]"));
             alter_table_queries.push(format!("ALTER TABLE [{table_id}] DROP COLUMN [{name}]"));
             alter_table_queries
                 .push(format!("ALTER TABLE [{table_id}] ADD COLUMN [{name}] {sql_type}"));
-            alter_table_queries.push(format!("UPDATE [{table_id}] SET [{name}] = {sql_value}"));
+            alter_table_queries.push(format!(
+                "UPDATE [{table_id}] SET [{name}] = (SELECT [{name}] FROM tmp_values_{name} \
+                 WHERE tmp_values_{name}.internal_id = [{table_id}].internal_id)"
+            ));
             alter_table_queries.push(format!("DROP TABLE tmp_values_{name}"));
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
let modify_column =
|alter_table_queries: &mut Vec<String>, name: &str, sql_type: &str, sql_value: &str| {
// SQLite doesn't support ALTER COLUMN directly, so we need to:
// 1. Create a temporary table to store the current values
// 2. Drop the old column & index
// 3. Create new column with new type/constraint
// 4. Copy values back & create new index
alter_table_queries.push(format!(
"CREATE TEMPORARY TABLE tmp_values_{name} AS SELECT internal_id, [{name}] FROM \
[{table_id}]"
));
alter_table_queries.push(format!("DROP INDEX IF EXISTS [idx_{table_id}_{name}]"));
alter_table_queries.push(format!("ALTER TABLE [{table_id}] DROP COLUMN [{name}]"));
alter_table_queries
.push(format!("ALTER TABLE [{table_id}] ADD COLUMN [{name}] {sql_type}"));
alter_table_queries.push(format!("UPDATE [{table_id}] SET [{name}] = {sql_value}"));
alter_table_queries.push(format!("DROP TABLE tmp_values_{name}"));
alter_table_queries.push(format!(
"CREATE INDEX IF NOT EXISTS [idx_{table_id}_{name}] ON [{table_id}] ([{name}]);"
));
};
let modify_column =
|alter_table_queries: &mut Vec<String>, name: &str, sql_type: &str, sql_value: &str| {
// SQLite doesn't support ALTER COLUMN directly, so we need to:
// 1. Create a temporary table to store the current values
// 2. Drop the old column & index
// 3. Create new column with new type/constraint
// 4. Copy values back & create new index
alter_table_queries.push(format!(
"CREATE TEMPORARY TABLE tmp_values_{name} AS SELECT internal_id, [{name}] FROM \
[{table_id}]"
));
alter_table_queries.push(format!("DROP INDEX IF EXISTS [idx_{table_id}_{name}]"));
alter_table_queries.push(format!("ALTER TABLE [{table_id}] DROP COLUMN [{name}]"));
alter_table_queries
.push(format!("ALTER TABLE [{table_id}] ADD COLUMN [{name}] {sql_type}"));
alter_table_queries.push(format!(
"UPDATE [{table_id}] SET [{name}] = (SELECT [{name}] FROM tmp_values_{name} \
WHERE tmp_values_{name}.internal_id = [{table_id}].internal_id)"
));
alter_table_queries.push(format!("DROP TABLE tmp_values_{name}"));
alter_table_queries.push(format!(
"CREATE INDEX IF NOT EXISTS [idx_{table_id}_{name}] ON [{table_id}] ([{name}]);"
));
};


match ty {
Ty::Struct(s) => {
let struct_diff =
if let Some(upgrade_diff) = upgrade_diff { upgrade_diff.as_struct() } else { None };

for member in &s.children {
if let Some(upgrade_diff) = upgrade_diff {
if !upgrade_diff
.as_struct()
.unwrap()
.children
.iter()
.any(|m| m.name == member.name)
{
let member_diff = if let Some(diff) = struct_diff {
if let Some(m) = diff.children.iter().find(|m| m.name == member.name) {
Some(&m.ty)

Check warning on line 880 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L879-L880 

Added lines #L879 - L880 were not covered by tests
} else {
// If the member is not in the diff, skip it
continue;
}
}
} else {
None
};

let mut new_path = path.to_vec();
new_path.push(member.name.clone());
Expand All @@ -872,23 +896,36 @@
alter_table_queries,
indices,
table_id,
None,
member_diff,
)?;
}
}
Ty::Tuple(tuple) => {
let tuple_diff =
if let Some(upgrade_diff) = upgrade_diff { upgrade_diff.as_tuple() } else { None };

Check warning on line 905 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L904-L905 

Added lines #L904 - L905 were not covered by tests

for (idx, member) in tuple.iter().enumerate() {
let mut new_path = path.to_vec();
new_path.push(idx.to_string());

let member_diff = if let Some(diff) = tuple_diff {
if let Some((_, m)) = diff.iter().enumerate().find(|(i, _)| *i == idx) {
Some(m)

Check warning on line 913 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L911-L913 

Added lines #L911 - L913 were not covered by tests
} else {
continue;

Check warning on line 915 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L915 

Added line #L915 was not covered by tests
}
} else {
None

Check warning on line 918 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L918 

Added line #L918 was not covered by tests
};

add_columns_recursive(
&new_path,
member,
columns,
alter_table_queries,
indices,
table_id,
None,
member_diff,

Check warning on line 928 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L928 

Added line #L928 was not covered by tests
)?;
}
}
Expand All @@ -899,17 +936,45 @@
add_column(&column_name, "TEXT");
}
Ty::Enum(e) => {
// The variant of the enum
let enum_diff =
if let Some(upgrade_diff) = upgrade_diff { upgrade_diff.as_enum() } else { None };

let column_name =
if column_prefix.is_empty() { "option".to_string() } else { column_prefix };

let all_options =
e.options.iter().map(|c| format!("'{}'", c.name)).collect::<Vec<_>>().join(", ");

let sql_type = format!("TEXT CHECK([{column_name}] IN ({all_options}))");
add_column(&column_name, &sql_type);
let sql_type = format!(
"TEXT CONSTRAINT [{column_name}_check] CHECK([{column_name}] IN ({all_options}))"
);
if enum_diff.is_some() {
// For upgrades, modify the existing option column to add the new options to the
// CHECK constraint We need to drop the old column and create a new
// one with the new CHECK constraint
modify_column(
alter_table_queries,
&column_name,
&sql_type,
&format!("[{column_name}]"),
);

Check warning on line 960 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L952-L960 

Added lines #L952 - L960 were not covered by tests
} else {
// For new tables, create the column directly
add_column(&column_name, &sql_type);
}
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codebase verification

🛠️ Refactor suggestion

Ohayo sensei! The review comment is spot-on about enum variant removal!

The current implementation is indeed unsafe when removing enum variants. The code modifies the CHECK constraint without validating if any rows contain the removed variants, which could lead to invalid data states. The suggested validation using RAISE(ABORT) is essential to fail fast and prevent data inconsistency.

🔗 Analysis chain

Ohayo sensei! Handle removed enum variants.

The current implementation only handles new variants but doesn't account for variants that might have been removed in the upgrade. This could lead to invalid data if variants are removed.

Run this script to check if any enums have removed variants:

Consider adding validation to handle removed variants:

             let sql_type = format!(
                 "TEXT CONSTRAINT [{column_name}_check] CHECK([{column_name}] IN ({all_options}))"
             );
             if enum_diff.is_some() {
+                // Validate existing data before modifying the constraint
+                alter_table_queries.push(format!(
+                    "SELECT CASE WHEN EXISTS (
+                        SELECT 1 FROM [{table_id}]
+                        WHERE [{column_name}] NOT IN ({all_options})
+                    )
+                    THEN RAISE(ABORT, 'Found invalid enum variants')
+                    END"
+                ));
+
                 modify_column(
                     alter_table_queries,
                     &column_name,
🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Description: Check for removed enum variants in the codebase

# Search for enum definitions and their variants
ast-grep --pattern 'enum $name {
  $$$
  $variant($$$)
  $$$
}'

Length of output: 65

Script:

#!/bin/bash
# Search for enum definitions in Rust
ast-grep --pattern 'enum $name { $$$ }'

# Search for enum-related tests and migration code
rg -l "enum.*diff|enum.*migration" --type rust

# Search for EnumDiff implementation
rg "impl.*EnumDiff" -A 5 --type rust

Length of output: 192

Script:

#!/bin/bash
# Find EnumDiff type definition and usage
rg "struct EnumDiff|type EnumDiff" -A 5 --type rust

# Look for enum migration tests
rg "test.*enum.*migration|test.*enum.*upgrade" -B 2 -A 10 --type rust

# Search for variant removal handling
rg "enum.*remove|variant.*remove|DROP.*CONSTRAINT" -A 5 --type rust

Length of output: 1181

Script:

#!/bin/bash
# Search for CHECK constraint and column modification logic
rg "CHECK.*IN|modify_column|alter_table.*enum" -B 2 -A 5 --type rust

# Look for SQL migration validation
rg "RAISE.*ABORT|SELECT.*EXISTS.*WHERE" -A 5 --type rust

# Find column modification implementation
rg "fn modify_column" -B 2 -A 10 --type rust

Length of output: 1868


for child in &e.options {
// If we have a diff, only process new variants that aren't in the original enum
let variant_diff = if let Some(diff) = enum_diff {
if let Some(v) = diff.options.iter().find(|v| v.name == child.name) {
Some(&v.ty)

Check warning on line 970 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L969-L970 

Added lines #L969 - L970 were not covered by tests
} else {
continue;

Check warning on line 972 in crates/torii/sqlite/src/lib.rs

View check run for this annotation

Codecov / codecov/patch

crates/torii/sqlite/src/lib.rs#L972 

Added line #L972 was not covered by tests
}
} else {
None
};

if let Ty::Tuple(tuple) = &child.ty {
if tuple.is_empty() {
continue;
Expand All @@ -926,7 +991,7 @@
alter_table_queries,
indices,
table_id,
None,
variant_diff,
)?;
}
}
Expand Down
Loading