ooni-probe-cli/internal/database/models.go

package database

import (
	"os"
	"path/filepath"
	"time"

	"github.com/ooni/probe-cli/nettests/summary"
	"github.com/pkg/errors"
	"upper.io/db.v3/lib/sqlbuilder"
)

// Network represents a network tested by the user
type Network struct {
	ID          int64  `db:"id"`
	NetworkName string `db:"network_name"`
	IP          string `db:"ip"`
	ASN         int    `db:"asn"`
	CountryCode string `db:"country_code"`
}

type URL struct {
	ID           int64  `db:"id"`
	URL          int64  `db:"url"`
	CategoryCode string `db:"category_code"`
	CountryCode  string `db:"country_code"`
}

// Measurement model
type Measurement struct {
	ID               int64     `db:"id"`
	TestName         string    `db:"test_name"`
	StartTime        time.Time `db:"start_time"`
	Runtime          float64   `db:"runtime"`    // Fractional number of seconds
	NetworkID        int64     `db:"network_id"` // Used to include a Network
	IsDone           bool      `db:"is_done"`
	IsUploaded       bool      `db:"is_uploaded"`
	IsFailed         string    `db:"is_failed"`
	FailureMsg       string    `db:"failure_msg"`
	IsUploadFailed   bool      `db:"is_upload_failed"`
	UploadFailureMsg string    `db:"upload_failure_msg"`
	IsRerun          bool      `db:"is_rerun"`
	ReportID         string    `db:"report_id"`
	URLID            string    `db:"url_id"` // Used to reference URL
	MeasurementID    int64     `db:"measurement_id"`
	IsAnomaly        bool      `db:"is_anomaly"`
	TestKeys         struct{}  `db:"test_keys"`
	ResultID         int64     `db:"result_id"`
	ReportFilePath   string    `db:"report_file_path"`
}

// Result model
type Result struct {
	ID             int64     `db:"id"`
	TestGroupName  string    `db:"test_group_name"`
	StartTime      time.Time `db:"start_time"`
	Runtime        float64   `db:"runtime"` // Runtime is expressed in fractional seconds
	IsViewed       bool      `db:"is_viewed"`
	IsDone         bool      `db:"is_done"`
	DataUsageUp    int64     `db:"data_usage_up"`
	DataUsageDown  int64     `db:"data_usage_down"`
	MeasurementDir string    `db:"measurement_dir"`
}

// Finished marks the result as done and sets the runtime
func (r *Result) Finished(sess sqlbuilder.Database, makeSummary summary.ResultSummaryFunc) error {
	if r.IsDone == true || r.Runtime != 0 {
		return errors.New("Result is already finished")
	}
	r.Runtime = time.Now().UTC().Sub(r.StartTime).Seconds()
	r.IsDone = true

	err := sess.Collection("results").Find("id", r.ID).Update(r)
	if err != nil {
		return errors.Wrap(err, "updating finished result")
	}
	return nil
}

// SetGeoIPInfo for the Measurement
func (m *Measurement) SetGeoIPInfo() error {
	return nil
}

// Failed writes the error string to the measurement
func (m *Measurement) Failed(sess sqlbuilder.Database, failure string) error {
	m.FailureMsg = failure
	err := sess.Collection("measurements").Find("id", m.ID).Update(m)
	if err != nil {
		return errors.Wrap(err, "updating measurement")
	}
	return nil
}

// Done marks the measurement as completed
func (m *Measurement) Done(sess sqlbuilder.Database) error {
	runtime := time.Now().UTC().Sub(m.StartTime)
	m.Runtime = runtime.Seconds()
	m.IsDone = true

	err := sess.Collection("measurements").Find("id", m.ID).Update(m)
	if err != nil {
		return errors.Wrap(err, "updating measurement")
	}
	return nil
}

// UploadFailed writes the error string for the upload failure to the measurement
func (m *Measurement) UploadFailed(sess sqlbuilder.Database, failure string) error {
	m.UploadFailureMsg = failure
	m.IsUploaded = false

	err := sess.Collection("measurements").Find("id", m.ID).Update(m)
	if err != nil {
		return errors.Wrap(err, "updating measurement")
	}
	return nil
}

// UploadSucceeded writes the error string for the upload failure to the measurement
func (m *Measurement) UploadSucceeded(sess sqlbuilder.Database) error {
	m.IsUploaded = true

	err := sess.Collection("measurements").Find("id", m.ID).Update(m)
	if err != nil {
		return errors.Wrap(err, "updating measurement")
	}
	return nil
}

// WriteSummary writes the summary to the measurement
func (m *Measurement) WriteSummary(sess sqlbuilder.Database, summary string) error {
	// XXX remove m.Summary = summary

	err := sess.Collection("measurements").Find("id", m.ID).Update(m)
	if err != nil {
		return errors.Wrap(err, "updating measurement")
	}
	return nil
}

// AddToResult adds a measurement to a result
func (m *Measurement) AddToResult(sess sqlbuilder.Database, result *Result) error {
	var err error

	m.ResultID = result.ID
	finalPath := filepath.Join(result.MeasurementDir,
		filepath.Base(m.ReportFilePath))

	// If the finalPath already exists, it means it has already been moved there.
	// This happens in multi input reports
	if _, err = os.Stat(finalPath); os.IsNotExist(err) {
		err = os.Rename(m.ReportFilePath, finalPath)
		if err != nil {
			return errors.Wrap(err, "moving report file")
		}
	}
	m.ReportFilePath = finalPath

	err = sess.Collection("measurements").Find("id", m.ID).Update(m)
	if err != nil {
		return errors.Wrap(err, "updating measurement")
	}
	return nil
}