chore: update db utils (#2177)

cmd/memos.go

@@ -42,11 +42,16 @@ var (
 		Run: func(_cmd *cobra.Command, _args []string) {
 			ctx, cancel := context.WithCancel(context.Background())
 			db := db.NewDB(profile)
-			if err := db.Open(ctx); err != nil {
+			if err := db.Open(); err != nil {
 				cancel()
 				log.Error("failed to open db", zap.Error(err))
 				return
 			}
+			if err := db.Migrate(ctx); err != nil {
+				cancel()
+				log.Error("failed to migrate db", zap.Error(err))
+				return
+			}
 
 			store := store.New(db.DBInstance, profile)
 			s, err := server.NewServer(ctx, profile, store)

cmd/mvrss.go

@@ -39,10 +39,14 @@ var (
 			}
 
 			db := db.NewDB(profile)
-			if err := db.Open(ctx); err != nil {
+			if err := db.Open(); err != nil {
 				fmt.Printf("failed to open db, error: %+v\n", err)
 				return
 			}
+			if err := db.Migrate(ctx); err != nil {
+				fmt.Printf("failed to migrate db, error: %+v\n", err)
+				return
+			}
 
 			s := store.New(db.DBInstance, profile)
 			resources, err := s.ListResources(ctx, &store.FindResource{})

cmd/setup.go

@@ -36,10 +36,14 @@ var (
 			}
 
 			db := db.NewDB(profile)
-			if err := db.Open(ctx); err != nil {
+			if err := db.Open(); err != nil {
 				fmt.Printf("failed to open db, error: %+v\n", err)
 				return
 			}
+			if err := db.Migrate(ctx); err != nil {
+				fmt.Printf("failed to migrate db, error: %+v\n", err)
+				return
+			}
 
 			store := store.New(db.DBInstance, profile)
 			if err := ExecuteSetup(ctx, store, hostUsername, hostPassword); err != nil {

plugin/cron/cron.go

+// Package cron implements a crontab-like service to execute and schedule
+// repeative tasks/jobs.
+//
+// Example:
+//
+//	c := cron.New()
+//	c.MustAdd("dailyReport", "0 0 * * *", func() { ... })
+//	c.Start()
+package cron
+
+import (
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+)
+
+type job struct {
+	schedule *Schedule
+	run      func()
+}
+
+// Cron is a crontab-like struct for tasks/jobs scheduling.
+type Cron struct {
+	sync.RWMutex
+
+	interval time.Duration
+	timezone *time.Location
+	ticker   *time.Ticker
+	jobs     map[string]*job
+}
+
+// New create a new Cron struct with default tick interval of 1 minute
+// and timezone in UTC.
+//
+// You can change the default tick interval with Cron.SetInterval().
+// You can change the default timezone with Cron.SetTimezone().
+func New() *Cron {
+	return &Cron{
+		interval: 1 * time.Minute,
+		timezone: time.UTC,
+		jobs:     map[string]*job{},
+	}
+}
+
+// SetInterval changes the current cron tick interval
+// (it usually should be >= 1 minute).
+func (c *Cron) SetInterval(d time.Duration) {
+	// update interval
+	c.Lock()
+	wasStarted := c.ticker != nil
+	c.interval = d
+	c.Unlock()
+
+	// restart the ticker
+	if wasStarted {
+		c.Start()
+	}
+}
+
+// SetTimezone changes the current cron tick timezone.
+func (c *Cron) SetTimezone(l *time.Location) {
+	c.Lock()
+	defer c.Unlock()
+
+	c.timezone = l
+}
+
+// MustAdd is similar to Add() but panic on failure.
+func (c *Cron) MustAdd(jobID string, cronExpr string, run func()) {
+	if err := c.Add(jobID, cronExpr, run); err != nil {
+		panic(err)
+	}
+}
+
+// Add registers a single cron job.
+//
+// If there is already a job with the provided id, then the old job
+// will be replaced with the new one.
+//
+// cronExpr is a regular cron expression, eg. "0 */3 * * *" (aka. at minute 0 past every 3rd hour).
+// Check cron.NewSchedule() for the supported tokens.
+func (c *Cron) Add(jobID string, cronExpr string, run func()) error {
+	if run == nil {
+		return errors.New("failed to add new cron job: run must be non-nil function")
+	}
+
+	c.Lock()
+	defer c.Unlock()
+
+	schedule, err := NewSchedule(cronExpr)
+	if err != nil {
+		return fmt.Errorf("failed to add new cron job: %w", err)
+	}
+
+	c.jobs[jobID] = &job{
+		schedule: schedule,
+		run:      run,
+	}
+
+	return nil
+}
+
+// Remove removes a single cron job by its id.
+func (c *Cron) Remove(jobID string) {
+	c.Lock()
+	defer c.Unlock()
+
+	delete(c.jobs, jobID)
+}
+
+// RemoveAll removes all registered cron jobs.
+func (c *Cron) RemoveAll() {
+	c.Lock()
+	defer c.Unlock()
+
+	c.jobs = map[string]*job{}
+}
+
+// Total returns the current total number of registered cron jobs.
+func (c *Cron) Total() int {
+	c.RLock()
+	defer c.RUnlock()
+
+	return len(c.jobs)
+}
+
+// Stop stops the current cron ticker (if not already).
+//
+// You can resume the ticker by calling Start().
+func (c *Cron) Stop() {
+	c.Lock()
+	defer c.Unlock()
+
+	if c.ticker == nil {
+		return // already stopped
+	}
+
+	c.ticker.Stop()
+	c.ticker = nil
+}
+
+// Start starts the cron ticker.
+//
+// Calling Start() on already started cron will restart the ticker.
+func (c *Cron) Start() {
+	c.Stop()
+
+	c.Lock()
+	c.ticker = time.NewTicker(c.interval)
+	c.Unlock()
+
+	go func() {
+		for t := range c.ticker.C {
+			c.runDue(t)
+		}
+	}()
+}
+
+// HasStarted checks whether the current Cron ticker has been started.
+func (c *Cron) HasStarted() bool {
+	c.RLock()
+	defer c.RUnlock()
+
+	return c.ticker != nil
+}
+
+// runDue runs all registered jobs that are scheduled for the provided time.
+func (c *Cron) runDue(t time.Time) {
+	c.RLock()
+	defer c.RUnlock()
+
+	moment := NewMoment(t.In(c.timezone))
+
+	for _, j := range c.jobs {
+		if j.schedule.IsDue(moment) {
+			go j.run()
+		}
+	}
+}

plugin/cron/cron_test.go

+package cron
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCronNew(t *testing.T) {
+	c := New()
+
+	expectedInterval := 1 * time.Minute
+	if c.interval != expectedInterval {
+		t.Fatalf("Expected default interval %v, got %v", expectedInterval, c.interval)
+	}
+
+	expectedTimezone := time.UTC
+	if c.timezone.String() != expectedTimezone.String() {
+		t.Fatalf("Expected default timezone %v, got %v", expectedTimezone, c.timezone)
+	}
+
+	if len(c.jobs) != 0 {
+		t.Fatalf("Expected no jobs by default, got \n%v", c.jobs)
+	}
+
+	if c.ticker != nil {
+		t.Fatal("Expected the ticker NOT to be initialized")
+	}
+}
+
+func TestCronSetInterval(t *testing.T) {
+	c := New()
+
+	interval := 2 * time.Minute
+
+	c.SetInterval(interval)
+
+	if c.interval != interval {
+		t.Fatalf("Expected interval %v, got %v", interval, c.interval)
+	}
+}
+
+func TestCronSetTimezone(t *testing.T) {
+	c := New()
+
+	timezone, _ := time.LoadLocation("Asia/Tokyo")
+
+	c.SetTimezone(timezone)
+
+	if c.timezone.String() != timezone.String() {
+		t.Fatalf("Expected timezone %v, got %v", timezone, c.timezone)
+	}
+}
+
+func TestCronAddAndRemove(t *testing.T) {
+	c := New()
+
+	if err := c.Add("test0", "* * * * *", nil); err == nil {
+		t.Fatal("Expected nil function error")
+	}
+
+	if err := c.Add("test1", "invalid", func() {}); err == nil {
+		t.Fatal("Expected invalid cron expression error")
+	}
+
+	if err := c.Add("test2", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := c.Add("test3", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := c.Add("test4", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	// overwrite test2
+	if err := c.Add("test2", "1 2 3 4 5", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := c.Add("test5", "1 2 3 4 5", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	// mock job deletion
+	c.Remove("test4")
+
+	// try to remove non-existing (should be no-op)
+	c.Remove("missing")
+
+	// check job keys
+	{
+		expectedKeys := []string{"test3", "test2", "test5"}
+
+		if v := len(c.jobs); v != len(expectedKeys) {
+			t.Fatalf("Expected %d jobs, got %d", len(expectedKeys), v)
+		}
+
+		for _, k := range expectedKeys {
+			if c.jobs[k] == nil {
+				t.Fatalf("Expected job with key %s, got nil", k)
+			}
+		}
+	}
+
+	// check the jobs schedule
+	{
+		expectedSchedules := map[string]string{
+			"test2": `{"minutes":{"1":{}},"hours":{"2":{}},"days":{"3":{}},"months":{"4":{}},"daysOfWeek":{"5":{}}}`,
+			"test3": `{"minutes":{"0":{},"1":{},"10":{},"11":{},"12":{},"13":{},"14":{},"15":{},"16":{},"17":{},"18":{},"19":{},"2":{},"20":{},"21":{},"22":{},"23":{},"24":{},"25":{},"26":{},"27":{},"28":{},"29":{},"3":{},"30":{},"31":{},"32":{},"33":{},"34":{},"35":{},"36":{},"37":{},"38":{},"39":{},"4":{},"40":{},"41":{},"42":{},"43":{},"44":{},"45":{},"46":{},"47":{},"48":{},"49":{},"5":{},"50":{},"51":{},"52":{},"53":{},"54":{},"55":{},"56":{},"57":{},"58":{},"59":{},"6":{},"7":{},"8":{},"9":{}},"hours":{"0":{},"1":{},"10":{},"11":{},"12":{},"13":{},"14":{},"15":{},"16":{},"17":{},"18":{},"19":{},"2":{},"20":{},"21":{},"22":{},"23":{},"3":{},"4":{},"5":{},"6":{},"7":{},"8":{},"9":{}},"days":{"1":{},"10":{},"11":{},"12":{},"13":{},"14":{},"15":{},"16":{},"17":{},"18":{},"19":{},"2":{},"20":{},"21":{},"22":{},"23":{},"24":{},"25":{},"26":{},"27":{},"28":{},"29":{},"3":{},"30":{},"31":{},"4":{},"5":{},"6":{},"7":{},"8":{},"9":{}},"months":{"1":{},"10":{},"11":{},"12":{},"2":{},"3":{},"4":{},"5":{},"6":{},"7":{},"8":{},"9":{}},"daysOfWeek":{"0":{},"1":{},"2":{},"3":{},"4":{},"5":{},"6":{}}}`,
+			"test5": `{"minutes":{"1":{}},"hours":{"2":{}},"days":{"3":{}},"months":{"4":{}},"daysOfWeek":{"5":{}}}`,
+		}
+		for k, v := range expectedSchedules {
+			raw, err := json.Marshal(c.jobs[k].schedule)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if string(raw) != v {
+				t.Fatalf("Expected %q schedule \n%s, \ngot \n%s", k, v, raw)
+			}
+		}
+	}
+}
+
+func TestCronMustAdd(t *testing.T) {
+	c := New()
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("test1 didn't panic")
+		}
+	}()
+
+	c.MustAdd("test1", "* * * * *", nil)
+
+	c.MustAdd("test2", "* * * * *", func() {})
+
+	if _, ok := c.jobs["test2"]; !ok {
+		t.Fatal("Couldn't find job test2")
+	}
+}
+
+func TestCronRemoveAll(t *testing.T) {
+	c := New()
+
+	if err := c.Add("test1", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := c.Add("test2", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := c.Add("test3", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if v := len(c.jobs); v != 3 {
+		t.Fatalf("Expected %d jobs, got %d", 3, v)
+	}
+
+	c.RemoveAll()
+
+	if v := len(c.jobs); v != 0 {
+		t.Fatalf("Expected %d jobs, got %d", 0, v)
+	}
+}
+
+func TestCronTotal(t *testing.T) {
+	c := New()
+
+	if v := c.Total(); v != 0 {
+		t.Fatalf("Expected 0 jobs, got %v", v)
+	}
+
+	if err := c.Add("test1", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := c.Add("test2", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	// overwrite
+	if err := c.Add("test1", "* * * * *", func() {}); err != nil {
+		t.Fatal(err)
+	}
+
+	if v := c.Total(); v != 2 {
+		t.Fatalf("Expected 2 jobs, got %v", v)
+	}
+}
+
+func TestCronStartStop(t *testing.T) {
+	c := New()
+
+	c.SetInterval(1 * time.Second)
+
+	test1 := 0
+	test2 := 0
+
+	err := c.Add("test1", "* * * * *", func() {
+		test1++
+	})
+	require.NoError(t, err)
+
+	err = c.Add("test2", "* * * * *", func() {
+		test2++
+	})
+	require.NoError(t, err)
+
+	expectedCalls := 3
+
+	// call twice Start to check if the previous ticker will be reseted
+	c.Start()
+	c.Start()
+
+	time.Sleep(3250 * time.Millisecond)
+
+	// call twice Stop to ensure that the second stop is no-op
+	c.Stop()
+	c.Stop()
+
+	if test1 != expectedCalls {
+		t.Fatalf("Expected %d test1, got %d", expectedCalls, test1)
+	}
+	if test2 != expectedCalls {
+		t.Fatalf("Expected %d test2, got %d", expectedCalls, test2)
+	}
+
+	// resume for ~5 seconds
+	c.Start()
+	time.Sleep(5250 * time.Millisecond)
+	c.Stop()
+
+	expectedCalls += 5
+
+	if test1 != expectedCalls {
+		t.Fatalf("Expected %d test1, got %d", expectedCalls, test1)
+	}
+	if test2 != expectedCalls {
+		t.Fatalf("Expected %d test2, got %d", expectedCalls, test2)
+	}
+}

plugin/cron/schedule.go

+package cron
+
+import (
+	"errors"
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+)
+
+// Moment represents a parsed single time moment.
+type Moment struct {
+	Minute    int `json:"minute"`
+	Hour      int `json:"hour"`
+	Day       int `json:"day"`
+	Month     int `json:"month"`
+	DayOfWeek int `json:"dayOfWeek"`
+}
+
+// NewMoment creates a new Moment from the specified time.
+func NewMoment(t time.Time) *Moment {
+	return &Moment{
+		Minute:    t.Minute(),
+		Hour:      t.Hour(),
+		Day:       t.Day(),
+		Month:     int(t.Month()),
+		DayOfWeek: int(t.Weekday()),
+	}
+}
+
+// Schedule stores parsed information for each time component when a cron job should run.
+type Schedule struct {
+	Minutes    map[int]struct{} `json:"minutes"`
+	Hours      map[int]struct{} `json:"hours"`
+	Days       map[int]struct{} `json:"days"`
+	Months     map[int]struct{} `json:"months"`
+	DaysOfWeek map[int]struct{} `json:"daysOfWeek"`
+}
+
+// IsDue checks whether the provided Moment satisfies the current Schedule.
+func (s *Schedule) IsDue(m *Moment) bool {
+	if _, ok := s.Minutes[m.Minute]; !ok {
+		return false
+	}
+
+	if _, ok := s.Hours[m.Hour]; !ok {
+		return false
+	}
+
+	if _, ok := s.Days[m.Day]; !ok {
+		return false
+	}
+
+	if _, ok := s.DaysOfWeek[m.DayOfWeek]; !ok {
+		return false
+	}
+
+	if _, ok := s.Months[m.Month]; !ok {
+		return false
+	}
+
+	return true
+}
+
+// NewSchedule creates a new Schedule from a cron expression.
+//
+// A cron expression is consisted of 5 segments separated by space,
+// representing: minute, hour, day of the month, month and day of the week.
+//
+// Each segment could be in the following formats:
+//   - wildcard: *
+//   - range:    1-30
+//   - step:     */n or 1-30/n
+//   - list:     1,2,3,10-20/n
+func NewSchedule(cronExpr string) (*Schedule, error) {
+	segments := strings.Split(cronExpr, " ")
+	if len(segments) != 5 {
+		return nil, errors.New("invalid cron expression - must have exactly 5 space separated segments")
+	}
+
+	minutes, err := parseCronSegment(segments[0], 0, 59)
+	if err != nil {
+		return nil, err
+	}
+
+	hours, err := parseCronSegment(segments[1], 0, 23)
+	if err != nil {
+		return nil, err
+	}
+
+	days, err := parseCronSegment(segments[2], 1, 31)
+	if err != nil {
+		return nil, err
+	}
+
+	months, err := parseCronSegment(segments[3], 1, 12)
+	if err != nil {
+		return nil, err
+	}
+
+	daysOfWeek, err := parseCronSegment(segments[4], 0, 6)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Schedule{
+		Minutes:    minutes,
+		Hours:      hours,
+		Days:       days,
+		Months:     months,
+		DaysOfWeek: daysOfWeek,
+	}, nil
+}
+
+// parseCronSegment parses a single cron expression segment and
+// returns its time schedule slots.
+func parseCronSegment(segment string, min int, max int) (map[int]struct{}, error) {
+	slots := map[int]struct{}{}
+
+	list := strings.Split(segment, ",")
+	for _, p := range list {
+		stepParts := strings.Split(p, "/")
+
+		// step (*/n, 1-30/n)
+		var step int
+		switch len(stepParts) {
+		case 1:
+			step = 1
+		case 2:
+			parsedStep, err := strconv.Atoi(stepParts[1])
+			if err != nil {
+				return nil, err
+			}
+			if parsedStep < 1 || parsedStep > max {
+				return nil, fmt.Errorf("invalid segment step boundary - the step must be between 1 and the %d", max)
+			}
+			step = parsedStep
+		default:
+			return nil, errors.New("invalid segment step format - must be in the format */n or 1-30/n")
+		}
+
+		// find the min and max range of the segment part
+		var rangeMin, rangeMax int
+		if stepParts[0] == "*" {
+			rangeMin = min
+			rangeMax = max
+		} else {
+			// single digit (1) or range (1-30)
+			rangeParts := strings.Split(stepParts[0], "-")
+			switch len(rangeParts) {
+			case 1:
+				if step != 1 {
+					return nil, errors.New("invalid segement step - step > 1 could be used only with the wildcard or range format")
+				}
+				parsed, err := strconv.Atoi(rangeParts[0])
+				if err != nil {
+					return nil, err
+				}
+				if parsed < min || parsed > max {
+					return nil, errors.New("invalid segment value - must be between the min and max of the segment")
+				}
+				rangeMin = parsed
+				rangeMax = rangeMin
+			case 2:
+				parsedMin, err := strconv.Atoi(rangeParts[0])
+				if err != nil {
+					return nil, err
+				}
+				if parsedMin < min || parsedMin > max {
+					return nil, fmt.Errorf("invalid segment range minimum - must be between %d and %d", min, max)
+				}
+				rangeMin = parsedMin
+
+				parsedMax, err := strconv.Atoi(rangeParts[1])
+				if err != nil {
+					return nil, err
+				}
+				if parsedMax < parsedMin || parsedMax > max {
+					return nil, fmt.Errorf("invalid segment range maximum - must be between %d and %d", rangeMin, max)
+				}
+				rangeMax = parsedMax
+			default:
+				return nil, errors.New("invalid segment range format - the range must have 1 or 2 parts")
+			}
+		}
+
+		// fill the slots
+		for i := rangeMin; i <= rangeMax; i += step {
+			slots[i] = struct{}{}
+		}
+	}
+
+	return slots, nil
+}

store/db/db.go

@@ -36,7 +36,10 @@ func NewDB(profile *profile.Profile) *DB {
 	return db
 }
 
-func (db *DB) Open(ctx context.Context) (err error) {
+// Open opens a database specified by its database driver name and a
+// driver-specific data source name, usually consisting of at least a
+// database name and connection information.
+func (db *DB) Open() error {
 	// Ensure a DSN is set before attempting to open the database.
 	if db.profile.DSN == "" {
 		return fmt.Errorf("dsn required")
@@ -61,7 +64,11 @@ func (db *DB) Open(ctx context.Context) (err error) {
 		return fmt.Errorf("failed to open db with dsn: %s, err: %w", db.profile.DSN, err)
 	}
 	db.DBInstance = sqliteDB
+	return nil
+}
 
+// Migrate applies the latest schema to the database.
+func (db *DB) Migrate(ctx context.Context) error {
 	if db.profile.Mode == "prod" {
 		_, err := os.Stat(db.profile.DSN)
 		if err != nil {

test/server/server.go

@@ -32,9 +32,12 @@ type TestingServer struct {
 func NewTestingServer(ctx context.Context, t *testing.T) (*TestingServer, error) {
 	profile := test.GetTestingProfile(t)
 	db := db.NewDB(profile)
-	if err := db.Open(ctx); err != nil {
+	if err := db.Open(); err != nil {
 		return nil, errors.Wrap(err, "failed to open db")
 	}
+	if err := db.Migrate(ctx); err != nil {
+		return nil, errors.Wrap(err, "failed to migrate db")
+	}
 
 	store := store.New(db.DBInstance, profile)
 	server, err := server.NewServer(ctx, profile, store)

test/store/store.go

@@ -16,9 +16,12 @@ import (
 func NewTestingStore(ctx context.Context, t *testing.T) *store.Store {
 	profile := test.GetTestingProfile(t)
 	db := db.NewDB(profile)
-	if err := db.Open(ctx); err != nil {
+	if err := db.Open(); err != nil {
 		fmt.Printf("failed to open db, error: %+v\n", err)
 	}
+	if err := db.Migrate(ctx); err != nil {
+		fmt.Printf("failed to migrate db, error: %+v\n", err)
+	}
 
 	store := store.New(db.DBInstance, profile)
 	return store