pocket-id/backend/internal/bootstrap/db_bootstrap.go

package bootstrap

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/url"
	"strings"
	"time"

	_ "github.com/golang-migrate/migrate/v4/source/github"
	sqlinstrument "github.com/italypaleale/go-sql-utils/instrument"
	postgresinstrument "github.com/italypaleale/go-sql-utils/instrument/postgres"
	sqliteinstrument "github.com/italypaleale/go-sql-utils/instrument/sqlite"
	sqlitekit "github.com/italypaleale/go-sql-utils/sqlite"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/jackc/pgx/v5/stdlib"
	"github.com/libtnb/sqlite"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	gormLogger "gorm.io/gorm/logger"
	gormMetrics "gorm.io/plugin/opentelemetry/metrics"

	"github.com/pocket-id/pocket-id/backend/internal/common"
	"github.com/pocket-id/pocket-id/backend/internal/utils"
	sqliteutil "github.com/pocket-id/pocket-id/backend/internal/utils/sqlite"
)

func NewDatabase(ctx context.Context) (db *gorm.DB, pg *pgxpool.Pool, err error) {
	db, pg, err = ConnectDatabase(ctx)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to connect to database: %w", err)
	}
	sqlDb, err := db.DB()
	if err != nil {
		return nil, nil, fmt.Errorf("failed to get sql.DB: %w", err)
	}

	// Run migrations
	err = utils.MigrateDatabase(ctx, sqlDb)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to run migrations: %w", err)
	}

	return db, pg, nil
}

func ConnectDatabase(ctx context.Context) (db *gorm.DB, pg *pgxpool.Pool, err error) {
	var dialector gorm.Dialector

	// Choose the correct database provider
	switch common.EnvConfig.DbProvider {
	case common.DbProviderSqlite:
		if common.EnvConfig.DbConnectionString == "" {
			return nil, nil, errors.New("missing required env var 'DB_CONNECTION_STRING' for SQLite database")
		}

		sqliteutil.RegisterSqliteFunctions()

		// The connector validates the connection string and performs the filesystem setup SQLite needs: it creates the database and temporary directories
		// It also warns when the database lives on a networked filesystem, which is unsupported
		connector, err := sqlitekit.NewConnector(sqlitekit.ConnectOpts{
			ConnString: addSqliteDatetimeParams(common.EnvConfig.DbConnectionString),
			Logger:     slog.Default(),
		})
		if err != nil {
			return nil, nil, err
		}

		// We open the connection ourselves, rather than letting Gorm do it, so it goes through the instrumented driver
		// It also caps in-memory databases to a single connection, which they need to see the whole data
		sqliteDB, err := sqliteinstrument.Open(connector, sqlInstrumentOptions())
		if err != nil {
			return nil, nil, fmt.Errorf("failed to open SQLite database: %w", err)
		}

		dialector = sqlite.New(sqlite.Config{Conn: sqliteDB})
	case common.DbProviderPostgres:
		if common.EnvConfig.DbConnectionString == "" {
			return nil, nil, errors.New("missing required env var 'DB_CONNECTION_STRING' for Postgres database")
		}

		// We need a pgxpool object for francis, so we open this as a pgxpool...
		poolCfg, err := pgxpool.ParseConfig(common.EnvConfig.DbConnectionString)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to parse Postgres connection string: %w", err)
		}

		// ...with the instrumented tracer attached, chaining any tracer the connection string may have configured
		poolCfg.ConnConfig.Tracer = postgresinstrument.NewTracer(sqlInstrumentOptions(), poolCfg.ConnConfig.Tracer)

		pg, err = pgxpool.NewWithConfig(ctx, poolCfg)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to create Postgres pool: %w", err)
		}

		// Test it with a ping
		pingCtx, pingCancel := context.WithTimeout(ctx, 10*time.Second)
		defer pingCancel()
		err = pg.Ping(pingCtx)
		if err != nil {
			pg.Close()
			return nil, nil, fmt.Errorf("failed to ping Postgres database: %w", err)
		}

		// ...then create the dialector by adapting it to *sql.DB
		dialector = postgres.New(postgres.Config{
			Conn: stdlib.OpenDBFromPool(pg),
		})
	default:
		return nil, nil, fmt.Errorf("unsupported database provider: %s", common.EnvConfig.DbProvider)
	}

	// Try connecting up to 3 times
	for i := 1; i <= 3; i++ {
		db, err = gorm.Open(dialector, &gorm.Config{
			TranslateError: true,
			// Disable logging in Gorm because the driver itself is instrumented
			Logger: gormLogger.Discard,
		})
		if err == nil {
			slog.Info("Connected to database", slog.String("provider", string(common.EnvConfig.DbProvider)))

			conn, err := db.DB()
			if err != nil {
				if pg != nil {
					pg.Close()
				}
				return nil, nil, fmt.Errorf("failed to get *sql.DB connection from Gorm: %w", err)
			}

			// Report the metrics for the connection pool
			// Gorm's OpenTelemetry plugin is not used for this: it would wrap every statement span the instrumentation already emits in a second span
			gormMetrics.ReportDBStatsMetrics(conn)

			return db, pg, nil
		}

		// If we're here, the connection failed
		slog.Warn("Failed to connect to database, will retry in 3s", slog.Int("attempt", i), slog.String("provider", string(common.EnvConfig.DbProvider)), slog.Any("error", err))
		time.Sleep(3 * time.Second)
	}

	slog.Error("Failed to connect to database after 3 attempts", slog.String("provider", string(common.EnvConfig.DbProvider)), slog.Any("error", err))

	if pg != nil {
		pg.Close()
	}

	return nil, nil, err
}

// sqlInstrumentOptions returns the instrumentation applied to the database connection, shared by both providers
// This enables tracing in addition to logs
func sqlInstrumentOptions() *sqlinstrument.Options {
	return &sqlinstrument.Options{
		Log: slog.Default().With("scope", "sql"),
		// Logging every statement is only useful while debugging, and the instrumentation drops the records anyway unless the logger is at debug level
		QueryLog: common.EnvConfig.LogLevel == "debug",
		// Query parameter values can contain sensitive data, so they are optional and disabled by default
		IncludeParameters: common.EnvConfig.LogQueryArgs,
		// Slow statements are worth a warning at any log level
		SlowThreshold: 250 * time.Millisecond,
	}
}

// addSqliteDatetimeParams adds the datetime parameters to a SQLite connection string, leaving any the user set explicitly alone.
func addSqliteDatetimeParams(connString string) string {
	// sqliteDatetimeParams are the DSN parameters the Gorm SQLite driver injects when it opens the connection itself.
	// Pocket ID opens the connection instead, to instrument it, so they have to be set here or modernc.org/sqlite stops returning time.Time for datetime columns.
	// See injectDSNParams in github.com/libtnb/sqlite.
	var sqliteDatetimeParams = map[string]string{
		"_texttotime":  "1",
		"_inttotime":   "1",
		"_time_format": "sqlite",
	}

	path, rawQuery, _ := strings.Cut(connString, "?")
	if path == "" {
		// Return the connection string so the driver reports an error
		return connString
	}

	qs, err := url.ParseQuery(rawQuery)
	if err != nil {
		// Return the connection string so the driver reports an error
		return connString
	}

	for k, v := range sqliteDatetimeParams {
		if len(qs[k]) == 0 {
			qs.Set(k, v)
		}
	}

	return path + "?" + qs.Encode()
}