From 4da801445d0523ce287c00600d1c1fd3f5af330f Mon Sep 17 00:00:00 2001 From: Ben Rosen <34146575+berosen@users.noreply.github.com> Date: Thu, 21 Oct 2021 17:26:45 -0600 Subject: [PATCH] feat: Identity Column Support (#726) Adds support for a column identity/autoincrement for snowflake tables. Only supports adding an identity to a column for the follow operations * `CREATE TABLE` * `ALTER TABLE ADD COLUMN` ## Test Plan * [x] acceptance tests * [x] unit tests ## References * https://docs.snowflake.com/en/sql-reference/sql/create-table.html#optional-parameters (see autoincrement/identity section) * #538 --- docs/resources/table.md | 21 +++ .../resources/snowflake_table/resource.tf | 11 ++ pkg/resources/table.go | 69 +++++++- pkg/resources/table_acceptance_test.go | 156 ++++++++++++++++++ pkg/snowflake/table.go | 63 ++++++- pkg/snowflake/table_test.go | 41 ++++- 6 files changed, 353 insertions(+), 8 deletions(-) diff --git a/docs/resources/table.md b/docs/resources/table.md index c72e5518a2..d018898c8c 100644 --- a/docs/resources/table.md +++ b/docs/resources/table.md @@ -44,6 +44,17 @@ resource "snowflake_table" "table" { } } + column { + name = "identity" + type = "NUMBER(38,0)" + nullable = true + + identity { + start_num = 1 + step_num = 3 + } + } + column { name = "data" type = "text" @@ -104,6 +115,7 @@ Optional: - **comment** (String) Column comment - **default** (Block List, Max: 1) Defines the column default value; note due to limitations of Snowflake's ALTER TABLE ADD/MODIFY COLUMN updates to default will not be applied (see [below for nested schema](#nestedblock--column--default)) +- **identity** (Block List, Max: 1) Defines the identity start/step values for a column. **Note** Identity/default are mutually exclusive. (see [below for nested schema](#nestedblock--column--identity)) - **nullable** (Boolean) Whether this column can contain null values. **Note**: Depending on your Snowflake version, the default value will not suffice if this column is used in a primary key constraint. @@ -116,6 +128,15 @@ Optional: - **sequence** (String) The default sequence to use for the column + +### Nested Schema for `column.identity` + +Optional: + +- **start_num** (Number) The number to start incrementing at. +- **step_num** (Number) Step size to increment by. + + ### Nested Schema for `primary_key` diff --git a/examples/resources/snowflake_table/resource.tf b/examples/resources/snowflake_table/resource.tf index bdae8d099d..7c58d68c55 100644 --- a/examples/resources/snowflake_table/resource.tf +++ b/examples/resources/snowflake_table/resource.tf @@ -29,6 +29,17 @@ resource "snowflake_table" "table" { } } + column { + name = "identity" + type = "NUMBER(38,0)" + nullable = true + + identity { + start_num = 1 + step_num = 3 + } + } + column { name = "data" type = "text" diff --git a/pkg/resources/table.go b/pkg/resources/table.go index 2756854964..3d9b2fb863 100644 --- a/pkg/resources/table.go +++ b/pkg/resources/table.go @@ -95,6 +95,32 @@ var tableSchema = map[string]*schema.Schema{ }, }, }, + /*Note: Identity and default are mutually exclusive. From what I can tell we can't enforce this here + the snowflake query will error so we can defer enforcement to there. + */ + "identity": { + Type: schema.TypeList, + Optional: true, + Description: "Defines the identity start/step values for a column. **Note** Identity/default are mutually exclusive.", + MinItems: 1, + MaxItems: 1, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "start_num": { + Type: schema.TypeInt, + Optional: true, + Description: "The number to start incrementing at.", + Default: 1, + }, + "step_num": { + Type: schema.TypeInt, + Optional: true, + Description: "Step size to increment by.", + Default: 1, + }, + }, + }, + }, "comment": { Type: schema.TypeString, Optional: true, @@ -251,11 +277,23 @@ func (cd *columnDefault) _type() string { return "unknown" } +type columnIdentity struct { + startNum int + stepNum int +} + +func (identity *columnIdentity) toSnowflakeColumnIdentity() *snowflake.ColumnIdentity { + snowIdentity := snowflake.ColumnIdentity{} + return snowIdentity.WithStartNum(identity.startNum).WithStep(identity.stepNum) + +} + type column struct { name string dataType string nullable bool _default *columnDefault + identity *columnIdentity comment string } @@ -266,6 +304,10 @@ func (c column) toSnowflakeColumn() snowflake.Column { sC = sC.WithDefault(c._default.toSnowflakeColumnDefault()) } + if c.identity != nil { + sC = sC.WithIdentity(c.identity.toSnowflakeColumnIdentity()) + } + return *sC.WithName(c.name). WithType(c.dataType). WithNullable(c.nullable). @@ -349,20 +391,38 @@ func getColumnDefault(def map[string]interface{}) *columnDefault { return nil } +func getColumnIdentity(identity map[string]interface{}) *columnIdentity { + if len(identity) > 0 { + + startNum := identity["start_num"].(int) + stepNum := identity["step_num"].(int) + return &columnIdentity{startNum, stepNum} + } + + return nil +} + func getColumn(from interface{}) (to column) { c := from.(map[string]interface{}) var cd *columnDefault + var id *columnIdentity _default := c["default"].([]interface{}) + identity := c["identity"].([]interface{}) + if len(_default) == 1 { cd = getColumnDefault(_default[0].(map[string]interface{})) } + if len(identity) == 1 { + id = getColumnIdentity(identity[0].(map[string]interface{})) + } return column{ name: c["name"].(string), dataType: c["type"].(string), nullable: c["nullable"].(bool), _default: cd, + identity: id, comment: c["comment"].(string), } } @@ -573,14 +633,17 @@ func UpdateTable(d *schema.ResourceData, meta interface{}) error { } for _, cA := range added { var q string - if cA._default == nil { - q = builder.AddColumn(cA.name, cA.dataType, cA.nullable, nil, cA.comment) + + if cA.identity == nil && cA._default == nil { + q = builder.AddColumn(cA.name, cA.dataType, cA.nullable, nil, nil, cA.comment) + } else if cA.identity != nil { + q = builder.AddColumn(cA.name, cA.dataType, cA.nullable, nil, cA.identity.toSnowflakeColumnIdentity(), cA.comment) } else { if cA._default._type() != "constant" { return fmt.Errorf("Failed to add column %v => Only adding a column as a constant is supported by Snowflake", cA.name) } - q = builder.AddColumn(cA.name, cA.dataType, cA.nullable, cA._default.toSnowflakeColumnDefault(), cA.comment) + q = builder.AddColumn(cA.name, cA.dataType, cA.nullable, cA._default.toSnowflakeColumnDefault(), nil, cA.comment) } err := snowflake.Exec(db, q) diff --git a/pkg/resources/table_acceptance_test.go b/pkg/resources/table_acceptance_test.go index 5a3c331aad..1b39218c80 100644 --- a/pkg/resources/table_acceptance_test.go +++ b/pkg/resources/table_acceptance_test.go @@ -1052,3 +1052,159 @@ resource "snowflake_table" "test_table" { ` return fmt.Sprintf(s, name, tagName, tag2Name) } + +func TestAcc_TableIdentity(t *testing.T) { + accName := strings.ToUpper(acctest.RandStringFromCharSet(10, acctest.CharSetAlpha)) + + resource.ParallelTest(t, resource.TestCase{ + Providers: providers(), + Steps: []resource.TestStep{ + { + Config: tableColumnWithIdentityDefault(accName), + Check: resource.ComposeTestCheckFunc( + resource.TestCheckResourceAttr("snowflake_table.test_table", "name", accName), + resource.TestCheckResourceAttr("snowflake_table.test_table", "database", accName), + resource.TestCheckResourceAttr("snowflake_table.test_table", "schema", accName), + resource.TestCheckResourceAttr("snowflake_table.test_table", "data_retention_days", "1"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "change_tracking", "false"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "comment", "Terraform acceptance test"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.#", "3"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.name", "column1"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.type", "NUMBER(38,0)"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.0.type.default.0.expression"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.0.type.default.0.sequence"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.name", "column2"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.type", "TIMESTAMP_NTZ(9)"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.1.type.default.0.constant"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.1.type.default.0.sequence"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.2.name", "column3"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.2.type", "NUMBER(38,0)"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.2.type.default.0.constant"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.2.type.default.0.expression"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.2.identity.0.start_num", "1"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.2.identity.0.step_num", "1"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "primary_key"), + ), + }, + { + Config: tableColumnWithIdentity(accName), + Check: resource.ComposeTestCheckFunc( + resource.TestCheckResourceAttr("snowflake_table.test_table", "name", accName), + resource.TestCheckResourceAttr("snowflake_table.test_table", "database", accName), + resource.TestCheckResourceAttr("snowflake_table.test_table", "schema", accName), + resource.TestCheckResourceAttr("snowflake_table.test_table", "data_retention_days", "1"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "change_tracking", "false"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "comment", "Terraform acceptance test"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.#", "3"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.name", "column1"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.type", "NUMBER(38,0)"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.0.type.default.0.expression"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.0.type.default.0.sequence"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.name", "column2"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.type", "TIMESTAMP_NTZ(9)"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.1.type.default.0.constant"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.1.type.default.0.sequence"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.2.type.default.0.constant"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.2.type.default.0.expression"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.0.identity.0.start_num"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "column.0.identity.0.step_num"), + // we've dropped the previous identity column and making sure that adding a new column as an identity works + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.2.identity.0.start_num", "2"), + resource.TestCheckResourceAttr("snowflake_table.test_table", "column.2.identity.0.step_num", "4"), + resource.TestCheckNoResourceAttr("snowflake_table.test_table", "primary_key"), + ), + }, + }, + }) +} + +func tableColumnWithIdentityDefault(name string) string { + s := ` +resource "snowflake_database" "test_database" { + name = "%s" + comment = "Terraform acceptance test" +} + +resource "snowflake_schema" "test_schema" { + name = "%s" + database = snowflake_database.test_database.name + comment = "Terraform acceptance test" +} + +resource "snowflake_sequence" "test_seq" { + database = snowflake_database.test_database.name + schema = snowflake_schema.test_schema.name + name = "%s" +} + +resource "snowflake_table" "test_table" { + database = snowflake_database.test_database.name + schema = snowflake_schema.test_schema.name + name = "%s" + comment = "Terraform acceptance test" + + column { + name = "column1" + type = "NUMBER(38,0)" + } + column { + name = "column2" + type = "TIMESTAMP_NTZ(9)" + } + column { + name = "column3" + type = "NUMBER(38,0)" + identity { + } + } +} +` + return fmt.Sprintf(s, name, name, name, name) +} + +func tableColumnWithIdentity(name string) string { + s := ` +resource "snowflake_database" "test_database" { + name = "%s" + comment = "Terraform acceptance test" +} + +resource "snowflake_schema" "test_schema" { + name = "%s" + database = snowflake_database.test_database.name + comment = "Terraform acceptance test" +} + +resource "snowflake_sequence" "test_seq" { + database = snowflake_database.test_database.name + schema = snowflake_schema.test_schema.name + name = "%s" +} + +resource "snowflake_table" "test_table" { + database = snowflake_database.test_database.name + schema = snowflake_schema.test_schema.name + name = "%s" + comment = "Terraform acceptance test" + + column { + name = "column1" + type = "NUMBER(38,0)" + } + column { + name = "column2" + type = "TIMESTAMP_NTZ(9)" + } + + column { + name = "column4" + type = "NUMBER(38,0)" + identity { + start_num = 2 + step_num = 4 + } + } +} +` + return fmt.Sprintf(s, name, name, name, name) +} diff --git a/pkg/snowflake/table.go b/pkg/snowflake/table.go index d7b4ca405c..ae533fd1a4 100644 --- a/pkg/snowflake/table.go +++ b/pkg/snowflake/table.go @@ -43,6 +43,22 @@ type ColumnDefault struct { expression string } +type ColumnIdentity struct { + startNum int + stepNum int +} + +func (id *ColumnIdentity) WithStartNum(start int) *ColumnIdentity { + id.startNum = start + return id + +} + +func (id *ColumnIdentity) WithStep(step int) *ColumnIdentity { + id.stepNum = step + return id +} + func NewColumnDefaultWithConstant(constant string) *ColumnDefault { return &ColumnDefault{ _type: columnDefaultTypeConstant, @@ -98,7 +114,8 @@ type Column struct { _type string // type is reserved nullable bool _default *ColumnDefault // default is reserved - comment string // pointer as value is nullable + identity *ColumnIdentity + comment string // pointer as value is nullable } // WithName set the column name @@ -130,6 +147,11 @@ func (c *Column) WithComment(comment string) *Column { return c } +func (c *Column) WithIdentity(id *ColumnIdentity) *Column { + c.identity = id + return c +} + func (c *Column) getColumnDefinition(withInlineConstraints bool, withComment bool) string { if c == nil { @@ -148,6 +170,10 @@ func (c *Column) getColumnDefinition(withInlineConstraints bool, withComment boo colDef.WriteString(fmt.Sprintf(` DEFAULT %v`, c._default.String(c._type))) } + if c.identity != nil { + colDef.WriteString(fmt.Sprintf(` IDENTITY(%v, %v)`, c.identity.startNum, c.identity.stepNum)) + } + if withComment { colDef.WriteString(fmt.Sprintf(` COMMENT '%v'`, EscapeString(c.comment))) } @@ -211,6 +237,7 @@ func NewColumns(tds []tableDescription) Columns { _type: td.Type.String, nullable: td.IsNullable(), _default: td.ColumnDefault(), + identity: td.ColumnIdentity(), comment: td.Comment.String, }) } @@ -240,6 +267,12 @@ func (c Columns) Flatten() []interface{} { flat["default"] = []interface{}{def} } + if col.identity != nil { + id := map[string]interface{}{} + id["start_num"] = col.identity.startNum + id["step_num"] = col.identity.stepNum + flat["identity"] = []interface{}{id} + } flattened = append(flattened, flat) } return flattened @@ -507,12 +540,13 @@ func (tb *TableBuilder) ChangeChangeTracking(changeTracking bool) string { } // AddColumn returns the SQL query that will add a new column to the table. -func (tb *TableBuilder) AddColumn(name string, dataType string, nullable bool, _default *ColumnDefault, comment string) string { +func (tb *TableBuilder) AddColumn(name string, dataType string, nullable bool, _default *ColumnDefault, identity *ColumnIdentity, comment string) string { col := Column{ name: name, _type: dataType, nullable: nullable, _default: _default, + identity: identity, comment: comment, } return fmt.Sprintf(`ALTER TABLE %s ADD COLUMN %s`, tb.QualifiedName(), col.getColumnDefinition(true, true)) @@ -647,9 +681,34 @@ func (td *tableDescription) ColumnDefault() *ColumnDefault { return NewColumnDefaultWithConstant(UnescapeSnowflakeString(td.Default.String)) } + if td.ColumnIdentity() != nil { + /* + Identity/autoincrement information is stored in the same column as default information. We want to handle the identity seperate so will return nil + here if identity information is present. Default/identity are mutually exclusive + */ + return nil + } + return NewColumnDefaultWithConstant(td.Default.String) } +func (td *tableDescription) ColumnIdentity() *ColumnIdentity { + // if autoincrement is used this is reflected back IDENTITY START 1 INCREMENT 1 + if !td.Default.Valid { + return nil + } + if strings.Contains(td.Default.String, "IDENTITY") { + + split := strings.Split(td.Default.String, " ") + start, _ := strconv.Atoi(split[2]) + step, _ := strconv.Atoi(split[4]) + + return &ColumnIdentity{start, step} + + } + return nil +} + type primaryKeyDescription struct { ColumnName sql.NullString `db:"column_name"` KeySequence sql.NullString `db:"key_sequence"` diff --git a/pkg/snowflake/table_test.go b/pkg/snowflake/table_test.go index 460b6e7623..32918d4f64 100644 --- a/pkg/snowflake/table_test.go +++ b/pkg/snowflake/table_test.go @@ -81,6 +81,35 @@ func TestTableCreate(t *testing.T) { r.Equal(s.Create(), `CREATE TABLE "test_db"."test_schema"."test_table" ("column1" OBJECT COMMENT '', "column2" VARCHAR COMMENT 'only populated when data is available', "column3" NUMBER(38,0) NOT NULL DEFAULT "test_db"."test_schema"."test_seq".NEXTVAL COMMENT '', "column4" VARCHAR NOT NULL DEFAULT 'test default''s' COMMENT '', "column5" TIMESTAMP_NTZ NOT NULL DEFAULT CURRENT_TIMESTAMP() COMMENT '' ,CONSTRAINT "MY_KEY" PRIMARY KEY("column1")) COMMENT = 'Test Comment' CLUSTER BY LINEAR(column1) DATA_RETENTION_TIME_IN_DAYS = 10 CHANGE_TRACKING = true WITH TAG ("test_db"."test_schema"."tag" = "value", "test_db"."test_schema"."tag2" = "value2")`) } +func TestTableCreateIdentity(t *testing.T) { + r := require.New(t) + s := Table("test_table", "test_db", "test_schema") + cols := []Column{ + { + name: "column1", + _type: "OBJECT", + nullable: true, + }, + { + name: "column2", + _type: "VARCHAR", + nullable: true, + comment: "only populated when data is available", + }, + { + name: "column3", + _type: "NUMBER(38,0)", + nullable: false, + identity: &ColumnIdentity{2, 5}, + }, + } + + s.WithColumns(Columns(cols)) + r.Equal(s.QualifiedName(), `"test_db"."test_schema"."test_table"`) + + r.Equal(`CREATE TABLE "test_db"."test_schema"."test_table" ("column1" OBJECT COMMENT '', "column2" VARCHAR COMMENT 'only populated when data is available', "column3" NUMBER(38,0) NOT NULL IDENTITY(2, 5) COMMENT '') DATA_RETENTION_TIME_IN_DAYS = 0 CHANGE_TRACKING = false`, s.Create()) +} + func TestTableChangeComment(t *testing.T) { r := require.New(t) s := Table("test_table", "test_db", "test_schema") @@ -96,19 +125,25 @@ func TestTableRemoveComment(t *testing.T) { func TestTableAddColumn(t *testing.T) { r := require.New(t) s := Table("test_table", "test_db", "test_schema") - r.Equal(s.AddColumn("new_column", "VARIANT", true, nil, ""), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" VARIANT COMMENT ''`) + r.Equal(s.AddColumn("new_column", "VARIANT", true, nil, nil, ""), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" VARIANT COMMENT ''`) } func TestTableAddColumnWithComment(t *testing.T) { r := require.New(t) s := Table("test_table", "test_db", "test_schema") - r.Equal(s.AddColumn("new_column", "VARIANT", true, nil, "some comment"), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" VARIANT COMMENT 'some comment'`) + r.Equal(s.AddColumn("new_column", "VARIANT", true, nil, nil, "some comment"), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" VARIANT COMMENT 'some comment'`) } func TestTableAddColumnWithDefault(t *testing.T) { r := require.New(t) s := Table("test_table", "test_db", "test_schema") - r.Equal(s.AddColumn("new_column", "NUMBER(38,0)", true, NewColumnDefaultWithConstant("1"), ""), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" NUMBER(38,0) DEFAULT 1 COMMENT ''`) + r.Equal(s.AddColumn("new_column", "NUMBER(38,0)", true, NewColumnDefaultWithConstant("1"), nil, ""), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" NUMBER(38,0) DEFAULT 1 COMMENT ''`) +} + +func TestTableAddColumnWithIdentity(t *testing.T) { + r := require.New(t) + s := Table("test_table", "test_db", "test_schema") + r.Equal(s.AddColumn("new_column", "NUMBER(38,0)", true, nil, &ColumnIdentity{1, 4}, ""), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" NUMBER(38,0) IDENTITY(1, 4) COMMENT ''`) } func TestTableDropColumn(t *testing.T) {