diff --git a/core/schema.rs b/core/schema.rs index 415d5f4e1..210830b81 100644 --- a/core/schema.rs +++ b/core/schema.rs @@ -915,6 +915,12 @@ impl Index { if col.unique { // Unique columns in Table should always be named let col_name = col.name.as_ref().unwrap(); + if has_primary_key_index + && table.primary_key_columns.len() == 1 + && &table.primary_key_columns.first().as_ref().unwrap().0 == col_name { + // skip unique columns that are satisfied with pk constraint + return None; + } let (index_name, root_page) = auto_indices.next().expect("number of auto_indices in schema should be same number of indices calculated"); Some(Index { name: normalize_ident(index_name.as_str()), @@ -955,34 +961,50 @@ impl Index { } if let Some(unique_sets) = table.unique_sets.as_ref() { - let unique_set_indices = unique_sets.iter().map(|set| { - let (index_name, root_page) = auto_indices.next().expect( + let unique_set_indices = unique_sets + .iter() + .filter(|set| { + if has_primary_key_index + && table.primary_key_columns.len() == set.len() + && table + .primary_key_columns + .iter() + .all(|col| set.contains(col)) + { + // skip unique columns that are satisfied with pk constraint + return false; + } else { + true + } + }) + .map(|set| { + let (index_name, root_page) = auto_indices.next().expect( "number of auto_indices in schema should be same number of indices calculated", ); - let index_cols = set.iter().map(|(col_name, order)| { - let Some((pos_in_table, _)) = table.get_column(col_name) else { - // This is clearly an invariant that should be maintained, so a panic seems more correct here - panic!( - "Column {} is in index {} but not found in table {}", - col_name, index_name, table.name - ); - }; - IndexColumn { - name: normalize_ident(col_name), - order: *order, - pos_in_table, + let index_cols = set.iter().map(|(col_name, order)| { + let Some((pos_in_table, _)) = table.get_column(col_name) else { + // This is clearly an invariant that should be maintained, so a panic seems more correct here + panic!( + "Column {} is in index {} but not found in table {}", + col_name, index_name, table.name + ); + }; + IndexColumn { + name: normalize_ident(col_name), + order: *order, + pos_in_table, + } + }); + Index { + name: normalize_ident(index_name.as_str()), + table_name: table.name.clone(), + root_page, + columns: index_cols.collect(), + unique: true, + ephemeral: false, } }); - Index { - name: normalize_ident(index_name.as_str()), - table_name: table.name.clone(), - root_page, - columns: index_cols.collect(), - unique: true, - ephemeral: false, - } - }); indices.extend(unique_set_indices); } @@ -1509,4 +1531,85 @@ mod tests { Ok(()) } + + #[test] + fn test_automatic_index_primary_key_is_unique() -> Result<()> { + let sql = r#"CREATE TABLE t1 (a primary key unique);"#; + let table = BTreeTable::from_sql(sql, 0)?; + let mut index = Index::automatic_from_primary_key_and_unique( + &table, + vec![("sqlite_autoindex_t1_1".to_string(), 2)], + )?; + + assert!(index.len() == 1); + let index = index.pop().unwrap(); + + assert_eq!(index.name, "sqlite_autoindex_t1_1"); + assert_eq!(index.table_name, "t1"); + assert_eq!(index.root_page, 2); + assert!(index.unique); + assert_eq!(index.columns.len(), 1); + assert_eq!(index.columns[0].name, "a"); + assert!(matches!(index.columns[0].order, SortOrder::Asc)); + + Ok(()) + } + + #[test] + fn test_automatic_index_primary_key_is_unique_and_composite() -> Result<()> { + let sql = r#"CREATE TABLE t1 (a, b, PRIMARY KEY(a, b), UNIQUE(a, b));"#; + let table = BTreeTable::from_sql(sql, 0)?; + let mut index = Index::automatic_from_primary_key_and_unique( + &table, + vec![("sqlite_autoindex_t1_1".to_string(), 2)], + )?; + + assert!(index.len() == 1); + let index = index.pop().unwrap(); + + assert_eq!(index.name, "sqlite_autoindex_t1_1"); + assert_eq!(index.table_name, "t1"); + assert_eq!(index.root_page, 2); + assert!(index.unique); + assert_eq!(index.columns.len(), 2); + assert_eq!(index.columns[0].name, "a"); + assert_eq!(index.columns[1].name, "b"); + assert!(matches!(index.columns[0].order, SortOrder::Asc)); + + Ok(()) + } + + #[test] + fn test_automatic_index_unique_and_a_pk() -> Result<()> { + let sql = r#"CREATE TABLE t1 (a NUMERIC UNIQUE UNIQUE, b TEXT PRIMARY KEY)"#; + let table = BTreeTable::from_sql(sql, 0)?; + let mut indexes = Index::automatic_from_primary_key_and_unique( + &table, + vec![ + ("sqlite_autoindex_t1_1".to_string(), 2), + ("sqlite_autoindex_t1_2".to_string(), 3), + ], + )?; + + assert!(indexes.len() == 2); + let index = indexes.pop().unwrap(); + assert_eq!(index.name, "sqlite_autoindex_t1_2"); + assert_eq!(index.table_name, "t1"); + assert_eq!(index.root_page, 3); + assert!(index.unique); + assert_eq!(index.columns.len(), 1); + assert_eq!(index.columns[0].name, "a"); + assert!(matches!(index.columns[0].order, SortOrder::Asc)); + + let index = indexes.pop().unwrap(); + assert_eq!(index.name, "sqlite_autoindex_t1_1"); + assert_eq!(index.table_name, "t1"); + assert_eq!(index.root_page, 2); + assert!(index.unique); + assert_eq!(index.columns.len(), 1); + assert_eq!(index.columns[0].name, "b"); + assert!(matches!(index.columns[0].order, SortOrder::Asc)); + + Ok(()) + } } diff --git a/core/translate/schema.rs b/core/translate/schema.rs index 7c0b3d7cc..4942a9564 100644 --- a/core/translate/schema.rs +++ b/core/translate/schema.rs @@ -275,7 +275,6 @@ fn check_automatic_pk_index_required( options, } => { let mut primary_key_definition = None; - let mut unique_def_count = 0usize; // Used to dedup named unique constraints let mut unique_sets = vec![]; @@ -286,6 +285,9 @@ fn check_automatic_pk_index_required( columns: pk_cols, .. } = &constraint.constraint { + if primary_key_definition.is_some() { + bail_parse_error!("table {} has more than one primary key", tbl_name); + } let primary_key_column_results: Vec> = pk_cols .iter() .map(|col| match &col.expr { @@ -312,22 +314,28 @@ fn check_automatic_pk_index_required( bail_parse_error!("No such column: {}", column_name); } - if matches!( - primary_key_definition, - Some(PrimaryKeyDefinitionType::Simple { .. }) - ) { - primary_key_definition = Some(PrimaryKeyDefinitionType::Composite); - continue; - } - if primary_key_definition.is_none() { - let column_def = column_def.unwrap(); - let typename = - column_def.col_type.as_ref().map(|t| t.name.as_str()); - let is_descending = pk_info.is_descending; - primary_key_definition = Some(PrimaryKeyDefinitionType::Simple { - typename, - is_descending, - }); + match &mut primary_key_definition { + Some(PrimaryKeyDefinitionType::Simple { column, .. }) => { + let mut columns = HashSet::new(); + columns.insert(column.clone()); + primary_key_definition = + Some(PrimaryKeyDefinitionType::Composite { columns }); + } + Some(PrimaryKeyDefinitionType::Composite { columns }) => { + columns.insert(column_name.clone()); + } + None => { + let column_def: &ast::ColumnDefinition = column_def.unwrap(); + let typename = + column_def.col_type.as_ref().map(|t| t.name.as_str()); + let is_descending = pk_info.is_descending; + primary_key_definition = + Some(PrimaryKeyDefinitionType::Simple { + typename, + is_descending, + column: column_name.clone(), + }); + } } } } else if let ast::TableConstraint::Unique { @@ -354,7 +362,7 @@ fn check_automatic_pk_index_required( } } - // Check column constraints for PRIMARY KEY + // Check column constraints for PRIMARY KEY and UNIQUE for (_, col_def) in columns.iter() { for constraint in &col_def.constraints { if matches!( @@ -368,9 +376,12 @@ fn check_automatic_pk_index_required( primary_key_definition = Some(PrimaryKeyDefinitionType::Simple { typename, is_descending: false, + column: col_def.col_name.0.clone(), }); } else if matches!(constraint.constraint, ast::ColumnConstraint::Unique(..)) { - unique_def_count += 1; + let mut single_set = HashSet::new(); + single_set.insert(col_def.col_name.0.clone()); + unique_sets.push(single_set); } } } @@ -381,36 +392,40 @@ fn check_automatic_pk_index_required( } unique_sets.dedup(); - let mut total_indices = unique_def_count + unique_sets.len(); // Check if we need an automatic index + let mut pk_is_unique = false; let auto_index_pk = if let Some(primary_key_definition) = &primary_key_definition { match primary_key_definition { PrimaryKeyDefinitionType::Simple { typename, is_descending, + column, } => { + pk_is_unique = unique_sets + .iter() + .any(|set| set.len() == 1 && set.contains(column)); let is_integer = typename.is_some() && typename.unwrap().eq_ignore_ascii_case("INTEGER"); // Should match on any case of INTEGER !is_integer || *is_descending } - PrimaryKeyDefinitionType::Composite => true, + PrimaryKeyDefinitionType::Composite { columns } => { + pk_is_unique = unique_sets.iter().any(|set| set == columns); + true + } } } else { false }; - if auto_index_pk { + let mut total_indices = unique_sets.len(); + // if pk needs and index, but we already found out we primary key is unique, we only need a single index since constraint pk == unique + if auto_index_pk && !pk_is_unique { total_indices += 1; } if total_indices > 0 { - if total_indices == 1 { - let index_reg = program.alloc_register(); - Ok(Some(index_reg..index_reg + 1)) - } else { - let index_start_reg = program.alloc_registers(total_indices + 1); - Ok(Some(index_start_reg..index_start_reg + total_indices + 1)) - } + let index_start_reg = program.alloc_registers(total_indices); + Ok(Some(index_start_reg..index_start_reg + total_indices)) } else { Ok(None) } @@ -423,15 +438,19 @@ fn check_automatic_pk_index_required( enum PrimaryKeyDefinitionType<'a> { Simple { + column: String, typename: Option<&'a str>, is_descending: bool, }, - Composite, + Composite { + columns: HashSet, + }, } struct TableFormatter<'a> { body: &'a ast::CreateTableBody, } + impl Display for TableFormatter<'_> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { self.body.to_fmt(f)