readyset-server: add support for ALTER TABLE ADD KEY

This commit adds support for the ALTER TABLE ADD KEY statement to
not require a re-snapshot of the table.
When we restart readyset we essentially re-run the snapshot for all
tables. Then we check if the body of the table we have stored in the
registry is the same as the body from the DDL we are trying to apply.
If it is, we don't need to re-snapshot the table.
The change on this commit ensures that we update the in-memory copy of
the DDL we have stored and also persists a new copy of the table body
with the new key in the registry, to avoid re-snapshotting the table
when we restart readyset.

Fixes: REA-5753
Fixes: #1521

Release-Note-Core: Avoid re-snapshotting tables when user executes
  an ALTER TABLE ADD KEY statement.

Change-Id: If42bc9b9ba78179fe22e7f6cb8bb030f90e1e713
Reviewed-on: https://gerrit.readyset.name/c/readyset/+/9601
Reviewed-by: Michael Zink <michael.z@readyset.io>
Tested-by: Buildkite CI

Author: altmannmarcelo

readyset-server/src/controller/sql/mod.rs

@@ -16,9 +16,9 @@ use readyset_errors::{
     ReadySetResult,
 };
 use readyset_sql::ast::{
-    self, CompoundSelectStatement, CreateTableBody, CreateTableOption, FieldDefinitionExpr,
-    NonReplicatedRelation, NotReplicatedReason, Relation, SelectSpecification, SelectStatement,
-    SqlIdentifier, SqlType, TableExpr,
+    self, AlterTableDefinition, CompoundSelectStatement, CreateTableBody, CreateTableOption,
+    FieldDefinitionExpr, NonReplicatedRelation, NotReplicatedReason, Relation, SelectSpecification,
+    SelectStatement, SqlIdentifier, SqlType, TableExpr, TableKey,
 };
 use readyset_sql::DialectDisplay;
 use readyset_sql_passes::alias_removal::TableAliasRewrite;
@@ -292,7 +292,7 @@ impl SqlIncorporator {
                             // Table already exists, so check if it has been changed.
                             if current_body != &body {
                                 // Table has changed. Drop and recreate.
-                                trace!(
+                                info!(
                                     table = %cts.table.display_unquoted(),
                                     "table exists and has changed. Dropping and recreating..."
                                 );
@@ -408,9 +408,13 @@ impl SqlIncorporator {
                 Change::CreateCache(cc) => {
                     self.add_query(cc.name, *cc.statement, cc.always, &schema_search_path, mig)?;
                 }
-                Change::AlterTable(_) => {
-                    // The only ALTER TABLE changes that can end up here (currently) are ones that
-                    // aren't relevant to ReadySet, so we can just ignore them.
+                ref change @ Change::AlterTable(ref alter_table_body) => {
+                    // for PG we don't replicate create index. Also Foreign key is not retrieved
+                    // as part of SHOW CREATE TABLE so we don't need to modify the base schema
+                    if dialect == Dialect::DEFAULT_MYSQL {
+                        let mut table = alter_table_body.table.clone();
+                        self.modify_base_schema(&mut table, &schema_search_path, change)?;
+                    }
                 }
                 Change::CreateType { mut name, ty } => {
                     if let Some(first_schema) = schema_search_path.first() {
@@ -969,6 +973,100 @@ impl SqlIncorporator {
         }
     }
 
+    pub(super) fn propagate_default_search_path(
+        &self,
+        name: &mut Relation,
+        schema_search_path: &[SqlIdentifier],
+    ) -> ReadySetResult<()> {
+        if name.schema.is_none() {
+            let Some(first_schema) = schema_search_path.first() else {
+                return Err(invalid_query_err!("No schema search path provided"));
+            };
+            name.schema = Some(first_schema.clone());
+        }
+        Ok(())
+    }
+
+    /// Modify the base schema of a table. This needs to be aligned with [`Change::requires_resnapshot`]
+    /// This will update the base_schemas and the registry.
+    pub(super) fn modify_base_schema(
+        &mut self,
+        name: &mut Relation,
+        schema_search_path: &[SqlIdentifier],
+        change: &Change,
+    ) -> ReadySetResult<()> {
+        self.propagate_default_search_path(name, schema_search_path)?;
+        if change.requires_resnapshot() {
+            return Err(invalid_query_err!("Unsupported alter table type"));
+        }
+        let base_schema = self
+            .base_schemas
+            .get_mut(name)
+            .ok_or_else(|| invalid_query_err!("Table not found"))?;
+
+        if let Change::AlterTable(alter_table_statement) = change {
+            alter_table_statement
+                .definitions
+                .iter()
+                .try_for_each(|definition| {
+                    definition.iter().try_for_each(|alter_table_body| {
+                        match *alter_table_body {
+                            // Extensive list of operations to adjust the base schema.
+                            AlterTableDefinition::AddKey(ref key @ TableKey::Key { .. })
+                            | AlterTableDefinition::AddKey(ref key @ TableKey::ForeignKey { .. })
+                            | AlterTableDefinition::AddKey(
+                                ref key @ TableKey::CheckConstraint { .. },
+                            )
+                            | AlterTableDefinition::AddKey(
+                                ref key @ TableKey::FulltextKey { .. },
+                            ) => {
+                                let mut key = key.clone();
+                                // propagate default search path to the new key. Otherwise we might get different table body on restarts.
+                                key.get_columns_mut()
+                                    .iter_mut()
+                                    .filter(|col| col.table.is_none())
+                                    .for_each(|col| col.table = Some(name.clone()));
+                                base_schema
+                                    .statement
+                                    .add_key_with_name_mysql(key.clone(), name.clone());
+
+                                if let Some(RecipeExpr::Table { ref body, .. }) =
+                                    self.registry.get(name)
+                                {
+                                    let mut new_body = body.clone();
+                                    new_body.add_key_with_name_mysql(key.clone(), name.clone());
+                                    let _ = self.registry.update_query(
+                                        name,
+                                        RecipeExpr::Table {
+                                            name: name.clone(),
+                                            body: new_body,
+                                            pg_meta: base_schema.pg_meta.clone(),
+                                        },
+                                    );
+                                }
+                                Ok(())
+                            }
+                            // These operations will change the underlying table, and should have been caught earlier
+                            AlterTableDefinition::AddColumn(_)
+                            | AlterTableDefinition::AddKey(TableKey::PrimaryKey { .. })
+                            | AlterTableDefinition::AddKey(TableKey::UniqueKey { .. })
+                            | AlterTableDefinition::AlterColumn { .. }
+                            | AlterTableDefinition::DropColumn { .. }
+                            | AlterTableDefinition::ChangeColumn { .. }
+                            | AlterTableDefinition::RenameColumn { .. }
+                            | AlterTableDefinition::DropConstraint { .. }
+                            | AlterTableDefinition::ReplicaIdentity(_)
+                            | AlterTableDefinition::DropForeignKey { .. } => {
+                                Err(invalid_query_err!("Unsupported ALTER TABLE type"))
+                            }
+                        }
+                    })
+                })?;
+        }
+
+        Ok(())
+    }
+
     fn add_base_via_mir(
         &mut self,
         name: Relation,

readyset-server/src/controller/sql/registry.rs

@@ -423,6 +423,22 @@ impl ExprRegistry {
         Ok(true)
     }
 
+    /// Update the [`RecipeExpr`] associated with the given name (or alias).
+    pub(super) fn update_query(
+        &mut self,
+        name: &Relation,
+        expression: RecipeExpr,
+    ) -> ReadySetResult<()> {
+        let query_id = self.aliases.get(name).ok_or_else(|| {
+            ReadySetError::RecipeInvariantViolated(format!(
+                "Query {} does not exist",
+                name.display_unquoted(),
+            ))
+        })?;
+        self.expressions.insert(*query_id, expression);
+        Ok(())
+    }
+
     /// Retrieves the [`RecipeExpr`] associated with the given name or alias.
     /// If no query is found, returns `None`.
     pub(super) fn get(&self, alias: &Relation) -> Option<&RecipeExpr> {

readyset-sql/src/ast/common.rs

@@ -158,6 +158,16 @@ impl fmt::Display for NullsDistinct {
     }
 }
 
+#[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize, Arbitrary)]
+pub enum TableKeyType {
+    PrimaryKey,
+    UniqueKey,
+    FulltextKey,
+    Key,
+    ForeignKey,
+    CheckConstraint,
+}
+
 #[derive(Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize, Arbitrary)]
 pub enum TableKey {
     PrimaryKey {
@@ -357,6 +367,7 @@ impl TableKey {
             _ => None,
         }
     }
+
     /// Returns a mutable reference to the target columns of the foreign key
     pub fn target_columns_mut(&mut self) -> Option<&mut Vec<Column>> {
         match self {
@@ -396,6 +407,11 @@ impl TableKey {
         matches!(self, TableKey::ForeignKey { .. })
     }
 
+    /// Check if the key is of type Key
+    pub fn is_key(&self) -> bool {
+        matches!(self, TableKey::Key { .. })
+    }
+
     /// Get the columns that the key is defined on
     pub fn get_columns(&self) -> &[Column] {
         match self {
@@ -419,6 +435,18 @@ impl TableKey {
             TableKey::CheckConstraint { .. } => &mut [],
         }
     }
+
+    /// Order the column keys following the rules of MySQL
+    pub fn priority(&self) -> u8 {
+        match self {
+            TableKey::PrimaryKey { .. } => 0,
+            TableKey::UniqueKey { .. } => 1,
+            TableKey::Key { .. } => 2,
+            TableKey::FulltextKey { .. } => 3,
+            TableKey::ForeignKey { .. } => 4,
+            TableKey::CheckConstraint { .. } => 5,
+        }
+    }
 }
 
 impl DialectDisplay for TableKey {
@@ -573,6 +601,19 @@ impl DialectDisplay for TableKey {
     }
 }
 
+impl PartialOrd for TableKey {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for TableKey {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        // Compare by priority
+        self.priority().cmp(&other.priority())
+    }
+}
+
 #[derive(Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Serialize, Deserialize, Arbitrary)]
 #[allow(clippy::large_enum_variant)] // NOTE(aspen): do we actually care about this?
 #[derive(Default)]

readyset-sql/src/ast/create.rs

@@ -145,6 +145,7 @@ impl CreateTableBody {
     pub fn add_key(&mut self, key: TableKey) {
         if let Some(keys) = &mut self.keys {
             keys.push(key);
+            keys.sort();
         } else {
             self.keys = Some(vec![key]);
         }
@@ -153,15 +154,16 @@ impl CreateTableBody {
     /// Add a new key and generate a name for it if it doesn't have one
     /// The column name uses the first column as index name, if that name is already in use,
     /// we append a number to the name.
-    pub fn add_key_with_name(&mut self, key: TableKey, table: Relation) {
+    pub fn add_key_with_name_mysql(&mut self, key: TableKey, table: Relation) {
         let mut new_key = key;
 
         // check index name for non-foreign keys
         if !new_key.is_foreign_key() && new_key.index_name().is_none() {
             let first_column = new_key.get_columns().first().unwrap();
             let base_index_name = first_column.name.as_str().to_string();
             let mut index_name = base_index_name.clone();
-            let mut i = 0;
+            // [make_unique_key_name](https://github.com/mysql/mysql-server/blob/mysql-8.0.42/sql/sql_table.cc#L10386)
+            let mut i = 2;
             while self.is_index_name_in_use(&index_name) {
                 index_name = format!("{base_index_name}_{i}");
                 i += 1;
@@ -171,31 +173,47 @@ impl CreateTableBody {
 
         // special case for foreign keys
         if new_key.is_foreign_key() {
-            if new_key.constraint_name().is_none() {
-                let table_name = table.name.clone();
-                let base_constraint_name = format!("{table_name}_ibfk");
-                let mut i = 1;
-                let mut constraint_name = format!("{base_constraint_name}_{i}");
-                while self.is_constraint_name_in_use(&constraint_name) {
-                    i += 1;
-                    constraint_name = format!("{base_constraint_name}_{i}");
-                }
-                new_key.set_constraint_name(constraint_name.into());
-            }
+            self.handle_add_foreign_key_mysql(&mut new_key, table);
+        }
+        self.add_key(new_key);
+    }
 
-            // target_table
-            if let Some(target_table) = new_key.target_table_mut() {
-                target_table.set_schema(table.schema.clone());
+    /// MySQL logic to add a new foreign key to the table
+    fn handle_add_foreign_key_mysql(&mut self, new_key: &mut TableKey, table: Relation) {
+        if new_key.constraint_name().is_none() {
+            let table_name = table.name.clone();
+            let base_constraint_name = format!("{table_name}_ibfk");
+            let mut i = 1;
+            let mut constraint_name = format!("{base_constraint_name}_{i}");
+            while self.is_constraint_name_in_use(&constraint_name) {
+                i += 1;
+                constraint_name = format!("{base_constraint_name}_{i}");
             }
+            new_key.set_constraint_name(constraint_name.into());
+        }
 
-            // target_columns
-            if let Some(target_columns) = new_key.target_columns_mut() {
-                for column in target_columns {
-                    column.table = Some(table.clone());
-                }
+        // target_table
+        if let Some(target_table) = new_key.target_table_mut() {
+            target_table.set_schema(table.schema.clone());
+        }
+
+        // target_columns
+        if let Some(target_columns) = new_key.target_columns_mut() {
+            for column in target_columns {
+                column.table = Some(table.clone());
             }
         }
-        self.add_key(new_key);
+        let columns = new_key.get_columns();
+        if !self.has_key_on_columns(columns) {
+            // add a new key of type Key on target columns
+            let target_col_new_key = TableKey::Key {
+                index_name: new_key.constraint_name().clone(),
+                columns: columns.to_vec(),
+                index_type: None,
+                constraint_name: None,
+            };
+            self.add_key(target_col_new_key);
+        }
     }
 
     /// Returns the primary key of the table, if one exists.
@@ -213,6 +231,18 @@ impl CreateTableBody {
             .iter()
             .find(|key| matches!(key, TableKey::UniqueKey { .. }))
     }
+
+    /// Check if there is an index of type Key covering a list of columns
+    pub fn has_key_on_columns(&self, columns: &[Column]) -> bool {
+        self.keys.as_ref().unwrap_or(&Vec::new()).iter().any(|key| {
+            key.is_key()
+                && key.get_columns().iter().all(|column| {
+                    columns
+                        .iter()
+                        .any(|c| c.name.as_str() == column.name.as_str())
+                })
+        })
+    }
 }
 
 impl DialectDisplay for CreateTableBody {