Fix upserting documents with custom fields

pkg/models/document.go

@@ -80,59 +80,8 @@ const (
 
 // BeforeSave is a hook used to find associations before saving.
 func (d *Document) BeforeSave(tx *gorm.DB) error {
-	// Get approvers.
-	var approvers []*User
-	for _, a := range d.Approvers {
-		if err := a.Get(tx); err != nil {
-			return fmt.Errorf("error getting approver: %w", err)
-		}
-		approvers = append(approvers, a)
-	}
-	d.Approvers = approvers
-
-	// Get contributors.
-	var contributors []*User
-	for _, c := range d.Contributors {
-		if err := c.FirstOrCreate(tx); err != nil {
-			return fmt.Errorf("error getting contributor: %w", err)
-		}
-		contributors = append(contributors, c)
-	}
-	d.Contributors = contributors
-
-	// Get custom fields.
-	var customFields []*DocumentCustomField
-	for _, c := range d.CustomFields {
-		if err := c.DocumentTypeCustomField.Get(tx); err != nil {
-			return fmt.Errorf("error getting document type custom field: %w", err)
-		}
-		c.DocumentTypeCustomFieldID = c.DocumentTypeCustomField.DocumentType.ID
-		customFields = append(customFields, c)
-	}
-	d.CustomFields = customFields
-
-	// Get document type.
-	dt := d.DocumentType
-	if err := dt.Get(tx); err != nil {
-		return fmt.Errorf("error getting document type: %w", err)
-	}
-	d.DocumentType = dt
-	d.DocumentTypeID = dt.ID
-
-	// Get owner.
-	if d.Owner != nil && d.Owner.EmailAddress != "" {
-		if err := d.Owner.Get(tx); err != nil {
-			return fmt.Errorf("error getting owner: %w", err)
-		}
-		d.OwnerID = &d.Owner.ID
-	}
-
-	// Get product.
-	if d.Product.Name != "" {
-		if err := d.Product.Get(tx); err != nil {
-			return fmt.Errorf("error getting product: %w", err)
-		}
-		d.ProductID = d.Product.ID
+	if err := d.getAssociations(tx); err != nil {
+		return fmt.Errorf("error getting associations: %w", err)
 	}
 
 	return nil
@@ -217,10 +166,19 @@ func (d *Document) Get(db *gorm.DB) error {
 		return err
 	}
 
-	return db.
+	if err := db.
 		Where(Document{GoogleFileID: d.GoogleFileID}).
 		Preload(clause.Associations).
-		First(&d).Error
+		First(&d).
+		Error; err != nil {
+		return err
+	}
+
+	if err := d.getAssociations(db); err != nil {
+		return fmt.Errorf("error getting associations: %w", err)
+	}
+
+	return nil
 }
 
 // GetLatestProductNumber gets the latest document number for a product.
@@ -311,7 +269,7 @@ func (d *Document) Upsert(db *gorm.DB) error {
 		}
 
 		if err := d.Get(tx); err != nil {
-			return fmt.Errorf("error getting the document after upsert")
+			return fmt.Errorf("error getting the document after upsert: %w", err)
 		}
 
 		return nil
@@ -351,6 +309,81 @@ func (d *Document) createAssocations(db *gorm.DB) error {
 	return nil
 }
 
+// getAssociations gets associations.
+func (d *Document) getAssociations(db *gorm.DB) error {
+	// Get approvers.
+	var approvers []*User
+	for _, a := range d.Approvers {
+		if err := a.Get(db); err != nil {
+			return fmt.Errorf("error getting approver: %w", err)
+		}
+		approvers = append(approvers, a)
+	}
+	d.Approvers = approvers
+
+	// Get contributors.
+	var contributors []*User
+	for _, c := range d.Contributors {
+		if err := c.FirstOrCreate(db); err != nil {
+			return fmt.Errorf("error getting contributor: %w", err)
+		}
+		contributors = append(contributors, c)
+	}
+	d.Contributors = contributors
+
+	// Get custom fields.
+	var customFields []*DocumentCustomField
+	for _, c := range d.CustomFields {
+		// If we already know the document type custom field ID, get the rest of its
+		// data.
+		if c.DocumentTypeCustomFieldID != 0 {
+			if err := db.
+				Model(DocumentTypeCustomField{
+					Model: gorm.Model{
+						ID: c.DocumentTypeCustomFieldID,
+					},
+				}).
+				First(&c.DocumentTypeCustomField).
+				Error; err != nil {
+				return fmt.Errorf("error getting document type custom field: %w", err)
+			}
+		}
+		c.DocumentTypeCustomField.DocumentType.Name = d.DocumentType.Name
+		if err := c.DocumentTypeCustomField.Get(db); err != nil {
+			return fmt.Errorf("error getting document type custom field: %w", err)
+		}
+		c.DocumentTypeCustomFieldID = c.DocumentTypeCustomField.DocumentType.ID
+		customFields = append(customFields, c)
+	}
+	d.CustomFields = customFields
+
+	// Get document type.
+	dt := d.DocumentType
+	if err := dt.Get(db); err != nil {
+		return fmt.Errorf("error getting document type: %w", err)
+	}
+	d.DocumentType = dt
+	d.DocumentTypeID = dt.ID
+
+	// Get owner.
+	if d.Owner != nil && d.Owner.EmailAddress != "" {
+		if err := d.Owner.Get(db); err != nil {
+			return fmt.Errorf("error getting owner: %w", err)
+		}
+		d.OwnerID = &d.Owner.ID
+	}
+
+	// Get product.
+	if d.Product.Name != "" {
+		if err := d.Product.Get(db); err != nil {
+			return fmt.Errorf("error getting product: %w", err)
+		}
+		d.ProductID = d.Product.ID
+	}
+
+	return nil
+}
+
 // replaceAssocations replaces assocations for a document.
 func (d *Document) replaceAssocations(db *gorm.DB) error {
 	// Replace approvers.

pkg/models/document_test.go

@@ -549,7 +549,6 @@ func TestDocumentModel(t *testing.T) {
 	*/
 
 	t.Run("Upsert Summary", func(t *testing.T) {
-		// assert, require := assert.New(t), require.New(t)
 		db, tearDownTest := setupTest(t, dsn)
 		defer tearDownTest(t)
 
@@ -627,6 +626,103 @@ func TestDocumentModel(t *testing.T) {
 			assert.Equal("summary2", d.Summary)
 		})
 	})
+
+	t.Run("Upsert a document with custom fields", func(t *testing.T) {
+		db, tearDownTest := setupTest(t, dsn)
+		defer tearDownTest(t)
+
+		t.Run("Create a document type", func(t *testing.T) {
+			_, require := assert.New(t), require.New(t)
+			dt := DocumentType{
+				Name:     "DT1",
+				LongName: "DocumentType1",
+			}
+			err := dt.FirstOrCreate(db)
+			require.NoError(err)
+		})
+
+		t.Run("Create a document type custom field",
+			func(t *testing.T) {
+				_, require := assert.New(t), require.New(t)
+
+				d := DocumentTypeCustomField{
+					Name: "CustomStringField",
+					DocumentType: DocumentType{
+						Name: "DT1",
+					},
+					Type: StringDocumentTypeCustomFieldType,
+				}
+				err := d.Upsert(db)
+				require.NoError(err)
+			})
+
+		t.Run("Create a product", func(t *testing.T) {
+			_, require := assert.New(t), require.New(t)
+			p := Product{
+				Name:         "Product1",
+				Abbreviation: "P1",
+			}
+			err := p.FirstOrCreate(db)
+			require.NoError(err)
+		})
+
+		t.Run("Create a document using Upsert", func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+			d := Document{
+				GoogleFileID: "fileID1",
+				Approvers: []*User{
+					{
+						EmailAddress: "[email protected]",
+					},
+					{
+						EmailAddress: "[email protected]",
+					},
+				},
+				CustomFields: []*DocumentCustomField{
+					{
+						DocumentTypeCustomField: DocumentTypeCustomField{
+							Name: "CustomStringField",
+							DocumentType: DocumentType{
+								Name: "DT1",
+							},
+						},
+						Value: "string value 1",
+					},
+				},
+				DocumentType: DocumentType{
+					Name: "DT1",
+				},
+				Product: Product{
+					Name: "Product1",
+				},
+			}
+			err := d.Upsert(db)
+			require.NoError(err)
+			assert.EqualValues(1, d.ID)
+			require.Len(d.CustomFields, 1)
+			assert.Equal("CustomStringField",
+				d.CustomFields[0].DocumentTypeCustomField.Name)
+			assert.Equal("DT1",
+				d.CustomFields[0].DocumentTypeCustomField.DocumentType.Name)
+			assert.Equal("string value 1", d.CustomFields[0].Value)
+		})
+
+		t.Run("Get the document", func(t *testing.T) {
+			assert, require := assert.New(t), require.New(t)
+			d := Document{
+				GoogleFileID: "fileID1",
+			}
+			err := d.Get(db)
+			require.NoError(err)
+			assert.EqualValues(1, d.ID)
+			require.Len(d.CustomFields, 1)
+			assert.Equal("CustomStringField",
+				d.CustomFields[0].DocumentTypeCustomField.Name)
+			assert.Equal("DT1",
+				d.CustomFields[0].DocumentTypeCustomField.DocumentType.Name)
+			assert.Equal("string value 1", d.CustomFields[0].Value)
+		})
+	})
 }
 
 func TestGetLatestProductNumber(t *testing.T) {