Generate uuid in go rather than in db (flyteorg#58)

Signed-off-by: Yee Hing Tong <[email protected]>

datacalog/go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/golang/protobuf v1.4.3
 	github.com/jackc/pgconn v1.8.1
 	github.com/mitchellh/mapstructure v1.4.1
+	github.com/satori/go.uuid v1.2.0
 	github.com/spf13/cobra v1.1.1
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.7.0

datacalog/pkg/repositories/gormimpl/dataset_test.go

@@ -34,6 +34,7 @@ func getTestDataset() models.Dataset {
 			Domain:  "testDomain",
 			Name:    "testName",
 			Version: "testVersion",
+			UUID:    "test-uuid",
 		},
 		SerializedMetadata: []byte{1, 2, 3},
 		PartitionKeys: []models.PartitionKey{
@@ -87,13 +88,14 @@ func TestCreateDatasetNoPartitions(t *testing.T) {
 
 	// Only match on queries that append expected filters
 	GlobalMock.NewMock().WithQuery(
-		`INSERT INTO "datasets" ("created_at","updated_at","deleted_at","project","name","domain","version","serialized_metadata") VALUES ($1,$2,$3,$4,$5,$6,$7,$8)`).WithCallback(
+		`INSERT INTO "datasets" ("created_at","updated_at","deleted_at","project","name","domain","version","uuid","serialized_metadata") VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)`).WithCallback(
 		func(s string, values []driver.NamedValue) {
 			assert.EqualValues(t, dataset.Project, values[3].Value)
 			assert.EqualValues(t, dataset.Name, values[4].Value)
 			assert.EqualValues(t, dataset.Domain, values[5].Value)
 			assert.EqualValues(t, dataset.Version, values[6].Value)
-			assert.EqualValues(t, dataset.SerializedMetadata, values[7].Value)
+			assert.EqualValues(t, dataset.UUID, values[7].Value)
+			assert.EqualValues(t, dataset.SerializedMetadata, values[8].Value)
 			datasetCreated = true
 		},
 	)
@@ -115,13 +117,14 @@ func TestCreateDataset(t *testing.T) {
 
 	// Only match on queries that append expected filters
 	GlobalMock.NewMock().WithQuery(
-		`INSERT INTO "datasets" ("created_at","updated_at","deleted_at","project","name","domain","version","serialized_metadata") VALUES ($1,$2,$3,$4,$5,$6,$7,$8)`).WithCallback(
+		`INSERT INTO "datasets" ("created_at","updated_at","deleted_at","project","name","domain","version","uuid","serialized_metadata") VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)`).WithCallback(
 		func(s string, values []driver.NamedValue) {
 			assert.EqualValues(t, dataset.Project, values[3].Value)
 			assert.EqualValues(t, dataset.Name, values[4].Value)
 			assert.EqualValues(t, dataset.Domain, values[5].Value)
 			assert.EqualValues(t, dataset.Version, values[6].Value)
-			assert.EqualValues(t, dataset.SerializedMetadata, values[7].Value)
+			assert.EqualValues(t, dataset.UUID, values[7].Value)
+			assert.EqualValues(t, dataset.SerializedMetadata, values[8].Value)
 			datasetCreated = true
 		},
 	).WithReply([]map[string]interface{}{{"dataset_uuid": getDatasetUUID()}})
@@ -159,7 +162,7 @@ func TestGetDataset(t *testing.T) {
 	GlobalMock.Logging = true
 
 	// Only match on queries that append expected filters
-	GlobalMock.NewMock().WithQuery(`SELECT * FROM "datasets" WHERE "datasets"."project" = $1 AND "datasets"."name" = $2 AND "datasets"."domain" = $3 AND "datasets"."version" = $4 ORDER BY "datasets"."created_at" LIMIT 1`).WithReply(expectedDatasetResponse)
+	GlobalMock.NewMock().WithQuery(`SELECT * FROM "datasets" WHERE "datasets"."project" = $1 AND "datasets"."name" = $2 AND "datasets"."domain" = $3 AND "datasets"."version" = $4 AND "datasets"."uuid" = $5 ORDER BY "datasets"."created_at" LIMIT 1`).WithReply(expectedDatasetResponse)
 
 	expectedPartitionKeyResponse := make([]map[string]interface{}, 0)
 	samplePartitionKey := make(map[string]interface{})
@@ -237,7 +240,7 @@ func TestCreateDatasetAlreadyExists(t *testing.T) {
 
 	// Only match on queries that append expected filters
 	GlobalMock.NewMock().WithQuery(
-		`INSERT INTO "datasets" ("created_at","updated_at","deleted_at","project","name","domain","version","serialized_metadata") VALUES ($1,$2,$3,$4,$5,$6,$7,$8)`).WithError(
+		`INSERT INTO "datasets" ("created_at","updated_at","deleted_at","project","name","domain","version","uuid","serialized_metadata") VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)`).WithError(
 		getAlreadyExistsErr(),
 	)
 

datacalog/pkg/repositories/handle.go

@@ -69,10 +69,6 @@ func (h *DBHandle) CreateDB(dbName string) error {
 }
 
 func (h *DBHandle) Migrate(ctx context.Context) error {
-	if h.db.Config.Dialector.Name() == config.Postgres {
-		logger.Infof(ctx, "Creating postgres extension uuid-ossp if it does not exist")
-		h.db.Exec("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\"")
-	}
 	if err := h.db.AutoMigrate(&models.Dataset{}); err != nil {
 		return err
 	}

datacalog/pkg/repositories/models/dataset.go

@@ -1,11 +1,16 @@
 package models
 
+import (
+	uuid "github.com/satori/go.uuid"
+	"gorm.io/gorm"
+)
+
 type DatasetKey struct {
 	Project string `gorm:"primary_key;"`                          // part of pkey, no index needed as it is first column in the pkey
 	Name    string `gorm:"primary_key;index:dataset_name_idx"`    // part of pkey and has separate index for filtering
 	Domain  string `gorm:"primary_key;index:dataset_domain_idx"`  // part of pkey and has separate index for filtering
 	Version string `gorm:"primary_key;index:dataset_version_idx"` // part of pkey and has separate index for filtering
-	UUID    string `gorm:"type:uuid;unique;default:uuid_generate_v4()"`
+	UUID    string `gorm:"type:uuid;unique;"`
 }
 
 type Dataset struct {
@@ -20,3 +25,12 @@ type PartitionKey struct {
 	DatasetUUID string `gorm:"type:uuid;primary_key"`
 	Name        string `gorm:"primary_key"`
 }
+
+// BeforeCreate so that we set the UUID in golang rather than from a DB function call
+func (dataset *Dataset) BeforeCreate(tx *gorm.DB) error {
+	if dataset.UUID == "" {
+		generated := uuid.NewV4()
+		tx.Model(dataset).Update("UUID", generated)
+	}
+	return nil
+}