From 2a38b842f71f7dc49391af8835869fd66cb3d2fe Mon Sep 17 00:00:00 2001
From: Will Norris <will@tailscale.com>
Date: Thu, 16 Jun 2022 09:20:09 -0700
Subject: [PATCH] cmd/golink: enable sqlite and data migration

Add flags to specify sqlite database and migration from files to sqlite.
Migration is only supported in one direction (and flag name makes that
clear). If both sqlite and linkdir flags are provided, sqlite is
preferred.  This has no change in behavior if neither of the new flags
are provided.

Change-Id: Ie13116ec3b80835b4389d228583e7a1476c00bfb
---
 build.sh  |  2 +-
 golink.go | 95 ++++++++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 77 insertions(+), 20 deletions(-)

diff --git a/build.sh b/build.sh
index da8ec0d..4cd66fb 100755
--- a/build.sh
+++ b/build.sh
@@ -12,4 +12,4 @@ go run github.com/tailscale/mkctr \
   --tags="latest" \
   --repos="registry.fly.io/tsgo" \
   --push \
-  /tsgo -verbose -linkdir=/root/linkdata -mkdir
+  /tsgo -verbose -sqlitedb /root/golinks.db
diff --git a/golink.go b/golink.go
index 6f5f48f..0d530bd 100644
--- a/golink.go
+++ b/golink.go
@@ -29,10 +29,12 @@ import (
 )
 
 var (
-	verbose = flag.Bool("verbose", false, "be verbose")
-	linkDir = flag.String("linkdir", "", "the directory to store one JSON file per go/ shortlink")
-	dev     = flag.String("dev-listen", "", "if non-empty, listen on this addr and run in dev mode; auto-set linkDir if empty and don't use tsnet")
-	doMkdir = flag.Bool("mkdir", false, "whether to make --linkdir at start")
+	verbose    = flag.Bool("verbose", false, "be verbose")
+	linkDir    = flag.String("linkdir", "", "the directory to store one JSON file per go/ shortlink")
+	sqlitefile = flag.String("sqlitedb", "", "path of SQLite database to store links")
+	migrate    = flag.Bool("migrate-to-sqlite", false, "migrate link data from file storage to sqlite")
+	dev        = flag.String("dev-listen", "", "if non-empty, listen on this addr and run in dev mode; auto-set linkDir if empty and don't use tsnet")
+	doMkdir    = flag.Bool("mkdir", false, "whether to make --linkdir at start")
 )
 
 var stats struct {
@@ -57,23 +59,10 @@ var localClient *tailscale.LocalClient
 func main() {
 	flag.Parse()
 
-	if *linkDir == "" {
-		if devMode() {
-			var err error
-			*linkDir, err = ioutil.TempDir("", "golink_dev_*")
-			if err != nil {
-				log.Fatal(err)
-			}
-			log.Printf("Dev mode temp dir: %s", *linkDir)
-		} else {
-			log.Fatalf("--linkdir is required")
-		}
-	}
-
 	var err error
-	db, err = NewFileDB(*linkDir, *doMkdir)
+	db, err = setupDB()
 	if err != nil {
-		log.Fatalf("NewFileDB(%q): %v", *linkDir, err)
+		log.Fatalf("setting up database: %v", err)
 	}
 
 	if err := restoreLastSnapshot(); err != nil {
@@ -120,6 +109,74 @@ func main() {
 	}
 }
 
+// setupDB returns a DB used for link storage based on CLI flags and migrates
+// data if requested.  If flags are provided for both sqlite and file-base
+// storage, sqlite is preferred.
+func setupDB() (DB, error) {
+	if *sqlitefile == "" && *linkDir == "" && !devMode() {
+		return nil, errors.New("must specify linkdir or sqlitedb")
+	}
+
+	var sqliteDB *SQLiteDB
+	if *sqlitefile != "" {
+		var err error
+		if sqliteDB, err = NewSQLiteDB(*sqlitefile); err != nil {
+			return nil, fmt.Errorf("NewSQLiteDB(%q): %w", *sqlitefile, err)
+		}
+		if !*migrate {
+			// not migrating data, so return early
+			return sqliteDB, nil
+		}
+	}
+
+	if *linkDir == "" && devMode() {
+		var err error
+		*linkDir, err = ioutil.TempDir("", "golink_dev_*")
+		if err != nil {
+			return nil, err
+		}
+		log.Printf("Dev mode temp dir: %s", *linkDir)
+	}
+
+	var fileDB *FileDB
+	if *linkDir != "" {
+		var err error
+		if fileDB, err = NewFileDB(*linkDir, *doMkdir); err != nil {
+			return nil, fmt.Errorf("NewFileDB(%q): %w", *linkDir, err)
+		}
+	}
+
+	if *migrate {
+		if sqliteDB == nil || fileDB == nil {
+			return nil, errors.New("migrate-to-sqlite requires both linkdir and sqlitedb to be specified")
+		}
+
+		links, err := fileDB.LoadAll()
+		if err != nil {
+			return nil, err
+		}
+		for _, link := range links {
+			if err := sqliteDB.Save(link); err != nil {
+				return nil, err
+			}
+		}
+
+		stats, err := fileDB.LoadStats()
+		if err != nil {
+			return nil, err
+		}
+		if err := sqliteDB.SaveStats(stats); err != nil {
+			return nil, err
+		}
+	}
+
+	if sqliteDB != nil {
+		return sqliteDB, nil
+	} else {
+		return fileDB, nil
+	}
+}
+
 // homeTmpl is the template used by the http://go/ index page where you can
 // create or edit links.
 var homeTmpl *template.Template