edge/pkg/storage/driver/sqlite/document_store.go
William Petit 8e574c299b
All checks were successful
arcad/edge/pipeline/pr-master This commit looks good
feat(storage): rpc based implementation
2023-09-28 12:36:30 -06:00

363 lines
7.9 KiB
Go

package sqlite
import (
"context"
"database/sql"
"fmt"
"math"
"time"
"forge.cadoles.com/arcad/edge/pkg/storage"
"forge.cadoles.com/arcad/edge/pkg/storage/filter"
filterSQL "forge.cadoles.com/arcad/edge/pkg/storage/filter/sql"
"github.com/pkg/errors"
"gitlab.com/wpetit/goweb/logger"
_ "modernc.org/sqlite"
)
type DocumentStore struct {
getDB GetDBFunc
}
// Delete implements storage.DocumentStore
func (s *DocumentStore) Delete(ctx context.Context, collection string, id storage.DocumentID) error {
err := s.withTx(ctx, func(tx *sql.Tx) error {
query := `
DELETE FROM documents
WHERE collection = $1 AND id = $2
`
_, err := tx.ExecContext(ctx, query, collection, string(id))
if err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return errors.WithStack(err)
}
return nil
}
// Get implements storage.DocumentStore
func (s *DocumentStore) Get(ctx context.Context, collection string, id storage.DocumentID) (storage.Document, error) {
var document storage.Document
err := s.withTx(ctx, func(tx *sql.Tx) error {
query := `
SELECT id, data, created_at, updated_at
FROM documents
WHERE collection = $1 AND id = $2
`
row := tx.QueryRowContext(ctx, query, collection, string(id))
var (
createdAt time.Time
updatedAt time.Time
data JSONMap
)
err := row.Scan(&id, &data, &createdAt, &updatedAt)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return errors.WithStack(storage.ErrDocumentNotFound)
}
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
document = storage.Document(data)
document[storage.DocumentAttrID] = id
document[storage.DocumentAttrCreatedAt] = createdAt
document[storage.DocumentAttrUpdatedAt] = updatedAt
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return document, nil
}
// Query implements storage.DocumentStore
func (s *DocumentStore) Query(ctx context.Context, collection string, filter *filter.Filter, funcs ...storage.QueryOptionFunc) ([]storage.Document, error) {
opts := &storage.QueryOptions{}
for _, fn := range funcs {
fn(opts)
}
var documents []storage.Document
err := s.withTx(ctx, func(tx *sql.Tx) error {
criteria := "1 = 1"
args := make([]any, 0)
var err error
if filter != nil {
criteria, args, err = filterSQL.ToSQL(
filter.Root(),
filterSQL.WithPreparedParameter("$", 2),
filterSQL.WithTransform(transformOperator),
filterSQL.WithKeyTransform(func(key string) string {
return fmt.Sprintf("json_extract(data, '$.%s')", key)
}),
)
if err != nil {
return errors.WithStack(err)
}
}
query := `
SELECT id, data, created_at, updated_at
FROM documents
WHERE collection = $1 AND (` + criteria + `)
`
args = append([]interface{}{collection}, args...)
if opts.OrderBy != nil {
direction := storage.OrderDirectionAsc
if opts.OrderDirection != nil {
direction = *opts.OrderDirection
}
query, args = withOrderByClause(query, args, *opts.OrderBy, direction)
}
if opts.Offset != nil || opts.Limit != nil {
offset := 0
if opts.Offset != nil {
offset = *opts.Offset
}
limit := math.MaxInt
if opts.Limit != nil {
limit = *opts.Limit
}
query, args = withLimitOffsetClause(query, args, limit, offset)
}
logger.Debug(
ctx, "executing query",
logger.F("query", query),
logger.F("args", args),
)
rows, err := tx.QueryContext(ctx, query, args...)
if err != nil {
return errors.WithStack(err)
}
defer func() {
if err := rows.Close(); err != nil {
logger.Error(ctx, "could not close rows", logger.E(errors.WithStack(err)))
}
}()
documents = make([]storage.Document, 0)
for rows.Next() {
var (
id storage.DocumentID
createdAt time.Time
updatedAt time.Time
data JSONMap
)
if err := rows.Scan(&id, &data, &createdAt, &updatedAt); err != nil {
return errors.WithStack(err)
}
document := storage.Document(data)
document[storage.DocumentAttrID] = id
document[storage.DocumentAttrCreatedAt] = createdAt
document[storage.DocumentAttrUpdatedAt] = updatedAt
documents = append(documents, document)
}
if err := rows.Err(); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return documents, nil
}
// Upsert implements storage.DocumentStore
func (s *DocumentStore) Upsert(ctx context.Context, collection string, document storage.Document) (storage.Document, error) {
var upsertedDocument storage.Document
err := s.withTx(ctx, func(tx *sql.Tx) error {
query := `
INSERT INTO documents (id, collection, data, created_at, updated_at)
VALUES($1, $2, $3, $4, $4)
ON CONFLICT (id, collection) DO UPDATE SET
data = $3, updated_at = $4
RETURNING "id", "data", "created_at", "updated_at"
`
now := time.Now().UTC()
id, exists := document.ID()
if !exists || id == "" {
id = storage.NewDocumentID()
}
args := []any{id, collection, JSONMap(document), now, now}
logger.Debug(
ctx, "executing query",
logger.F("query", query),
logger.F("args", args),
)
row := tx.QueryRowContext(ctx, query, args...)
var (
createdAt time.Time
updatedAt time.Time
data JSONMap
)
err := row.Scan(&id, &data, &createdAt, &updatedAt)
if err != nil {
return errors.WithStack(err)
}
if err := row.Err(); err != nil {
return errors.WithStack(err)
}
upsertedDocument = storage.Document(data)
upsertedDocument[storage.DocumentAttrID] = id
upsertedDocument[storage.DocumentAttrCreatedAt] = createdAt
upsertedDocument[storage.DocumentAttrUpdatedAt] = updatedAt
return nil
})
if err != nil {
return nil, errors.WithStack(err)
}
return upsertedDocument, nil
}
func (s *DocumentStore) withTx(ctx context.Context, fn func(tx *sql.Tx) error) error {
var db *sql.DB
db, err := s.getDB(ctx)
if err != nil {
return errors.WithStack(err)
}
if err := WithTx(ctx, db, fn); err != nil {
return errors.WithStack(err)
}
return nil
}
func ensureDocumentTables(ctx context.Context, db *sql.DB) error {
err := WithTx(ctx, db, func(tx *sql.Tx) error {
query := `
CREATE TABLE IF NOT EXISTS documents (
id TEXT PRIMARY KEY,
collection TEXT NOT NULL,
data TEXT,
created_at TIMESTAMP NOT NULL,
updated_at TIMESTAMP NOT NULL,
UNIQUE(id, collection) ON CONFLICT REPLACE
);
`
if _, err := tx.ExecContext(ctx, query); err != nil {
return errors.WithStack(err)
}
query = `
CREATE INDEX IF NOT EXISTS collection_idx ON documents (collection);
`
if _, err := tx.ExecContext(ctx, query); err != nil {
return errors.WithStack(err)
}
return nil
})
if err != nil {
return errors.WithStack(err)
}
return nil
}
func withOrderByClause(query string, args []any, orderBy string, orderDirection storage.OrderDirection) (string, []any) {
direction := "ASC"
if orderDirection == storage.OrderDirectionDesc {
direction = "DESC"
}
var column string
switch orderBy {
case storage.DocumentAttrID:
column = "id"
case storage.DocumentAttrCreatedAt:
column = "created_at"
case storage.DocumentAttrUpdatedAt:
column = "updated_at"
default:
column = fmt.Sprintf("json_extract(data, '$.' || $%d)", len(args)+1)
args = append(args, orderBy)
}
query += fmt.Sprintf(` ORDER BY %s %s`, column, direction)
return query, args
}
func withLimitOffsetClause(query string, args []any, limit int, offset int) (string, []any) {
query += fmt.Sprintf(` LIMIT $%d OFFSET $%d`, len(args)+1, len(args)+2)
args = append(args, limit, offset)
return query, args
}
func NewDocumentStore(path string) *DocumentStore {
getDB := NewGetDBFunc(path, ensureDocumentTables)
return &DocumentStore{
getDB: getDB,
}
}
func NewDocumentStoreWithDB(db *sql.DB) *DocumentStore {
getDB := NewGetDBFuncFromDB(db, ensureDocumentTables)
return &DocumentStore{
getDB: getDB,
}
}
var _ storage.DocumentStore = &DocumentStore{}