Fix encoding snowflake ID for object identifiers

Snowflake-Labs · Dec 12, 2023 · 52b0695 · 52b0695
1 parent 2d26046
commit 52b0695
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 24 deletions.
diff --git a/pkg/helpers/helpers.go b/pkg/helpers/helpers.go
@@ -65,10 +65,18 @@ func EncodeSnowflakeID(attributes ...interface{}) string {
 	// is attribute already an object identifier?
 	if len(attributes) == 1 {
 		if id, ok := attributes[0].(sdk.ObjectIdentifier); ok {
-			// remove quotes and replace dots with pipes
-			parts := strings.Split(id.FullyQualifiedName(), ".")
-			for i, part := range parts {
-				parts[i] = strings.Trim(part, `"`)
+			parts := make([]string, 0)
+			switch id.(type) {
+			case sdk.AccountObjectIdentifier:
+				parts = append(parts, id.(sdk.AccountObjectIdentifier).Name())
+			case sdk.DatabaseObjectIdentifier:
+				parts = append(parts, id.(sdk.DatabaseObjectIdentifier).DatabaseName(), id.(sdk.DatabaseObjectIdentifier).Name())
+			case sdk.SchemaObjectIdentifier:
+				parts = append(parts, id.(sdk.SchemaObjectIdentifier).DatabaseName(), id.(sdk.SchemaObjectIdentifier).SchemaName(), id.(sdk.SchemaObjectIdentifier).Name())
+			case sdk.TableColumnIdentifier:
+				parts = append(parts, id.(sdk.TableColumnIdentifier).DatabaseName(), id.(sdk.TableColumnIdentifier).SchemaName(), id.(sdk.TableColumnIdentifier).TableName(), id.(sdk.TableColumnIdentifier).Name())
+			default:
+				log.Panicf("Unsupported object identifier: %v", id)
 			}
 			return strings.Join(parts, IDDelimiter)
 		}

diff --git a/pkg/helpers/helpers_test.go b/pkg/helpers/helpers_test.go
@@ -80,40 +80,40 @@ func TestEncodeSnowflakeID(t *testing.T) {
 		expectedEncodedID string
 	}{
 		"encodes account object identifier": {
-			identifier:        sdk.NewAccountObjectIdentifier("account"),
-			expectedEncodedID: `account`,
+			identifier:        sdk.NewAccountObjectIdentifier("database"),
+			expectedEncodedID: `database`,
 		},
 		"encodes quoted account object identifier": {
-			identifier:        sdk.NewAccountObjectIdentifier("\"account\""),
-			expectedEncodedID: `account`,
+			identifier:        sdk.NewAccountObjectIdentifier("\"database\""),
+			expectedEncodedID: `database`,
 		},
 		"encodes account object identifier with a dot": {
-			identifier:        sdk.NewAccountObjectIdentifier("acc.ount"),
-			expectedEncodedID: `acc.ount`,
+			identifier:        sdk.NewAccountObjectIdentifier("data.base"),
+			expectedEncodedID: `data.base`,
 		},
 		"encodes database object identifier": {
-			identifier:        sdk.NewDatabaseObjectIdentifier("account", "database"),
-			expectedEncodedID: `account|database`,
+			identifier:        sdk.NewDatabaseObjectIdentifier("database", "schema"),
+			expectedEncodedID: `database|schema`,
 		},
 		"encodes quoted database object identifier": {
-			identifier:        sdk.NewDatabaseObjectIdentifier("\"account\"", "\"database\""),
-			expectedEncodedID: `account|database`,
+			identifier:        sdk.NewDatabaseObjectIdentifier("\"database\"", "\"schema\""),
+			expectedEncodedID: `database|schema`,
 		},
 		"encodes database object identifier with dots": {
-			identifier:        sdk.NewDatabaseObjectIdentifier("acc.ount", "data.base"),
-			expectedEncodedID: `acc.ount|data.base`,
+			identifier:        sdk.NewDatabaseObjectIdentifier("data.base", "sche.ma"),
+			expectedEncodedID: `data.base|sche.ma`,
 		},
 		"encodes schema object identifier": {
-			identifier:        sdk.NewSchemaObjectIdentifier("account", "database", "schema"),
-			expectedEncodedID: `account|database|schema`,
+			identifier:        sdk.NewSchemaObjectIdentifier("database", "schema", "table"),
+			expectedEncodedID: `database|schema|table`,
 		},
 		"encodes quoted schema object identifier": {
-			identifier:        sdk.NewSchemaObjectIdentifier("\"account\"", "\"database\"", "\"schema\""),
-			expectedEncodedID: `account|database|schema`,
+			identifier:        sdk.NewSchemaObjectIdentifier("\"database\"", "\"schema\"", "\"table\""),
+			expectedEncodedID: `database|schema|table`,
 		},
 		"encodes schema object identifier with dots": {
-			identifier:        sdk.NewSchemaObjectIdentifier("acc.ount", "data.base", "sche.ma"),
-			expectedEncodedID: `acc.ount|data.base|sche.ma`,
+			identifier:        sdk.NewSchemaObjectIdentifier("data.base", "sche.ma", "tab.le"),
+			expectedEncodedID: `data.base|sche.ma|tab.le`,
 		},
 		"encodes table column identifier": {
 			identifier:        sdk.NewTableColumnIdentifier("database", "schema", "table", "column"),
@@ -135,4 +135,20 @@ func TestEncodeSnowflakeID(t *testing.T) {
 			require.Equal(t, tc.expectedEncodedID, encodedID)
 		})
 	}
+
+	t.Run("panics for unsupported object identifier", func(t *testing.T) {
+		require.Panics(t, func() {
+			EncodeSnowflakeID(unsupportedObjectIdentifier{})
+		})
+	})
+}
+
+type unsupportedObjectIdentifier struct{}
+
+func (i unsupportedObjectIdentifier) Name() string {
+	return "name"
+}
+
+func (i unsupportedObjectIdentifier) FullyQualifiedName() string {
+	return "fully qualified name"
 }
diff --git a/pkg/sdk/identifier_helpers.go b/pkg/sdk/identifier_helpers.go
@@ -121,7 +121,9 @@ type AccountObjectIdentifier struct {
 }
 
 func NewAccountObjectIdentifier(name string) AccountObjectIdentifier {
-	return AccountObjectIdentifier{name: name}
+	return AccountObjectIdentifier{
+		name: strings.Trim(name, `"`),
+	}
 }
 
 func NewAccountObjectIdentifierFromFullyQualifiedName(fullyQualifiedName string) AccountObjectIdentifier {
@@ -268,7 +270,12 @@ type TableColumnIdentifier struct {
 }
 
 func NewTableColumnIdentifier(databaseName, schemaName, tableName, columnName string) TableColumnIdentifier {
-	return TableColumnIdentifier{databaseName: databaseName, schemaName: schemaName, tableName: tableName, columnName: columnName}
+	return TableColumnIdentifier{
+		databaseName: strings.Trim(databaseName, `"`),
+		schemaName:   strings.Trim(schemaName, `"`),
+		tableName:    strings.Trim(tableName, `"`),
+		columnName:   strings.Trim(columnName, `"`),
+	}
 }
 
 func NewTableColumnIdentifierFromFullyQualifiedName(fullyQualifiedName string) TableColumnIdentifier {