[Feature] Add support for constraints in databricks_sql_table resource

catalog/resource_sql_table.go

@@ -32,6 +32,15 @@ type SqlColumnInfo struct {
 	TypeJson string         `json:"type_json,omitempty" tf:"computed"`
 }
 
+type ConstraintInfo struct {
+	Name          string   `json:"name"`
+	Type          string   `json:"type"`
+	KeyColumns    []string `json:"key_columns"`
+	ParentTable   string   `json:"parent_table,omitempty"`
+	ParentColumns []string `json:"parent_columns,omitempty"`
+	Rely          bool     `json:"rely,omitempty" tf:"default:false"`
+}
+
 type TypeJson struct {
 	Metadata map[string]any `json:"metadata,omitempty"`
 }
@@ -51,6 +60,7 @@ type SqlTableInfo struct {
 	ColumnInfos           []SqlColumnInfo   `json:"columns,omitempty" tf:"alias:column,computed"`
 	Partitions            []string          `json:"partitions,omitempty" tf:"force_new"`
 	ClusterKeys           []string          `json:"cluster_keys,omitempty"`
+	Constraints           []ConstraintInfo  `json:"constraints,omitempty" tf:"alias:constraint"`
 	StorageLocation       string            `json:"storage_location,omitempty" tf:"suppress_diff"`
 	StorageCredentialName string            `json:"storage_credential_name,omitempty" tf:"force_new"`
 	ViewDefinition        string            `json:"view_definition,omitempty"`
@@ -89,6 +99,9 @@ func (ti SqlTableInfo) CustomizeSchema(s *common.CustomizableSchema) *common.Cus
 	s.SchemaPath("column", "type").SetCustomSuppressDiff(func(k, old, new string, d *schema.ResourceData) bool {
 		return getColumnType(old) == getColumnType(new)
 	})
+	s.SchemaPath("constraint", "type").SetCustomSuppressDiff(func(k, old, new string, d *schema.ResourceData) bool {
+		return getColumnType(old) == getColumnType(new)
+	})
 	return s
 }
 
@@ -242,6 +255,39 @@ func (ti *SqlTableInfo) serializeColumnInfos() string {
 	return strings.Join(columnFragments[:], ", ") // id INT NOT NULL, name STRING, age INT
 }
 
+func serializePrimaryKeyConstraint(constraint ConstraintInfo) string {
+	constraint_clause := fmt.Sprintf("CONSTRAINT %s PRIMARY KEY(%s)", constraint.getWrappedConstraintName(), constraint.getWrappedKeyColumnNames())
+	if constraint.Rely {
+		constraint_clause += " RELY"
+	}
+	return constraint_clause
+}
+
+func serializeForeignKeyConstraint(constraint ConstraintInfo) string {
+	constraint_clause := fmt.Sprintf("CONSTRAINT %s FOREIGN KEY(%s) REFERENCES %s", constraint.getWrappedConstraintName(), constraint.getWrappedKeyColumnNames(), constraint.ParentTable)
+	if len(constraint.ParentColumns) > 0 {
+		constraint_clause += fmt.Sprintf("(%s)", constraint.getWrappedParentColumnNames())
+	}
+	if constraint.Rely {
+		constraint_clause += " RELY"
+	}
+	return constraint_clause
+}
+
+func (ti *SqlTableInfo) serializeConstraints() string {
+	constraintFragments := make([]string, len(ti.Constraints))
+
+	for i, constraint := range ti.Constraints {
+		if constraint.Type == "PRIMARY KEY" {
+			constraintFragments[i] = serializePrimaryKeyConstraint(constraint)
+		} else if constraint.Type == "FOREIGN KEY" {
+			constraintFragments[i] = serializeForeignKeyConstraint(constraint)
+		}
+	}
+
+	return strings.Join(constraintFragments[:], ", ") // CONSTRAINT `pk`` PRIMARY KEY (`id`, `nickname`), CONSTRAINT `fk` FOREIGN KEY (`player_id`) REFERENCES players
+}
+
 func (ti *SqlTableInfo) serializeProperties() string {
 	propsMap := make([]string, 0, len(ti.Properties))
 	for key, value := range ti.Properties {
@@ -290,7 +336,11 @@ func (ti *SqlTableInfo) buildTableCreateStatement() string {
 	statements = append(statements, fmt.Sprintf("CREATE %s%s %s", externalFragment, createType, ti.SQLFullName()))
 
 	if len(ti.ColumnInfos) > 0 {
-		statements = append(statements, fmt.Sprintf(" (%s)", ti.serializeColumnInfos()))
+		columnInfosClause := ti.serializeColumnInfos()
+		if len(ti.Constraints) > 0 {
+			columnInfosClause += fmt.Sprintf(", %s", ti.serializeConstraints())
+		}
+		statements = append(statements, fmt.Sprintf(" (%s)", columnInfosClause))
 	}
 
 	if !isView {
@@ -342,6 +392,21 @@ func (ti *SqlTableInfo) getWrappedClusterKeys() string {
 	return "`" + strings.Join(ti.ClusterKeys, "`,`") + "`"
 }
 
+// Wrapping the constraint name with backticks to avoid special character messing things up.
+func (ci ConstraintInfo) getWrappedConstraintName() string {
+	return fmt.Sprintf("`%s`", ci.Name)
+}
+
+// Wrapping constraint column names with backticks to avoid special character messing things up.
+func (ci ConstraintInfo) getWrappedKeyColumnNames() string {
+	return "`" + strings.Join(ci.KeyColumns, "`,`") + "`"
+}
+
+// Wrapping parent column name with backticks to avoid special character messing things up.
+func (ci ConstraintInfo) getWrappedParentColumnNames() string {
+	return "`" + strings.Join(ci.ParentColumns, "`,`") + "`"
+}
+
 func (ti *SqlTableInfo) getStatementsForColumnDiffs(oldti *SqlTableInfo, statements []string, typestring string) []string {
 	if len(ti.ColumnInfos) != len(oldti.ColumnInfos) {
 		statements = ti.addOrRemoveColumnStatements(oldti, statements, typestring)
@@ -413,6 +478,43 @@ func (ti *SqlTableInfo) alterExistingColumnStatements(oldti *SqlTableInfo, state
 	return statements
 }
 
+func (ti *SqlTableInfo) addOrRemoveConstraintStatements(oldti *SqlTableInfo, statements []string, typestring string) []string {
+	nameToOldConstraint := make(map[string]ConstraintInfo)
+	nameToNewConstraint := make(map[string]ConstraintInfo)
+	for _, ci := range oldti.Constraints {
+		nameToOldConstraint[ci.Name] = ci
+	}
+	for _, newCi := range ti.Constraints {
+		nameToNewConstraint[newCi.Name] = newCi
+	}
+
+	removeConstraintStatements := make([]string, 0)
+
+	for name, oldCi := range nameToOldConstraint {
+		if _, exists := nameToNewConstraint[name]; !exists {
+			// Remove old constraint if old constraint is no longer found in the config.
+			removeConstraintStatements = append(removeConstraintStatements, oldCi.getWrappedConstraintName())
+		}
+	}
+	for _, removeStatement := range removeConstraintStatements {
+		statements = append(statements, fmt.Sprintf("ALTER %s %s DROP CONSTRAINT IF EXISTS %s", typestring, ti.SQLFullName(), removeStatement))
+	}
+
+	for _, newCi := range ti.Constraints {
+		if _, exists := nameToOldConstraint[newCi.Name]; !exists {
+			var newConstraintStatement string
+			if newCi.Type == "PRIMARY KEY" {
+				newConstraintStatement = serializePrimaryKeyConstraint(newCi)
+			} else if newCi.Type == "FOREIGN KEY" {
+				newConstraintStatement = serializeForeignKeyConstraint(newCi)
+			}
+			statements = append(statements, fmt.Sprintf("ALTER %s %s ADD %s", typestring, ti.SQLFullName(), newConstraintStatement))
+		}
+	}
+
+	return statements
+}
+
 func (ti *SqlTableInfo) diff(oldti *SqlTableInfo) ([]string, error) {
 	statements := make([]string, 0)
 	typestring := ti.getTableTypeString()
@@ -454,6 +556,7 @@ func (ti *SqlTableInfo) diff(oldti *SqlTableInfo) ([]string, error) {
 	}
 
 	statements = ti.getStatementsForColumnDiffs(oldti, statements, typestring)
+	statements = ti.addOrRemoveConstraintStatements(oldti, statements, typestring)
 
 	return statements, nil
 }

catalog/resource_sql_table_test.go

@@ -141,6 +141,87 @@ func TestResourceSqlTableCreateStatement_Partition(t *testing.T) {
 	assert.Contains(t, stmt, "PARTITIONED BY (baz, bazz)")
 }
 
+func TestResourceSqlTableCreateStatement_Constraints(t *testing.T) {
+	ti := &SqlTableInfo{
+		Name:                  "bar",
+		CatalogName:           "main",
+		SchemaName:            "foo",
+		TableType:             "EXTERNAL",
+		DataSourceFormat:      "DELTA",
+		StorageLocation:       "s3://ext-main/foo/bar1",
+		StorageCredentialName: "somecred",
+		Comment:               "terraform managed",
+		ColumnInfos: []SqlColumnInfo{
+			{
+				Name: "id",
+				Type: "int",
+			},
+			{
+				Name: "external_id",
+				Type: "int",
+			},
+			{
+				Name: "external_name",
+				Type: "string",
+			},
+		},
+		Constraints: []ConstraintInfo{
+			{
+				Name:       "pk",
+				Type:       "PRIMARY KEY",
+				KeyColumns: []string{"id"},
+				Rely:       true,
+			},
+			{
+				Name:        "fk",
+				Type:        "FOREIGN KEY",
+				KeyColumns:  []string{"external_id", "external_name"},
+				ParentTable: "some_table",
+			},
+		},
+	}
+	stmt := ti.buildTableCreateStatement()
+	assert.Contains(t, stmt, "CREATE EXTERNAL TABLE `main`.`foo`.`bar`")
+	assert.Contains(t, stmt, "USING DELTA")
+	assert.Contains(t, stmt, "(`id` int NOT NULL, `external_id` int NOT NULL, `external_name` string NOT NULL")
+	assert.Contains(t, stmt, "CONSTRAINT `pk` PRIMARY KEY(`id`) RELY")
+	assert.Contains(t, stmt, "CONSTRAINT `fk` FOREIGN KEY(`external_id`,`external_name`) REFERENCES some_table)")
+	assert.Contains(t, stmt, "LOCATION 's3://ext-main/foo/bar1' WITH (CREDENTIAL `somecred`)")
+	assert.Contains(t, stmt, "COMMENT 'terraform managed'")
+}
+
+func TestResourceSqlTableCreateStatement_ConstraintsWithoutColumns(t *testing.T) {
+	ti := &SqlTableInfo{
+		Name:                  "bar",
+		CatalogName:           "main",
+		SchemaName:            "foo",
+		TableType:             "EXTERNAL",
+		DataSourceFormat:      "DELTA",
+		StorageLocation:       "s3://ext-main/foo/bar1",
+		StorageCredentialName: "somecred",
+		Comment:               "terraform managed",
+		Constraints: []ConstraintInfo{
+			{
+				Name:       "pk",
+				Type:       "PRIMARY KEY",
+				KeyColumns: []string{"id"},
+			},
+			{
+				Name:        "fk",
+				Type:        "FOREIGN KEY",
+				KeyColumns:  []string{"external_id", "external_name"},
+				ParentTable: "some_table",
+			},
+		},
+	}
+	stmt := ti.buildTableCreateStatement()
+	assert.Contains(t, stmt, "CREATE EXTERNAL TABLE `main`.`foo`.`bar`")
+	assert.Contains(t, stmt, "USING DELTA")
+	assert.NotContains(t, stmt, "CONSTRAINT")
+	assert.Contains(t, stmt, "LOCATION 's3://ext-main/foo/bar1' WITH (CREDENTIAL `somecred`)")
+	assert.Contains(t, stmt, "COMMENT 'terraform managed'")
+}
+
 func TestResourceSqlTableCreateStatement_Liquid(t *testing.T) {
 	ti := &SqlTableInfo{
 		Name:                  "bar",
@@ -198,7 +279,7 @@ func TestResourceSqlTableCreateTable(t *testing.T) {
 		table_type         = "MANAGED"
 		data_source_format = "DELTA"
 		storage_location   = "abfss:container@account/somepath"
-	  
+
 		column {
 		  name      = "id"
 		  type      = "int"
@@ -251,7 +332,7 @@ func TestResourceSqlTableCreateTableWithOwner(t *testing.T) {
 		table_type         = "MANAGED"
 		data_source_format = "DELTA"
 		storage_location   = "abfss:container@account/somepath"
-	  
+
 		column {
 		  name      = "id"
 		  type      = "int"
@@ -326,7 +407,7 @@ func TestResourceSqlTableCreateTable_Error(t *testing.T) {
 		table_type         = "MANAGED"
 		data_source_format = "DELTA"
 		storage_location   = "abfss:container@account/somepath"
-	  
+
 		column {
 		  name      = "id"
 		  type      = "int"
@@ -1310,7 +1391,7 @@ func TestResourceSqlTableCreateTable_ExistingSQLWarehouse(t *testing.T) {
 		data_source_format = "DELTA"
 		storage_location   = "abfss://container@account/somepath"
 		warehouse_id       = "existingwarehouse"
-	  
+
 		column {
 		  name      = "id"
 		  type      = "int"
@@ -1380,7 +1461,7 @@ func TestResourceSqlTableCreateTableWithIdentityColumn_ExistingSQLWarehouse(t *t
 		data_source_format = "DELTA"
 		storage_location   = "abfss://container@account/somepath"
 		warehouse_id       = "existingwarehouse"
-	  
+
 		column {
 		  name      	= "id"
 		  type 			= "bigint"
@@ -1476,8 +1557,8 @@ func TestResourceSqlTableReadTableWithIdentityColumn_ExistingSQLWarehouse(t *tes
 		data_source_format = "DELTA"
 		storage_location   = "abfss://container@account/somepath"
 		warehouse_id       = "existingwarehouse"
-	  
-	
+
+
 		comment = "this table is managed by terraform"
 		`,
 		Fixtures: []qa.HTTPFixture{
@@ -1542,7 +1623,7 @@ func TestResourceSqlTableCreateTable_OnlyManagedProperties(t *testing.T) {
 		table_type         = "MANAGED"
 		data_source_format = "DELTA"
 		warehouse_id       = "existingwarehouse"
-	  
+
 		column {
 		  name      = "id"
 		  type      = "int"

docs/resources/sql_table.md

@@ -149,6 +149,35 @@ resource "databricks_sql_table" "thing" {
 }
 ```
 
+## Use Constraints
+
+```hcl
+resource "databricks_sql_table" "table_with_constraints" {
+  provider           = databricks.workspace
+  name               = "constraints_table"
+  catalog_name       = databricks_catalog.sandbox.name
+  schema_name        = databricks_schema.things.name
+  table_type         = "MANAGED"
+  data_source_format = "DELTA"
+  storage_location   = ""
+  column {
+    name     = "id"
+    type     = "int"
+  }
+  column {
+    name    = "name"
+    type    = "string"
+    comment = "name of thing"
+  }
+  constraint {
+    name        = "pk"
+    type        = "PRIMARY KEY"
+    key_columns = ["id"]
+  }
+  comment = "this table contains PRIMARY KEY constraint on `id` column"
+}
+```
+
 ## Argument Reference
 
 The following arguments are supported:
@@ -173,14 +202,25 @@ The following arguments are supported:
 ### `column` configuration block
 
 For table columns
-Currently, changing the column definitions for a table will require dropping and re-creating the table
 
 * `name` - User-visible name of column
 * `type` - Column type spec (with metadata) as SQL text. Not supported for `VIEW` table_type.
 * `identity` - (Optional) Whether field is an identity column. Can be `default`, `always` or unset. It is unset by default.
 * `comment` - (Optional) User-supplied free-form text.
 * `nullable` - (Optional) Whether field is nullable (Default: `true`)
 
+### `constraint` configuration block
+
+For table constraints
+Only `PRIMARY KEY` and `FOREIGN KEY` constraints are supported right now.
+
+* `name` - User-visible name of the constraint
+* `type` - Type of the constraint. Supported values: `PRIMARY KEY`, `FOREIGN KEY`
+* `key_columns` - List of the columns in the defined table that should be used for the constraint.
+* `parent_table` - (Optional) Required for `FOREIGN KEY` constraint. Full name of the parent table.
+* `parent_columns` - (Optional) Optional for `FOREIGN KEY` constraint. Should be used if column names in the parent table differ.
+* `rely` - (Optional) Set to `true` if `RELY` option should be used for the constraint
+
 ## Attribute Reference
 
 In addition to all arguments above, the following attributes are exported: