Refactor some database code and added delete function

This commit is contained in:
David J. Allen 2023-09-29 09:24:18 -06:00
parent 54dde2dda6
commit 2a6ffd16bb

View file

@ -8,12 +8,7 @@ import (
"github.com/jmoiron/sqlx" "github.com/jmoiron/sqlx"
) )
func InsertProbeResults(path string, states *[]magellan.BMCProbeResult) error { func CreateProbeResultsIfNotExists(path string) (*sqlx.DB, error) {
if states == nil {
return fmt.Errorf("states == nil")
}
// create database if it doesn't already exist
schema := ` schema := `
CREATE TABLE IF NOT EXISTS magellan_scanned_ports ( CREATE TABLE IF NOT EXISTS magellan_scanned_ports (
host TEXT NOT NULL, host TEXT NOT NULL,
@ -25,9 +20,22 @@ func InsertProbeResults(path string, states *[]magellan.BMCProbeResult) error {
` `
db, err := sqlx.Open("sqlite3", path) db, err := sqlx.Open("sqlite3", path)
if err != nil { if err != nil {
return fmt.Errorf("could not open database: %v", err) return nil, fmt.Errorf("could not open database: %v", err)
} }
db.MustExec(schema) db.MustExec(schema)
return db, nil
}
func InsertProbeResults(path string, states *[]magellan.BMCProbeResult) error {
if states == nil {
return fmt.Errorf("states == nil")
}
// create database if it doesn't already exist
db, err := CreateProbeResultsIfNotExists(path)
if err != nil {
return err
}
// insert all probe states into db // insert all probe states into db
tx := db.MustBegin() tx := db.MustBegin()
@ -46,6 +54,30 @@ func InsertProbeResults(path string, states *[]magellan.BMCProbeResult) error {
return nil return nil
} }
func DeleteProbeResults(path string, results *[]magellan.BMCProbeResult) error {
if results == nil {
return fmt.Errorf("no probe results found")
}
db, err := sqlx.Open("sqlite3", path)
if err != nil {
return fmt.Errorf("could not open database: %v", err)
}
tx := db.MustBegin()
for _, state := range *results {
sql := `DELETE FROM magellan_scanned_ports WHERE host = :host, port = :port;`
_, err := tx.NamedExec(sql, &state)
if err != nil {
fmt.Printf("could not execute transaction: %v\n", err)
}
}
err = tx.Commit()
if err != nil {
return fmt.Errorf("could not commit transaction: %v", err)
}
return nil
}
func GetProbeResults(path string) ([]magellan.BMCProbeResult, error) { func GetProbeResults(path string) ([]magellan.BMCProbeResult, error) {
db, err := sqlx.Open("sqlite3", path) db, err := sqlx.Open("sqlite3", path)
if err != nil { if err != nil {