sql_server: add additional upstream constraint support (#33991)

Adds support for unique constraints and isolates constraint detection
from table discovery. This also fixes an issue where a column appears
multiple times in the source description if it is a member of multiple
unique/PK constraints.


### Motivation

Implements MaterializeInc/database-issues#9723
 
### Tips for reviewer

Functional changes:
- `src/sql-server-util/src/inspect.rs` (changes SQL Server queries:
simplify table desc query, break out constraint query)
- `src/sql/src/pure/sql_server.rs` (building out
SourceExportStatementDetails)

### Checklist

- [ ] This PR has adequate test coverage / QA involvement has been duly
considered. ([trigger-ci for additional test/nightly
runs](https://trigger-ci.dev.materialize.com/))
- [ ] This PR has an associated up-to-date [design
doc](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/README.md),
is a design doc
([template](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/00000000_template.md)),
or is sufficiently small to not require a design.
  <!-- Reference the design in the description. -->
- [ ] If this PR evolves [an existing `$T ⇔ Proto$T`
mapping](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/command-and-response-binary-encoding.md)
(possibly in a backwards-incompatible way), then it is tagged with a
`T-proto` label.
- [ ] If this PR will require changes to cloud orchestration or tests,
there is a companion cloud PR to account for those changes that is
tagged with the release-blocker label
([example](MaterializeInc/cloud#5021)).
<!-- Ask in #team-cloud on Slack if you need help preparing the cloud
PR. -->
- [ ] If this PR includes major [user-facing behavior
changes](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/guide-changes.md#what-changes-require-a-release-note),
I have pinged the relevant PM to schedule a changelog post.

Author: martykulma

Cargo.lock

@@ -240,7 +240,7 @@ def validate(self) -> Testdrive:
             # > SELECT regexp_match(create_sql, 'TEXT COLUMNS = \\((.*?)\\)')[1] FROM (SHOW CREATE SOURCE sql_server_source_tableA{self.suffix});
             # "\"f4\""
 
-            # Confirm that the primary key information has been propagated from Pg
+            # Confirm that the primary key information has been propagated from SQL Server
             > SELECT key FROM (SHOW INDEXES ON sql_server_source_tableA{self.suffix});
             {{f1,f2}}
 

misc/python/materialize/checks/all_checks/sql_server_cdc.py

@@ -66,6 +66,7 @@ mz-segment = { path = "../segment" }
 mz-service = { path = "../service" }
 mz-sql = { path = "../sql" }
 mz-sql-parser = { path = "../sql-parser" }
+mz-sql-server-util = { path = "../sql-server-util"}
 mz-ssh-util = { path = "../ssh-util" }
 mz-storage-client = { path = "../storage-client" }
 mz-storage-types = { path = "../storage-types" }

src/adapter/Cargo.toml

@@ -148,6 +148,7 @@ pub(crate) async fn migrate(
         // Migration functions may also take `tx` as input to stage
         // arbitrary changes to the catalog.
         ast_rewrite_create_sink_partition_strategy(stmt)?;
+        ast_rewrite_sql_server_constraints(stmt)?;
         Ok(())
     })?;
 
@@ -853,3 +854,90 @@ fn ast_rewrite_create_sink_partition_strategy(
         .retain(|op| op.name != CreateSinkOptionName::PartitionStrategy);
     Ok(())
 }
+
+// Migrate SQL Server constraint information from the columns to dedicated constraints field.
+fn ast_rewrite_sql_server_constraints(stmt: &mut Statement<Raw>) -> Result<(), anyhow::Error> {
+    use mz_sql::ast::{
+        CreateSubsourceOptionName, TableFromSourceOptionName, Value, WithOptionValue,
+    };
+    use mz_sql_server_util::desc::{SqlServerTableConstraint, SqlServerTableConstraintType};
+    use mz_storage_types::sources::ProtoSourceExportStatementDetails;
+    use mz_storage_types::sources::proto_source_export_statement_details::Kind;
+
+    let deets: Option<&mut String> = match stmt {
+        Statement::CreateSubsource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
+            if matches!(option.name, CreateSubsourceOptionName::Details)
+                && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
+            {
+                Some(details)
+            } else {
+                None
+            }
+        }),
+        Statement::CreateTableFromSource(stmt) => stmt.with_options.iter_mut().find_map(|option| {
+            if matches!(option.name, TableFromSourceOptionName::Details)
+                && let Some(WithOptionValue::Value(Value::String(ref mut details))) = option.value
+            {
+                Some(details)
+            } else {
+                None
+            }
+        }),
+        _ => None,
+    };
+    let Some(deets) = deets else {
+        return Ok(());
+    };
+
+    let current_value = hex::decode(&mut *deets)?;
+    let current_value = ProtoSourceExportStatementDetails::decode(&*current_value)?;
+
+    // avoid further work if this isn't SQL Server
+    if !matches!(current_value.kind, Some(Kind::SqlServer(_))) {
+        return Ok(());
+    };
+
+    let SourceExportStatementDetails::SqlServer {
+        mut table,
+        capture_instance,
+        initial_lsn,
+    } = SourceExportStatementDetails::from_proto(current_value)?
+    else {
+        unreachable!("statement details must exist for SQL Server");
+    };
+
+    // Migration has already occured or did not need to happen.
+    if !table.constraints.is_empty() {
+        return Ok(());
+    }
+
+    // Relocates the primary key constraint information from the individual columns to the
+    // constraints field. This ensures that the columns no longer hold constraint information.
+    let mut migrated_constraints: BTreeMap<_, Vec<_>> = BTreeMap::new();
+    for col in table.columns.iter_mut() {
+        if let Some(constraint_name) = col.primary_key_constraint.take() {
+            migrated_constraints
+                .entry(constraint_name)
+                .or_default()
+                .push(col.name.to_string());
+        }
+    }
+
+    table.constraints = migrated_constraints
+        .into_iter()
+        .map(|(constraint_name, column_names)| SqlServerTableConstraint {
+            constraint_name: constraint_name.to_string(),
+            constraint_type: SqlServerTableConstraintType::PrimaryKey,
+            column_names,
+        })
+        .collect();
+
+    let new_value = SourceExportStatementDetails::SqlServer {
+        table,
+        capture_instance,
+        initial_lsn,
+    };
+    *deets = hex::encode(new_value.into_proto().encode_to_vec());
+
+    Ok(())
+}

src/adapter/src/catalog/migrate.rs

@@ -20,6 +20,16 @@ message ProtoSqlServerTableDesc {
   string name = 1;
   string schema_name = 2;
   repeated ProtoSqlServerColumnDesc columns = 3;
+  repeated ProtoSqlServerTableConstraint constraints = 5;
+}
+
+message ProtoSqlServerTableConstraint {
+  string constraint_name = 1;
+  oneof constraint_type {
+    google.protobuf.Empty primary_key = 2;
+    google.protobuf.Empty unique = 3;
+  }
+  repeated string column_names = 4;
 }
 
 message ProtoSqlServerColumnDesc {

src/sql-server-util/src/desc.proto

@@ -38,6 +38,7 @@ use serde::{Deserialize, Serialize};
 use std::collections::BTreeSet;
 use std::sync::Arc;
 
+use crate::desc::proto_sql_server_table_constraint::ConstraintType;
 use crate::{SqlServerDecodeError, SqlServerError};
 
 include!(concat!(env!("OUT_DIR"), "/mz_sql_server_util.rs"));
@@ -59,24 +60,34 @@ pub struct SqlServerTableDesc {
     pub name: Arc<str>,
     /// Columns for the table.
     pub columns: Box<[SqlServerColumnDesc]>,
+    /// Constraints for the table.
+    pub constraints: Vec<SqlServerTableConstraint>,
 }
 
 impl SqlServerTableDesc {
     /// Creating a [`SqlServerTableDesc`] from a [`SqlServerTableRaw`] description.
     ///
     /// Note: Not all columns from SQL Server can be ingested into Materialize. To determine if a
     /// column is supported see [`SqlServerColumnDesc::decode_type`].
-    pub fn new(raw: SqlServerTableRaw) -> Self {
+    pub fn new(
+        raw: SqlServerTableRaw,
+        raw_constraints: Vec<SqlServerTableConstraintRaw>,
+    ) -> Result<Self, SqlServerError> {
         let columns: Box<[_]> = raw
             .columns
             .into_iter()
             .map(SqlServerColumnDesc::new)
             .collect();
-        SqlServerTableDesc {
+        let constraints = raw_constraints
+            .into_iter()
+            .map(SqlServerTableConstraint::try_from)
+            .collect::<Result<Vec<_>, _>>()?;
+        Ok(SqlServerTableDesc {
             schema_name: raw.schema_name,
             name: raw.name,
             columns,
-        }
+            constraints,
+        })
     }
 
     /// Returns the [`SqlServerQualifiedTableName`] for this [`SqlServerTableDesc`].
@@ -121,6 +132,7 @@ impl RustType<ProtoSqlServerTableDesc> for SqlServerTableDesc {
             name: self.name.to_string(),
             schema_name: self.schema_name.to_string(),
             columns: self.columns.iter().map(|c| c.into_proto()).collect(),
+            constraints: self.constraints.iter().map(|c| c.into_proto()).collect(),
         }
     }
 
@@ -130,10 +142,99 @@ impl RustType<ProtoSqlServerTableDesc> for SqlServerTableDesc {
             .into_iter()
             .map(|c| c.into_rust())
             .collect::<Result<_, _>>()?;
+        let constraints = proto
+            .constraints
+            .into_iter()
+            .map(|c| c.into_rust())
+            .collect::<Result<_, _>>()?;
         Ok(SqlServerTableDesc {
             schema_name: proto.schema_name.into(),
             name: proto.name.into(),
             columns,
+            constraints,
+        })
+    }
+}
+
+/// SQL Server table constraint type (e.g. PRIMARY KEY, UNIQUE, etc.)
+/// See <https://learn.microsoft.com/en-us/sql/relational-databases/system-information-schema-views/table-constraints-transact-sql?view=sql-server-ver17>
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Arbitrary)]
+pub enum SqlServerTableConstraintType {
+    PrimaryKey,
+    Unique,
+}
+
+impl TryFrom<String> for SqlServerTableConstraintType {
+    type Error = SqlServerError;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        match value.as_str() {
+            "PRIMARY KEY" => Ok(Self::PrimaryKey),
+            "UNIQUE" => Ok(Self::Unique),
+            name => Err(SqlServerError::InvalidData {
+                column_name: "constraint_type".into(),
+                error: format!("Unknown constraint type: {name}"),
+            }),
+        }
+    }
+}
+
+impl RustType<proto_sql_server_table_constraint::ConstraintType> for SqlServerTableConstraintType {
+    fn into_proto(&self) -> proto_sql_server_table_constraint::ConstraintType {
+        match self {
+            SqlServerTableConstraintType::PrimaryKey => ConstraintType::PrimaryKey(()),
+            SqlServerTableConstraintType::Unique => ConstraintType::Unique(()),
+        }
+    }
+
+    fn from_proto(
+        proto: proto_sql_server_table_constraint::ConstraintType,
+    ) -> Result<Self, mz_proto::TryFromProtoError> {
+        Ok(match proto {
+            ConstraintType::PrimaryKey(_) => SqlServerTableConstraintType::PrimaryKey,
+            ConstraintType::Unique(_) => SqlServerTableConstraintType::Unique,
+        })
+    }
+}
+
+/// SQL Server table constraint.
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Arbitrary)]
+pub struct SqlServerTableConstraint {
+    pub constraint_name: String,
+    pub constraint_type: SqlServerTableConstraintType,
+    pub column_names: Vec<String>,
+}
+
+impl TryFrom<SqlServerTableConstraintRaw> for SqlServerTableConstraint {
+    type Error = SqlServerError;
+
+    fn try_from(value: SqlServerTableConstraintRaw) -> Result<Self, Self::Error> {
+        Ok(SqlServerTableConstraint {
+            constraint_name: value.constraint_name,
+            constraint_type: value.constraint_type.try_into()?,
+            column_names: value.columns,
+        })
+    }
+}
+
+impl RustType<ProtoSqlServerTableConstraint> for SqlServerTableConstraint {
+    fn into_proto(&self) -> ProtoSqlServerTableConstraint {
+        ProtoSqlServerTableConstraint {
+            constraint_name: self.constraint_name.clone(),
+            constraint_type: Some(self.constraint_type.into_proto()),
+            column_names: self.column_names.clone(),
+        }
+    }
+
+    fn from_proto(
+        proto: ProtoSqlServerTableConstraint,
+    ) -> Result<Self, mz_proto::TryFromProtoError> {
+        Ok(SqlServerTableConstraint {
+            constraint_name: proto.constraint_name,
+            constraint_type: proto
+                .constraint_type
+                .into_rust_if_some("ProtoSqlServerTableConstraint::constraint_type")?,
+            column_names: proto.column_names,
         })
     }
 }
@@ -189,7 +290,8 @@ pub struct SqlServerColumnDesc {
     /// Note: This type might differ from the `decode_type`, e.g. a user can
     /// specify `TEXT COLUMNS` to decode columns as text.
     pub column_type: Option<SqlColumnType>,
-    /// If this column is part of the primary key for the table, and the name of the constraint.
+    /// This field is deprecated and will be removed in a future version.  This exists only for the
+    /// purpose of migrating from old representations.
     pub primary_key_constraint: Option<Arc<str>>,
     /// Rust type we should parse the data from a [`tiberius::Row`] as.
     pub decode_type: SqlServerColumnDecodeType,
@@ -223,7 +325,7 @@ impl SqlServerColumnDesc {
         };
         SqlServerColumnDesc {
             name: Arc::clone(&raw.name),
-            primary_key_constraint: raw.primary_key_constraint.clone(),
+            primary_key_constraint: None,
             column_type,
             decode_type,
             raw_type: Arc::clone(&raw.data_type),
@@ -528,8 +630,6 @@ pub struct SqlServerColumnRaw {
     pub data_type: Arc<str>,
     /// Whether or not the column is nullable.
     pub is_nullable: bool,
-    /// If the column is part of the primary key for the table, and the name of the constraint.
-    pub primary_key_constraint: Option<Arc<str>>,
     /// Maximum length (in bytes) of the column.
     ///
     /// For `varchar(max)`, `nvarchar(max)`, `varbinary(max)`, or `xml` this will be `-1`. For
@@ -548,6 +648,14 @@ pub struct SqlServerColumnRaw {
     pub is_computed: bool,
 }
 
+/// Raw metadata for a table constraint.
+#[derive(Clone, Debug)]
+pub struct SqlServerTableConstraintRaw {
+    pub constraint_name: String,
+    pub constraint_type: String,
+    pub columns: Vec<String>,
+}
+
 /// Rust type that we should use when reading a column from SQL Server.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Arbitrary)]
 pub enum SqlServerColumnDecodeType {
@@ -1049,7 +1157,6 @@ mod tests {
                 name: name.into(),
                 data_type: data_type.into(),
                 is_nullable: false,
-                primary_key_constraint: None,
                 max_length: 0,
                 precision: 0,
                 scale: 0,
@@ -1147,7 +1254,7 @@ mod tests {
             }),
             columns: sql_server_columns.into(),
         };
-        let sql_server_desc = SqlServerTableDesc::new(sql_server_desc);
+        let sql_server_desc = SqlServerTableDesc::new(sql_server_desc, vec![]).unwrap();
 
         let max_length = Some(VarCharMaxLength::try_from(16).unwrap());
         let relation_desc = RelationDesc::builder()
@@ -1231,7 +1338,7 @@ mod tests {
                 }),
                 columns: columns.into(),
             };
-            let mut sql_server_desc = SqlServerTableDesc::new(sql_server_desc);
+            let mut sql_server_desc = SqlServerTableDesc::new(sql_server_desc, vec![]).unwrap();
             sql_server_desc.apply_text_columns(&BTreeSet::from(["a"]));
 
             // We should support decoding every datatype to a string.