fix: PK Addition Order (#607)

bbernays · web-flow · commit eff40e76ae65 · 2023-01-14T12:36:52.000Z
#### Summary

This will only add the default PKs once the `Transformers` have run and added any and all PKs


&lt;!--
Explain what problem this PR addresses
--&gt;

---
diff --git a/plugins/source/plugin.go b/plugins/source/plugin.go
@@ -112,11 +112,11 @@ func NewPlugin(name string, version string, tables []*schema.Table, newExecution
 		metrics:            &Metrics{TableClient: make(map[string]map[string]*TableClientMetrics)},
 		caser:              caser.New(),
 	}
-	addInternalColumns(p.tables)
 	setParents(p.tables, nil)
 	if err := transformTables(p.tables); err != nil {
 		panic(err)
 	}
+	addInternalColumns(p.tables)
 	if err := p.validate(); err != nil {
 		panic(err)
 	}
diff --git a/plugins/source/plugin_test.go b/plugins/source/plugin_test.go
@@ -6,9 +6,11 @@ import (
 
 	"github.com/cloudquery/plugin-sdk/schema"
 	"github.com/cloudquery/plugin-sdk/specs"
+	"github.com/cloudquery/plugin-sdk/transformers"
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/uuid"
 	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/assert"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -316,3 +318,39 @@ func TestIgnoredColumns(t *testing.T) {
 		},
 	}})
 }
+
+var testTable struct {
+	PrimaryKey   string
+	SecondaryKey string
+	TertiaryKey  string
+	Quaternary   string
+}
+
+func TestNewPluginPrimaryKeys(t *testing.T) {
+	testTransforms := []struct {
+		transformerOptions []transformers.StructTransformerOption
+		resultKeys         []string
+	}{
+		{
+			transformerOptions: []transformers.StructTransformerOption{transformers.WithPrimaryKeys("PrimaryKey")},
+			resultKeys:         []string{"primary_key"},
+		},
+		{
+			transformerOptions: []transformers.StructTransformerOption{},
+			resultKeys:         []string{"_cq_id"},
+		},
+	}
+	for _, tc := range testTransforms {
+		tables := []*schema.Table{
+			{
+				Name: "test_table",
+				Transform: transformers.TransformWithStruct(
+					&testTable, tc.transformerOptions...,
+				),
+			},
+		}
+
+		plugin := NewPlugin("testSourcePlugin", "1.0.0", tables, newTestExecutionClient)
+		assert.Equal(t, tc.resultKeys, plugin.tables[0].PrimaryKeys())
+	}
+}

Original file line number	Diff line number	Diff line change
`@@ -112,11 +112,11 @@ func NewPlugin(name string, version string, tables []*schema.Table, newExecution`
`112`	`112`	`metrics: &Metrics{TableClient: make(map[string]map[string]*TableClientMetrics)},`
`113`	`113`	`caser: caser.New(),`
`114`	`114`	`}`
`115`		`- addInternalColumns(p.tables)`
`116`	`115`	`setParents(p.tables, nil)`
`117`	`116`	`if err := transformTables(p.tables); err != nil {`
`118`	`117`	`panic(err)`
`119`	`118`	`}`
	`119`	`+ addInternalColumns(p.tables)`
`120`	`120`	`if err := p.validate(); err != nil {`
`121`	`121`	`panic(err)`
`122`	`122`	`}`