golang-saas-starter-kit/internal/mid/serverless.go

package mid

import (
	"context"
	"net/http"
	"time"

	"geeks-accelerator/oss/saas-starter-kit/internal/platform/web"
	"github.com/jmoiron/sqlx"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

// ctxServerlessKey represents the type of value for the context key.
type ctxServerlessKey int

// Key is used to store/retrieve a Serverless value from a context.Context.
const ServerlessKey ctxServerlessKey = 1

type (
	// WaitForDbResumedConfig defines the config for WaitForDbResumed middleware.
	WaitForDbResumedConfig struct {
		RedirectConfig

		// Database handle to be used to ensure its online.
		DB *sqlx.DB

		// WaitHandler defines the handler to render for the user to when the database is being resumed.
		WaitHandler web.Handler
	}
)

// WaitForDbResumed returns an middleware with for ensuring an serverless database is resumed.
func WaitForDbResumed(config WaitForDbResumedConfig) web.Middleware {

	if config.Skipper == nil {
		config.Skipper = DefaultSkipper
	}
	if config.Code == 0 {
		config.Code = DefaultRedirectConfig.Code
	}

	verifyDb := func() error {
		// When the database is paused, Postgres will return the error, "Canceling statement due to user request"

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
		defer cancel()

		_, err := config.DB.ExecContext(ctx, "SELECT NULL")
		if err != nil {
			return err
		}

		return nil
	}

	// This is the actual middleware function to be executed.
	f := func(after web.Handler) web.Handler {

		h := func(ctx context.Context, w http.ResponseWriter, r *http.Request, params map[string]string) error {
			span, ctx := tracer.StartSpanFromContext(ctx, "internal.mid.serverless")
			defer span.Finish()

			if config.Skipper(ctx, w, r, params) {
				return after(ctx, w, r, params)
			}

			if err := verifyDb(); err != nil {
				ctx = context.WithValue(ctx, ServerlessKey, err)
				return config.WaitHandler(ctx, w, r, params)
			}

			return after(ctx, w, r, params)
		}

		return h
	}

	return f
}