memos/cmd/memos/main.go

package main

import (
	"context"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"

	"github.com/usememos/memos/internal/profile"
	"github.com/usememos/memos/internal/version"
	"github.com/usememos/memos/internal/webhook"
	"github.com/usememos/memos/server"
	"github.com/usememos/memos/store"
	"github.com/usememos/memos/store/db"
)

func initSlogDefault() {
	level, err := parseSlogLevel(viper.GetString("log-level"))
	if err != nil {
		slog.Warn("invalid log-level value, defaulting to info", "error", err)
	}
	slog.SetDefault(newLogger(level, os.Stderr))
}

var (
	rootCmd = &cobra.Command{
		Use:   "memos",
		Short: `An open source, lightweight note-taking service. Easily capture and share your great thoughts.`,
		Run: func(_ *cobra.Command, _ []string) {
			instanceProfile := &profile.Profile{
				Demo:        viper.GetBool("demo"),
				Addr:        viper.GetString("addr"),
				Port:        viper.GetInt("port"),
				UNIXSock:    viper.GetString("unix-sock"),
				Data:        viper.GetString("data"),
				Driver:      viper.GetString("driver"),
				DSN:         viper.GetString("dsn"),
				InstanceURL: viper.GetString("instance-url"),
			}
			instanceProfile.Version = version.GetCurrentVersion()
			instanceProfile.Commit = version.Commit
			webhook.AllowPrivateIPs = viper.GetBool("allow-private-webhooks")

			if err := instanceProfile.Validate(); err != nil {
				slog.Error("failed to validate profile", "error", err)
				return
			}

			ctx, cancel := context.WithCancel(context.Background())
			dbDriver, err := db.NewDBDriver(instanceProfile)
			if err != nil {
				cancel()
				slog.Error("failed to create db driver", "error", err)
				return
			}

			storeInstance := store.New(dbDriver, instanceProfile)
			if err := storeInstance.Migrate(ctx); err != nil {
				cancel()
				slog.Error("failed to migrate", "error", err)
				return
			}

			s, err := server.NewServer(ctx, instanceProfile, storeInstance)
			if err != nil {
				cancel()
				slog.Error("failed to create server", "error", err)
				return
			}

			c := make(chan os.Signal, 1)
			// Trigger graceful shutdown on SIGINT or SIGTERM.
			// The default signal sent by the `kill` command is SIGTERM,
			// which is taken as the graceful shutdown signal for many systems, eg., Kubernetes, Gunicorn.
			signal.Notify(c, os.Interrupt, syscall.SIGTERM)

			if err := s.Start(ctx); err != nil {
				if err != http.ErrServerClosed {
					slog.Error("failed to start server", "error", err)
					cancel()
					return
				}
			}

			printGreetings(instanceProfile)

			go func() {
				<-c
				s.Shutdown(ctx)
				cancel()
			}()

			// Wait for CTRL-C.
			<-ctx.Done()
		},
	}
	versionCmd = &cobra.Command{
		Use:   "version",
		Short: "Print the current Memos version",
		Run: func(_ *cobra.Command, _ []string) {
			fmt.Println(version.GetCurrentVersion())
		},
	}
)

func init() {
	cobra.OnInitialize(initSlogDefault)

	viper.SetDefault("demo", false)
	viper.SetDefault("driver", "sqlite")
	viper.SetDefault("port", 8081)

	rootCmd.PersistentFlags().Bool("demo", false, "enable demo mode")
	rootCmd.PersistentFlags().String("addr", "", "address of server")
	rootCmd.PersistentFlags().Int("port", 8081, "port of server")
	rootCmd.PersistentFlags().String("unix-sock", "", "path to the unix socket, overrides --addr and --port")
	rootCmd.PersistentFlags().String("data", "", "data directory")
	rootCmd.PersistentFlags().String("driver", "sqlite", "database driver")
	rootCmd.PersistentFlags().String("dsn", "", "database source name(aka. DSN)")
	rootCmd.PersistentFlags().String("instance-url", "", "the url of your memos instance")
	rootCmd.PersistentFlags().Bool("allow-private-webhooks", false, "allow webhook URLs to resolve to private/reserved IP addresses")
	rootCmd.PersistentFlags().String("log-level", "info", "log verbosity level (debug, info, warn, error)")

	if err := viper.BindPFlag("demo", rootCmd.PersistentFlags().Lookup("demo")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("addr", rootCmd.PersistentFlags().Lookup("addr")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("port", rootCmd.PersistentFlags().Lookup("port")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("unix-sock", rootCmd.PersistentFlags().Lookup("unix-sock")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("data", rootCmd.PersistentFlags().Lookup("data")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("driver", rootCmd.PersistentFlags().Lookup("driver")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("dsn", rootCmd.PersistentFlags().Lookup("dsn")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("instance-url", rootCmd.PersistentFlags().Lookup("instance-url")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("allow-private-webhooks", rootCmd.PersistentFlags().Lookup("allow-private-webhooks")); err != nil {
		panic(err)
	}
	if err := viper.BindPFlag("log-level", rootCmd.PersistentFlags().Lookup("log-level")); err != nil {
		panic(err)
	}

	viper.SetEnvPrefix("memos")
	viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
	viper.AutomaticEnv()

	rootCmd.AddCommand(versionCmd)
}

func printGreetings(profile *profile.Profile) {
	fmt.Printf("Memos %s started successfully!\n", profile.Version)

	if profile.Demo {
		fmt.Fprint(os.Stderr, "Development mode is enabled\n")
		if profile.DSN != "" {
			fmt.Fprintf(os.Stderr, "Database: %s\n", profile.DSN)
		}
	}

	// Server information
	fmt.Printf("Data directory: %s\n", profile.Data)
	fmt.Printf("Database driver: %s\n", profile.Driver)

	// Connection information
	if len(profile.UNIXSock) == 0 {
		if len(profile.Addr) == 0 {
			fmt.Printf("Server running on port %d\n", profile.Port)
			fmt.Printf("Access your memos at: http://localhost:%d\n", profile.Port)
		} else {
			fmt.Printf("Server running on %s:%d\n", profile.Addr, profile.Port)
			fmt.Printf("Access your memos at: http://%s:%d\n", profile.Addr, profile.Port)
		}
	} else {
		fmt.Printf("Server running on unix socket: %s\n", profile.UNIXSock)
	}

	// Access mode is derived from instance_url: set = public, unset = private.
	accessMode := "private"
	if profile.AllowAnonymous() {
		accessMode = "public"
	}
	fmt.Printf("Access mode: %s\n", accessMode)

	fmt.Println()
	fmt.Printf("Documentation: %s\n", "https://usememos.com")
	fmt.Printf("Source code: %s\n", "https://github.com/usememos/memos")
	fmt.Println("\nHappy note-taking!")
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}