From f234e49f1288fe006334879d875e63c2e10986f7 Mon Sep 17 00:00:00 2001
From: Travis Ralston <travpc@gmail.com>
Date: Fri, 8 Sep 2023 17:27:47 -0600
Subject: [PATCH] Make "machine ID" easier to supply on import

---
 cmd/gdpr_import/main.go                 |  4 ++--
 cmd/import_synapse/main.go              |  3 ++-
 common/import_cmdline/ask_machine_id.go | 27 +++++++++++++++++++++++++
 util/ids/snowflake.go                   | 15 ++++++++++++++
 4 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 common/import_cmdline/ask_machine_id.go

diff --git a/cmd/gdpr_import/main.go b/cmd/gdpr_import/main.go
index 9ac4f720..25d248a2 100644
--- a/cmd/gdpr_import/main.go
+++ b/cmd/gdpr_import/main.go
@@ -1,7 +1,6 @@
 package main
 
 import (
-	"errors"
 	"flag"
 	"os"
 	"path"
@@ -10,6 +9,7 @@ import (
 	"github.com/turt2live/matrix-media-repo/archival/v2archive"
 	"github.com/turt2live/matrix-media-repo/common/assets"
 	"github.com/turt2live/matrix-media-repo/common/config"
+	"github.com/turt2live/matrix-media-repo/common/import_cmdline"
 	"github.com/turt2live/matrix-media-repo/common/logging"
 	"github.com/turt2live/matrix-media-repo/common/rcontext"
 	"github.com/turt2live/matrix-media-repo/common/runtime"
@@ -37,7 +37,7 @@ func main() {
 	assets.SetupMigrations(*migrationsPath)
 
 	if ids.GetMachineId() == 0 {
-		panic(errors.New("expected custom machine ID for import process (unsafe to import as Machine 0)"))
+		import_cmdline.AskMachineId()
 	}
 
 	var err error
diff --git a/cmd/import_synapse/main.go b/cmd/import_synapse/main.go
index ef086d1e..7d59f60d 100644
--- a/cmd/import_synapse/main.go
+++ b/cmd/import_synapse/main.go
@@ -15,6 +15,7 @@ import (
 	"github.com/sirupsen/logrus"
 	"github.com/turt2live/matrix-media-repo/common/assets"
 	"github.com/turt2live/matrix-media-repo/common/config"
+	"github.com/turt2live/matrix-media-repo/common/import_cmdline"
 	"github.com/turt2live/matrix-media-repo/common/logging"
 	"github.com/turt2live/matrix-media-repo/common/rcontext"
 	"github.com/turt2live/matrix-media-repo/common/runtime"
@@ -52,7 +53,7 @@ func main() {
 	assets.SetupMigrations(*migrationsPath)
 
 	if ids.GetMachineId() == 0 {
-		panic(errors.New("expected custom machine ID for import process (unsafe to import as Machine 0)"))
+		import_cmdline.AskMachineId()
 	}
 
 	var realPsqlPassword string
diff --git a/common/import_cmdline/ask_machine_id.go b/common/import_cmdline/ask_machine_id.go
new file mode 100644
index 00000000..069c4587
--- /dev/null
+++ b/common/import_cmdline/ask_machine_id.go
@@ -0,0 +1,27 @@
+package import_cmdline
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/turt2live/matrix-media-repo/util/ids"
+	"golang.org/x/term"
+)
+
+func AskMachineId() {
+	fmt.Println("The importer runs as a MMR worker and needs to have a dedicated MACHINE_ID. See https://docs.t2bot.io/matrix-media-repo/deployment/horizontal_scaling.html for details on what a MACHINE_ID is.")
+	if !term.IsTerminal(int(os.Stdin.Fd())) {
+		fmt.Println("Please specify a MACHINE_ID environment variable.")
+		os.Exit(2)
+		return // for good measure
+	}
+	fmt.Println("If you don't use horizontal scaling, you can use '1' as the machine ID. Otherwise, please enter an unused machine ID in your environment.")
+	fmt.Printf("Machine ID: ")
+	var machineId int64
+	if _, err := fmt.Scanf("%d", &machineId); err != nil {
+		panic(err)
+	}
+	if err := ids.SetMachineId(machineId); err != nil {
+		panic(err)
+	}
+}
diff --git a/util/ids/snowflake.go b/util/ids/snowflake.go
index 73f531ac..d942a712 100644
--- a/util/ids/snowflake.go
+++ b/util/ids/snowflake.go
@@ -1,6 +1,7 @@
 package ids
 
 import (
+	"errors"
 	"os"
 	"strconv"
 
@@ -30,3 +31,17 @@ func makeSnowflake() (*snowflake.Node, error) {
 	sfnode = node
 	return sfnode, nil
 }
+
+func SetMachineId(id int64) error {
+	if err := os.Setenv("MACHINE_ID", strconv.FormatInt(id, 10)); err != nil {
+		return err
+	}
+	sfnode = nil
+	if GetMachineId() != id {
+		return errors.New("unexpected error setting machine ID")
+	}
+	if _, err := makeSnowflake(); err != nil {
+		return err
+	}
+	return nil
+}