edge/cmd/storage-server/command/run.go

package command

import (
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/keegancsmith/rpc"
	"gitlab.com/wpetit/goweb/logger"

	"forge.cadoles.com/arcad/edge/pkg/storage/rpc/server"
	"forge.cadoles.com/arcad/edge/pkg/storage/sqlite"
	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/pkg/errors"
	"github.com/urfave/cli/v2"
)

func Run() *cli.Command {
	return &cli.Command{
		Name:  "run",
		Usage: "Run server",
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    "address",
				Aliases: []string{"addr"},
				Value:   ":3001",
			},
			&cli.StringFlag{
				Name:  "data-dir",
				Value: "./data",
			},
			&cli.StringFlag{
				Name:  "dsn-query-string",
				Value: fmt.Sprintf("_pragma=foreign_keys(1)&_pragma=busy_timeout=%d", (60 * time.Second).Milliseconds()),
			},
		},
		Action: func(ctx *cli.Context) error {
			addr := ctx.String("address")
			dataDir := ctx.String("data-dir")
			dsnQueryString := ctx.String("dsn-query-string")

			router := chi.NewRouter()

			router.Use(middleware.RealIP)
			router.Use(middleware.Logger)

			router.Handle("/blobstore", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				tenant := r.URL.Query().Get("tenant")
				if tenant == "" {
					http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)

					return
				}

				server, err := getBlobStoreStoreServer(dataDir, tenant, dsnQueryString)
				if err != nil {
					logger.Error(r.Context(), "could not retrieve blob store server", logger.E(errors.WithStack(err)), logger.F("tenant", tenant))
					http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)

					return
				}

				server.ServeHTTP(w, r)

			}))

			router.Handle("/documentstore", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				tenant := r.URL.Query().Get("tenant")
				if tenant == "" {
					http.Error(w, http.StatusText(http.StatusBadRequest), http.StatusBadRequest)

					return
				}

				server, err := getDocumentStoreServer(dataDir, tenant, dsnQueryString)
				if err != nil {
					logger.Error(r.Context(), "could not retrieve document store server", logger.E(errors.WithStack(err)), logger.F("tenant", tenant))
					http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)

					return
				}

				server.ServeHTTP(w, r)
			}))

			if err := http.ListenAndServe(addr, router); err != nil {
				return errors.WithStack(err)
			}

			return nil
		},
	}
}

var documentStoreTenants sync.Map

func getDocumentStoreServer(dataDir, tenant, dsnQueryString string) (*rpc.Server, error) {
	dir := filepath.Join(dataDir, tenant)
	if err := os.MkdirAll(dir, os.FileMode(0750)); err != nil {
		return nil, errors.WithStack(err)
	}

	file := filepath.Join(dir, "documentstore.sqlite")
	dsn := fmt.Sprintf("%s?%s", file, dsnQueryString)

	documentStore := sqlite.NewDocumentStore(dsn)
	documentStoreServer := server.NewDocumentStoreServer(documentStore)

	rawDocumentStoreServer, _ := documentStoreTenants.LoadOrStore(tenant, documentStoreServer)

	documentStoreServer, ok := rawDocumentStoreServer.(*rpc.Server)
	if !ok {
		return nil, errors.Errorf("unexpected document store server value of type '%T'", rawDocumentStoreServer)
	}

	return documentStoreServer, nil
}

var blobStoreTenants sync.Map

func getBlobStoreStoreServer(dataDir, tenant, dsnQueryString string) (*rpc.Server, error) {
	dir := filepath.Join(dataDir, tenant)
	if err := os.MkdirAll(dir, os.FileMode(0750)); err != nil {
		return nil, errors.WithStack(err)
	}

	file := filepath.Join(dir, "blobstore.sqlite")
	dsn := fmt.Sprintf("%s?%s", file, dsnQueryString)

	blobStore := sqlite.NewBlobStore(dsn)
	blobStoreServer := server.NewBlobStoreServer(blobStore)

	rawBlobStoreServer, _ := documentStoreTenants.LoadOrStore(tenant, blobStoreServer)

	blobStoreServer, ok := rawBlobStoreServer.(*rpc.Server)
	if !ok {
		return nil, errors.Errorf("unexpected document store server value of type '%T'", rawBlobStoreServer)
	}

	return blobStoreServer, nil
}