86 lines
2.4 KiB
Go
86 lines
2.4 KiB
Go
package sqlitestorage
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"time"
|
|
|
|
"git.maronato.dev/maronato/goshort/internal/config"
|
|
"git.maronato.dev/maronato/goshort/internal/errs"
|
|
bunstorage "git.maronato.dev/maronato/goshort/internal/storage/bun"
|
|
"github.com/uptrace/bun"
|
|
"github.com/uptrace/bun/dialect/sqlitedialect"
|
|
_ "modernc.org/sqlite" // Import the SQLite driver.
|
|
)
|
|
|
|
// NewSQLiteStorage creates a new SQLite storage.
|
|
func NewSQLiteStorage(_ context.Context, cfg *config.Config) *bunstorage.BunStorage {
|
|
if cfg.DBType == config.DBTypeMemory {
|
|
cfg.DBURL = "file::memory:?cache=shared&mode=memory&_foreign_keys=1"
|
|
} else {
|
|
cfg.DBURL += "?_pragma=journal_mode=WAL&_pragma=foreign_keys=1"
|
|
}
|
|
|
|
sqldb, err := sql.Open("sqlite", cfg.DBURL)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
db := bun.NewDB(sqldb, sqlitedialect.New())
|
|
|
|
// Use Bun as the storage implementation.
|
|
strg := bunstorage.NewBunStorage(cfg, db)
|
|
|
|
// Config vacuuming of the database every hour.
|
|
ticker := time.NewTicker(time.Hour)
|
|
tickerCtx, tickerCancel := context.WithCancel(context.Background())
|
|
// Start vacuuming the database when the storage is started.
|
|
strg.RegisterStartHook(func(ctx context.Context, activeDB *bun.DB) error {
|
|
// Vacuum now
|
|
_, err := activeDB.NewRaw("VACUUM").Exec(ctx)
|
|
if err != nil {
|
|
return errs.Errorf("failed to vacuum database: %w", err)
|
|
}
|
|
|
|
// Vacuum every hour, until the context is canceled.
|
|
go func() {
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
_, err := activeDB.NewRaw("VACUUM").Exec(ctx)
|
|
if err != nil {
|
|
panic(errs.Errorf("failed to vacuum database: %w", err))
|
|
}
|
|
case <-ctx.Done():
|
|
return
|
|
case <-tickerCtx.Done():
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
return nil
|
|
})
|
|
// So some cleaning up and optimization when the storage is stopped.
|
|
strg.RegisterStopHook(func(ctx context.Context, activeDB *bun.DB) error {
|
|
// Stop the ticker and cancel the ticker context.
|
|
ticker.Stop()
|
|
tickerCancel()
|
|
// Run PRAGMA optimize to optimize the database.
|
|
// By this point, the passed context is likely cancelled, so we use the
|
|
// background context since otherwise the query will fail.
|
|
// This will delay the shutdown of the storage, but it'll be for a good cause.
|
|
// The user can force the shutdown by using SIGINT or SIGTERM.
|
|
noCancelCtx := context.WithoutCancel(ctx)
|
|
|
|
_, err := activeDB.NewRaw("PRAGMA optimize").Exec(noCancelCtx)
|
|
if err != nil {
|
|
return errs.Errorf("failed to optimize database: %w", err)
|
|
}
|
|
|
|
return nil
|
|
})
|
|
|
|
return strg
|
|
}
|