Skip to content

Commit

Permalink
[Feature] Add support for Identity Column in databricks_sql_table (#…
Browse files Browse the repository at this point in the history 
…4035)

## Changes
Add support for Identity Column in `databricks_sql_table`

## Tests

- [X] `make test` run locally
- [X] relevant change in `docs/` folder
- [X] covered with integration tests in `internal/acceptance`
- [X] relevant acceptance tests are passing
- [X] using Go SDK

---------

Co-authored-by: Miles Yucht <miles@databricks.com>
  • Loading branch information
@hectorcast-db @mgyucht
hectorcast-db and mgyucht authored Sep 30, 2024
1 parent ae06c79 commit 1cfc531
Show file tree
Hide file tree
Showing 4 changed files with 363 additions and 5 deletions.
70 changes: 65 additions & 5 deletions catalog/resource_sql_table.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package catalog

import (
"context"
"encoding/json"
"fmt"
"log"
"reflect"
Expand All @@ -22,12 +23,24 @@ import (
var MaxSqlExecWaitTimeout = 50

type SqlColumnInfo struct {
Name string `json:"name"`
Type string `json:"type_text,omitempty" tf:"alias:type,computed"`
Comment string `json:"comment,omitempty"`
Nullable bool `json:"nullable,omitempty" tf:"default:true"`
Name string `json:"name"`
Type string `json:"type_text,omitempty" tf:"alias:type,computed"`
Identity IdentityColumn `json:"identity,omitempty"`
Comment string `json:"comment,omitempty"`
Nullable bool `json:"nullable,omitempty" tf:"default:true"`
TypeJson string `json:"type_json,omitempty" tf:"computed"`
}

type TypeJson struct {
Metadata map[string]any `json:"metadata,omitempty"`
}

type IdentityColumn string

const IdentityColumnNone IdentityColumn = ""
const IdentityColumnAlways IdentityColumn = "always"
const IdentityColumnDefault IdentityColumn = "default"

type SqlTableInfo struct {
Name string `json:"name"`
CatalogName string `json:"catalog_name" tf:"force_new"`
Expand Down Expand Up @@ -108,6 +121,28 @@ func parseComment(s string) string {
return strings.ReplaceAll(strings.ReplaceAll(s, `\'`, `'`), `'`, `\'`)
}

func reconstructIdentity(c *SqlColumnInfo) (IdentityColumn, error) {
if c.TypeJson == "" {
return IdentityColumnNone, nil
}
var typeJson TypeJson
err := json.Unmarshal([]byte(c.TypeJson), &typeJson)
if err != nil {
return IdentityColumnNone, err
}
if _, ok := typeJson.Metadata["delta.identity.start"]; !ok {
return IdentityColumnNone, nil
}
explicit, ok := typeJson.Metadata["delta.identity.allowExplicitInsert"]
if !ok {
return IdentityColumnNone, nil
}
if explicit.(bool) {
return IdentityColumnDefault, nil
}
return IdentityColumnAlways, nil
}

func (ti *SqlTableInfo) initCluster(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) (err error) {
defaultClusterName := "terraform-sql-table"
clustersAPI := clusters.NewClustersAPI(ctx, c)
Expand Down Expand Up @@ -171,7 +206,22 @@ func (ti *SqlTableInfo) getOrCreateCluster(clusterName string, clustersAPI clust
return aclCluster.ClusterID, nil
}

func (ci *SqlColumnInfo) getColumnType() string {
var colType string
switch ci.Identity {
case IdentityColumnAlways:
colType = fmt.Sprintf("%s GENERATED ALWAYS AS IDENTITY", ci.Type)
case IdentityColumnDefault:
colType = fmt.Sprintf("%s GENERATED BY DEFAULT AS IDENTITY", ci.Type)
default:
colType = ci.Type
}
return colType
}

func (ti *SqlTableInfo) serializeColumnInfo(col SqlColumnInfo) string {
var colType = col.getColumnType()

notNull := ""
if !col.Nullable {
notNull = " NOT NULL"
Expand All @@ -181,7 +231,7 @@ func (ti *SqlTableInfo) serializeColumnInfo(col SqlColumnInfo) string {
if col.Comment != "" {
comment = fmt.Sprintf(" COMMENT '%s'", parseComment(col.Comment))
}
return fmt.Sprintf("%s %s%s%s", col.getWrappedColumnName(), col.Type, notNull, comment) // id INT NOT NULL COMMENT 'something'
return fmt.Sprintf("%s %s%s%s", col.getWrappedColumnName(), colType, notNull, comment) // id INT NOT NULL COMMENT 'something'
}

func (ti *SqlTableInfo) serializeColumnInfos() string {
Expand Down Expand Up @@ -502,6 +552,9 @@ func assertNoColumnTypeDiff(oldCols []interface{}, newColumnInfos []SqlColumnInf
if getColumnType(oldColMap["type"].(string)) != getColumnType(newColumnInfos[i].Type) {
return fmt.Errorf("changing the 'type' of an existing column is not supported")
}
if oldColMap["identity"].(string) != string(newColumnInfos[i].Identity) {
return fmt.Errorf("changing the 'identity' type of an existing column is not supported")
}
}
return nil
}
Expand Down Expand Up @@ -602,6 +655,13 @@ func ResourceSqlTable() common.Resource {
if err != nil {
return err
}
for i := range ti.ColumnInfos {
c := &ti.ColumnInfos[i]
c.Identity, err = reconstructIdentity(c)
if err != nil {
return err
}
}
return common.StructToData(ti, tableSchema, d)
},
Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
Expand Down
193 changes: 193 additions & 0 deletions catalog/resource_sql_table_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,36 @@ func TestResourceSqlTableCreateStatement_External(t *testing.T) {
assert.Contains(t, stmt, "COMMENT 'terraform managed'")
}

func TestResourceSqlTableCreateStatement_IdentityColumn(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: "bigint",
Identity: "default",
},
{
Name: "name",
Comment: "a comment",
},
},
}
stmt := ti.buildTableCreateStatement()
assert.Contains(t, stmt, "CREATE EXTERNAL TABLE `main`.`foo`.`bar`")
assert.Contains(t, stmt, "USING DELTA")
assert.Contains(t, stmt, "(`id` bigint GENERATED BY DEFAULT AS IDENTITY NOT NULL, `name` NOT NULL COMMENT 'a comment')")
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",
Expand Down Expand Up @@ -1334,6 +1364,169 @@ func TestResourceSqlTableCreateTable_ExistingSQLWarehouse(t *testing.T) {
assert.NoError(t, err)
}

func TestResourceSqlTableCreateTableWithIdentityColumn_ExistingSQLWarehouse(t *testing.T) {
qa.ResourceFixture{
CommandMock: func(commandStr string) common.CommandResults {
return common.CommandResults{
ResultType: "",
Data: nil,
}
},
HCL: `
name = "bar"
catalog_name = "main"
schema_name = "foo"
table_type = "MANAGED"
data_source_format = "DELTA"
storage_location = "abfss://container@account/somepath"
warehouse_id = "existingwarehouse"
column {
name = "id"
type = "bigint"
identity = "default"
}
column {
name = "name"
type = "string"
comment = "name of thing"
}
column {
name = "number"
type = "bigint"
identity = "always"
}
comment = "this table is managed by terraform"
`,
Fixtures: []qa.HTTPFixture{
{
Method: "POST",
Resource: "/api/2.0/sql/statements/",
ExpectedRequest: sql.ExecuteStatementRequest{
Statement: "CREATE TABLE `main`.`foo`.`bar` (`id` bigint GENERATED BY DEFAULT AS IDENTITY, `name` string COMMENT 'name of thing', `number` bigint GENERATED ALWAYS AS IDENTITY)\nUSING DELTA\nCOMMENT 'this table is managed by terraform'\nLOCATION 'abfss://container@account/somepath';",
WaitTimeout: "50s",
WarehouseId: "existingwarehouse",
OnWaitTimeout: sql.ExecuteStatementRequestOnWaitTimeoutCancel,
},
Response: sql.StatementResponse{
StatementId: "statement1",
Status: &sql.StatementStatus{
State: "SUCCEEDED",
},
},
},
{
Method: "GET",
Resource: "/api/2.1/unity-catalog/tables/main.foo.bar",
Response: SqlTableInfo{
Name: "bar",
CatalogName: "main",
SchemaName: "foo",
TableType: "EXTERNAL",
DataSourceFormat: "DELTA",
StorageLocation: "s3://ext-main/foo/bar1",
StorageCredentialName: "somecred",
Comment: "terraform managed",
Properties: map[string]string{
"one": "two",
"three": "four",
},
ColumnInfos: []SqlColumnInfo{
{
Name: "id",
Type: "bigint",
TypeJson: "{\"type\":\"bigint\",\"nullable\":true, \"metadata\":{\"delta.identity.start\":1,\"delta.identity.allowExplicitInsert\":true}}",
},
{
Name: "name",
Type: "string",
Comment: "name of thing",
},
{
Name: "number",
Type: "bigint",
TypeJson: "{\"type\":\"bigint\",\"nullable\":true, \"metadata\":{\"delta.identity.start\":1,\"delta.identity.allowExplicitInsert\":false}}",
},
},
},
},
},
Create: true,
Resource: ResourceSqlTable(),
}.ApplyAndExpectData(t, map[string]any{
"column.0.identity": "default",
"column.1.identity": "",
"column.2.identity": "always",
})
}

func TestResourceSqlTableReadTableWithIdentityColumn_ExistingSQLWarehouse(t *testing.T) {
qa.ResourceFixture{
CommandMock: func(commandStr string) common.CommandResults {
return common.CommandResults{
ResultType: "",
Data: nil,
}
},
HCL: `
name = "bar"
catalog_name = "main"
schema_name = "foo"
table_type = "MANAGED"
data_source_format = "DELTA"
storage_location = "abfss://container@account/somepath"
warehouse_id = "existingwarehouse"
comment = "this table is managed by terraform"
`,
Fixtures: []qa.HTTPFixture{
{
Method: "GET",
Resource: "/api/2.1/unity-catalog/tables/main.foo.bar",
Response: SqlTableInfo{
Name: "bar",
CatalogName: "main",
SchemaName: "foo",
TableType: "EXTERNAL",
DataSourceFormat: "DELTA",
StorageLocation: "s3://ext-main/foo/bar1",
StorageCredentialName: "somecred",
Comment: "terraform managed",
Properties: map[string]string{
"one": "two",
"three": "four",
},
ColumnInfos: []SqlColumnInfo{
{
Name: "id",
Type: "bigint",
TypeJson: "{\"type\":\"bigint\",\"nullable\":true, \"metadata\":{\"delta.identity.start\":1,\"delta.identity.allowExplicitInsert\":false}}",
},
{
Name: "name",
Type: "string",
Comment: "name of thing",
},
{
Name: "number",
Type: "bigint",
TypeJson: "{\"type\":\"bigint\",\"nullable\":true, \"metadata\":{\"delta.identity.start\":1,\"delta.identity.allowExplicitInsert\":true}}",
},
},
},
},
},
ID: "main.foo.bar",
Read: true,
Resource: ResourceSqlTable(),
}.ApplyAndExpectData(t, map[string]any{
"column.0.identity": "always",
"column.1.identity": "",
"column.2.identity": "default",
})
}

func TestResourceSqlTableCreateTable_OnlyManagedProperties(t *testing.T) {
qa.ResourceFixture{
CommandMock: func(commandStr string) common.CommandResults {
Expand Down
41 changes: 41 additions & 0 deletions docs/resources/sql_table.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,46 @@ resource "databricks_sql_table" "thing_view" {
}
```

## Use an Identity Column

```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"
identity = "default"
}
column {
name = "name"
type = "string"
comment = "name of thing"
}
comment = "this table is managed by terraform"
}
```

## Argument Reference

The following arguments are supported:
Expand Down Expand Up @@ -137,6 +177,7 @@ Currently, changing the column definitions for a table will require dropping and

* `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`)

Expand Down
Loading

0 comments on commit 1cfc531

Please sign in to comment.