dbstore: rename Tx to Begin, add Tx which accepts callback and handles commit/rollback

1 files changed, 104 insertions, 86 deletions

diff --git a/dbstore/dbstore.go b/dbstore/dbstore.go
index 21a80a3..ccddd72 100644
--- a/dbstore/dbstore.go
+++ b/dbstore/dbstore.go
@@ -83,11 +83,35 @@ var addCpeDictionaryQueryIds = []string {
   "cpe/insert-ref",
 }
 
-// create new transaction
-func (me DbStore) Tx(ctx context.Context, queryIds []string) (Tx, error) {
+// Begin new transaction and create prepared statements.
+func (me DbStore) Begin(ctx context.Context, queryIds []string) (Tx, error) {
   return newTx(ctx, me.db, queryIds)
 }
 
+// Create a transaction, pass it to callback, then commit the transaction
+// if the callback returns success and rollback the transaction if the
+// callback returns an error.
+func (me DbStore) Tx(ctx context.Context, queryIds []string, fn func(Tx) error) error {
+  // create transaction
+  tx, err := me.Begin(ctx, queryIds)
+  if err != nil {
+    return err
+  }
+
+  if err := fn(tx); err != nil {
+    // rollback
+    if rb_err := tx.Rollback(); rb_err != nil {
+      return rb_err
+    }
+
+    // return error
+    return err
+  } else {
+    // commit transaction
+    return tx.Commit()
+  }
+}
+
 // Execute query and invoke callback with each row of result.
 func (me DbStore) Query(
   ctx context.Context,
@@ -139,44 +163,41 @@ func (me DbStore) AddCpeDictionary(ctx context.Context, dict cpedict.Dictionary)
     return err
   }
 
-  tx, err := me.Tx(ctx, addCpeDictionaryQueryIds)
-  if err != nil {
-    return err
-  }
-
-  // add items
-  for _, item := range(dict.Items) {
-    // add cpe
-    rs, err := tx.Exec(ctx, "cpe/insert", item.CpeUri, item.Cpe23Item.Name)
-    if err != nil {
-      return err
-    }
-
-    // get last row ID
-    id, err := rs.LastInsertId()
-    if err != nil {
-      return err
-    }
-
-    // add titles
-    for _, title := range(item.Titles) {
-      _, err := tx.Exec(ctx, "cpe/insert-title", id, title.Lang, title.Text)
+  return me.Tx(ctx, addCpeDictionaryQueryIds, func(tx Tx) error {
+    // add items
+    for _, item := range(dict.Items) {
+      // add cpe
+      rs, err := tx.Exec(ctx, "cpe/insert", item.CpeUri, item.Cpe23Item.Name)
       if err != nil {
         return err
       }
-    }
 
-    // add refs
-    for _, ref := range(item.References) {
-      _, err := tx.Exec(ctx, "cpe/insert-ref", id, ref.Href, ref.Text)
+      // get last row ID
+      id, err := rs.LastInsertId()
       if err != nil {
         return err
       }
+
+      // add titles
+      for _, title := range(item.Titles) {
+        _, err := tx.Exec(ctx, "cpe/insert-title", id, title.Lang, title.Text)
+        if err != nil {
+          return err
+        }
+      }
+
+      // add refs
+      for _, ref := range(item.References) {
+        _, err := tx.Exec(ctx, "cpe/insert-ref", id, ref.Href, ref.Text)
+        if err != nil {
+          return err
+        }
+      }
     }
-  }
 
-  // commit changes, return result
-  return tx.Commit()
+    // return success
+    return nil
+  })
 }
 
 // search CPEs
@@ -227,74 +248,71 @@ func (me DbStore) AddCpeMatches(ctx context.Context, matches cpematch.Matches) e
   }
 
   // begin transaction
-  tx, err := me.Tx(ctx, addCpeMatchesQueryIds)
-  if err != nil {
-    return err
-  }
-
-  // add matches
-  for _, m := range(matches.Matches) {
-    // add cpe
-    rs, err := tx.Exec(ctx, "cpe-match/insert", m.Cpe23Uri, m.Cpe22Uri)
-    if err != nil {
-      return err
-    }
-
-    // get last row ID
-    id, err := rs.LastInsertId()
-    if err != nil {
-      return err
-    }
-
-    // add vulnerable
-    if m.Vulnerable != nil {
-      _, err := tx.Exec(ctx, "cpe-match/insert-vulnerable", id, *m.Vulnerable)
-      if  err != nil {
+  return me.Tx(ctx, addCpeMatchesQueryIds, func(tx Tx) error {
+    // add matches
+    for _, m := range(matches.Matches) {
+      // add cpe
+      rs, err := tx.Exec(ctx, "cpe-match/insert", m.Cpe23Uri, m.Cpe22Uri)
+      if err != nil {
         return err
       }
-    }
 
-    // add version minimum
-    if m.VersionStartIncluding != "" && m.VersionStartExcluding != "" {
-      return fmt.Errorf("cannot specify both VersionStartIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionStartIncluding, m.VersionStartExcluding)
-    } else if m.VersionStartIncluding != "" {
-      _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, true, m.VersionStartIncluding)
-      if  err != nil {
+      // get last row ID
+      id, err := rs.LastInsertId()
+      if err != nil {
         return err
       }
-    } else if m.VersionStartExcluding != "" {
-      _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, false, m.VersionStartExcluding)
-      if  err != nil {
-        return err
+
+      // add vulnerable
+      if m.Vulnerable != nil {
+        _, err := tx.Exec(ctx, "cpe-match/insert-vulnerable", id, *m.Vulnerable)
+        if  err != nil {
+          return err
+        }
       }
-    }
 
-    // add version maximum
-    if m.VersionEndIncluding != "" && m.VersionEndExcluding != "" {
-      return fmt.Errorf("cannot specify both VersionEndIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionEndIncluding, m.VersionEndExcluding)
-    } else if m.VersionEndIncluding != "" {
-      _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, true, m.VersionEndIncluding)
-      if  err != nil {
-        return err
+      // add version minimum
+      if m.VersionStartIncluding != "" && m.VersionStartExcluding != "" {
+        return fmt.Errorf("cannot specify both VersionStartIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionStartIncluding, m.VersionStartExcluding)
+      } else if m.VersionStartIncluding != "" {
+        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, true, m.VersionStartIncluding)
+        if  err != nil {
+          return err
+        }
+      } else if m.VersionStartExcluding != "" {
+        _, err := tx.Exec(ctx, "cpe-match/insert-version-min", id, false, m.VersionStartExcluding)
+        if  err != nil {
+          return err
+        }
       }
-    } else if m.VersionEndExcluding != "" {
-      _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, false, m.VersionEndExcluding)
-      if  err != nil {
-        return err
+
+      // add version maximum
+      if m.VersionEndIncluding != "" && m.VersionEndExcluding != "" {
+        return fmt.Errorf("cannot specify both VersionEndIncluding = \"%s\", VersionEndIncluding \"%s\"", m.VersionEndIncluding, m.VersionEndExcluding)
+      } else if m.VersionEndIncluding != "" {
+        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, true, m.VersionEndIncluding)
+        if  err != nil {
+          return err
+        }
+      } else if m.VersionEndExcluding != "" {
+        _, err := tx.Exec(ctx, "cpe-match/insert-version-max", id, false, m.VersionEndExcluding)
+        if  err != nil {
+          return err
+        }
       }
-    }
 
-    // add names
-    for _, name := range(m.Names) {
-      _, err := tx.Exec(ctx, "cpe-match/insert-name", id, name.Cpe23Uri, name.Cpe22Uri)
-      if err != nil {
-        return err
+      // add names
+      for _, name := range(m.Names) {
+        _, err := tx.Exec(ctx, "cpe-match/insert-name", id, name.Cpe23Uri, name.Cpe22Uri)
+        if err != nil {
+          return err
+        }
       }
     }
-  }
 
-  // commit changes, return result
-  return tx.Commit()
+    // return success
+    return nil
+  })
 }
 
 // search CPE matches