fix: support different tag association queries for COLUMN object types (

#1380) * support different tag association queries for COLUMN object types * refactor after testing * fix logic for getting table and column name * fix logic for when non-FQ name is used * db name is required by the resource, updating tests/logic accordingly * fix sql syntax to use MODIFY rather than ALTER on column * fix show logic to support column-level tagging
Snowflake-Labs · Feb 1, 2023 · 546d0a1 · 546d0a1
1 parent 00fc00c
commit 546d0a1
Show file tree

Hide file tree

Showing 2 changed files with 93 additions and 4 deletions.
diff --git a/pkg/snowflake/tag_association.go b/pkg/snowflake/tag_association.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"log"
+	"strings"
 
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/validation"
 	"github.com/jmoiron/sqlx"
@@ -49,14 +50,25 @@ func (tb *TagAssociationBuilder) GetTagDatabase() string {
 
 // GetTagName returns the value of the tag name of TagAssociationBuilder.
 func (tb *TagAssociationBuilder) GetTagName() string {
-	return tb.schemaName
+	return tb.tagName
 }
 
 // GetTagSchema returns the value of the tag schema of TagAssociationBuilder.
 func (tb *TagAssociationBuilder) GetTagSchema() string {
 	return tb.schemaName
 }
 
+func (tb *TagAssociationBuilder) GetTableAndColumnName() (string, string) {
+	if strings.ToUpper(tb.objectType) != "COLUMN" {
+		return tb.objectIdentifier, ""
+	} else {
+		splObjIdentifier := strings.Split(tb.objectIdentifier, ".")
+		tableName := strings.ReplaceAll(splObjIdentifier[2], "\"", "")
+		columnName := strings.ReplaceAll(splObjIdentifier[3], "\"", "")
+		return fmt.Sprintf(`"%s"."%s"."%s"`, tb.databaseName, tb.schemaName, tableName), columnName
+	}
+}
+
 // TagAssociation returns a pointer to a Builder that abstracts the DDL operations for a tag sssociation.
 //
 // Supported DDL operations are:
@@ -76,17 +88,33 @@ func NewTagAssociationBuilder(tagID string) *TagAssociationBuilder {
 
 // Create returns the SQL query that will set the tag on an object.
 func (tb *TagAssociationBuilder) Create() string {
-	return fmt.Sprintf(`ALTER %v %v SET TAG "%v"."%v"."%v" = '%v'`, tb.objectType, tb.objectIdentifier, tb.databaseName, tb.schemaName, tb.tagName, EscapeString(tb.tagValue))
+	if strings.ToUpper(tb.objectType) == "COLUMN" {
+		tableName, columnName := tb.GetTableAndColumnName()
+		return fmt.Sprintf(`ALTER TABLE %v MODIFY COLUMN %v SET TAG "%v"."%v"."%v" = '%v'`, tableName, columnName, tb.databaseName, tb.schemaName, tb.tagName, EscapeString(tb.tagValue))
+	} else {
+		return fmt.Sprintf(`ALTER %v %v SET TAG "%v"."%v"."%v" = '%v'`, tb.objectType, tb.objectIdentifier, tb.databaseName, tb.schemaName, tb.tagName, EscapeString(tb.tagValue))
+	}
 }
 
 // Drop returns the SQL query that will remove a tag from an object.
 func (tb *TagAssociationBuilder) Drop() string {
-	return fmt.Sprintf(`ALTER %v %v UNSET TAG "%v"."%v"."%v"`, tb.objectType, tb.objectIdentifier, tb.databaseName, tb.schemaName, tb.tagName)
+	if strings.ToUpper(tb.objectType) == "COLUMN" {
+		tableName, columnName := tb.GetTableAndColumnName()
+		return fmt.Sprintf(`ALTER TABLE %v MODIFY COLUMN %v UNSET TAG "%v"."%v"."%v"`, tableName, columnName, tb.databaseName, tb.schemaName, tb.tagName)
+	} else {
+		return fmt.Sprintf(`ALTER %v %v UNSET TAG "%v"."%v"."%v"`, tb.objectType, tb.objectIdentifier, tb.databaseName, tb.schemaName, tb.tagName)
+	}
 }
 
 // Show returns the SQL query that will show the current tag value on an object.
 func (tb *TagAssociationBuilder) Show() string {
-	return fmt.Sprintf(`SELECT SYSTEM$GET_TAG('"%v"."%v"."%v"', '%v', '%v') TAG_VALUE WHERE TAG_VALUE IS NOT NULL`, tb.databaseName, tb.schemaName, tb.tagName, tb.objectIdentifier, tb.objectType)
+	if strings.ToUpper(tb.objectType) == "COLUMN" {
+		fqTableName, columnName := tb.GetTableAndColumnName()
+		fqColumnName := fmt.Sprintf(`%v."%v"`, fqTableName, columnName)
+		return fmt.Sprintf(`SELECT SYSTEM$GET_TAG('"%v"."%v"."%v"', '%v', '%v') TAG_VALUE WHERE TAG_VALUE IS NOT NULL`, tb.databaseName, tb.schemaName, tb.tagName, fqColumnName, tb.objectType)
+	} else {
+		return fmt.Sprintf(`SELECT SYSTEM$GET_TAG('"%v"."%v"."%v"', '%v', '%v') TAG_VALUE WHERE TAG_VALUE IS NOT NULL`, tb.databaseName, tb.schemaName, tb.tagName, tb.objectIdentifier, tb.objectType)
+	}
 }
 
 func ScanTagAssociation(row *sqlx.Row) (*TagAssociation, error) {

diff --git a/pkg/snowflake/tag_association_test.go b/pkg/snowflake/tag_association_test.go
@@ -0,0 +1,61 @@
+package snowflake
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type TagAssociationTest struct {
+	Builder        *TagAssociationBuilder
+	ExpectedCreate string
+	ExpectedDrop   string
+	ExpectedShow   string
+}
+
+func TestTagAssociation(t *testing.T) {
+	tests := []TagAssociationTest{
+		{
+			Builder:        NewTagAssociationBuilder("test_db|test_schema|sensitive").WithObjectIdentifier(`"test_db"."test_schema"."test_table"`).WithObjectType("TABLE").WithTagValue("true"),
+			ExpectedCreate: `ALTER TABLE "test_db"."test_schema"."test_table" SET TAG "test_db"."test_schema"."sensitive" = 'true'`,
+			ExpectedDrop:   `ALTER TABLE "test_db"."test_schema"."test_table" UNSET TAG "test_db"."test_schema"."sensitive"`,
+			ExpectedShow:   `SELECT SYSTEM$GET_TAG('"test_db"."test_schema"."sensitive"', '"test_db"."test_schema"."test_table"', 'TABLE') TAG_VALUE WHERE TAG_VALUE IS NOT NULL`,
+		},
+		{
+			Builder:        NewTagAssociationBuilder("test_db|test_schema|sensitive").WithObjectIdentifier(`"test_db"."test_schema"."test_table.important"`).WithObjectType("COLUMN").WithTagValue("true"),
+			ExpectedCreate: `ALTER TABLE "test_db"."test_schema"."test_table" MODIFY COLUMN important SET TAG "test_db"."test_schema"."sensitive" = 'true'`,
+			ExpectedDrop:   `ALTER TABLE "test_db"."test_schema"."test_table" MODIFY COLUMN important UNSET TAG "test_db"."test_schema"."sensitive"`,
+			ExpectedShow:   `SELECT SYSTEM$GET_TAG('"test_db"."test_schema"."sensitive"', '"test_db"."test_schema"."test_table"."important"', 'COLUMN') TAG_VALUE WHERE TAG_VALUE IS NOT NULL`,
+		},
+		{
+			Builder:        NewTagAssociationBuilder("OPERATION_DB|SECURITY|PII_2").WithObjectIdentifier(`"OPERATION_DB"."SECURITY"."test_table.important"`).WithObjectType("COLUMN").WithTagValue("true"),
+			ExpectedCreate: `ALTER TABLE "OPERATION_DB"."SECURITY"."test_table" MODIFY COLUMN important SET TAG "OPERATION_DB"."SECURITY"."PII_2" = 'true'`,
+			ExpectedDrop:   `ALTER TABLE "OPERATION_DB"."SECURITY"."test_table" MODIFY COLUMN important UNSET TAG "OPERATION_DB"."SECURITY"."PII_2"`,
+			ExpectedShow:   `SELECT SYSTEM$GET_TAG('"OPERATION_DB"."SECURITY"."PII_2"', '"OPERATION_DB"."SECURITY"."test_table"."important"', 'COLUMN') TAG_VALUE WHERE TAG_VALUE IS NOT NULL`,
+		},
+	}
+	for _, testCase := range tests {
+		r := require.New(t)
+		r.Equal(testCase.ExpectedCreate, testCase.Builder.Create())
+		r.Equal(testCase.ExpectedDrop, testCase.Builder.Drop())
+		r.Equal(testCase.ExpectedShow, testCase.Builder.Show())
+	}
+}
+
+type TableColumnNameTest struct {
+	Builder                               *TagAssociationBuilder
+	expectedTableName, expectedColumnName string
+}
+
+func TestTableColumnName(t *testing.T) {
+	tests := []TableColumnNameTest{
+		{NewTagAssociationBuilder("a|b|sensitive").WithObjectIdentifier(`"a"."b"."c"`).WithObjectType("TABLE"), `"a"."b"."c"`, ""},
+		{NewTagAssociationBuilder("db|schema|sensitive").WithObjectIdentifier(`"db"."schema"."table.column"`).WithObjectType("COLUMN"), `"db"."schema"."table"`, "column"},
+	}
+	for _, testCase := range tests {
+		r := require.New(t)
+		tableName, columnName := testCase.Builder.GetTableAndColumnName()
+		r.Equal(testCase.expectedTableName, tableName)
+		r.Equal(testCase.expectedColumnName, columnName)
+	}
+}