feat: introduce database type (#982)

See https://github.com/ooni/probe/issues/2352

Co-authored-by: decfox <decfox@github.com>
This commit is contained in:
DecFox
2022-11-16 20:21:41 +05:30
committed by GitHub
parent 6b01264373
commit 28aabe0947
14 changed files with 144 additions and 94 deletions
+2 -3
View File
@@ -7,7 +7,6 @@ import (
"github.com/apex/log"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/cli/root"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/output"
"github.com/ooni/probe-cli/v3/internal/database"
)
func init() {
@@ -20,7 +19,7 @@ func init() {
return err
}
if *resultID > 0 {
measurements, err := database.ListMeasurements(probeCLI.DB(), *resultID)
measurements, err := probeCLI.DB().ListMeasurements(*resultID)
if err != nil {
log.WithError(err).Error("failed to list measurements")
return err
@@ -63,7 +62,7 @@ func init() {
}
output.MeasurementSummary(msmtSummary)
} else {
doneResults, incompleteResults, err := database.ListResults(probeCLI.DB())
doneResults, incompleteResults, err := probeCLI.DB().ListResults()
if err != nil {
log.WithError(err).Error("failed to list results")
return err
+1 -1
View File
@@ -25,7 +25,7 @@ func init() {
log.WithError(err).Error("failed to close the DB")
return err
}
if *force == true {
if *force {
os.RemoveAll(ctx.Home())
log.Infof("Deleted %s", ctx.Home())
} else {
+6 -6
View File
@@ -12,7 +12,7 @@ import (
"github.com/upper/db/v4"
)
func deleteAll(sess db.Session, skipInteractive bool) error {
func deleteAll(d *database.Database, skipInteractive bool) error {
if skipInteractive == false {
answer := ""
confirm := &survey.Select{
@@ -25,21 +25,21 @@ func deleteAll(sess db.Session, skipInteractive bool) error {
return errors.New("canceled by user")
}
}
doneResults, incompleteResults, err := database.ListResults(sess)
doneResults, incompleteResults, err := d.ListResults()
if err != nil {
log.WithError(err).Error("failed to list results")
return err
}
cnt := 0
for _, result := range incompleteResults {
err = database.DeleteResult(sess, result.Result.ID)
err = d.DeleteResult(result.Result.ID)
if err == db.ErrNoMoreRows {
log.WithError(err).Errorf("failed to delete result #%d", result.Result.ID)
}
cnt++
}
for _, result := range doneResults {
err = database.DeleteResult(sess, result.Result.ID)
err = d.DeleteResult(result.Result.ID)
if err == db.ErrNoMoreRows {
log.WithError(err).Errorf("failed to delete result #%d", result.Result.ID)
}
@@ -68,7 +68,7 @@ func init() {
}
if *yes == true {
err = database.DeleteResult(ctx.DB(), *resultID)
err = ctx.DB().DeleteResult(*resultID)
if err == db.ErrNoMoreRows {
return errors.New("result not found")
}
@@ -84,7 +84,7 @@ func init() {
if answer == "false" {
return errors.New("canceled by user")
}
err = database.DeleteResult(ctx.DB(), *resultID)
err = ctx.DB().DeleteResult(*resultID)
if err == db.ErrNoMoreRows {
return errors.New("result not found")
}
+1 -2
View File
@@ -5,7 +5,6 @@ import (
"github.com/apex/log"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/cli/root"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/output"
"github.com/ooni/probe-cli/v3/internal/database"
)
func init() {
@@ -17,7 +16,7 @@ func init() {
log.WithError(err).Error("failed to initialize root context")
return err
}
msmt, err := database.GetMeasurementJSON(ctx.DB(), *msmtID)
msmt, err := ctx.DB().GetMeasurementJSON(*msmtID)
if err != nil {
log.Errorf("error: %v", err)
return err
+1 -1
View File
@@ -25,7 +25,7 @@ func (n DNSCheck) lookupURLs(ctl *Controller) ([]string, error) {
if err != nil {
return nil, err
}
return ctl.BuildAndSetInputIdxMap(ctl.Probe.DB(), testlist)
return ctl.BuildAndSetInputIdxMap(testlist)
}
// Run starts the nettest.
+13 -13
View File
@@ -14,7 +14,6 @@ import (
engine "github.com/ooni/probe-cli/v3/internal/engine"
"github.com/ooni/probe-cli/v3/internal/model"
"github.com/pkg/errors"
"github.com/upper/db/v4"
)
// Nettest interface. Every Nettest should implement this.
@@ -90,14 +89,13 @@ type Controller struct {
// - on success, a list of strings containing URLs to test;
//
// - on failure, an error.
func (c *Controller) BuildAndSetInputIdxMap(
sess db.Session, testlist []model.OOAPIURLInfo) ([]string, error) {
func (c *Controller) BuildAndSetInputIdxMap(testlist []model.OOAPIURLInfo) ([]string, error) {
var urls []string
urlIDMap := make(map[int64]int64)
for idx, url := range testlist {
log.Debugf("Going over URL %d", idx)
urlID, err := database.CreateOrUpdateURL(
sess, url.URL, url.CategoryCode, url.CountryCode,
urlID, err := c.Probe.DB().CreateOrUpdateURL(
url.URL, url.CategoryCode, url.CountryCode,
)
if err != nil {
log.Error("failed to add to the URL table")
@@ -124,6 +122,7 @@ func (c *Controller) SetNettestIndex(i, n int) {
// This function will continue to run in most cases but will
// immediately halt if something's wrong with the file system.
func (c *Controller) Run(builder model.ExperimentBuilder, inputs []string) error {
db := c.Probe.DB()
// This will configure the controller as handler for the callbacks
// called by ooni/probe-engine/experiment.Experiment.
builder.SetCallbacks(model.ExperimentCallbacks(c))
@@ -168,6 +167,7 @@ func (c *Controller) Run(builder model.ExperimentBuilder, inputs []string) error
log.Debug("disabling maxRuntime with user-provided input")
maxRuntime = 0
}
sess := db.Session()
start := time.Now()
c.ntStartTime = start
for idx, input := range inputs {
@@ -187,8 +187,8 @@ func (c *Controller) Run(builder model.ExperimentBuilder, inputs []string) error
urlID = sql.NullInt64{Int64: c.inputIdxMap[idx64], Valid: true}
}
msmt, err := database.CreateMeasurement(
c.Probe.DB(), reportID, exp.Name(), c.res.MeasurementDir, idx, resultID, urlID,
msmt, err := db.CreateMeasurement(
reportID, exp.Name(), c.res.MeasurementDir, idx, resultID, urlID,
)
if err != nil {
return errors.Wrap(err, "failed to create measurement")
@@ -201,7 +201,7 @@ func (c *Controller) Run(builder model.ExperimentBuilder, inputs []string) error
measurement, err := exp.MeasureWithContext(context.Background(), input)
if err != nil {
log.WithError(err).Debug(color.RedString("failure.measurement"))
if err := c.msmts[idx64].Failed(c.Probe.DB(), err.Error()); err != nil {
if err := c.msmts[idx64].Failed(sess, err.Error()); err != nil {
return errors.Wrap(err, "failed to mark measurement as failed")
}
// Since https://github.com/ooni/probe-cli/pull/527, the Measure
@@ -221,10 +221,10 @@ func (c *Controller) Run(builder model.ExperimentBuilder, inputs []string) error
// bit of a spew in the logs, perhaps, but stopping seems less efficient.
if err := exp.SubmitAndUpdateMeasurementContext(context.Background(), measurement); err != nil {
log.Debug(color.RedString("failure.measurement_submission"))
if err := c.msmts[idx64].UploadFailed(c.Probe.DB(), err.Error()); err != nil {
if err := c.msmts[idx64].UploadFailed(sess, err.Error()); err != nil {
return errors.Wrap(err, "failed to mark upload as failed")
}
} else if err := c.msmts[idx64].UploadSucceeded(c.Probe.DB()); err != nil {
} else if err := c.msmts[idx64].UploadSucceeded(sess); err != nil {
return errors.Wrap(err, "failed to mark upload as succeeded")
} else {
// Everything went OK, don't save to disk
@@ -238,7 +238,7 @@ func (c *Controller) Run(builder model.ExperimentBuilder, inputs []string) error
}
}
if err := c.msmts[idx64].Done(c.Probe.DB()); err != nil {
if err := c.msmts[idx64].Done(sess); err != nil {
return errors.Wrap(err, "failed to mark measurement as done")
}
@@ -253,11 +253,11 @@ func (c *Controller) Run(builder model.ExperimentBuilder, inputs []string) error
continue
}
log.Debugf("Fetching: %d %v", idx, c.msmts[idx64])
if err := database.AddTestKeys(c.Probe.DB(), c.msmts[idx64], tk); err != nil {
if err := db.AddTestKeys(c.msmts[idx64], tk); err != nil {
return errors.Wrap(err, "failed to add test keys to summary")
}
}
database.UpdateUploadedStatus(c.Probe.DB(), c.res)
db.UpdateUploadedStatus(c.res)
log.Debugf("status.end")
return nil
}
@@ -8,7 +8,6 @@ import (
"testing"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/ooni"
"github.com/ooni/probe-cli/v3/internal/database"
"github.com/ooni/probe-cli/v3/internal/model"
)
@@ -53,11 +52,12 @@ func TestRun(t *testing.T) {
if err != nil {
t.Fatal(err)
}
network, err := database.CreateNetwork(probe.DB(), sess)
db := probe.DB()
network, err := db.CreateNetwork(sess)
if err != nil {
t.Fatal(err)
}
res, err := database.CreateResult(probe.DB(), probe.Home(), "middlebox", network.ID)
res, err := db.CreateResult(probe.Home(), "middlebox", network.ID)
if err != nil {
t.Fatal(err)
}
+6 -6
View File
@@ -8,7 +8,6 @@ import (
"github.com/apex/log"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/ooni"
"github.com/ooni/probe-cli/v3/internal/database"
"github.com/ooni/probe-cli/v3/internal/model"
"github.com/pkg/errors"
)
@@ -72,7 +71,8 @@ func RunGroup(config RunGroupConfig) error {
log.WithError(err).Error("Failed to lookup the location of the probe")
return err
}
network, err := database.CreateNetwork(config.Probe.DB(), sess)
db := config.Probe.DB()
network, err := db.CreateNetwork(sess)
if err != nil {
log.WithError(err).Error("Failed to create the network row")
return err
@@ -89,8 +89,8 @@ func RunGroup(config RunGroupConfig) error {
}
log.Debugf("Running test group %s", group.Label)
result, err := database.CreateResult(
config.Probe.DB(), config.Probe.Home(), config.GroupName, network.ID)
result, err := db.CreateResult(
config.Probe.Home(), config.GroupName, network.ID)
if err != nil {
log.Errorf("DB result error: %s", err)
return err
@@ -131,8 +131,8 @@ func RunGroup(config RunGroupConfig) error {
if err != nil {
os.Remove(result.MeasurementDir)
}
if err = result.Finished(config.Probe.DB()); err != nil {
dbSess := db.Session()
if err = result.Finished(dbSess); err != nil {
return err
}
return nil
@@ -25,7 +25,7 @@ func (n STUNReachability) lookupURLs(ctl *Controller) ([]string, error) {
if err != nil {
return nil, err
}
return ctl.BuildAndSetInputIdxMap(ctl.Probe.DB(), testlist)
return ctl.BuildAndSetInputIdxMap(testlist)
}
// Run starts the nettest.
@@ -31,7 +31,7 @@ func (n WebConnectivity) lookupURLs(ctl *Controller, categories []string) ([]str
if err != nil {
return nil, err
}
return ctl.BuildAndSetInputIdxMap(ctl.Probe.DB(), testlist)
return ctl.BuildAndSetInputIdxMap(testlist)
}
// WebConnectivity test implementation
+4 -5
View File
@@ -19,7 +19,6 @@ import (
"github.com/ooni/probe-cli/v3/internal/legacy/assetsdir"
"github.com/ooni/probe-cli/v3/internal/model"
"github.com/pkg/errors"
"github.com/upper/db/v4"
)
// DefaultSoftwareName is the default software name.
@@ -33,7 +32,7 @@ var logger = log.WithFields(log.Fields{
// ProbeCLI is the OONI Probe CLI context.
type ProbeCLI interface {
Config() *config.Config
DB() db.Session
DB() *database.Database
IsBatch() bool
Home() string
TempDir() string
@@ -53,7 +52,7 @@ type ProbeEngine interface {
// Probe contains the ooniprobe CLI context.
type Probe struct {
config *config.Config
db db.Session
db *database.Database
isBatch bool
home string
@@ -86,7 +85,7 @@ func (p *Probe) Config() *config.Config {
}
// DB returns the database we're using
func (p *Probe) DB() db.Session {
func (p *Probe) DB() *database.Database {
return p.db
}
@@ -180,7 +179,7 @@ func (p *Probe) Init(softwareName, softwareVersion, proxy string) error {
p.dbPath = utils.DBDir(p.home, "main")
log.Debugf("Connecting to database sqlite3://%s", p.dbPath)
db, err := database.Connect(p.dbPath)
db, err := database.Open(p.dbPath)
if err != nil {
return err
}
+3 -3
View File
@@ -8,8 +8,8 @@ import (
"github.com/apex/log"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/config"
"github.com/ooni/probe-cli/v3/cmd/ooniprobe/internal/ooni"
"github.com/ooni/probe-cli/v3/internal/database"
"github.com/ooni/probe-cli/v3/internal/model"
"github.com/upper/db/v4"
)
// FakeOutput allows to fake the output package.
@@ -28,7 +28,7 @@ func (fo *FakeOutput) SectionTitle(s string) {
// FakeProbeCLI fakes ooni.ProbeCLI
type FakeProbeCLI struct {
FakeConfig *config.Config
FakeDB db.Session
FakeDB *database.Database
FakeIsBatch bool
FakeHome string
FakeTempDir string
@@ -42,7 +42,7 @@ func (cli *FakeProbeCLI) Config() *config.Config {
}
// DB implements ProbeCLI.DB
func (cli *FakeProbeCLI) DB() db.Session {
func (cli *FakeProbeCLI) DB() *database.Database {
return cli.FakeDB
}