Repository: chanzuckerberg/terraform-provider-snowflake
Pull request #599 (2021-07-08 21:20:31)

feat: table constraints

<!-- summary of changes -->

* Adds null constraint to snowflake tables
* Adds specifying a primary key constraint to snowflake tables

## Test Plan
<!-- detail ways in which this PR has been tested or needs to be tested -->
* [x] acceptance tests

## References
<!-- issues documentation links, etc  -->

* #538
* https://docs.snowflake.com/en/sql-reference/constraints-overview.html

alldoami: /ok-to-test sha=9f43de1
github-actions[bot]: Integration tests success for [9f43de1](https://github.com/chanzuckerberg/terraform-provider-snowflake/actions/runs/1023846446)
alldoami: @berosen if you can update your branch, we'll be good to go!
berosen: > @berosen if you can update your branch, we'll be good to go!

Sure thing! Thank you!
alldoami: /ok-to-test sha=f380225

Patch:
diff --git a/docs/resources/table.md b/docs/resources/table.md
index e01baf39..9e0b74c3 100644
--- a/docs/resources/table.md
+++ b/docs/resources/table.md
@@ -21,19 +21,27 @@ resource snowflake_table table {
   cluster_by = ["to_date(DATE)"]
   
   column {
-    name = "id"
-    type = "int"
+    name     = "id"
+    type     = "int"
+    nullable = true
   }
 
   column {
-    name = "data"
-    type = "text"
+    name     = "data"
+    type     = "text"
+    nullable = false
   }
 
   column {
     name = "DATE"
     type = "TIMESTAMP_NTZ(9)"
   }
+
+  primary_key {
+    name = "my_key"
+    keys = ["data"]
+
+  }
 }
 ```
 
@@ -49,9 +57,10 @@ resource snowflake_table table {
 
 ### Optional
 
-- **cluster_by** (List of String) A list of one of more table columns/expressions to be used as clustering key(s) for the table
+- **cluster_by** (List of String) A list of one or more table columns/expressions to be used as clustering key(s) for the table
 - **comment** (String) Specifies a comment for the table.
 - **id** (String) The ID of this resource.
+- **primary_key** (Block List, Max: 1) Definitions of primary key constraint to create on table (see [below for nested schema](#nestedblock--primary_key))
 
 ### Read-Only
 
@@ -65,6 +74,22 @@ Required:
 - **name** (String) Column name
 - **type** (String) Column type, e.g. VARIANT
 
+Optional:
+
+- **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.
+
+
+<a id="nestedblock--primary_key"></a>
+### Nested Schema for `primary_key`
+
+Required:
+
+- **keys** (List of String) Columns to use in primary key
+
+Optional:
+
+- **name** (String) Name of constraint
+
 ## Import
 
 Import is supported using the following syntax:
diff --git a/examples/resources/snowflake_table/resource.tf b/examples/resources/snowflake_table/resource.tf
index 5c46412b..0578d91f 100644
--- a/examples/resources/snowflake_table/resource.tf
+++ b/examples/resources/snowflake_table/resource.tf
@@ -6,17 +6,25 @@ resource snowflake_table table {
   cluster_by = ["to_date(DATE)"]
   
   column {
-    name = "id"
-    type = "int"
+    name     = "id"
+    type     = "int"
+    nullable = true
   }
 
   column {
-    name = "data"
-    type = "text"
+    name     = "data"
+    type     = "text"
+    nullable = false
   }
 
   column {
     name = "DATE"
     type = "TIMESTAMP_NTZ(9)"
   }
+
+  primary_key {
+    name = "my_key"
+    keys = ["data"]
+
+  }
 }
diff --git a/pkg/resources/table.go b/pkg/resources/table.go
index 7f1dc824..f7a18979 100644
--- a/pkg/resources/table.go
+++ b/pkg/resources/table.go
@@ -40,7 +40,7 @@ var tableSchema = map[string]*schema.Schema{
 		Type:        schema.TypeList,
 		Elem:        &schema.Schema{Type: schema.TypeString},
 		Optional:    true,
-		Description: "A list of one of more table columns/expressions to be used as clustering key(s) for the table",
+		Description: "A list of one or more table columns/expressions to be used as clustering key(s) for the table",
 	},
 	"column": {
 		Type:        schema.TypeList,
@@ -59,6 +59,12 @@ var tableSchema = map[string]*schema.Schema{
 					Required:    true,
 					Description: "Column type, e.g. VARIANT",
 				},
+				"nullable": {
+					Type:        schema.TypeBool,
+					Optional:    true,
+					Default:     true,
+					Description: "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.",
+				},
 			},
 		},
 	},
@@ -72,6 +78,29 @@ var tableSchema = map[string]*schema.Schema{
 		Computed:    true,
 		Description: "Name of the role that owns the table.",
 	},
+	"primary_key": {
+		Type:        schema.TypeList,
+		Optional:    true,
+		MaxItems:    1,
+		Description: "Definitions of primary key constraint to create on table",
+		Elem: &schema.Resource{
+			Schema: map[string]*schema.Schema{
+				"name": {
+					Type:        schema.TypeString,
+					Optional:    true,
+					Description: "Name of constraint",
+				},
+				"keys": {
+					Type: schema.TypeList,
+					Elem: &schema.Schema{
+						Type: schema.TypeString,
+					},
+					Required:    true,
+					Description: "Columns to use in primary key",
+				},
+			},
+		},
+	},
 }
 
 func Table() *schema.Resource {
@@ -137,11 +166,12 @@ func tableIDFromString(stringID string) (*tableID, error) {
 type column struct {
 	name     string
 	dataType string
+	nullable bool
 }
 
 func (c column) toSnowflakeColumn() snowflake.Column {
 	sC := snowflake.Column{}
-	return *sC.WithName(c.name).WithType(c.dataType)
+	return *sC.WithName(c.name).WithType(c.dataType).WithNullable(c.nullable)
 }
 
 type columns []column
@@ -171,20 +201,34 @@ func (old columns) getNewIn(new columns) (added columns) {
 	return
 }
 
-func (old columns) getChangedTypes(new columns) (changed columns) {
-	changed = columns{}
+type changedColumns []changedColumn
+
+type changedColumn struct {
+	newColumn             column //our new column
+	changedDataType       bool
+	changedNullConstraint bool
+}
+
+func (old columns) getChangedColumnProperties(new columns) (changed changedColumns) {
+	changed = changedColumns{}
 	for _, cO := range old {
 		for _, cN := range new {
+			changeColumn := changedColumn{cN, false, false}
 			if cO.name == cN.name && cO.dataType != cN.dataType {
-				changed = append(changed, cN)
+				changeColumn.changedDataType = true
 			}
+			if cO.name == cN.name && cO.nullable != cN.nullable {
+				changeColumn.changedNullConstraint = true
+			}
+
+			changed = append(changed, changeColumn)
 		}
 	}
 	return
 }
 
-func (old columns) diffs(new columns) (removed columns, added columns, changed columns) {
-	return old.getNewIn(new), new.getNewIn(old), old.getChangedTypes(new)
+func (old columns) diffs(new columns) (removed columns, added columns, changed changedColumns) {
+	return old.getNewIn(new), new.getNewIn(old), old.getChangedColumnProperties(new)
 }
 
 func getColumn(from interface{}) (to column) {
@@ -192,6 +236,7 @@ func getColumn(from interface{}) (to column) {
 	return column{
 		name:     c["name"].(string),
 		dataType: c["type"].(string),
+		nullable: c["nullable"].(bool),
 	}
 }
 
@@ -204,6 +249,29 @@ func getColumns(from interface{}) (to columns) {
 	return to
 }
 
+type primarykey struct {
+	name string
+	keys []string
+}
+
+func getPrimaryKey(from interface{}) (to primarykey) {
+	pk := from.([]interface{})
+	to = primarykey{}
+	if len(pk) > 0 {
+		pkDetails := pk[0].(map[string]interface{})
+		to.name = pkDetails["name"].(string)
+		to.keys = expandStringList(pkDetails["keys"].([]interface{}))
+		return to
+	}
+	return to
+}
+
+func (pk primarykey) toSnowflakePrimaryKey() snowflake.PrimaryKey {
+	snowPk := snowflake.PrimaryKey{}
+	return *snowPk.WithName(pk.name).WithKeys(pk.keys)
+
+}
+
 // CreateTable implements schema.CreateFunc
 func CreateTable(d *schema.ResourceData, meta interface{}) error {
 	db := meta.(*sql.DB)
@@ -212,6 +280,7 @@ func CreateTable(d *schema.ResourceData, meta interface{}) error {
 	name := d.Get("name").(string)
 
 	columns := getColumns(d.Get("column").([]interface{}))
+
 	builder := snowflake.TableWithColumnDefinitions(name, database, schema, columns.toSnowflakeColumns())
 
 	// Set optionals
@@ -223,6 +292,11 @@ func CreateTable(d *schema.ResourceData, meta interface{}) error {
 		builder.WithClustering(expandStringList(v.([]interface{})))
 	}
 
+	if v, ok := d.GetOk("primary_key"); ok {
+		pk := getPrimaryKey(v.([]interface{}))
+		builder.WithPrimaryKey(pk.toSnowflakePrimaryKey())
+	}
+
 	stmt := builder.Create()
 	err := snowflake.Exec(db, stmt)
 	if err != nil {
@@ -275,15 +349,26 @@ func ReadTable(d *schema.ResourceData, meta interface{}) error {
 		return err
 	}
 
+	showPkrows, err := snowflake.Query(db, builder.ShowPrimaryKeys())
+	if err != nil {
+		return err
+	}
+
+	pkDescription, err := snowflake.ScanPrimaryKeyDescription(showPkrows)
+	if err != nil {
+		return err
+	}
+
 	// Set the relevant data in the state
 	toSet := map[string]interface{}{
-		"name":       table.TableName.String,
-		"owner":      table.Owner.String,
-		"database":   tableID.DatabaseName,
-		"schema":     tableID.SchemaName,
-		"comment":    table.Comment.String,
-		"column":     snowflake.NewColumns(tableDescription).Flatten(),
-		"cluster_by": snowflake.ClusterStatementToList(table.ClusterBy.String),
+		"name":        table.TableName.String,
+		"owner":       table.Owner.String,
+		"database":    tableID.DatabaseName,
+		"schema":      tableID.SchemaName,
+		"comment":     table.Comment.String,
+		"column":      snowflake.NewColumns(tableDescription).Flatten(),
+		"cluster_by":  snowflake.ClusterStatementToList(table.ClusterBy.String),
+		"primary_key": snowflake.FlattenTablePrimaryKey(pkDescription),
 	}
 
 	for key, val := range toSet {
@@ -345,17 +430,56 @@ func UpdateTable(d *schema.ResourceData, meta interface{}) error {
 			}
 		}
 		for _, cA := range added {
-			q := builder.AddColumn(cA.name, cA.dataType)
+			q := builder.AddColumn(cA.name, cA.dataType, cA.nullable)
 			err := snowflake.Exec(db, q)
 			if err != nil {
 				return errors.Wrapf(err, "error adding column on %v", d.Id())
 			}
 		}
 		for _, cA := range changed {
-			q := builder.ChangeColumnType(cA.name, cA.dataType)
+
+			if cA.changedDataType {
+
+				q := builder.ChangeColumnType(cA.newColumn.name, cA.newColumn.dataType)
+				err := snowflake.Exec(db, q)
+				if err != nil {
+					return errors.Wrapf(err, "error changing property on %v", d.Id())
+
+				}
+			}
+			if cA.changedNullConstraint {
+
+				q := builder.ChangeNullConstraint(cA.newColumn.name, cA.newColumn.nullable)
+				err := snowflake.Exec(db, q)
+				if err != nil {
+					return errors.Wrapf(err, "error changing property on %v", d.Id())
+
+				}
+			}
+
+		}
+	}
+	if d.HasChange("primary_key") {
+		opk, npk := d.GetChange("primary_key")
+
+		newpk := getPrimaryKey(npk)
+		oldpk := getPrimaryKey(opk)
+
+		if len(oldpk.keys) > 0 || len(newpk.keys) == 0 {
+			//drop our pk if there was an old primary key, or pk has been removed
+			q := builder.DropPrimaryKey()
+			err := snowflake.Exec(db, q)
+			if err != nil {
+				return errors.Wrapf(err, "error changing primary key first on %v", d.Id())
+			}
+		}
+
+		if len(newpk.keys) > 0 {
+			// add our new pk
+			q := builder.ChangePrimaryKey(newpk.toSnowflakePrimaryKey())
 			err := snowflake.Exec(db, q)
 			if err != nil {
-				return errors.Wrapf(err, "error changing column type on %v", d.Id())
+				return errors.Wrapf(err, "error changing property on %v", d.Id())
 			}
 		}
 	}
diff --git a/pkg/snowflake/table.go b/pkg/snowflake/table.go
index 24d3d782..16b76ee4 100644
--- a/pkg/snowflake/table.go
+++ b/pkg/snowflake/table.go
@@ -3,14 +3,32 @@ package snowflake
 import (
 	"database/sql"
 	"fmt"
+	"sort"
+	"strconv"
 	"strings"
 
 	"github.com/jmoiron/sqlx"
 )
 
+type PrimaryKey struct {
+	name string
+	keys []string
+}
+
+func (pk *PrimaryKey) WithName(name string) *PrimaryKey {
+	pk.name = name
+	return pk
+}
+
+func (pk *PrimaryKey) WithKeys(keys []string) *PrimaryKey {
+	pk.keys = keys
+	return pk
+}
+
 type Column struct {
-	name  string
-	_type string // type is reserved
+	name     string
+	_type    string // type is reserved
+	nullable bool
 }
 
 func (c *Column) WithName(name string) *Column {
@@ -22,11 +40,66 @@ func (c *Column) WithType(t string) *Column {
 	return c
 }
 
-func (c *Column) getColumnDefinition() string {
+func (c *Column) WithNullable(nullable bool) *Column {
+	c.nullable = nullable
+	return c
+}
+
+func (c *Column) getColumnDefinition(withInlineConstraints bool) string {
+
 	if c == nil {
 		return ""
 	}
-	return fmt.Sprintf(`"%v" %v`, EscapeString(c.name), EscapeString(c._type))
+	var colDef strings.Builder
+	colDef.WriteString(fmt.Sprintf(`"%v" %v`, EscapeString(c.name), EscapeString(c._type)))
+	if withInlineConstraints {
+		if !c.nullable {
+			colDef.WriteString(` NOT NULL`)
+		}
+	}
+
+	return colDef.String()
+
+}
+
+func FlattenTablePrimaryKey(pkds []primaryKeyDescription) []interface{} {
+	flattened := []interface{}{}
+	if len(pkds) == 0 {
+		return flattened
+	}
+
+	sort.SliceStable(pkds, func(i, j int) bool {
+		num1, _ := strconv.Atoi(pkds[i].KeySequence.String)
+		num2, _ := strconv.Atoi(pkds[j].KeySequence.String)
+		return num1 < num2
+	})
+	//sort our keys on the key sequence
+
+	flat := map[string]interface{}{}
+	var keys []string
+	var name string
+	var nameSet bool
+
+	for _, pk := range pkds {
+		//set as empty string, sys_constraint means it was an unnnamed constraint
+		if strings.Contains(pk.ConstraintName.String, "SYS_CONSTRAINT") && !nameSet {
+			name = ""
+			nameSet = true
+		}
+		if !nameSet {
+			name = pk.ConstraintName.String
+			nameSet = true
+		}
+
+		keys = append(keys, pk.ColumnName.String)
+
+	}
+
+	flat["name"] = name
+	flat["keys"] = keys
+	flattened = append(flattened, flat)
+	return flattened
+
 }
 
 type Columns []Column
@@ -39,8 +112,9 @@ func NewColumns(tds []tableDescription) Columns {
 			continue
 		}
 		cs = append(cs, Column{
-			name:  td.Name.String,
-			_type: td.Type.String,
+			name:     td.Name.String,
+			_type:    td.Type.String,
+			nullable: td.IsNullable(),
 		})
 	}
 	return Columns(cs)
@@ -52,17 +126,18 @@ func (c Columns) Flatten() []interface{} {
 		flat := map[string]interface{}{}
 		flat["name"] = col.name
 		flat["type"] = col._type
+		flat["nullable"] = col.nullable
 
 		flattened = append(flattened, flat)
 	}
 	return flattened
 }
 
-func (c Columns) getColumnDefinitions() string {
+func (c Columns) getColumnDefinitions(withInlineConstraints bool) string {
 	// TODO(el): verify Snowflake reflects column order back in desc table calls
 	columnDefinitions := []string{}
 	for _, column := range c {
-		columnDefinitions = append(columnDefinitions, column.getColumnDefinition())
+		columnDefinitions = append(columnDefinitions, column.getColumnDefinition(withInlineConstraints))
 	}
 
 	// NOTE: intentionally blank leading space
@@ -71,12 +146,13 @@ func (c Columns) getColumnDefinitions() string {
 
 // TableBuilder abstracts the creation of SQL queries for a Snowflake schema
 type TableBuilder struct {
-	name      string
-	db        string
-	schema    string
-	columns   Columns
-	comment   string
-	clusterBy []string
+	name       string
+	db         string
+	schema     string
+	columns    Columns
+	comment    string
+	clusterBy  []string
+	primaryKey PrimaryKey
 }
 
 // QualifiedName prepends the db and schema if set and escapes everything nicely
@@ -118,10 +194,55 @@ func (tb *TableBuilder) WithClustering(c []string) *TableBuilder {
 	return tb
 }
 
+func (tb *TableBuilder) WithPrimaryKey(pk PrimaryKey) *TableBuilder {
+	tb.primaryKey = pk
+	return tb
+}
+
 //Function to get clustering definition
 func (tb *TableBuilder) GetClusterKeyString() string {
 
-	return fmt.Sprint(strings.Join(tb.clusterBy[:], ", "))
+	return JoinStringList(tb.clusterBy[:], ", ")
+}
+
+func JoinStringList(instrings []string, delimiter string) string {
+
+	return fmt.Sprint(strings.Join(instrings[:], delimiter))
+
+}
+
+func quoteStringList(instrings []string) []string {
+	var clean []string
+	for _, word := range instrings {
+		quoted := fmt.Sprintf(`"%s"`, word)
+		clean = append(clean, quoted)
+
+	}
+	return clean
+
+}
+
+func (tb *TableBuilder) getCreateStatementBody() string {
+	var q strings.Builder
+
+	colDef := tb.columns.getColumnDefinitions(true)
+
+	if len(tb.primaryKey.keys) > 0 {
+		colDef = strings.TrimSuffix(colDef, ")") //strip trailing
+		q.WriteString(colDef)
+		if tb.primaryKey.name != "" {
+			q.WriteString(fmt.Sprintf(` ,CONSTRAINT "%v" PRIMARY KEY(%v)`, tb.primaryKey.name, JoinStringList(quoteStringList(tb.primaryKey.keys), ",")))
+
+		} else {
+			q.WriteString(fmt.Sprintf(` ,PRIMARY KEY(%v)`, JoinStringList(quoteStringList(tb.primaryKey.keys), ",")))
+		}
+
+		q.WriteString(")") // add closing
+	} else {
+		q.WriteString(colDef)
+	}
+
+	return q.String()
 }
 
 //function to take the literal snowflake cluster statement returned from SHOW TABLES and convert it to a list of keys.
@@ -178,7 +299,7 @@ func TableWithColumnDefinitions(name, db, schema string, columns Columns) *Table
 func (tb *TableBuilder) Create() string {
 	q := strings.Builder{}
 	q.WriteString(fmt.Sprintf(`CREATE TABLE %v`, tb.QualifiedName()))
-	q.WriteString(tb.columns.getColumnDefinitions())
+	q.WriteString(tb.getCreateStatementBody())
 
 	if tb.comment != "" {
 		q.WriteString(fmt.Sprintf(` COMMENT = '%v'`, EscapeString(tb.comment)))
@@ -204,12 +325,13 @@ func (tb *TableBuilder) ChangeComment(c string) string {
 }
 
 // AddColumn returns the SQL query that will add a new column to the table.
-func (tb *TableBuilder) AddColumn(name string, dataType string) string {
+func (tb *TableBuilder) AddColumn(name string, dataType string, nullable bool) string {
 	col := Column{
-		name:  name,
-		_type: dataType,
+		name:     name,
+		_type:    dataType,
+		nullable: nullable,
 	}
-	return fmt.Sprintf(`ALTER TABLE %s ADD COLUMN %s`, tb.QualifiedName(), col.getColumnDefinition())
+	return fmt.Sprintf(`ALTER TABLE %s ADD COLUMN %s`, tb.QualifiedName(), col.getColumnDefinition(true))
 }
 
 // DropColumn returns the SQL query that will add a new column to the table.
@@ -223,7 +345,7 @@ func (tb *TableBuilder) ChangeColumnType(name string, dataType string) string {
 		name:  name,
 		_type: dataType,
 	}
-	return fmt.Sprintf(`ALTER TABLE %s MODIFY COLUMN %s`, tb.QualifiedName(), col.getColumnDefinition())
+	return fmt.Sprintf(`ALTER TABLE %s MODIFY COLUMN %s`, tb.QualifiedName(), col.getColumnDefinition(false))
 }
 
 // RemoveComment returns the SQL query that will remove the comment on the table.
@@ -231,6 +353,28 @@ func (tb *TableBuilder) RemoveComment() string {
 	return fmt.Sprintf(`ALTER TABLE %v UNSET COMMENT`, tb.QualifiedName())
 }
 
+// Return sql to set/unset null constraint on column
+func (tb *TableBuilder) ChangeNullConstraint(name string, nullable bool) string {
+	if nullable {
+		return fmt.Sprintf(`ALTER TABLE %s MODIFY COLUMN "%s" DROP NOT NULL`, tb.QualifiedName(), name)
+	} else {
+		return fmt.Sprintf(`ALTER TABLE %s MODIFY COLUMN "%s" SET NOT NULL`, tb.QualifiedName(), name)
+	}
+}
+
+func (tb *TableBuilder) ChangePrimaryKey(newPk PrimaryKey) string {
+	tb.WithPrimaryKey(newPk)
+	pks := JoinStringList(quoteStringList(newPk.keys), ", ")
+	if tb.primaryKey.name != "" {
+		return fmt.Sprintf(`ALTER TABLE %s ADD CONSTRAINT "%v" PRIMARY KEY(%v)`, tb.QualifiedName(), tb.primaryKey.name, pks)
+	}
+	return fmt.Sprintf(`ALTER TABLE %s ADD PRIMARY KEY(%v)`, tb.QualifiedName(), pks)
+}
+
+func (tb *TableBuilder) DropPrimaryKey() string {
+	return fmt.Sprintf(`ALTER TABLE %s DROP PRIMARY KEY`, tb.QualifiedName())
+}
+
 // RemoveClustering returns the SQL query that will remove data clustering from the table
 func (tb *TableBuilder) DropClustering() string {
 	return fmt.Sprintf(`ALTER TABLE %v DROP CLUSTERING KEY`, tb.QualifiedName())
@@ -250,6 +394,10 @@ func (tb *TableBuilder) ShowColumns() string {
 	return fmt.Sprintf(`DESC TABLE %s`, tb.QualifiedName())
 }
 
+func (tb *TableBuilder) ShowPrimaryKeys() string {
+	return fmt.Sprintf(`SHOW PRIMARY KEYS IN TABLE %s`, tb.QualifiedName())
+}
+
 type table struct {
 	CreatedOn           sql.NullString `db:"created_on"`
 	TableName           sql.NullString `db:"name"`
@@ -273,9 +421,24 @@ func ScanTable(row *sqlx.Row) (*table, error) {
 }
 
 type tableDescription struct {
-	Name sql.NullString `db:"name"`
-	Type sql.NullString `db:"type"`
-	Kind sql.NullString `db:"kind"`
+	Name     sql.NullString `db:"name"`
+	Type     sql.NullString `db:"type"`
+	Kind     sql.NullString `db:"kind"`
+	Nullable sql.NullString `db:"null?"`
+}
+
+func (td *tableDescription) IsNullable() bool {
+	if td.Nullable.String == "Y" {
+		return true
+	} else {
+		return false
+	}
+}
+
+type primaryKeyDescription struct {
+	ColumnName     sql.NullString `db:"column_name"`
+	KeySequence    sql.NullString `db:"key_sequence"`
+	ConstraintName sql.NullString `db:"constraint_name"`
 }
 
 func ScanTableDescription(rows *sqlx.Rows) ([]tableDescription, error) {
@@ -290,3 +453,16 @@ func ScanTableDescription(rows *sqlx.Rows) ([]tableDescription, error) {
 	}
 	return tds, rows.Err()
 }
+
+func ScanPrimaryKeyDescription(rows *sqlx.Rows) ([]primaryKeyDescription, error) {
+	pkds := []primaryKeyDescription{}
+	for rows.Next() {
+		pk := primaryKeyDescription{}
+		err := rows.StructScan(&pk)
+		if err != nil {
+			return nil, err
+		}
+		pkds = append(pkds, pk)
+	}
+	return pkds, rows.Err()
+}


Test patch:
diff --git a/pkg/resources/table_acceptance_test.go b/pkg/resources/table_acceptance_test.go
index 12d31d97..4570b7a0 100644
--- a/pkg/resources/table_acceptance_test.go
+++ b/pkg/resources/table_acceptance_test.go
@@ -29,6 +29,7 @@ func TestAcc_Table(t *testing.T) {
 					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.type", "VARIANT"),
 					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.name", "column2"),
 					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.type", "VARCHAR(16)"),
+					resource.TestCheckNoResourceAttr("snowflake_table.test_table", "primary_key"),
 				),
 			},
 			{
@@ -44,6 +45,7 @@ func TestAcc_Table(t *testing.T) {
 					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.name", "column3"),
 					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.type", "FLOAT"),
 					resource.TestCheckNoResourceAttr("snowflake_table.test_table", "cluster_by"),
+					resource.TestCheckNoResourceAttr("snowflake_table.test_table", "primary_key"),
 				),
 			},
 			{
@@ -92,6 +94,63 @@ func TestAcc_Table(t *testing.T) {
 					resource.TestCheckResourceAttr("snowflake_table.test_table2", "cluster_by.0", "\"col2\""),
 				),
 			},
+			{
+				Config: tableConfig6(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", "comment", "Terraform acceptance test"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.#", "2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.name", "column2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.type", "VARCHAR(16777216)"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.nullable", "true"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.name", "column3"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.type", "FLOAT"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.nullable", "false"),
+					resource.TestCheckNoResourceAttr("snowflake_table.test_table", "cluster_by"),
+					resource.TestCheckNoResourceAttr("snowflake_table.test_table", "primary_key"),
+				),
+			},
+			{
+				Config: tableConfig7(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", "comment", "Terraform acceptance test"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.#", "2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.name", "column2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.type", "VARCHAR(16777216)"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.nullable", "true"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.name", "column3"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.type", "FLOAT"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.nullable", "false"),
+					resource.TestCheckNoResourceAttr("snowflake_table.test_table", "cluster_by"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "primary_key.0.keys.0", "column2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "primary_key.0.name", ""),
+				),
+			},
+			{
+				Config: tableConfig8(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", "comment", "Terraform acceptance test"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.#", "2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.name", "column2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.type", "VARCHAR(16777216)"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.0.nullable", "true"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.name", "column3"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.type", "FLOAT"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "column.1.nullable", "false"),
+					resource.TestCheckNoResourceAttr("snowflake_table.test_table", "cluster_by"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "primary_key.0.keys.0", "column2"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "primary_key.0.keys.1", "column3"),
+					resource.TestCheckResourceAttr("snowflake_table.test_table", "primary_key.0.name", "new_name"),
+				),
+			},
 		},
 	})
 }
@@ -253,3 +312,107 @@ resource "snowflake_table" "test_table2" {
 `
 	return fmt.Sprintf(s, name, name, table2Name)
 }
+
+func tableConfig6(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_table" "test_table" {
+	database = snowflake_database.test_database.name
+	schema   = snowflake_schema.test_schema.name
+	name     = "%s"
+	comment  = "Terraform acceptance test"
+	column {
+		name = "column2"
+		type = "VARCHAR(16777216)"
+	}
+	column {
+		name = "column3"
+		type = "FLOAT"
+		nullable = false
+	}
+}
+`
+	return fmt.Sprintf(s, name, name, name)
+}
+
+func tableConfig7(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_table" "test_table" {
+	database = snowflake_database.test_database.name
+	schema   = snowflake_schema.test_schema.name
+	name     = "%s"
+	comment  = "Terraform acceptance test"
+	column {
+		name = "column2"
+		type = "VARCHAR(16777216)"
+	}
+	column {
+		name = "column3"
+		type = "FLOAT"
+		nullable = false
+	}
+	primary_key {
+		name = ""
+		keys = ["column2"]
+	}
+}
+`
+	return fmt.Sprintf(s, name, name, name)
+}
+
+func tableConfig8(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_table" "test_table" {
+	database = snowflake_database.test_database.name
+	schema   = snowflake_schema.test_schema.name
+	name     = "%s"
+	comment  = "Terraform acceptance test"
+	column {
+		name = "column2"
+		type = "VARCHAR(16777216)"
+	}
+	column {
+		name = "column3"
+		type = "FLOAT"
+		nullable = false
+	}
+	primary_key {
+		name = "new_name"
+		keys = ["column2","column3"]
+	}
+}
+`
+	return fmt.Sprintf(s, name, name, name)
+}
diff --git a/pkg/resources/table_test.go b/pkg/resources/table_test.go
index 65ed8270..ddb6f0dc 100644
--- a/pkg/resources/table_test.go
+++ b/pkg/resources/table_test.go
@@ -22,16 +22,17 @@ func TestTableCreate(t *testing.T) {
 	r := require.New(t)
 
 	in := map[string]interface{}{
-		"name":     "good_name",
-		"database": "database_name",
-		"schema":   "schema_name",
-		"comment":  "great comment",
-		"column":   []interface{}{map[string]interface{}{"name": "column1", "type": "OBJECT"}, map[string]interface{}{"name": "column2", "type": "VARCHAR"}},
+		"name":        "good_name",
+		"database":    "database_name",
+		"schema":      "schema_name",
+		"comment":     "great comment",
+		"column":      []interface{}{map[string]interface{}{"name": "column1", "type": "OBJECT"}, map[string]interface{}{"name": "column2", "type": "VARCHAR", "nullable": false}},
+		"primary_key": []interface{}{map[string]interface{}{"name": "MY_KEY", "keys": []interface{}{"column1"}}},
 	}
 	d := table(t, "database_name|schema_name|good_name", in)
 
 	WithMockDb(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
-		mock.ExpectExec(`CREATE TABLE "database_name"."schema_name"."good_name" \("column1" OBJECT, "column2" VARCHAR\) COMMENT = 'great comment'`).WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectExec(`CREATE TABLE "database_name"."schema_name"."good_name" \("column1" OBJECT, "column2" VARCHAR NOT NULL ,CONSTRAINT "MY_KEY" PRIMARY KEY\("column1"\)\) COMMENT = 'great comment'`).WillReturnResult(sqlmock.NewResult(1, 1))
 		expectTableRead(mock)
 		err := resources.CreateTable(d, db)
 		r.NoError(err)
@@ -43,11 +44,16 @@ func expectTableRead(mock sqlmock.Sqlmock) {
 	rows := sqlmock.NewRows([]string{"name", "type", "kind", "null?", "default", "primary key", "unique key", "check", "expression", "comment"}).AddRow("good_name", "VARCHAR()", "COLUMN", "Y", "NULL", "NULL", "N", "N", "NULL", "mock comment")
 	mock.ExpectQuery(`SHOW TABLES LIKE 'good_name' IN SCHEMA "database_name"."schema_name"`).WillReturnRows(rows)
 
-	describeRows := sqlmock.NewRows([]string{"name", "type", "kind"}).
-		AddRow("column1", "OBJECT", "COLUMN").
-		AddRow("column2", "VARCHAR", "COLUMN")
+	describeRows := sqlmock.NewRows([]string{"name", "type", "kind", "null?"}).
+		AddRow("column1", "OBJECT", "COLUMN", "Y").
+		AddRow("column2", "VARCHAR", "COLUMN", "N")
 
 	mock.ExpectQuery(`DESC TABLE "database_name"."schema_name"."good_name"`).WillReturnRows(describeRows)
+
+	pkRows := sqlmock.NewRows([]string{"column_name", "key_sequence", "constraint_name"}).AddRow("column1", "1", "MY_PK")
+
+	mock.ExpectQuery(`SHOW PRIMARY KEYS IN TABLE "database_name"."schema_name"."good_name"`).WillReturnRows(pkRows)
+
 }
 
 func TestTableRead(t *testing.T) {
diff --git a/pkg/snowflake/table_test.go b/pkg/snowflake/table_test.go
index 72513218..fa2d1320 100644
--- a/pkg/snowflake/table_test.go
+++ b/pkg/snowflake/table_test.go
@@ -11,12 +11,14 @@ func TestTableCreate(t *testing.T) {
 	s := Table("test_table", "test_db", "test_schema")
 	cols := []Column{
 		{
-			name:  "column1",
-			_type: "OBJECT",
+			name:     "column1",
+			_type:    "OBJECT",
+			nullable: true,
 		},
 		{
-			name:  "column2",
-			_type: "VARCHAR",
+			name:     "column2",
+			_type:    "VARCHAR",
+			nullable: true,
 		},
 	}
 
@@ -31,6 +33,8 @@ func TestTableCreate(t *testing.T) {
 	s.WithClustering([]string{"column1"})
 	r.Equal(s.Create(), `CREATE TABLE "test_db"."test_schema"."test_table" ("column1" OBJECT, "column2" VARCHAR) COMMENT = 'Test Comment' CLUSTER BY LINEAR(column1)`)
 
+	s.WithPrimaryKey(PrimaryKey{name: "MY_KEY", keys: []string{"column1"}})
+	r.Equal(s.Create(), `CREATE TABLE "test_db"."test_schema"."test_table" ("column1" OBJECT, "column2" VARCHAR ,CONSTRAINT "MY_KEY" PRIMARY KEY("column1")) COMMENT = 'Test Comment' CLUSTER BY LINEAR(column1)`)
 }
 
 func TestTableChangeComment(t *testing.T) {
@@ -48,7 +52,7 @@ 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"), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" VARIANT`)
+	r.Equal(s.AddColumn("new_column", "VARIANT", true), `ALTER TABLE "test_db"."test_schema"."test_table" ADD COLUMN "new_column" VARIANT`)
 }
 
 func TestTableDropColumn(t *testing.T) {
@@ -86,3 +90,27 @@ func TestTableShow(t *testing.T) {
 	s := Table("test_table", "test_db", "test_schema")
 	r.Equal(s.Show(), `SHOW TABLES LIKE 'test_table' IN SCHEMA "test_db"."test_schema"`)
 }
+
+func TestTableShowPrimaryKeys(t *testing.T) {
+	r := require.New(t)
+	s := Table("test_table", "test_db", "test_schema")
+	r.Equal(s.ShowPrimaryKeys(), `SHOW PRIMARY KEYS IN TABLE "test_db"."test_schema"."test_table"`)
+}
+
+func TestTableDropPrimaryKeys(t *testing.T) {
+	r := require.New(t)
+	s := Table("test_table", "test_db", "test_schema")
+	r.Equal(s.DropPrimaryKey(), `ALTER TABLE "test_db"."test_schema"."test_table" DROP PRIMARY KEY`)
+}
+
+func TestTableChangePrimaryKeysWithConstraintName(t *testing.T) {
+	r := require.New(t)
+	s := Table("test_table", "test_db", "test_schema")
+	r.Equal(s.ChangePrimaryKey(PrimaryKey{name: "MY_KEY", keys: []string{"column1", "column2"}}), `ALTER TABLE "test_db"."test_schema"."test_table" ADD CONSTRAINT "MY_KEY" PRIMARY KEY("column1", "column2")`)
+}
+
+func TestTableChangePrimaryKeysWithoutConstraintName(t *testing.T) {
+	r := require.New(t)
+	s := Table("test_table", "test_db", "test_schema")
+	r.Equal(s.ChangePrimaryKey(PrimaryKey{name: "", keys: []string{"column1", "column2"}}), `ALTER TABLE "test_db"."test_schema"."test_table" ADD PRIMARY KEY("column1", "column2")`)
+}