Merge 'Fix autoindex of primary key marked as unique' from Pere Diaz Bou

Primary keys that are marked as unique constraint, do not need to have
separate indexes, one is enough. In the case primary key is integer,
therefore rowid alias, we still need an index to satisfy unique
constraint.

Closes #1494

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(())
+    }
 }

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<Result<PrimaryKeyColumnInfo>> = 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<String>,
+    },
 }
 
 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)