refactor: add reusable function that runs callbacks in a transaction (p…

…rotobom#31) * refactor: add reusable function that runs callbacks in a transaction refactor: - optionally pass annotations through protobom StoreOptions Signed-off-by: Jonathan Howard <[email protected]> * fix: return no results if no document IDs are passed in Signed-off-by: Jonathan Howard <[email protected]> * fix: return no results if no document has specified annotation values Signed-off-by: Jonathan Howard <[email protected]> * refactor: clear annotations on backend options Signed-off-by: Jonathan Howard <[email protected]> --------- Signed-off-by: Jonathan Howard <[email protected]>
jhoward-lm · Sep 10, 2024 · 9c3fdd7 · 9c3fdd7
1 parent 7ea764c
commit 9c3fdd7
Show file tree

Hide file tree

Showing 4 changed files with 374 additions and 404 deletions.
diff --git a/backends/ent/annotations.go b/backends/ent/annotations.go
@@ -16,39 +16,6 @@ import (
 	"github.com/protobom/storage/internal/backends/ent/predicate"
 )
 
-func (backend *Backend) createAnnotations(data ...*ent.Annotation) error {
-	tx, err := backend.txClient()
-	if err != nil {
-		return err
-	}
-
-	builders := []*ent.AnnotationCreate{}
-
-	for idx := range data {
-		builder := tx.Annotation.Create().
-			SetDocumentID(data[idx].DocumentID).
-			SetName(data[idx].Name).
-			SetValue(data[idx].Value).
-			SetIsUnique(data[idx].IsUnique)
-
-		builders = append(builders, builder)
-	}
-
-	err = tx.Annotation.CreateBulk(builders...).
-		OnConflict().
-		UpdateNewValues().
-		Exec(backend.ctx)
-	if err != nil && !ent.IsConstraintError(err) {
-		return rollback(tx, fmt.Errorf("creating annotations: %w", err))
-	}
-
-	if err := tx.Commit(); err != nil {
-		return rollback(tx, err)
-	}
-
-	return nil
-}
-
 // AddAnnotations applies multiple named annotation values to a single document.
 func (backend *Backend) AddAnnotations(documentID, name string, values ...string) error {
 	data := ent.Annotations{}
@@ -60,7 +27,7 @@ func (backend *Backend) AddAnnotations(documentID, name string, values ...string
 		})
 	}
 
-	return backend.createAnnotations(data...)
+	return backend.withTx(backend.saveAnnotations(data...))
 }
 
 // AddAnnotationToDocuments applies a single named annotation value to multiple documents.
@@ -74,7 +41,7 @@ func (backend *Backend) AddAnnotationToDocuments(name, value string, documentIDs
 		})
 	}
 
-	return backend.createAnnotations(data...)
+	return backend.withTx(backend.saveAnnotations(data...))
 }
 
 // ClearAnnotations removes all annotations from the specified documents.
@@ -83,21 +50,14 @@ func (backend *Backend) ClearAnnotations(documentIDs ...string) error {
 		return nil
 	}
 
-	tx, err := backend.txClient()
-	if err != nil {
-		return err
-	}
-
-	_, err = tx.Annotation.Delete().Where(annotation.DocumentIDIn(documentIDs...)).Exec(backend.ctx)
-	if err != nil {
-		return rollback(tx, fmt.Errorf("clearing annotations: %w", err))
-	}
-
-	if err := tx.Commit(); err != nil {
-		return rollback(tx, err)
-	}
+	return backend.withTx(func(tx *ent.Tx) error {
+		_, err := tx.Annotation.Delete().Where(annotation.DocumentIDIn(documentIDs...)).Exec(backend.ctx)
+		if err != nil {
+			return fmt.Errorf("clearing annotations: %w", err)
+		}
 
-	return nil
+		return nil
+	})
 }
 
 // GetDocumentAnnotations gets all annotations for the specified
@@ -143,6 +103,10 @@ func (backend *Backend) GetDocumentsByAnnotation(name string, values ...string)
 		return nil, fmt.Errorf("querying documents table: %w", err)
 	}
 
+	if len(ids) == 0 {
+		return []*sbom.Document{}, nil
+	}
+
 	return backend.GetDocumentsByID(ids...)
 }
 
@@ -173,29 +137,23 @@ func (backend *Backend) GetDocumentUniqueAnnotation(documentID, name string) (st
 // RemoveAnnotations removes all annotations with the specified name from
 // the document, limited to a set of annotation values if specified.
 func (backend *Backend) RemoveAnnotations(documentID, name string, values ...string) error {
-	tx, err := backend.txClient()
-	if err != nil {
-		return err
-	}
-
-	predicates := []predicate.Annotation{
-		annotation.DocumentIDEQ(documentID),
-		annotation.NameEQ(name),
-	}
-
-	if len(values) > 0 {
-		predicates = append(predicates, annotation.ValueIn(values...))
-	}
-
-	if _, err := tx.Annotation.Delete().Where(predicates...).Exec(backend.ctx); err != nil {
-		return rollback(tx, fmt.Errorf("removing annotations: %w", err))
-	}
-
-	if err := tx.Commit(); err != nil {
-		return rollback(tx, err)
-	}
-
-	return nil
+	return backend.withTx(
+		func(tx *ent.Tx) error {
+			predicates := []predicate.Annotation{
+				annotation.DocumentIDEQ(documentID),
+				annotation.NameEQ(name),
+			}
+
+			if len(values) > 0 {
+				predicates = append(predicates, annotation.ValueIn(values...))
+			}
+
+			if _, err := tx.Annotation.Delete().Where(predicates...).Exec(backend.ctx); err != nil {
+				return fmt.Errorf("removing annotations: %w", err)
+			}
+
+			return nil
+		})
 }
 
 // SetAnnotations explicitly sets the named annotations for the specified document.
@@ -209,10 +167,12 @@ func (backend *Backend) SetAnnotations(documentID, name string, values ...string
 
 // SetUniqueAnnotation sets a named annotation value that is unique to the specified document.
 func (backend *Backend) SetUniqueAnnotation(documentID, name, value string) error {
-	return backend.createAnnotations(&ent.Annotation{
-		DocumentID: documentID,
-		Name:       name,
-		Value:      value,
-		IsUnique:   true,
-	})
+	return backend.withTx(
+		backend.saveAnnotations(&ent.Annotation{
+			DocumentID: documentID,
+			Name:       name,
+			Value:      value,
+			IsUnique:   true,
+		}),
+	)
 }
diff --git a/backends/ent/backend.go b/backends/ent/backend.go
@@ -81,6 +81,22 @@ func (backend *Backend) Debug() *Backend {
 	return backend
 }
 
+func (backend *Backend) WithAnnotation(name, value string, unique bool) *Backend {
+	backend.Options.Annotations = append(backend.Options.Annotations, &Annotation{
+		Name:     name,
+		Value:    value,
+		IsUnique: unique,
+	})
+
+	return backend
+}
+
+func (backend *Backend) WithAnnotations(annotations Annotations) *Backend {
+	backend.Options.Annotations = append(backend.Options.Annotations, annotations...)
+
+	return backend
+}
+
 func (backend *Backend) WithBackendOptions(opts *BackendOptions) *Backend {
 	backend.Options = opts
 
@@ -93,17 +109,31 @@ func (backend *Backend) WithDatabaseFile(file string) *Backend {
 	return backend
 }
 
-func (backend *Backend) txClient() (*ent.Tx, error) {
+func (backend *Backend) withTx(fns ...txFunc) error {
 	if backend.client == nil {
-		return nil, fmt.Errorf("%w", errUninitializedClient)
+		return fmt.Errorf("%w", errUninitializedClient)
 	}
 
 	tx, err := backend.client.Tx(backend.ctx)
 	if err != nil {
-		return nil, fmt.Errorf("creating transactional client: %w", err)
+		return fmt.Errorf("creating transactional client: %w", err)
 	}
 
 	backend.ctx = ent.NewTxContext(backend.ctx, tx)
 
-	return tx, nil
+	for _, fn := range fns {
+		if err := fn(tx); err != nil {
+			if rollbackErr := tx.Rollback(); rollbackErr != nil {
+				err = fmt.Errorf("%w: rolling back transaction: %w", err, rollbackErr)
+			}
+
+			return err
+		}
+	}
+
+	if err := tx.Commit(); err != nil {
+		return fmt.Errorf("committing transaction: %w", err)
+	}
+
+	return nil
 }
diff --git a/backends/ent/options.go b/backends/ent/options.go
@@ -7,6 +7,8 @@ package ent
 
 import (
 	"errors"
+
+	"github.com/protobom/storage/internal/backends/ent"
 )
 
 // Enable SQLite foreign key support.
@@ -18,11 +20,20 @@ var (
 )
 
 type (
+	// Annotation is the model entity for the Annotation schema.
+	Annotation = ent.Annotation
+
+	// Annotations is a parsable slice of Annotation.
+	Annotations = ent.Annotations
+
 	// BackendOptions contains options specific to the protobom ent backend.
 	BackendOptions struct {
 		// DatabaseFile is the file path of the SQLite database to be created.
 		DatabaseFile string
 
+		// Annotations is a slice of annotations to apply to stored document.
+		Annotations
+
 		// Debug configures the ent client to output all SQL statements during execution.
 		Debug bool
 	}