github.com/khulnasoft-lab/tunnel-db@v0.0.0-20231117205118-74e1113bd007/pkg/db/vulnid.go (about)

     1  package db
     2  
     3  import (
     4  	bolt "go.etcd.io/bbolt"
     5  	"golang.org/x/xerrors"
     6  )
     7  
     8  const (
     9  	vulnerabilityIDBucket = "vulnerability-id"
    10  )
    11  
    12  func (dbc Config) PutVulnerabilityID(tx *bolt.Tx, vulnID string) error {
    13  	bucket, err := tx.CreateBucketIfNotExists([]byte(vulnerabilityIDBucket))
    14  	if err != nil {
    15  		return xerrors.Errorf("failed to create %s bucket: %w", vulnerabilityIDBucket, err)
    16  	}
    17  	return bucket.Put([]byte(vulnID), []byte("{}"))
    18  }
    19  
    20  func (dbc Config) ForEachVulnerabilityID(f func(tx *bolt.Tx, vulnID string) error) error {
    21  	err := db.Batch(func(tx *bolt.Tx) error {
    22  		bucket := tx.Bucket([]byte(vulnerabilityIDBucket))
    23  		if bucket == nil {
    24  			return xerrors.Errorf("no such bucket: %s", vulnerabilityIDBucket)
    25  		}
    26  		err := bucket.ForEach(func(vulnID, _ []byte) error {
    27  			if err := f(tx, string(vulnID)); err != nil {
    28  				return xerrors.Errorf("something wrong: %w", err)
    29  			}
    30  			return nil
    31  		})
    32  		if err != nil {
    33  			return xerrors.Errorf("error in for each: %w", err)
    34  		}
    35  		return nil
    36  	})
    37  	return err
    38  }
    39  
    40  func (dbc Config) DeleteVulnerabilityIDBucket() error {
    41  	return dbc.deleteBucket(vulnerabilityIDBucket)
    42  }