[Feature] Add support for Primary and Foreign Key Constraints in databricks_sql_table

catalog/resource_sql_table.go

@@ -41,21 +41,73 @@ const IdentityColumnNone IdentityColumn = ""
 const IdentityColumnAlways IdentityColumn = "always"
 const IdentityColumnDefault IdentityColumn = "default"
 
+type SqlKeyConstraintInfo struct {
+	SqlKeyConstraint SqlKeyConstraint
+}
+
+type SqlKeyConstraint interface {
+	getConstraint() string
+}
+
+type SqlPrimaryKeyConstraint struct {
+	PrimaryKey string `json:"primary_key"`
+	Rely       bool   `json:"rely,omitempty" tf:"default:false"`
+}
+
+type SqlForeignKeyConstraint struct {
+	ReferencedKey        string `json:"referenced_key"`
+	ReferencedCatalog    string `json:"referenced_catalog"`
+	ReferencedSchema     string `json:"referenced_schema"`
+	ReferencedTable      string `json:"referenced_table"`
+	ReferencedForeignKey string `json:"referenced_foreign_key"`
+}
+
+func (sqlKeyConstraint SqlPrimaryKeyConstraint) getConstraint() string {
+	var constraint = fmt.Sprintf("PRIMARY KEY (%s)", sqlKeyConstraint.PrimaryKey)
+	if sqlKeyConstraint.Rely {
+		constraint += " RELY"
+	}
+	return constraint
+}
+
+func (sqlKeyConstraint SqlForeignKeyConstraint) getConstraint() string {
+	return fmt.Sprintf(
+		"FOREIGN KEY (%s) REFERENCES %s.%s.%s(%s)",
+		sqlKeyConstraint.ReferencedKey,
+		sqlKeyConstraint.ReferencedCatalog,
+		sqlKeyConstraint.ReferencedSchema,
+		sqlKeyConstraint.ReferencedTable,
+		sqlKeyConstraint.ReferencedForeignKey)
+}
+
+func (ti *SqlTableInfo) serializeSqlKeyConstraintInfo(keyConstraint SqlKeyConstraintInfo) string {
+	return keyConstraint.SqlKeyConstraint.getConstraint()
+}
+
+func (ti *SqlTableInfo) serializeSqlKeyConstraintInfos() string {
+	keyConstraintFragments := make([]string, len(ti.KeyConstraintInfos))
+	for i, keyConstraint := range ti.KeyConstraintInfos {
+		keyConstraintFragments[i] = ti.serializeSqlKeyConstraintInfo(keyConstraint)
+	}
+	return strings.Join(keyConstraintFragments[:], ", ")
+}
+
 type SqlTableInfo struct {
-	Name                  string            `json:"name"`
-	CatalogName           string            `json:"catalog_name" tf:"force_new"`
-	SchemaName            string            `json:"schema_name" tf:"force_new"`
-	TableType             string            `json:"table_type" tf:"force_new"`
-	DataSourceFormat      string            `json:"data_source_format,omitempty" tf:"force_new"`
-	ColumnInfos           []SqlColumnInfo   `json:"columns,omitempty" tf:"alias:column,computed"`
-	Partitions            []string          `json:"partitions,omitempty" tf:"force_new"`
-	ClusterKeys           []string          `json:"cluster_keys,omitempty"`
-	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"`
-	Comment               string            `json:"comment,omitempty"`
-	Properties            map[string]string `json:"properties,omitempty"`
-	Options               map[string]string `json:"options,omitempty" tf:"force_new"`
+	Name                  string                 `json:"name"`
+	CatalogName           string                 `json:"catalog_name" tf:"force_new"`
+	SchemaName            string                 `json:"schema_name" tf:"force_new"`
+	TableType             string                 `json:"table_type" tf:"force_new"`
+	DataSourceFormat      string                 `json:"data_source_format,omitempty" tf:"force_new"`
+	ColumnInfos           []SqlColumnInfo        `json:"columns,omitempty" tf:"alias:column,computed"`
+	KeyConstraintInfos    []SqlKeyConstraintInfo `json:"key_constraints,omitempty" tf:"alias:key_constraint"`
+	Partitions            []string               `json:"partitions,omitempty" tf:"force_new"`
+	ClusterKeys           []string               `json:"cluster_keys,omitempty"`
+	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"`
+	Comment               string                 `json:"comment,omitempty"`
+	Properties            map[string]string      `json:"properties,omitempty"`
+	Options               map[string]string      `json:"options,omitempty" tf:"force_new"`
 	// EffectiveProperties includes both properties and options. Options are prefixed with `option.`.
 	EffectiveProperties map[string]string `json:"effective_properties" tf:"computed"`
 	ClusterID           string            `json:"cluster_id,omitempty" tf:"computed"`
@@ -293,6 +345,10 @@ func (ti *SqlTableInfo) buildTableCreateStatement() string {
 		statements = append(statements, fmt.Sprintf(" (%s)", ti.serializeColumnInfos()))
 	}
 
+	if len(ti.KeyConstraintInfos) > 0 {
+		statements = append(statements, fmt.Sprintf(" (%s)", ti.serializeSqlKeyConstraintInfos()))
+	}
+
 	if !isView {
 		if ti.DataSourceFormat != "" {
 			statements = append(statements, fmt.Sprintf("\nUSING %s", ti.DataSourceFormat)) // USING CSV

catalog/resource_sql_table_test.go

@@ -65,6 +65,83 @@ func TestResourceSqlTableCreateStatement_IdentityColumn(t *testing.T) {
 	assert.Contains(t, stmt, "COMMENT 'terraform managed'")
 }
 
+func TestResourceSqlTableCreateStatement_PrimaryKeyConstraint(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",
+			},
+			{
+				Name:    "name",
+				Comment: "a comment",
+			},
+		},
+		KeyConstraintInfos: []SqlKeyConstraintInfo{
+			{
+				SqlKeyConstraint: SqlPrimaryKeyConstraint{
+					PrimaryKey: "id",
+					Rely:       true,
+				},
+			},
+		},
+	}
+	stmt := ti.buildTableCreateStatement()
+	assert.Contains(t, stmt, "CREATE EXTERNAL TABLE `main`.`foo`.`bar`")
+	assert.Contains(t, stmt, "USING DELTA")
+	assert.Contains(t, stmt, "(`id`  NOT NULL, `name`  NOT NULL COMMENT 'a comment')")
+	assert.Contains(t, stmt, "(PRIMARY KEY (id) RELY)")
+	assert.Contains(t, stmt, "LOCATION 's3://ext-main/foo/bar1' WITH (CREDENTIAL `somecred`)")
+	assert.Contains(t, stmt, "COMMENT 'terraform managed'")
+}
+
+func TestResourceSqlTableCreateStatement_ForeignKeyConstraint(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",
+			},
+			{
+				Name:    "name",
+				Comment: "a comment",
+			},
+		},
+		KeyConstraintInfos: []SqlKeyConstraintInfo{
+			{
+				SqlKeyConstraint: SqlForeignKeyConstraint{
+					ReferencedKey:        "id",
+					ReferencedCatalog:    "bronze",
+					ReferencedSchema:     "biz",
+					ReferencedTable:      "transactions",
+					ReferencedForeignKey: "transactionId",
+				},
+			},
+		},
+	}
+	stmt := ti.buildTableCreateStatement()
+	assert.Contains(t, stmt, "CREATE EXTERNAL TABLE `main`.`foo`.`bar`")
+	assert.Contains(t, stmt, "USING DELTA")
+	assert.Contains(t, stmt, "(`id`  NOT NULL, `name`  NOT NULL COMMENT 'a comment')")
+	assert.Contains(t, stmt, "(FOREIGN KEY (id) REFERENCES bronze.biz.transactions(transactionId)")
+	assert.Contains(t, stmt, "LOCATION 's3://ext-main/foo/bar1' WITH (CREDENTIAL `somecred`)")
+	assert.Contains(t, stmt, "COMMENT 'terraform managed'")
+}
+
 func TestResourceSqlTableCreateStatement_View(t *testing.T) {
 	ti := &SqlTableInfo{
 		Name:                  "bar",

docs/resources/sql_table.md

@@ -149,6 +149,57 @@ resource "databricks_sql_table" "thing" {
 }
 ```
 
+## Use Primary and Foreign Keys
+
+```hcl
+resource "databricks_catalog" "sandbox" {
+  name    = "sandbox"
+  comment = "this catalog is managed by terraform"
+  properties = {
+    purpose = "testing"
+  }
+}
+resource "databricks_schema" "things" {
+  catalog_name = databricks_catalog.sandbox.id
+  name         = "things"
+  comment      = "this database is managed by terraform"
+  properties = {
+    kind = "various"
+  }
+}
+resource "databricks_sql_table" "thing" {
+  provider           = databricks.workspace
+  name               = "quickstart_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     = "bigint"
+  }
+  column {
+    name    = "name"
+    type    = "string"
+    comment = "name of thing"
+  }
+  key_constraint {
+    primary_key = "id"
+    rely = "true"
+  }
+  key_constraint {
+    referenced_key = "external_id"
+    referenced_catalog = "bronze"
+    referenced_schema = "biz"
+    referenced_table = "transactions"
+    referenced_foreign_key = "transactionId"
+  }
+  comment = "this table is managed by terraform"
+}
+```
+
+
 ## Argument Reference
 
 The following arguments are supported:
@@ -163,6 +214,7 @@ The following arguments are supported:
 * `cluster_id` - (Optional) All table CRUD operations must be executed on a running cluster or SQL warehouse. If a cluster_id is specified, it will be used to execute SQL commands to manage this table. If empty, a cluster will be created automatically with the name `terraform-sql-table`.
 * `warehouse_id` - (Optional) All table CRUD operations must be executed on a running cluster or SQL warehouse. If a `warehouse_id` is specified, that SQL warehouse will be used to execute SQL commands to manage this table. Conflicts with `cluster_id`.
 * `cluster_keys` - (Optional) a subset of columns to liquid cluster the table by. Conflicts with `partitions`.
+* `key_constraint` - (Optional) Constraints for Primary and Foreign Keys.
 * `storage_credential_name` - (Optional) For EXTERNAL Tables only: the name of storage credential to use. Change forces creation of a new resource.
 * `owner` - (Optional) Username/groupname/sp application_id of the schema owner.
 * `comment` - (Optional) User-supplied free-form text. Changing comment is not currently supported on `VIEW` table_type.
@@ -181,6 +233,19 @@ Currently, changing the column definitions for a table will require dropping and
 * `comment` - (Optional) User-supplied free-form text.
 * `nullable` - (Optional) Whether field is nullable (Default: `true`)
 
+### `key_constraint` configuration block
+
+For Primary Keys
+* `primary_key` - Column that will get the constraint.
+* `rely` - (Optional) Whether utilizing rely optimization. Can be `true` or `false`. Default to `false`.
+
+For Foreign Keys
+* `referenced_key` - Column that will get the constraint.
+* `referenced_catalog` - Catalog name of the remote table.
+* `referenced_schema` - Schema name of the remote table.
+* `referenced_table` - Remote table name.
+* `referenced_foreign_key` - Column name of the foreign key in the remote table that will get referenced by the Foreign Key.
+
 ## Attribute Reference
 
 In addition to all arguments above, the following attributes are exported:

internal/acceptance/sql_table_test.go

@@ -136,6 +136,98 @@ func TestUcAccResourceSqlTableWithIdentityColumn_Managed(t *testing.T) {
 	})
 }
 
+func TestUcAccResourceSqlTableWithPrimaryAndForeignKeyConstraints_Managed(t *testing.T) {
+	if os.Getenv("GOOGLE_CREDENTIALS") != "" {
+		skipf(t)("databricks_sql_table resource not available on GCP")
+	}
+	UnityWorkspaceLevel(t, Step{
+		Template: `
+		resource "databricks_schema" "this" {
+			name         = "{var.STICKY_RANDOM}"
+			catalog_name = "main"
+		}
+
+		resource "databricks_sql_table" "this" {
+			name               = "bar"
+			catalog_name       = "main"
+			schema_name        = databricks_schema.this.name
+			table_type         = "MANAGED"
+			properties         = {
+				this      = "that"
+				something = "else"
+			}
+
+			column {
+				name      = "id"
+				type      = "bigint"
+			}
+			column {
+				name      = "name"
+				type      = "string"
+			}
+			column {
+				name      = "externalId"
+				type      = "string"
+			}
+			key_constraint {
+				primary_key = "id"
+				rely = "true"
+			}
+			key_constraint {
+				referenced_key = "external_id"
+				referenced_catalog = "bronze"
+				referenced_schema = "biz"
+				referenced_table = "transactions"
+				referenced_foreign_key = "transactionId"
+			}
+			comment = "this table is managed by terraform"
+			owner = "account users"
+		}`,
+	}, Step{
+		Template: `
+		resource "databricks_schema" "this" {
+			name         = "{var.STICKY_RANDOM}"
+			catalog_name = "main"
+		}
+
+		resource "databricks_sql_table" "this" {
+			name               = "bar"
+			catalog_name       = "main"
+			schema_name        = databricks_schema.this.name
+			table_type         = "MANAGED"
+			properties         = {
+				that      = "this"
+				something = "else2"
+			}
+
+			column {
+				name      = "id"
+				type      = "bigint"
+			}
+			column {
+				name      = "name"
+				type      = "string"
+			}
+			column {
+				name      = "externalId"
+				type      = "string"
+			}
+			key_constraint {
+				primary_key = "id"
+				rely = "true"
+			}
+			key_constraint {
+				referenced_key = "external_id"
+				referenced_catalog = "bronze"
+				referenced_schema = "biz"
+				referenced_table = "transactions"
+				referenced_foreign_key = "transactionId"
+			}
+			comment = "this table is managed by terraform..."
+		}`,
+	})
+}
+
 func TestUcAccResourceSqlTable_External(t *testing.T) {
 	UnityWorkspaceLevel(t, Step{
 		Template: `