edge/pkg/storage/sqlite/blob_store.go

package sqlite

import (
	"context"
	"database/sql"

	"forge.cadoles.com/arcad/edge/pkg/storage"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/logger"
)

type BlobStore struct {
	getDB getDBFunc
}

// DeleteBucket implements storage.BlobStore
func (s *BlobStore) DeleteBucket(ctx context.Context, name string) error {
	err := s.withTx(ctx, func(tx *sql.Tx) error {
		query := `DELETE FROM blobs WHERE bucket = $1`
		_, err := tx.ExecContext(ctx, query, name)
		if err != nil {
			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// ListBuckets implements storage.BlobStore
func (s *BlobStore) ListBuckets(ctx context.Context) ([]string, error) {
	buckets := make([]string, 0)

	err := s.withTx(ctx, func(tx *sql.Tx) error {
		query := `SELECT DISTINCT bucket FROM blobs`
		rows, err := tx.QueryContext(ctx, query)
		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)))
			}
		}()

		for rows.Next() {
			var name string
			if err := rows.Scan(&name); err != nil {
				return errors.WithStack(err)
			}

			buckets = append(buckets, name)
		}

		if err := rows.Err(); err != nil {
			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return buckets, nil
}

// OpenBucket implements storage.BlobStore
func (s *BlobStore) OpenBucket(ctx context.Context, name string) (storage.BlobBucket, error) {
	return &BlobBucket{
		name:  name,
		getDB: s.getDB,
	}, nil
}

func ensureBlobTables(ctx context.Context, db *sql.DB) error {
	logger.Debug(ctx, "creating blobs table")

	err := withTx(ctx, db, func(tx *sql.Tx) error {
		query := `
			CREATE TABLE IF NOT EXISTS blobs (
				id TEXT,
				bucket TEXT,
				data BLOB,
				content_type TEXT NOT NULL,
				mod_time TIMESTAMP NOT NULL,
				size INTEGER,
				PRIMARY KEY (id, bucket)
			);
		`
		if _, err := tx.ExecContext(ctx, query); err != nil {
			return errors.WithStack(err)
		}

		return nil
	})
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *BlobStore) 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 NewBlobStore(dsn string) *BlobStore {
	getDB := newGetDBFunc(dsn, ensureBlobTables)

	return &BlobStore{getDB}
}

func NewBlobStoreWithDB(db *sql.DB) *BlobStore {
	getDB := newGetDBFuncFromDB(db, ensureBlobTables)

	return &BlobStore{getDB}
}

var _ storage.BlobStore = &BlobStore{}