fix: Adding double quotes to the column names (#1729)

* adding quotes to the column name * modifying the tests --------- Co-authored-by: sfc-gh-kmaurya <[email protected]>
Snowflake-Labs · Apr 18, 2023 · 791dd0b · 791dd0b
1 parent 66d88bd
commit 791dd0b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/pkg/snowflake/tag_association.go b/pkg/snowflake/tag_association.go
@@ -89,7 +89,7 @@ func NewTagAssociationBuilder(tagID string) *TagAssociationBuilder {
 func (tb *TagAssociationBuilder) Create() string {
 	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))
+		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))
 	}
 	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))
 }
@@ -98,7 +98,7 @@ func (tb *TagAssociationBuilder) Create() string {
 func (tb *TagAssociationBuilder) Drop() string {
 	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)
+		return fmt.Sprintf(`ALTER TABLE %v MODIFY COLUMN "%v" UNSET TAG "%v"."%v"."%v"`, tableName, columnName, tb.databaseName, tb.schemaName, tb.tagName)
 	}
 	return fmt.Sprintf(`ALTER %v %v UNSET TAG "%v"."%v"."%v"`, tb.objectType, tb.objectIdentifier, tb.databaseName, tb.schemaName, tb.tagName)
 }

diff --git a/pkg/snowflake/tag_association_test.go b/pkg/snowflake/tag_association_test.go
@@ -23,14 +23,14 @@ func TestTagAssociation(t *testing.T) {
 		},
 		{
 			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"`,
+			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"`,
+			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`,
 		},
 	}