96 lines
2.1 KiB
Go
96 lines
2.1 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"strings"
|
|
|
|
_ "modernc.org/sqlite"
|
|
)
|
|
|
|
type SQLiteCounterStore struct {
|
|
db *sql.DB
|
|
}
|
|
|
|
func NewSQLiteCounterStore(path string) (*SQLiteCounterStore, error) {
|
|
db, err := sql.Open("sqlite", path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
db.SetMaxOpenConns(1)
|
|
return &SQLiteCounterStore{db: db}, nil
|
|
}
|
|
|
|
func (s *SQLiteCounterStore) Init(ctx context.Context) error {
|
|
_, err := s.db.ExecContext(ctx, `
|
|
CREATE TABLE IF NOT EXISTS ticks (
|
|
slug TEXT PRIMARY KEY,
|
|
count INTEGER NOT NULL DEFAULT 0,
|
|
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
`)
|
|
return err
|
|
}
|
|
|
|
func (s *SQLiteCounterStore) Increment(ctx context.Context, slug string) (int64, error) {
|
|
if slug == "" {
|
|
return 0, fmt.Errorf("empty slug")
|
|
}
|
|
|
|
_, err := s.db.ExecContext(ctx, `
|
|
INSERT INTO ticks(slug, count) VALUES(?, 1)
|
|
ON CONFLICT(slug) DO UPDATE SET
|
|
count = count + 1,
|
|
updated_at = CURRENT_TIMESTAMP;
|
|
`, slug)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return s.Get(ctx, slug)
|
|
}
|
|
|
|
func (s *SQLiteCounterStore) Get(ctx context.Context, slug string) (int64, error) {
|
|
var count int64
|
|
err := s.db.QueryRowContext(ctx, `SELECT count FROM ticks WHERE slug = ?`, slug).Scan(&count)
|
|
if err == sql.ErrNoRows {
|
|
return 0, nil
|
|
}
|
|
return count, err
|
|
}
|
|
|
|
func (s *SQLiteCounterStore) GetMany(ctx context.Context, slugs []string) (map[string]int64, error) {
|
|
out := make(map[string]int64, len(slugs))
|
|
if len(slugs) == 0 {
|
|
return out, nil
|
|
}
|
|
|
|
placeholders := strings.TrimRight(strings.Repeat("?,", len(slugs)), ",")
|
|
args := make([]any, len(slugs))
|
|
for i, slug := range slugs {
|
|
args[i] = slug
|
|
out[slug] = 0
|
|
}
|
|
|
|
rows, err := s.db.QueryContext(ctx, `SELECT slug, count FROM ticks WHERE slug IN (`+placeholders+`)`, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
for rows.Next() {
|
|
var slug string
|
|
var count int64
|
|
if err := rows.Scan(&slug, &count); err != nil {
|
|
return nil, err
|
|
}
|
|
out[slug] = count
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (s *SQLiteCounterStore) Close() error {
|
|
return s.db.Close()
|
|
}
|